Skip to content

adding memory and disk usage stats to bench tests #591

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MarkWolters wants to merge 1 commit into
base: main
Choose a base branch
from
Open

adding memory and disk usage stats to bench tests #591

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 61 additions & 7 deletions jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,13 @@ public class Grid {

private static int diagnostic_level;

/**
* Get the index build time for a dataset
*/
public static Double getIndexBuildTime(String datasetName) {
return indexBuildTimes.get(datasetName);
Copy link
Contributor

@jshook jshook Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps clarify the unit here, such as getIndexBuildTimeMillis or getIndexBuildTimeSeconds

Copy link
Contributor Author

@MarkWolters MarkWolters Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good point, I will add that.

}

static void runAll(DataSet ds,
List<Integer> mGrid,
List<Integer> efConstructionGrid,
Expand Down Expand Up @@ -158,13 +165,21 @@ static void runOneGraph(List<? extends Set<FeatureId>> featureSets,
DataSet ds,
Path testDirectory) throws IOException
{
// Capture initial memory and disk state
var diagnostics = new io.github.jbellis.jvector.example.benchmarks.diagnostics.BenchmarkDiagnostics(getDiagnosticLevel());
diagnostics.setMonitoredDirectory(testDirectory);
Copy link
Contributor

@jshook jshook Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does import not work for this symbol?

Copy link
Contributor Author

@MarkWolters MarkWolters Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do not recall why this is fully qualified but I will clean it up

diagnostics.capturePrePhaseSnapshot("Graph Build");

Map<Set<FeatureId>, ImmutableGraphIndex> indexes;
if (buildCompressor == null) {
indexes = buildInMemory(featureSets, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, testDirectory);
} else {
indexes = buildOnDisk(featureSets, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, testDirectory, buildCompressor);
}

// Capture post-build memory and disk state
diagnostics.capturePostPhaseSnapshot("Graph Build");

try {
for (var cpSupplier : compressionGrid) {
indexes.forEach((features, index) -> {
Expand All @@ -188,7 +203,7 @@ static void runOneGraph(List<? extends Set<FeatureId>> featureSets,
}

try (var cs = new ConfiguredSystem(ds, index, cv, featureSetForIndex)) {
testConfiguration(cs, topKGrid, usePruningGrid, M, efConstruction, neighborOverflow, addHierarchy, benchmarks);
testConfiguration(cs, topKGrid, usePruningGrid, M, efConstruction, neighborOverflow, addHierarchy, benchmarks, testDirectory);
} catch (Exception e) {
throw new RuntimeException(e);
}
Expand All @@ -197,6 +212,11 @@ static void runOneGraph(List<? extends Set<FeatureId>> featureSets,
for (var index : indexes.values()) {
index.close();
}

// Log final diagnostics summary
if (diagnostic_level > 0) {
diagnostics.logSummary();
}
} finally {
for (int n = 0; n < featureSets.size(); n++) {
Files.deleteIfExists(testDirectory.resolve("graph" + n));
Expand Down Expand Up @@ -432,13 +452,14 @@ private static void testConfiguration(ConfiguredSystem cs,
int efConstruction,
float neighborOverflow,
boolean addHierarchy,
Map<String, List<String>> benchmarkSpec) {
Map<String, List<String>> benchmarkSpec,
Path testDirectory) {
int queryRuns = 2;
System.out.format("Using %s:%n", cs.index);
// 1) Select benchmarks to run. Use .createDefault or .createEmpty (for other options)

var benchmarks = setupBenchmarks(benchmarkSpec);
QueryTester tester = new QueryTester(benchmarks);
QueryTester tester = new QueryTester(benchmarks, testDirectory, cs.ds.name);

// 2) Setup benchmark table for printing
for (var topK : topKGrid.keySet()) {
Expand Down Expand Up @@ -563,11 +584,22 @@ public static List<BenchResult> runAllAndCollectResults(
for (Function<DataSet, CompressorParameters> searchCompressor : compressionGrid) {
Path testDirectory = Files.createTempDirectory("bench");
try {
// Capture initial state
var diagnostics = new io.github.jbellis.jvector.example.benchmarks.diagnostics.BenchmarkDiagnostics(getDiagnosticLevel());
diagnostics.setMonitoredDirectory(testDirectory);
diagnostics.capturePrePhaseSnapshot("Build");

var compressor = getCompressor(buildCompressor, ds);
var searchCompressorObj = getCompressor(searchCompressor, ds);
CompressedVectors cvArg = (searchCompressorObj instanceof CompressedVectors) ? (CompressedVectors) searchCompressorObj : null;
var indexes = buildOnDisk(List.of(features), m, ef, neighborOverflow, addHierarchy, false, ds, testDirectory, compressor);
ImmutableGraphIndex index = indexes.get(features);

// Capture post-build state
diagnostics.capturePostPhaseSnapshot("Build");
var buildSnapshot = diagnostics.getLatestSystemSnapshot();
var buildDiskSnapshot = diagnostics.getLatestDiskSnapshot();

try (ConfiguredSystem cs = new ConfiguredSystem(ds, index, cvArg, features)) {
int queryRuns = 2;
List<QueryBenchmark> benchmarks = List.of(
Expand All @@ -578,7 +610,7 @@ public static List<BenchResult> runAllAndCollectResults(
CountBenchmark.createDefault(),
AccuracyBenchmark.createDefault()
);
QueryTester tester = new QueryTester(benchmarks);
QueryTester tester = new QueryTester(benchmarks, testDirectory, ds.name);
for (int topK : topKGrid.keySet()) {
for (boolean usePruning : usePruningGrid) {
for (double overquery : topKGrid.get(topK)) {
Expand All @@ -596,11 +628,33 @@ public static List<BenchResult> runAllAndCollectResults(
"overquery", overquery,
"usePruning", usePruning
);
// Collect all metrics including memory and disk usage
Map<String, Object> allMetrics = new HashMap<>();
for (Metric metric : metricsList) {
Map<String, Object> metrics = java.util.Map.of(metric.getHeader(), metric.getValue());
results.add(new BenchResult(ds.name, params, metrics));
allMetrics.put(metric.getHeader(), metric.getValue());
}

// Add build time if available
if (indexBuildTimes.containsKey(ds.name)) {
allMetrics.put("Index Build Time", indexBuildTimes.get(ds.name));
}

// Add memory metrics if available
if (buildSnapshot != null) {
allMetrics.put("Heap Memory Used (MB)", buildSnapshot.memoryStats.heapUsed / 1024.0 / 1024.0);
allMetrics.put("Heap Memory Max (MB)", buildSnapshot.memoryStats.heapMax / 1024.0 / 1024.0);
allMetrics.put("Off-Heap Direct (MB)", buildSnapshot.memoryStats.directBufferMemory / 1024.0 / 1024.0);
allMetrics.put("Off-Heap Mapped (MB)", buildSnapshot.memoryStats.mappedBufferMemory / 1024.0 / 1024.0);
allMetrics.put("Total Off-Heap (MB)", buildSnapshot.memoryStats.getTotalOffHeapMemory() / 1024.0 / 1024.0);
}

// Add disk metrics if available
if (buildDiskSnapshot != null) {
allMetrics.put("Disk Usage (MB)", buildDiskSnapshot.totalBytes / 1024.0 / 1024.0);
allMetrics.put("File Count", buildDiskSnapshot.fileCount);
}
results.add(new BenchResult(ds.name, params, Map.of("Index Build Time", indexBuildTimes.get(ds.name))));

results.add(new BenchResult(ds.name, params, allMetrics));
}
}
}
Expand Down
65 changes: 65 additions & 0 deletions jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/QueryTester.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,25 +16,49 @@

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

import java.nio.file.Path;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import io.github.jbellis.jvector.example.Grid;
import io.github.jbellis.jvector.example.Grid.ConfiguredSystem;
import io.github.jbellis.jvector.example.benchmarks.diagnostics.BenchmarkDiagnostics;

/**
* Orchestrates running a set of QueryBenchmark instances
* and collects their summary results.
*/
public class QueryTester {
private final List<QueryBenchmark> benchmarks;
private final Path monitoredDirectory;
private final String datasetName;

/**
* @param benchmarks the benchmarks to run, in the order provided
*/
public QueryTester(List<QueryBenchmark> benchmarks) {
this(benchmarks, null, null);
}

/**
* @param benchmarks the benchmarks to run, in the order provided
* @param monitoredDirectory optional directory to monitor for disk usage
*/
public QueryTester(List<QueryBenchmark> benchmarks, Path monitoredDirectory) {
this(benchmarks, monitoredDirectory, null);
}

/**
* @param benchmarks the benchmarks to run, in the order provided
* @param monitoredDirectory optional directory to monitor for disk usage
* @param datasetName optional dataset name for retrieving build time
*/
public QueryTester(List<QueryBenchmark> benchmarks, Path monitoredDirectory, String datasetName) {
this.benchmarks = benchmarks;
this.monitoredDirectory = monitoredDirectory;
this.datasetName = datasetName;
}

/**
Expand All @@ -56,11 +80,52 @@ public List<Metric> run(

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

// Capture memory and disk usage before running queries
// Use NONE level to suppress logging output that would break the table
var diagnostics = new BenchmarkDiagnostics(io.github.jbellis.jvector.example.benchmarks.diagnostics.DiagnosticLevel.NONE);
if (monitoredDirectory != null) {
diagnostics.setMonitoredDirectory(monitoredDirectory);
}
diagnostics.capturePrePhaseSnapshot("Query");

for (var benchmark : benchmarks) {
var metrics = benchmark.runBenchmark(cs, topK, rerankK, usePruning, queryRuns);
results.addAll(metrics);
}

// Capture memory and disk usage after running queries
diagnostics.capturePostPhaseSnapshot("Query");

// Add memory and disk metrics to results
var systemSnapshot = diagnostics.getLatestSystemSnapshot();
var diskSnapshot = diagnostics.getLatestDiskSnapshot();

if (systemSnapshot != null) {
// Max heap usage in MB
results.add(Metric.of("Max heap usage", ".1f",
systemSnapshot.memoryStats.heapUsed / (1024.0 * 1024.0)));

// Max off-heap usage (direct + mapped) in MB
results.add(Metric.of("Max offheap usage", ".1f",
systemSnapshot.memoryStats.getTotalOffHeapMemory() / (1024.0 * 1024.0)));
}

if (diskSnapshot != null) {
// Total file size in MB
results.add(Metric.of("Total file size", ".1f",
diskSnapshot.totalBytes / (1024.0 * 1024.0)));

// Number of files
results.add(Metric.of("Number of files", ".0f",
(double) diskSnapshot.fileCount));
}

// Add index build time if available
if (datasetName != null && Grid.getIndexBuildTime(datasetName) != null) {
results.add(Metric.of("Index build time", ".2f",
Grid.getIndexBuildTime(datasetName)));
}

return results;
}
}
Expand Down
Loading