Skip to content

[benchmark] Refactor diskann-inmem benchmarks - #1510

Open
Mark Hildebrand (hildebrandmw) wants to merge 22 commits into
mainfrom
mhildebr/inmem-benchmark
Open

Mark Hildebrand (hildebrandmw) wants to merge 22 commits into
mainfrom
mhildebr/inmem-benchmark

Conversation

@hildebrandmw

Copy link
Copy Markdown
Contributor

With diskann-inmem inching closer to replacing diskann-providers, we need it to replace the benchmarking roles that diskann-providers currently 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-benchmark and diskann-benchmark-runner:

  • Benchmark inputs used to not be automatically registered by benchmarks. This led to the pattern seen in diskann-benchmarks of 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 Raw vs Input split, which minimized the amount of validation parsing could do to what could be handled inline (like InputFile resolution).

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, build contains the index configuration (alpha, prune degrees, l_build, etc.) as its input. The parsed Config input directly contains the diskann::graph::Config that will be passed to the index. We don't have the awkward validation step + recreation step that is done for inmem1.

Then, build::Config is directly responsible for packaging the diskann::graph::Config with a DataProvider into an Engine, 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 how InplaceDeleteStrategy, 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 how diskann-inmem works. diskann-inmem uses a single Quantized Representation for 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:

  • I'm really tired of either writing out our really long crate names (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_runner becomes dbr, diskann_quantization becomes dq etc. You will see that all over this PR and I like it quite a bit.
  • I made aggressive use of exhaustive destructuring when moving data from DTOs and displaying. This has a nice property of alerting you when you add or change a field of all the bits that need updating.
  • I did not use the components in diskann_benchmark_core::search::graph or diskann_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 in diskann-inmem's benchmark through this process instead of expanding this refactor all over the place.

Future plans include:

  • Adding regression tests for diskann-inmem benchmarks. This is looonng overdue.
  • Filling out support for filtered and range search.
  • Getting to perf parity (or as close as I can get) with diskann-providers before officially deprecating the old providers.

@hildebrandmw
Mark Hildebrand (hildebrandmw) requested review from a team and a balanced review from Copilot October 24, 2026 00:08
@hildebrandmw Mark Hildebrand (hildebrandmw) changed the title [benchmark] Take diskann-inmem benchmarks more seriously [benchmark] Refactor diskann-inmem benchmarks Oct 24, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 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

Comment thread diskann-benchmark/src/index/inmem2/benchmark/knn.rs
Comment thread diskann-benchmark/src/index/inmem2/full.rs
Comment thread diskann-benchmark/src/index/inmem2/full.rs Outdated
@codecov-commenter

Codecov Comments Bot (codecov-commenter) commented Oct 24, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 68.67168% with 500 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.13%. Comparing base (8376a57) to head (20687e3).

Files with missing lines Patch % Lines
...ann-benchmark/src/index/inmem2/benchmark/bigann.rs 0.00% 363 Missing ⚠️
diskann-benchmark/src/index/inmem2/full.rs 60.00% 62 Missing ⚠️
...kann-benchmark/src/index/inmem2/benchmark/build.rs 89.88% 18 Missing ⚠️
...iskann-benchmark/src/index/inmem2/benchmark/knn.rs 96.73% 14 Missing ⚠️
...iskann-benchmark/src/index/inmem2/quantized/mod.rs 91.54% 12 Missing ⚠️
...ann-benchmark/src/index/inmem2/quantized/scalar.rs 84.21% 9 Missing ⚠️
...-benchmark/src/index/inmem2/quantized/spherical.rs 90.32% 6 Missing ⚠️
diskann-benchmark/src/index/streaming/stats.rs 0.00% 6 Missing ⚠️
diskann-benchmark/src/utils/datafiles.rs 54.54% 5 Missing ⚠️
...nn-benchmark/src/index/inmem2/quantized/product.rs 96.10% 3 Missing ⚠️
... and 1 more

❌ 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

Impacted file tree graph

@@            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     
Flag Coverage Δ
miri 92.13% <68.67%> (-0.04%) ⬇️
unittests 92.07% <68.67%> (-0.04%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
diskann-benchmark-core/src/search/api.rs 98.91% <ø> (ø)
diskann-benchmark-runner/src/utils/fmt.rs 98.85% <100.00%> (+0.05%) ⬆️
diskann-benchmark/src/index/build.rs 82.75% <100.00%> (-1.81%) ⬇️
diskann-benchmark/src/index/inmem2/mod.rs 100.00% <100.00%> (ø)
diskann-benchmark/src/index/mod.rs 100.00% <ø> (ø)
diskann-benchmark/src/utils/mod.rs 95.71% <ø> (ø)
diskann-benchmark/src/utils/progress.rs 100.00% <100.00%> (ø)
diskann-benchmark/src/utils/recall.rs 97.43% <ø> (ø)
diskann-benchmark/src/index/inmem2/common.rs 97.05% <97.05%> (ø)
...nn-benchmark/src/index/inmem2/quantized/product.rs 96.10% <96.10%> (ø)
... and 9 more

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants