Allow benchmarking larger-than-memory datasets

.gitignore

@@ -27,6 +27,7 @@ logging/
 ### JVM crashes
 hs_err_pid*
 replay_pid*
+java_pid*.hprof
 
 ### IntelliJ IDEA ###
 .idea

docs/benchmarking.md

@@ -115,3 +115,22 @@ To add a custom fvecs/ivecs dataset:
     ```
 
 For remote datasets, use `base_url` to specify where files should be downloaded from. The `${VAR}` and `${VAR:-default}` syntax is supported for environment variable expansion. See the example config for details.
+
+## BenchYAML configuration through system properties
+
+- `jvector.bench.index_cache_dir`: Set the directory to which the index is saved. It is recommended that this is a low-latency NVMe SSD. Defaults to `index_cache` (relative to the working directory).
+- `jvector.bench.work_dir_root_path`: Sets the directory to which the index is saved when not using the cache. A low-latency NVMe SSD is preferred. Defaults to the system temp directory.
+- `jvector.bench.dataset.mmap.enable`: Allow the dataset to be memory-mapped. See [Benchmarking in memory-constrained environments](#benchmarking-in-memory-constrained-environments). Defaults to `false`
+
+## Benchmarking in memory-constrained environments
+
+You may want to benchmark datasets that are too large to fit entirely in memory. This is the case if the system itself doesn't have enough memory, or if you want to measure the behaviour by artificially constraining the amount of memory available. By default, `BenchYAML` loads the entire dataset into memory, so this will cause an OOM. To avoid this, you can set the system property `jvector.bench.dataset.mmap.enable` to `true` which causes bench to read the dataset from a memory-mapped file instead of loading it upfront.
+
+Example:
+```sh
+java -cp ... -Xmx10g ... -Djvector.bench.dataset.mmap.enable=true io.github.jbellis.jvector.example.BenchYAML cohere-english-v3-10M
+
+# on Linux you can use cgroups to constrain total memory available to the program
+systemd-run --scope --user -p MemoryHigh=18G -p MemoryMax=20G -- \
+    java -cp ... -Xmx10g ... -Djvector.bench.dataset.mmap.enable=true io.github.jbellis.jvector.example.BenchYAML cohere-english-v3-10M
+```

docs/release notes/4.1.0/721.enhancement.md

@@ -0,0 +1,20 @@
+### Allow benchmarking larger-than-memory Datasets
+
+**Description**
+
+This PR series allows Jvector's `BenchYAML` benchmarking system to process datasets that don't fit in physical memory by memory-mapping the dataset's index vectors. See `docs/benchmarking.md` for the full `BenchYAML` documentation.
+
+Covers changes from #717, #718, #720, #721.
+
+**How to enable**
+
+Set the Java system property `jvector.bench.dataset.mmap.enable` to `true` while running `BenchYAML` normally.
+
+Examples:
+```sh
+java -cp ... -Xmx10g ... -Djvector.bench.dataset.mmap.enable=true io.github.jbellis.jvector.example.BenchYAML cohere-english-v3-10M
+
+# on Linux you can use cgroups to constrain total memory available to the program
+systemd-run --scope --user -p MemoryHigh=18G -p MemoryMax=20G -- \
+    java -cp ... -Xmx10g ... -Djvector.bench.dataset.mmap.enable=true io.github.jbellis.jvector.example.BenchYAML cohere-english-v3-10M
+```

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

@@ -17,6 +17,7 @@
 package io.github.jbellis.jvector.example;
 
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet;
+import io.github.jbellis.jvector.example.benchmarks.datasets.DataSetInfo;
 import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
 import io.github.jbellis.jvector.example.reporting.RunArtifacts;
 import io.github.jbellis.jvector.example.reporting.SearchReportingCatalog;
@@ -118,9 +119,17 @@ public static void main(String[] args) throws IOException {
         for (var config : allConfigs) {
 
             String datasetName = config.dataset;
-            DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow(
+            DataSetInfo dsInfo = DataSets.loadDataSet(datasetName).orElseThrow(
                     () -> new RuntimeException("Could not load dataset:" + datasetName)
-            ).getDataSet();
+            );
+            DataSet ds;
+            if (Boolean.getBoolean("jvector.bench.dataset.mmap.enable")) {
+                // Memory-maps index vectors instead of loading them
+                ds = dsInfo.getMappedDataSet();
+            } else {
+                // Loads all index vectors into memory
+                ds = dsInfo.getDataSet();
+            }
             // Register dataset info the first time we actually load the dataset for benchmarking
             artifacts.registerDataset(datasetName, ds);
 

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

@@ -78,7 +78,7 @@ public class Grid {
 
     private static final String pqCacheDir = "pq_cache";
 
-    private static final String indexCacheDir = "index_cache";
+    private static final String indexCacheDir = System.getProperty("jvector.bench.index_cache_dir", "index_cache");
 
     private static final String dirPrefix = "BenchGraphDir";
 
@@ -95,6 +95,14 @@ public static Double getIndexBuildTimeSeconds(String datasetName) {
         return indexBuildTimes.get(datasetName);
     }
 
+    private static Path createTmpWorkDir() throws IOException {
+        String maybeRootPath = System.getProperty("jvector.bench.work_dir_root_path");
+        if (maybeRootPath != null) {
+            return Files.createTempDirectory(Path.of(maybeRootPath), dirPrefix);
+        }
+        return Files.createTempDirectory(dirPrefix);
+    }
+
     static void runAll(DataSet ds,
                        boolean enableIndexCache,
                        List<Integer> mGrid,
@@ -113,7 +121,7 @@ static void runAll(DataSet ds,
         boolean success = false;
 
         // Always use a fresh temp directory for per-run artifacts
-        final Path workDir = Files.createTempDirectory(dirPrefix);
+        final Path workDir = createTmpWorkDir();
 
         // Initialize index cache (creates stable directory when enabled, cleans up stale temp files; no-op otherwise)
         final OnDiskGraphIndexCache cache =
@@ -512,12 +520,14 @@ private static BuilderWithSuppliers builderWithSuppliers(Set<FeatureId> features
         var builder = new RandomAccessOnDiskGraphIndexWriter.Builder(onHeapGraph, outPath);
         builder.withMapper(identityMapper);
 
+        var vv = floatVectors.threadLocalSupplier();
+
         Map<FeatureId, IntFunction<Feature.State>> suppliers = new EnumMap<>(FeatureId.class);
         for (var featureId : features) {
             switch (featureId) {
                 case INLINE_VECTORS:
                     builder.with(new InlineVectors(floatVectors.dimension()));
-                    suppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(floatVectors.getVector(ordinal)));
+                    suppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(vv.get().getVector(ordinal)));
                     break;
                 case FUSED_PQ:
                     if (pq == null) {
@@ -533,7 +543,7 @@ private static BuilderWithSuppliers builderWithSuppliers(Set<FeatureId> features
                             ? constructionMetrics.index("NVQ").timeCompute(() -> NVQuantization.compute(floatVectors, nSubVectors))
                             : NVQuantization.compute(floatVectors, nSubVectors);
                     builder.with(new NVQ(nvq));
-                    suppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(floatVectors.getVector(ordinal))));
+                    suppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(vv.get().getVector(ordinal))));
                     break;
                 default:
                     break;

jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java

@@ -24,4 +24,7 @@ public interface DataSetInfo extends DataSetProperties {
     ///
     /// @return the ready-to-use {@link InMemoryDataSet}
     public InMemoryDataSet getDataSet();
+
+    /// @return a {@link DataSet} whose index vectors are backed by memory-mapped files.
+    public DataSet getMappedDataSet();
 }

jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfoMFD.java

@@ -19,6 +19,8 @@
 import io.github.jbellis.jvector.example.util.SiftLoader;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 
+import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.util.Optional;
 
 /// A lightweight, lazy handle that separates *identifying* a dataset from *loading* its data.
@@ -136,4 +138,33 @@ public InMemoryDataSet getDataSet() {
         }
         return cached;
     }
+
+    /// Returns a {@link DataSet} whose index vectors are backed by memory-mapped files.
+    /// The query vectors and ground truth are still loaded into memory.
+    ///
+    /// The output is not cached, so multiple invocoations will cause duplicate loads
+    /// and create duplicate maps.
+    ///
+    /// @return the ready-to-use {@link DataSet}
+    public RavvDataSet getMappedDataSet() {
+        if (baseProperties.loadBehavior() == LoadBehavior.LEGACY_SCRUB) {
+            throw new UnsupportedOperationException(
+                "Can't map a dataset that wants to be scrubbed, load the full dataset with getDataSet()"
+            );
+        }
+        var maybeVsf = baseProperties.similarityFunction();
+        if (maybeVsf.isEmpty()) {
+            throw new RuntimeException("Need a similarity function to create dataset");
+        }
+        try {
+            return new RavvDataSet(
+                baseProperties.getName(),
+                maybeVsf.get(),
+                FvecRavv.of(dsFiles.getBaseFvecsPath()),
+                SiftLoader.readFvecs(dsFiles.getQueryFvecsPath().toString()),
+                SiftLoader.readIvecs(dsFiles.getGtIvecsPath().toString()));
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
 }

jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/RavvDataSet.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.example.benchmarks.datasets;
+
+import java.util.List;
+
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+public class RavvDataSet implements DataSet {
+    private final String name;
+    private final VectorSimilarityFunction similarityFunction;
+    private final RandomAccessVectorValues baseRavv;
+    private final List<VectorFloat<?>> queryVectors;
+    private final List<? extends List<Integer>> groundTruth;
+
+    public RavvDataSet(
+        String name,
+        VectorSimilarityFunction similarityFunction,
+        RandomAccessVectorValues baseRavv,
+        List<VectorFloat<?>> queryVectors,
+        List<? extends List<Integer>> groundTruth
+    ) {
+        if (queryVectors.isEmpty()) {
+            throw new IllegalArgumentException("Query vectors must not be empty");
+        }
+        if (groundTruth.isEmpty()) {
+            throw new IllegalArgumentException("Ground truth vectors must not be empty");
+        }
+
+        if (baseRavv.dimension() != queryVectors.get(0).length()) {
+            throw new IllegalArgumentException("Base and query vectors must have the same dimensionality");
+        }
+        if (queryVectors.size() != groundTruth.size()) {
+            throw new IllegalArgumentException("Query and ground truth lists must be the same size");
+        }
+
+        this.name = name;
+        this.similarityFunction = similarityFunction;
+        this.baseRavv = baseRavv;
+        this.queryVectors = queryVectors;
+        this.groundTruth = groundTruth;
+
+        System.out.format("%n%s: %d base and %d query vectors created, dimensions %d%n",
+                name, baseRavv.size(), queryVectors.size(), baseRavv.dimension());
+    }
+
+    @Override
+    public int getDimension() {
+        return baseRavv.dimension();
+    }
+
+    @Override
+    public RandomAccessVectorValues getBaseRavv() {
+        return baseRavv;
+    }
+
+    @Override
+    public String getName() {
+        return name;
+    }
+
+    @Override
+    public VectorSimilarityFunction getSimilarityFunction() {
+        return similarityFunction;
+    }
+
+    @Override
+    public List<VectorFloat<?>> getQueryVectors() {
+        return queryVectors;
+    }
+
+    @Override
+    public List<? extends List<Integer>> getGroundTruth() {
+        return groundTruth;
+    }
+}