Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
ANN-benchmarks: switch to use gbench (#1661)
Make the ANN benchmarks use the same google benchmark infrastructure as the prim benchmarks while keeping the functional changes minimal.
### Overview
- The command-line API largely stays the same, but enhanced with gbench-specific parameters, such as using regex to select algo configs, control the minimum run-time, and flexible reporting to console/files.
- There's just one executable `ANN_BENCH`, all of the algorithms are loaded as shared libraries. The CPU-only components do not require cuda at runtime (ANN_BENCH itself, hnswlib).
- Some dependencies are linked statically, it's possible to just copy the executable and the libs and run the benchmark on a linux machine with very few packages installed.
- Search benchmarks do not produce any output anymore, they use ground truth files to compute and report the recall in-place.
- Search/build parameters visible in the config files are passed as benchmark counters/labels/context.
- Extra functionality:
- `--data_prefix` to specify a custom path where the data sets are stored
- `--index_prefix` to specify a custom path where the index sets are stored
- `--override_kv=<key:value1:value2:...:valueN>` override one or more parameters of search/build for parameter-sweep benchmarks
__Breaking change__: the behavior of the ANN benchmark executables (library API is not touched). The executable CLI flags have changed, so the newer, adapted wrapper scripts won't work with the executables from the libraft-ann-bench-23.08 conda package.
### A primer
```bash
./cpp/build/ANN_BENCH \ # benchmark executable
--data_prefix=/datastore/my/local/data/path \ # override (prefix) path to local data
--benchmark_min_warmup_time=0.001 \ # spend some minimal time warming up
--benchmark_min_time=3s \ # run minimum 3 seconds on each case
--benchmark_out=ivf_pq.csv \ # duplicate output to this file
--benchmark_out_format=csv \ # the file output should be in CSV format
--benchmark_counters_tabular \ # the console output should be tabular
--benchmark_filter="raft_ivf_pq\..*" \ # use regex to filter benchmarks
--search \ # 'search' mode
--override_kv=k:1:10:100:200:500 \ # Parameter-sweep over the top-k value
--override_kv=n_queries:1:10:10000 \ # and the search batch size
--override_kv=smemLutDtype:"fp8" \ # Override a search parameter
cpp/bench/ann/conf/bigann-100M.json # specify the path to the config file
```
### Motivation
#### Eliminate huge bug-prone configs
The current config fixes the batch size and k to one value per-config, so the whole config needs to be copied to try multiple values. In the PR, both these parameters can be overwritten in the search parameters and/or via command line (`ANN_BENCH --override_kv=n_queries:1:100:1000 --override_kv=k:1:10:20:50:100:200:500:1000` would test all combinations in one go). Any of the build/search parameters can be overwritten at the same time.
#### Run the benchmarks and aggregate the data in the minimal environment
The new executable generates reports with QPS, Recall, and other metrics using gbench. Hence there's no need to copy back and forth dozens of result files and no need to install python environment for running or evaluating. A single CSV or JSON can be produced for all algorithms and run configurations per dataset+hardware pair.
#### Speedup the benchmarks
The current benchmark framework is extremely slow due to two factors:
- The dataset and the index need to be loaded for every test case, this takes orders of magnitude longer than the search test itself for large datasets. In my tests, the preparation phase for bigann-1B took ten minutes and the search could take anywhere between a few seconds and a minute.
- The benchmark always goes through the whole query dataset. That is, if the query set is 10K and the batch size is 1, the benchmark repeats 10K times (to produce the result file for evaluating the recall).
In the proposed solution, a user can set the desired time or number of iterations to run; the data is loaded only once and the index is cached between the search test cases. My subjective conservative estimate is the overall speedup of more than x100 for running a typical large-scale benchmark.
#### Better measurement of QPS
By default, the current benchmark reports the average execution time and does not warm-up iterations. As a result, the first test case on most of our plots is distorted (e.g. the first iteration of the first case takes about a second or two to run, and that significantly affects the average of the rest 999 ~100us iterations). `gbench` provides the `--benchmark_min_warmup_time` parameters to skip first one or few iterations, which solves the problem.
#### Extra context in the reports
The new benchmark executable uses gbench context to augment the report with some essential information: base and query set name, dimensionality, and size, distance metric, some CPU and GPU info, CUDA version. All this is appended directly to the generated CSV/JSON files, which makes the bookkeeping much easier.
In addition, a user may pass extra context via command line `--benchmark_context=<key>=<value>`; this could be e.g. the hostname, some ENV variables, etc.
#### Easier profiling
Thanks to flexible regex filtering and parameter overriding, now it's possible to specify a subset of cases and an exact number of times they should run. This makes the profiling using such tools as `nsys` and `ncu` much easier.
Authors:
- Artem M. Chirkin (https://github.com/achirkin)
Approvers:
- Corey J. Nolet (https://github.com/cjnolet)
URL: #1661- Loading branch information
