Skip to content

Commit 6ca3b5e

Browse files
MarkWoltersMarkWolters
authored
adding memory and disk usage stats to bench tests (#591)
* adding memory and disk usage stats to bench tests
1 parent a66fd91 commit 6ca3b5e

5 files changed

Lines changed: 708 additions & 21 deletions

File tree

‎jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java‎

Lines changed: 62 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import io.github.jbellis.jvector.example.benchmarks.QueryBenchmark;
2727
import io.github.jbellis.jvector.example.benchmarks.QueryTester;
2828
import io.github.jbellis.jvector.example.benchmarks.ThroughputBenchmark;
29+
import io.github.jbellis.jvector.example.benchmarks.diagnostics.BenchmarkDiagnostics;
2930
import io.github.jbellis.jvector.example.benchmarks.diagnostics.DiagnosticLevel;
3031
import io.github.jbellis.jvector.example.util.CompressorParameters;
3132
import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
@@ -86,6 +87,13 @@ public class Grid {
8687

8788
private static int diagnostic_level;
8889

90+
/**
91+
* Get the index build time for a dataset
92+
*/
93+
public static Double getIndexBuildTimeSeconds(String datasetName) {
94+
return indexBuildTimes.get(datasetName);
95+
}
96+
8997
static void runAll(DataSet ds,
9098
List<Integer> mGrid,
9199
List<Integer> efConstructionGrid,
@@ -157,13 +165,21 @@ static void runOneGraph(List<? extends Set<FeatureId>> featureSets,
157165
DataSet ds,
158166
Path testDirectory) throws IOException
159167
{
168+
// Capture initial memory and disk state
169+
var diagnostics = new BenchmarkDiagnostics(getDiagnosticLevel());
170+
diagnostics.setMonitoredDirectory(testDirectory);
171+
diagnostics.capturePrePhaseSnapshot("Graph Build");
172+
160173
Map<Set<FeatureId>, ImmutableGraphIndex> indexes;
161174
if (buildCompressor == null) {
162175
indexes = buildInMemory(featureSets, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, testDirectory);
163176
} else {
164177
indexes = buildOnDisk(featureSets, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, testDirectory, buildCompressor);
165178
}
166179

180+
// Capture post-build memory and disk state
181+
diagnostics.capturePostPhaseSnapshot("Graph Build");
182+
167183
try {
168184
for (var cpSupplier : compressionGrid) {
169185
indexes.forEach((features, index) -> {
@@ -187,7 +203,7 @@ static void runOneGraph(List<? extends Set<FeatureId>> featureSets,
187203
}
188204

189205
try (var cs = new ConfiguredSystem(ds, index, cv, featureSetForIndex)) {
190-
testConfiguration(cs, topKGrid, usePruningGrid, M, efConstruction, neighborOverflow, addHierarchy, benchmarks);
206+
testConfiguration(cs, topKGrid, usePruningGrid, M, efConstruction, neighborOverflow, addHierarchy, benchmarks, testDirectory);
191207
} catch (Exception e) {
192208
throw new RuntimeException(e);
193209
}
@@ -196,6 +212,11 @@ static void runOneGraph(List<? extends Set<FeatureId>> featureSets,
196212
for (var index : indexes.values()) {
197213
index.close();
198214
}
215+
216+
// Log final diagnostics summary
217+
if (diagnostic_level > 0) {
218+
diagnostics.logSummary();
219+
}
199220
} finally {
200221
for (int n = 0; n < featureSets.size(); n++) {
201222
Files.deleteIfExists(testDirectory.resolve("graph" + n));
@@ -431,13 +452,14 @@ private static void testConfiguration(ConfiguredSystem cs,
431452
int efConstruction,
432453
float neighborOverflow,
433454
boolean addHierarchy,
434-
Map<String, List<String>> benchmarkSpec) {
455+
Map<String, List<String>> benchmarkSpec,
456+
Path testDirectory) {
435457
int queryRuns = 2;
436458
System.out.format("Using %s:%n", cs.index);
437459
// 1) Select benchmarks to run. Use .createDefault or .createEmpty (for other options)
438460

439461
var benchmarks = setupBenchmarks(benchmarkSpec);
440-
QueryTester tester = new QueryTester(benchmarks);
462+
QueryTester tester = new QueryTester(benchmarks, testDirectory, cs.ds.getName());
441463

442464
// 2) Setup benchmark table for printing
443465
for (var topK : topKGrid.keySet()) {
@@ -562,11 +584,22 @@ public static List<BenchResult> runAllAndCollectResults(
562584
for (Function<DataSet, CompressorParameters> searchCompressor : compressionGrid) {
563585
Path testDirectory = Files.createTempDirectory("bench");
564586
try {
587+
// Capture initial state
588+
var diagnostics = new io.github.jbellis.jvector.example.benchmarks.diagnostics.BenchmarkDiagnostics(getDiagnosticLevel());
589+
diagnostics.setMonitoredDirectory(testDirectory);
590+
diagnostics.capturePrePhaseSnapshot("Build");
591+
565592
var compressor = getCompressor(buildCompressor, ds);
566593
var searchCompressorObj = getCompressor(searchCompressor, ds);
567594
CompressedVectors cvArg = (searchCompressorObj instanceof CompressedVectors) ? (CompressedVectors) searchCompressorObj : null;
568595
var indexes = buildOnDisk(List.of(features), m, ef, neighborOverflow, addHierarchy, false, ds, testDirectory, compressor);
569596
ImmutableGraphIndex index = indexes.get(features);
597+
598+
// Capture post-build state
599+
diagnostics.capturePostPhaseSnapshot("Build");
600+
var buildSnapshot = diagnostics.getLatestSystemSnapshot();
601+
var buildDiskSnapshot = diagnostics.getLatestDiskSnapshot();
602+
570603
try (ConfiguredSystem cs = new ConfiguredSystem(ds, index, cvArg, features)) {
571604
int queryRuns = 2;
572605
List<QueryBenchmark> benchmarks = List.of(
@@ -578,7 +611,7 @@ public static List<BenchResult> runAllAndCollectResults(
578611
AccuracyBenchmark.createDefault(),
579612
CountBenchmark.createDefault()
580613
);
581-
QueryTester tester = new QueryTester(benchmarks);
614+
QueryTester tester = new QueryTester(benchmarks, testDirectory, ds.getName());
582615
for (int topK : topKGrid.keySet()) {
583616
for (boolean usePruning : usePruningGrid) {
584617
for (double overquery : topKGrid.get(topK)) {
@@ -596,11 +629,33 @@ public static List<BenchResult> runAllAndCollectResults(
596629
"overquery", overquery,
597630
"usePruning", usePruning
598631
);
632+
// Collect all metrics including memory and disk usage
633+
Map<String, Object> allMetrics = new HashMap<>();
599634
for (Metric metric : metricsList) {
600-
Map<String, Object> metrics = java.util.Map.of(metric.getHeader(), metric.getValue());
601-
results.add(new BenchResult(ds.getName(), params, metrics));
635+
allMetrics.put(metric.getHeader(), metric.getValue());
636+
}
637+
638+
// Add build time if available
639+
if (indexBuildTimes.containsKey(ds.getName())) {
640+
allMetrics.put("Index Build Time", indexBuildTimes.get(ds.getName()));
641+
}
642+
643+
// Add memory metrics if available
644+
if (buildSnapshot != null) {
645+
allMetrics.put("Heap Memory Used (MB)", buildSnapshot.memoryStats.heapUsed / 1024.0 / 1024.0);
646+
allMetrics.put("Heap Memory Max (MB)", buildSnapshot.memoryStats.heapMax / 1024.0 / 1024.0);
647+
allMetrics.put("Off-Heap Direct (MB)", buildSnapshot.memoryStats.directBufferMemory / 1024.0 / 1024.0);
648+
allMetrics.put("Off-Heap Mapped (MB)", buildSnapshot.memoryStats.mappedBufferMemory / 1024.0 / 1024.0);
649+
allMetrics.put("Total Off-Heap (MB)", buildSnapshot.memoryStats.getTotalOffHeapMemory() / 1024.0 / 1024.0);
650+
}
651+
652+
// Add disk metrics if available
653+
if (buildDiskSnapshot != null) {
654+
allMetrics.put("Disk Usage (MB)", buildDiskSnapshot.totalBytes / 1024.0 / 1024.0);
655+
allMetrics.put("File Count", buildDiskSnapshot.fileCount);
602656
}
603-
results.add(new BenchResult(ds.getName(), params, Map.of("Index Build Time", indexBuildTimes.get(ds.getName()))));
657+
658+
results.add(new BenchResult(ds.getName(), params, allMetrics));
604659
}
605660
}
606661
}

‎jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/QueryTester.java‎

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,25 +16,48 @@
1616

1717
package io.github.jbellis.jvector.example.benchmarks;
1818

19+
import java.io.IOException;
20+
import java.nio.file.Path;
1921
import java.util.ArrayList;
20-
import java.util.LinkedHashMap;
2122
import java.util.List;
22-
import java.util.Map;
2323

24+
import io.github.jbellis.jvector.example.Grid;
2425
import io.github.jbellis.jvector.example.Grid.ConfiguredSystem;
26+
import io.github.jbellis.jvector.example.benchmarks.diagnostics.BenchmarkDiagnostics;
2527

2628
/**
2729
* Orchestrates running a set of QueryBenchmark instances
2830
* and collects their summary results.
2931
*/
3032
public class QueryTester {
3133
private final List<QueryBenchmark> benchmarks;
34+
private final Path monitoredDirectory;
35+
private final String datasetName;
3236

3337
/**
3438
* @param benchmarks the benchmarks to run, in the order provided
3539
*/
3640
public QueryTester(List<QueryBenchmark> benchmarks) {
41+
this(benchmarks, null, null);
42+
}
43+
44+
/**
45+
* @param benchmarks the benchmarks to run, in the order provided
46+
* @param monitoredDirectory optional directory to monitor for disk usage
47+
*/
48+
public QueryTester(List<QueryBenchmark> benchmarks, Path monitoredDirectory) {
49+
this(benchmarks, monitoredDirectory, null);
50+
}
51+
52+
/**
53+
* @param benchmarks the benchmarks to run, in the order provided
54+
* @param monitoredDirectory optional directory to monitor for disk usage
55+
* @param datasetName optional dataset name for retrieving build time
56+
*/
57+
public QueryTester(List<QueryBenchmark> benchmarks, Path monitoredDirectory, String datasetName) {
3758
this.benchmarks = benchmarks;
59+
this.monitoredDirectory = monitoredDirectory;
60+
this.datasetName = datasetName;
3861
}
3962

4063
/**
@@ -56,11 +79,56 @@ public List<Metric> run(
5679

5780
List<Metric> results = new ArrayList<>();
5881

82+
// Capture memory and disk usage before running queries
83+
// Use NONE level to suppress logging output that would break the table
84+
var diagnostics = new BenchmarkDiagnostics(io.github.jbellis.jvector.example.benchmarks.diagnostics.DiagnosticLevel.NONE);
85+
if (monitoredDirectory != null) {
86+
try {
87+
diagnostics.setMonitoredDirectory(monitoredDirectory);
88+
} catch (IOException e) {
89+
throw new RuntimeException(e);
90+
}
91+
}
92+
diagnostics.capturePrePhaseSnapshot("Query");
93+
5994
for (var benchmark : benchmarks) {
6095
var metrics = benchmark.runBenchmark(cs, topK, rerankK, usePruning, queryRuns);
6196
results.addAll(metrics);
6297
}
6398

99+
// Capture memory and disk usage after running queries
100+
diagnostics.capturePostPhaseSnapshot("Query");
101+
102+
// Add memory and disk metrics to results
103+
var systemSnapshot = diagnostics.getLatestSystemSnapshot();
104+
var diskSnapshot = diagnostics.getLatestDiskSnapshot();
105+
106+
if (systemSnapshot != null) {
107+
// Max heap usage in MB
108+
results.add(Metric.of("Max heap usage", ".1f",
109+
systemSnapshot.memoryStats.heapUsed / (1024.0 * 1024.0)));
110+
111+
// Max off-heap usage (direct + mapped) in MB
112+
results.add(Metric.of("Max offheap usage", ".1f",
113+
systemSnapshot.memoryStats.getTotalOffHeapMemory() / (1024.0 * 1024.0)));
114+
}
115+
116+
if (diskSnapshot != null) {
117+
// Total file size in MB
118+
results.add(Metric.of("Total file size", ".1f",
119+
diskSnapshot.totalBytes / (1024.0 * 1024.0)));
120+
121+
// Number of files
122+
results.add(Metric.of("Number of files", ".0f",
123+
(double) diskSnapshot.fileCount));
124+
}
125+
126+
// Add index build time if available
127+
if (datasetName != null && Grid.getIndexBuildTimeSeconds(datasetName) != null) {
128+
results.add(Metric.of("Index build time", ".2f",
129+
Grid.getIndexBuildTimeSeconds(datasetName)));
130+
}
131+
64132
return results;
65133
}
66134
}

0 commit comments

Comments
 (0)