Skip to content

fix for disk usage measurements #692

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

fix for disk usage measurements #692

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
47 changes: 34 additions & 13 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 @@ -230,7 +230,6 @@ static void runOneGraph(OnDiskGraphIndexCache cache,
// Prepare to collect index construction metrics for reporting....
var constructionMetrics = new ConstructionMetrics();

// TODO this does not capture disk usage for cached indexes. Need to update
// Capture initial memory and disk state
try (var diagnostics = new BenchmarkDiagnostics(getDiagnosticLevel())) {
diagnostics.startMonitoring("testDirectory", workDirectory);
Expand All @@ -253,6 +252,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache,
(buildCompressorObj == null) ? "None" : String.valueOf(buildCompressorObj);

Map<Set<FeatureId>, ImmutableGraphIndex> indexes = new HashMap<>();
Map<Set<FeatureId>, Long> indexFileSizes = new HashMap<>();
if (buildCompressorObj == null) {
indexes = buildInMemory(featureSets, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, workDirectory);
} else {
Expand All @@ -270,6 +270,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache,
if (cached.isPresent()) {
System.out.printf("%s: Using cached graph index for %s%n", key.datasetName, fs);
indexes.put(fs, cached.get());
try { indexFileSizes.put(fs, Files.size(cache.resolve(key).finalPath)); } catch (IOException ignored) {}
} else {
missing.add(fs);
if (cache.isEnabled()) {
Expand All @@ -287,9 +288,10 @@ static void runOneGraph(OnDiskGraphIndexCache cache,
if (!missing.isEmpty()) {
// At least one index needs to be built (b/c not in cache or cache is disabled)
// We pass the handles map so buildOnDisk knows exactly where to write
var newIndexes = buildOnDisk(missing, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph,
var result = buildOnDisk(missing, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph,
ds, outputDir, buildCompressorObj, handles, constructionMetrics);
indexes.putAll(newIndexes);
indexes.putAll(result.indexes);
indexFileSizes.putAll(result.fileSizes);
}
}

Expand All @@ -303,6 +305,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache,
try {
for (var cpSupplier : compressionGrid) {
indexes.forEach((features, index) -> {
constructionMetrics.indexFileSizeBytes = indexFileSizes.get(features);
final Set<FeatureId> featureSetForIndex = index instanceof OnDiskGraphIndex ? ((OnDiskGraphIndex) index).getFeatureSet() : Set.of();

CompressedVectors cv;
Expand Down Expand Up @@ -364,7 +367,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache,
}
}

private static Map<Set<FeatureId>, ImmutableGraphIndex> buildOnDisk(List<? extends Set<FeatureId>> featureSets,
private static BuildOnDiskResult buildOnDisk(List<? extends Set<FeatureId>> featureSets,
int M,
int efConstruction,
float neighborOverflow,
Expand Down Expand Up @@ -464,18 +467,20 @@ private static Map<Set<FeatureId>, ImmutableGraphIndex> buildOnDisk(List<? exten
entry.getValue().commit();
}

// open indexes
// open indexes and capture per-feature-set file sizes
Map<Set<FeatureId>, ImmutableGraphIndex> indexes = new HashMap<>();
Map<Set<FeatureId>, Long> fileSizes = new HashMap<>();
n = 0;
for (var features : featureSets) {
Path loadPath = handles.containsKey(features)
? handles.get(features).finalPath()
: outputDir.resolve("graph" + n++);

indexes.put(features, OnDiskGraphIndex.load(ReaderSupplierFactory.open(loadPath)));
fileSizes.put(features, Files.size(loadPath));
}

return indexes;
return new BuildOnDiskResult(indexes, fileSizes);
}

private static BuilderWithSuppliers builderWithSuppliers(Set<FeatureId> features,
Expand Down Expand Up @@ -845,6 +850,7 @@ public static List<BenchResult> runAllAndCollectResults(
diagnostics.startMonitoring("indexCache", Paths.get(indexCacheDir));
diagnostics.capturePrePhaseSnapshot("Build");
Map<Set<FeatureId>, ImmutableGraphIndex> indexes = new HashMap<>();
Map<Set<FeatureId>, Long> indexFileSizes = new HashMap<>();

var compressor = getCompressor(buildCompressor, ds);
var searchCompressorObj = getCompressor(searchCompressor, ds);
Expand Down Expand Up @@ -885,6 +891,7 @@ public static List<BenchResult> runAllAndCollectResults(
if (cached.isPresent()) {
System.out.printf("%s: Using cached graph index for %s%n", key.datasetName, fs);
indexes.put(fs, cached.get());
try { indexFileSizes.put(fs, Files.size(cache.resolve(key).finalPath)); } catch (IOException ignored) {}
} else {
missing.add(fs);
if (cache.isEnabled()) {
Expand All @@ -903,9 +910,10 @@ public static List<BenchResult> runAllAndCollectResults(
if (!missing.isEmpty()) {
// At least one index needs to be built (b/c not in cache or cache is disabled)
// We pass the handles map so buildOnDisk knows exactly where to write
var newIndexes = buildOnDisk(missing, m, ef, neighborOverflow, addHierarchy, refineFinalGraph,
var result = buildOnDisk(missing, m, ef, neighborOverflow, addHierarchy, refineFinalGraph,
ds, outputDir, compressor, handles, null);
indexes.putAll(newIndexes);
indexes.putAll(result.indexes);
indexFileSizes.putAll(result.fileSizes);
}

ImmutableGraphIndex index = indexes.get(features);
Expand All @@ -914,7 +922,6 @@ public static List<BenchResult> runAllAndCollectResults(
diagnostics.capturePostPhaseSnapshot("Build");
diagnostics.printDiskStatistics("Graph Index Build");
var buildSnapshot = diagnostics.getLatestSystemSnapshot();
DiskUsageMonitor.MultiDirectorySnapshot buildDiskSnapshot = diagnostics.getLatestDiskSnapshot();

try (ConfiguredSystem cs = new ConfiguredSystem(ds, index, cvArg, features)) {
int queryRuns = 2;
Expand Down Expand Up @@ -966,10 +973,10 @@ public static List<BenchResult> runAllAndCollectResults(
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.getTotalBytes() / 1024.0 / 1024.0);
allMetrics.put("File Count", buildDiskSnapshot.getTotalFileCount());
// Add per-feature-set index file size
Long fileSizeBytes = indexFileSizes.get(features);
if (fileSizeBytes != null) {
allMetrics.put("Disk Usage (MB)", fileSizeBytes / 1024.0 / 1024.0);
}

results.add(new BenchResult(ds.getName(), params, allMetrics));
Expand Down Expand Up @@ -1161,9 +1168,19 @@ private static String quantTypeOf(CompressorParameters cp) {
* <p>Call {@link #appendTo(List)} to emit these as {@link Metric} entries so they can be selected for
* console/CSV reporting.</p>
*/
private static final class BuildOnDiskResult {
final Map<Set<FeatureId>, ImmutableGraphIndex> indexes;
final Map<Set<FeatureId>, Long> fileSizes;
BuildOnDiskResult(Map<Set<FeatureId>, ImmutableGraphIndex> indexes, Map<Set<FeatureId>, Long> fileSizes) {
this.indexes = indexes;
this.fileSizes = fileSizes;
}
}

static final class ConstructionMetrics {
// null means “not applicable / not measured”
public Double indexBuildTimeS;
public Long indexFileSizeBytes;

// Index-construction phase (single pass per run)
private final Map<String, QuantStats> indexByType = new HashMap<>();
Expand Down Expand Up @@ -1193,6 +1210,10 @@ void appendTo(List<Metric> out) {
out.add(Metric.of("construction.index_build_time_s",
"Index build time (s)", ".2f", indexBuildTimeS));
}
if (indexFileSizeBytes != null) {
out.add(Metric.of("construction.index_file_size_mb",
"Index file size (MB)", ".2f", indexFileSizeBytes / 1024.0 / 1024.0));
}

// Index-construction quant timings
appendQuant(out, "construction.index_quant_time_s", "Idx", indexByType);
Expand Down
Loading