Repository navigation
Conversation
dataset_descriptor_host::state::get() skipped waiting for the descriptor initialization event when the caller's stream handle equaled the stream that ran the initialization. The default raft stream is cudaStreamPerThread, whose handle is identical in every host thread but refers to a different stream in each. When several threads search with the same raft::resources, a thread could pick up the cached descriptor and launch a search kernel before the initialization kernel queued on another thread's stream had run, reading an uninitialized descriptor (cudaErrorMisalignedAddress / cudaErrorIllegalAddress in NEIGHBORS_MG_TEST). Always wait on the ready event.
dataset_descriptor_init_with_cache() did the LRU get and set under separate locks. Threads that missed concurrently each created a descriptor, and the later insertion replaced the earlier one while other threads could already be using it. The last owner of the replaced descriptor then freed its device memory with cudaFreeAsync on its own stream, while a search kernel on another thread's stream could still be reading it. Serialize the lookup and insertion so that concurrent searches on one raft::resources share a single descriptor.
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
bdice
added
bug
|
/ok to test 733caab |
|
/ok to test cd97df7 |
bdice
left a comment
bdice
left a comment
There was a problem hiding this comment.
Workstream F of the flaky-CI sweep (#2783). NEIGHBORS_MG_TEST AnnMG/5 aborted with a CUDA error in raft::copy on 12.2.2, amd64, rockylinux8, l4, earliest-driver, oldest-deps. The sweep window has 1 hit; the triage comment found 10 between 2026-09-09 and 2026-10-01 (AnnMG/5 / AnnMG/8, mostly cudaErrorMisalignedAddress).
- Trigger: several OpenMP threads searched CAGRA concurrently with one rank's
raft::resources. The trigger in that test was the single-rank SHARDED loop (before #2704), which #2704 made sequential. The bug itself is still onmainfor any caller that shares a handle across threads. - Root cause: the cached dataset descriptor is initialized by a kernel on the first caller's stream.
state::get()skippedcudaStreamWaitEventwhen the stream handles compared equal (main).cudaStreamPerThreadhas the same handle in every thread but is a different stream in each, so a second thread could launch a search kernel that read an uninitialized descriptor. compute-sanitizer showed a misaligned read of a garbage dataset pointer. - Fix: always wait on the ready event, and make the descriptor cache's lookup-or-insert atomic.
- Evidence (#2796,
mainwith the pre-#2704 loop restored, 8 L4 jobs loopingAnnMG/5:/8:/14): 8/339 processes aborted before (run 38076210837) and 0/339 after (run 38081894648) (#2783 comment).
| 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)); |
There was a problem hiding this comment.
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).
| // Hold the lock across the lookup and the insertion. Otherwise threads that miss concurrently | ||
| // each create a descriptor and the later insertion replaces the earlier one, which may already | ||
| // be in use by another thread; the last owner then frees it on its own stream while a kernel on | ||
| // a different stream 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); |
There was a problem hiding this comment.
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?
| /** Serializes lookup-or-insert, so that concurrent searches share a single descriptor. */ | ||
| std::mutex mutex; |
There was a problem hiding this comment.
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?
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Summary
Merge Risk: ⚪ Minimal · up to No actionable merge-blocking issue is established for these changes. Merge after normal checks. Pre-merge checks |
  |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. (1 skipped: 1 … | Write docstrings for the functions missing them to satisfy the coverage threshold. |
✅ Passed checks (4 passed)
-
Check name Status Explanation Title check The title clearly and concisely identifies the main change: fixing a CAGRA dataset descriptor race across threads sharing a handle. Description check The description is directly related to the changeset. It explains both race conditions, the affected code paths, observed failures, and the intended fixes. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. (1 skipped: 1 unsupported.)
- Fix all pre-merge checks with AI
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
1 participant
CAGRA caches a dataset descriptor per
raft::resourcesand initializes it with a kernel on the first caller's stream.dataset_descriptor_host::state::get()only waited for that kernel when the caller's stream handle differed, but raft's defaultcudaStreamPerThreadhas the same handle and a different stream in every thread. So a second thread searching with the same handle could read an uninitialized descriptor, causingcudaErrorMisalignedAddress/cudaErrorIllegalAddressinNEIGHBORS_MG_TEST(AnnMG/5,AnnMG/8). CI hit this 10 times from 2026-09-09 to 2026-10-01 on L4, H100, RTX PRO 6000 and GH200 (9cudaErrorMisalignedAddress, 1cudaErrorIllegalAddress; see #2588). There have been no hits onmainsince #2704 removed the single-rank OpenMP search loop that triggered it (the 2026-10-01 hit is onrelease/26.10), but the race is still present for any threads that search concurrently with oneraft::resources, such as the multi-GPU ROUND_ROBIN mode.state::get().cudaFreeAsyncon the stream that initialized it, which is not ordered after kernels on other threads' streams.With concurrent searches on one handle, 8/339 test processes aborted before the fix and 0/339 after (reproducer: #2796). See #2783 (workstream F).