Make diskann-disk runtime-agnostic with async APIs#1184
Conversation
Previously diskann-disk created and owned its own tokio runtime: a create_runtime helper, a Runtime field on DiskIndexSearcher, and block_on calls inside build()/search(). This forced a runtime on callers, risked nested-runtime panics, and pulled in tokio's full feature set. Library: - DiskIndexBuilder::build() is now async; the redundant build_internal is inlined (BuildInterrupted handling preserved via an inner async block). Deleted build/builder/tokio.rs (create_runtime). - DiskIndexSearcher no longer holds a Runtime; new() drops the Option<Runtime> param; search()/search_internal() are async and await the inner futures directly. - Trimmed tokio features from ["full"] to ["rt","rt-multi-thread","macros","sync"]. - Tests: single-call search tests use #[tokio::test]; the two concurrent rayon tests share one runtime Handle; build helpers drive the async build via a local runtime. Consumers (diskann-tools, diskann-benchmark): - Search loops replace rayon for_each_in_pool with futures_util::buffer_unordered(num_threads) on a caller-owned runtime. - Build paths wrap build() in block_on on a multi-thread runtime. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1184 +/- ##
==========================================
- Coverage 89.47% 89.45% -0.02%
==========================================
Files 486 486
Lines 92161 92141 -20
==========================================
- Hits 82458 82429 -29
- Misses 9703 9712 +9
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
|
|
||
| let timer = Timer::new(); | ||
| disk_index.build()?; | ||
| let runtime = tokio::runtime::Builder::new_multi_thread() |
There was a problem hiding this comment.
For awareness — build_disk_index.rs and search_disk_index.rs are both dead code.
I've opened #1185 to remove them.
| load_query_result(params.storage_provider, params.truth_result_file_path); | ||
|
|
||
| let pool = create_thread_pool(params.thread_num.into_usize()).unwrap(); | ||
| let runtime = tokio::runtime::Builder::new_multi_thread().build().unwrap(); |
There was a problem hiding this comment.
The parallelism here comes from the rayon for_each_in_pool, and the search path never spawn or suspend, so the tokio workers stay idle.
new_current_thread() would be enough for search, since the runtime is only here to block_on.
| } | ||
|
|
||
| pub fn build(&mut self) -> ANNResult<()> { | ||
| let runtime = create_runtime(self.index_configuration.num_threads)?; |
There was a problem hiding this comment.
Worth noting: build parallelism is governed by two numbers — num_tasks (the tokio tasks spawned via JoinSet) and the runtime's worker-thread count. The old create_runtime(num_threads) derived both from num_threads, so they stayed in lockstep.
Now the runtime is caller-owned, so the two are independent: if the caller sizes worker threads below num_tasks, the build still runs but with reduced parallelism. Might be worth a note on build() about the expected worker-thread count.
| logger.log_checkpoint("index_loaded"); | ||
|
|
||
| let pool = create_thread_pool(search_params.num_threads)?; | ||
| let runtime = tokio::runtime::Builder::new_current_thread().build()?; |
There was a problem hiding this comment.
The CI benchmark shows QPS dropping ~70% on this path (run).
The cause looks like the current_thread runtime + buffer_unordered: the search futures never suspend (get_element is future::ready), so on a single worker thread they run one after another — buffer_unordered(num_threads) gives no real parallelism here.
There was a problem hiding this comment.
We may want to consider going benchmark core. It will manage the tokio parallelism.
4 participants
Summary
Makes
diskann-diskruntime-agnostic. Previously the crate created and owned its own tokio runtime (acreate_runtimehelper, aRuntimefield onDiskIndexSearcher, andblock_oncalls insidebuild()/search()). This forced a runtime on callers, risked nested-runtime panics, and pulled in tokio'sfullfeature set. The crate now exposesasyncAPIs and lets callers drive them on their own executor.Library changes (
diskann-disk)DiskIndexBuilder::build()is nowasync; the redundantbuild_internalis inlined (BuildInterruptedhandling preserved via an innerasyncblock). Deletedbuild/builder/tokio.rs(create_runtime).DiskIndexSearcherno longer holds aRuntime;new()drops theOption<Runtime>parameter;search()/search_internal()areasyncand await the inner futures directly.["full"]to["rt", "rt-multi-thread", "macros", "sync"].#[tokio::test]; the two concurrent rayon tests share one runtimeHandle; build helpers drive the async build via a local runtime.Consumer changes (
diskann-tools,diskann-benchmark)for_each_in_poolwithfutures_util::buffer_unordered(num_threads)on a caller-owned runtime.build()inblock_onon a multi-thread runtime.Validation
cargo check --workspaceclean.diskann-disk: 271 lib tests + 2 doc tests pass.cargo clippy -D warningsclean ondiskann-disk,diskann-tools, anddiskann-benchmark(disk-indexfeature).cargo fmt --allclean.Note on performance
The consumer search loops now use single-task
buffer_unorderedconcurrency rather than the previousnum_threads-wide rayon pool. This yields concurrency (interleaving at await points) but not CPU parallelism; for the CPU-bound search path this may reduce multi-core QPS in the benchmark. Restoring true parallelism would requireArc<searcher>+JoinSet/tokio::spawnwithSend + Sync + 'staticbounds on the generic search functions.🤖 Generated with GitHub Copilot CLI