Repository navigation
[benchmark] Refactor diskann-inmem benchmarks - #1510
Mark Hildebrand (hildebrandmw) wants to merge 22 commits into
Conversation
diskann-inmem benchmarks more seriouslydiskann-inmem benchmarks
There was a problem hiding this comment.
🟡 Changes recommended
The refactor introduces unresolved regressions in saved input metadata, supported datatype dispatch, and search-window validation.
3 open issues
What changed in this PR
Refactors diskann-inmem benchmarking into operation-focused modules, preparing it to replace the benchmarking role of diskann-providers.
Changes:
- Separates raw inputs from validated configurations for builds, KNN sweeps, and BigANN streaming.
- Shares quantized benchmark execution while isolating product, scalar, and spherical training.
- Adds structured statistics and shared formatting/progress helpers, and updates example inputs.
| File | Description |
|---|---|
diskann-benchmark/src/utils/recall.rs |
Enables recall-metric deserialization. |
diskann-benchmark/src/utils/progress.rs |
Extracts reusable build progress reporting. |
diskann-benchmark/src/utils/mod.rs |
Exposes progress utilities and nonzero constants. |
diskann-benchmark/src/utils/datafiles.rs |
Adds dataset type validation. |
diskann-benchmark/src/index/streaming/stats.rs |
Exposes timing and latency statistics. |
diskann-benchmark/src/index/mod.rs |
Updates feature-gated benchmark entries. |
diskann-benchmark/src/index/inmem2/quantized/spherical.rs |
Encapsulates spherical configuration and training. |
diskann-benchmark/src/index/inmem2/quantized/scalar.rs |
Encapsulates scalar configuration and training. |
diskann-benchmark/src/index/inmem2/quantized/product.rs |
Encapsulates product configuration and training. |
diskann-benchmark/src/index/inmem2/quantized/mod.rs |
Shares quantized build-and-search execution. |
diskann-benchmark/src/index/inmem2/mod.rs |
Organizes modules and benchmark registration. |
diskann-benchmark/src/index/inmem2/full.rs |
Connects full-precision inputs to benchmark operations. |
diskann-benchmark/src/index/inmem2/common.rs |
Shares dataset loading and parsing helpers. |
diskann-benchmark/src/index/inmem2/benchmark/mod.rs |
Defines combined build/search statistics. |
diskann-benchmark/src/index/inmem2/benchmark/knn.rs |
Implements search plans, execution, and aggregation. |
diskann-benchmark/src/index/inmem2/benchmark/build.rs |
Encapsulates build configuration, execution, and statistics. |
diskann-benchmark/src/index/inmem2/benchmark/bigann.rs |
Implements streaming runbook execution and reporting. |
diskann-benchmark/src/index/inmem2.rs |
Removes the monolithic benchmark implementation. |
diskann-benchmark/src/index/build.rs |
Uses the shared progress helper. |
diskann-benchmark/example/inmem2.json |
Migrates examples to the new input schemas. |
diskann-benchmark-runner/src/utils/fmt.rs |
Adds lazy debug-value formatting and tests. |
diskann-benchmark-core/src/search/api.rs |
Makes search runs cloneable. |
Review effort: Balanced
Codecov Report❌ Patch coverage is ❌ Your patch status has failed because the patch coverage (68.67%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #1510 +/- ##
==========================================
- Coverage 92.16% 92.13% -0.04%
==========================================
Files 592 602 +10
Lines 123582 124059 +477
==========================================
+ Hits 113899 114299 +400
- Misses 9683 9760 +77
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|


With
diskann-inmeminching closer to replacingdiskann-providers, we need it to replace the benchmarking roles thatdiskann-providerscurrently fulfills. The benchmark infra I slapped together earlier in development is not really sufficient. This PR overhauls it to be a bit more suitable for our long-term needs.The existing inmem benchmarks inherit a lot of structure from early designs of
diskann-benchmarkanddiskann-benchmark-runner:Benchmark inputs used to not be automatically registered by benchmarks. This led to the pattern seen in
diskann-benchmarksof all the inputs being dumped in one module while all the benchmarks are distributed throughout the crate.In my opinion, this makes benchmarks harder to encapsulate (both via feature gating and in the usual programming sense) since all the fields of a the benchmark inputs need to be more-or-less reachable from a different part of the crate.
Benchmark inputs also did not used to have a
RawvsInputsplit, which minimized the amount of validation parsing could do to what could be handled inline (likeInputFileresolution).With this PR, I'm exploring what we can do if we don't have these constraints. The idea is to have benchmark operations (here
build,knn-search, and bigann style streaming workloads) be more self-contained. The first thing to get out of the way is that I aggressively split benchmark inputs between raw DTOs and a parsed validation (even in some situations where there is no difference between the DTO and the non-DTO). This allows input parsing to be much more potent.For example,
buildcontains the index configuration (alpha, prune degrees,l_build, etc.) as its input. The parsedConfiginput directly contains thediskann::graph::Configthat will be passed to the index. We don't have the awkward validation step + recreation step that is done for inmem1.Then,
build::Configis directly responsible for packaging thediskann::graph::Configwith aDataProviderinto anEngine, which is then responsible for actually driving build operations. The point being that all build related stuff from input parsing, display, execution, and result aggregation are all contained within a single file. This makes it much more difficult for implementations to drift where only a subset of fields are used, rendering the rest silently impotent.With that in mind, the benchmark refactor has the following split:
inmem2/benchmark/*: Parsing, execution, and aggregation of benchmark operations.build.rs: One-shot index builds.knn.rs: Knn search sweeps and result aggregation.bigann.rs: BigANN style streaming workloads. Due to limitations with howInplaceDeleteStrategy, this is currently only sanely implementable for full-precision data. Part of the motivation of the inmem refactor is to lift this restriction.inmem2/full.rs: Input plumbing for full-precision one-shot builds and BigANN streaming. Note that the implementation is quite declarative.inmem2/quantization/*: This mirrors howdiskann-inmemworks.diskann-inmemuses a singleQuantizedRepresentationfor the quantizers, and the benchmarks do something similar. The quantizers themselves are mainly concerned with just building the quantizer from training data - everything else flows through common code.I'm not entirely sure how this will expand to handle situations where we have quantizer-specific parameters (e.g. the query layout for spherical), but I'm not too worried about how that will play out.
A few other comments about stylistic approaches I tried:
diskann_benchmark_runner,diskann_quantization, etc.) or importing all the things. I realized our crates have dissimilar enough names that the initialization is surprisingly unambiguous.diskann_benchmark_runnerbecomesdbr,diskann_quantizationbecomesdqetc. You will see that all over this PR and I like it quite a bit.diskann_benchmark_core::search::graphordiskann_benchmark_core::build::graph, instead using the search and build APIs manually. This is in anticipation of some future changes to how strategies work and how I plan to pass quantizer specific parameters. I think it makes more sense to have a little bit of duplication indiskann-inmem's benchmark through this process instead of expanding this refactor all over the place.Future plans include:
diskann-inmembenchmarks. This is looonng overdue.diskann-providersbefore officially deprecating the old providers.