diff --git a/.gitignore b/.gitignore index 70cede6f9..5460869f8 100644 --- a/.gitignore +++ b/.gitignore @@ -27,6 +27,7 @@ logging/ ### JVM crashes hs_err_pid* replay_pid* +java_pid*.hprof ### IntelliJ IDEA ### .idea diff --git a/docs/benchmarking.md b/docs/benchmarking.md index 79ddba82e..39d508d81 100644 --- a/docs/benchmarking.md +++ b/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 +``` diff --git a/docs/release notes/4.1.0/721.enhancement.md b/docs/release notes/4.1.0/721.enhancement.md new file mode 100644 index 000000000..3f7888e32 --- /dev/null +++ b/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 +``` diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java index e066a34dc..5a2a85733 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java +++ b/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); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index af5e54347..8f03188a4 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -79,7 +79,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"; @@ -96,6 +96,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 mGrid, @@ -114,7 +122,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 = @@ -495,12 +503,14 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features var builder = new OnDiskGraphIndexWriter.Builder(onHeapGraph, outPath); builder.withMapper(identityMapper); + var vv = floatVectors.threadLocalSupplier(); + Map> 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) { @@ -516,7 +526,7 @@ private static BuilderWithSuppliers builderWithSuppliers(Set 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; diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java index a18d63491..c219affe3 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.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); + } + } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/RavvDataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/RavvDataSet.java new file mode 100644 index 000000000..36c826a91 --- /dev/null +++ b/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> queryVectors; + private final List> groundTruth; + + public RavvDataSet( + String name, + VectorSimilarityFunction similarityFunction, + RandomAccessVectorValues baseRavv, + List> queryVectors, + List> 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> getQueryVectors() { + return queryVectors; + } + + @Override + public List> getGroundTruth() { + return groundTruth; + } +}