[WIP] Refactor raft-ann-bench: single CLI entrypoint, importable Python module, config validation, and testing #2182
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No breaking changes were made to the existing module
raft-ann-bench
.The proposal is for migrating functionalities of
raft-ann-bench
to the refactored module.Key features of the refactoring:
raft-ann-bench.run
and make it easier to write tests against. Currently themain()
andrun_build_and_search()
functions inpython/raft-ann-bench/src/raft-ann-bench/run/__main__.py
are rather difficult to test, as they include quite a bit of parameter parsing and config data processing logic. The refactor would break up the functions into smaller functions.Pydantic
to model the 4 types of config files. This makes config data parsing more robust, asPydantic
performs validation, and it helps remove quite a bit of data parsing code. The 4 types of config files modeled:datasets.yaml
asDatasetConfig
,algos.yaml
asAlgoLibConfig
, algo benchmark config (e.g.raft_cagra.yaml
) asAlgoConfig
, and lastly the benchmark executable's config (e.g. the deprecateddeep-image-96-inner.json
) asBenchConfig
.Typer
to organize the CLI commands in a single entrypoint. This would be helpful for incrementally migrating other functionalities, such asraft-ann-bench.plot
. Furthermore, documentation for the CLI can be auto-generated withtyper raft_ann_bench.cli utils docs --name raft-ann-bench --output raft_ann_benchmarks_cli.md
.cpp/bench/ann/src/common/benchmark.hpp
), and design the CLI to merely passthrough parameters to the ANN executable (e.g.raft-ann-bench run build
).raft-ann-bench
cannot be easily imported in Python; replacing it withraft_ann_bench
will work.Closes #2153
Closes #2151
Closes #2150
Closes #2149