Skip to content

Fix CAGRA dataset descriptor race across threads sharing a handle #2807

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bdice wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix CAGRA dataset descriptor race across threads sharing a handle #2807

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
456b98d
Fix CAGRA dataset descriptor race across threads sharing a handle
bdice Oct 10, 2026
733caab
Make the CAGRA descriptor cache lookup-or-insert atomic
bdice Oct 10, 2026
cd97df7
Correct the descriptor cache comment: the LRU cache does not deduplic…
bdice Oct 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions cpp/src/neighbors/detail/cagra/compute_distance.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,10 +253,13 @@ struct dataset_descriptor_host {
{
if (!ready.load(std::memory_order_acquire)) { eval(stream); }
// value is immutable at this point.
auto& [ptr, ready_stream] = std::get<ready_t>(value);
if (ready_stream != stream) {
RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0));
}
auto* ptr = std::get<0>(std::get<ready_t>(value));
// Always order `stream` after the initialization kernel. Comparing the stream handles is not
// enough: the per-thread default stream (`cudaStreamPerThread`, which is also raft's default
// stream) has the same handle in every host thread but refers to a different stream in each.
// When several threads share a raft::resources (and thus this cached descriptor), skipping
// the wait lets a search kernel read the descriptor before it is initialized.
RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0));
Comment on lines +256 to +262

@bdice bdice Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This is the fix (commit 1). ready_t stores the cuda::stream_ref that ran the init kernel (eval). The old code compared it with the caller's stream and skipped the wait when they were equal. That comparison is on handle values, and with raft's default cudaStreamPerThread every thread has the same handle value but its own stream. The wait is now unconditional. When the caller really is on the init stream, the event has already been recorded on it, so the extra cudaStreamWaitEvent adds no ordering. The filter payload cache already does it this way (cagra_filter_payload.hpp).

I checked this commit alone before CI: I patched the nightly 26.10 libcuvs.so so state::get always waits. On a local GPU, AnnMG/5:/8 went from 4/6 processes aborting to 0/6, and AnnMG/14 (ROUND_ROBIN) went from 34/200 failing tests to 0/200 (#2783 comment).

return ptr;
}
};
Expand Down
18 changes: 15 additions & 3 deletions cpp/src/neighbors/detail/cagra/factory.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@

#include <cuvs/neighbors/common.hpp>

#include <mutex>

namespace cuvs::neighbors::cagra::detail {

template <typename DataT,
Expand Down Expand Up @@ -147,6 +149,8 @@ struct store {
raft::cache::
lru<key, key_hash, std::equal_to<>, dataset_descriptor_host<DataT, IndexT, DistanceT>>
value{kDefaultSize};
/** Serializes lookup-or-insert, so that concurrent searches share a single descriptor. */
std::mutex mutex;
Comment on lines +152 to +153

@bdice bdice Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

The mutex lives in the per-raft::resources custom resource, next to the cache, so it only serializes searches that share a handle. Searches on different handles have separate caches and never contend.

Question for reviewers: One lifetime hazard is not changed here. When the last copy of a descriptor is dropped (for example after LRU eviction), ~state() frees it with cudaFreeAsync on the stream that ran the init kernel (compute_distance.hpp). That free is not ordered against a search kernel another thread may still have queued on a different stream with that descriptor. Is this worth a follow-up, or is it unreachable in practice with the 100-entry cache?

};

} // namespace descriptor_cache
Expand All @@ -170,9 +174,17 @@ auto dataset_descriptor_init_with_cache(const raft::resources& res,
-> dataset_descriptor_host<DataT, IndexT, DistanceT>
{
auto key = descriptor_cache::make_key(params, dataset, metric, dataset_norms);
auto& cache =
raft::resource::get_custom_resource<descriptor_cache::store<DataT, IndexT, DistanceT>>(res)
->value;
auto* store =
raft::resource::get_custom_resource<descriptor_cache::store<DataT, IndexT, DistanceT>>(res);
auto& cache = store->value;
// Hold the lock across the lookup and the insertion. Otherwise threads that miss concurrently
// each create a descriptor and insert it under the same key (the LRU cache does not deduplicate
// keys), so concurrent searches can end up with different descriptors. When the last owner of
// one releases it, its device memory is freed with `cudaFreeAsync` on the stream that
// initialized it, which is not ordered after a kernel on another thread's stream that can still
// be reading it. Only host-side work happens under the lock (the device-side initialization is
// deferred to the first `dev_ptr()` call).
std::lock_guard<std::mutex> lock(store->mutex);
Comment on lines +180 to +187

@bdice bdice Oct 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Commit 2, a second and narrower hazard. raft::cache::lru::get and set each take the cache's internal lock, but nothing holds a lock across them. Two threads that miss at the same time therefore both build a descriptor (each with its own device allocation and init kernel) and both insert it. This lock makes the miss path a single critical section. The work under it is host-only: the LRU lookup, dataset_descriptor_init (which builds the host object and creates a CUDA event), and the insertion. The device-side init still runs lazily in dev_ptr(). This commit was not tested on its own, but it was part of the after-fix CI run linked in the review body.

Question for reviewers: Reading raft::cache::lru::set, it doesn't check for an existing key. It takes the first free slot (or the LRU slot), so concurrent misses may leave two cached entries for one key rather than one replacing the other, as the code comment above says. If you agree, I'll reword the comment. Would you prefer a get_or_set on raft's lru instead of a second mutex in store?

dataset_descriptor_host<DataT, IndexT, DistanceT> desc;
if (!cache.get(key, &desc)) {
desc =
Expand Down
Loading