[STF] Make executable graph cache stream-affine - #11041
Conversation
Prevent concurrent scopes from reusing or updating an active executable by making cache entries stream-affine and skipping entries whose streams are busy.
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesThe graph cache now binds executable graphs to CUDA streams, reuses only idle matching entries, preserves busy entries, and serializes host access. Callers and tests pass streams explicitly. New tests cover stream affinity and concurrent scope construction. Stream-aware graph cache
Suggested reviewers: Merge Risk: 🟠 High · up to Concurrent use can update a cached executable before another caller launches it, causing the wrong graph to run, while destroyed streams may later be accessed during cache cleanup and trigger undefined behavior or crashes. The PR is not merge-ready until executable ownership and stream lifetime are protected. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuh (1)
165-189: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftimportant: Keep a cache entry reserved until its launch is enqueued.
cudaGraphExecUpdatemodifies the executable in place. If a sharedasync_resources_handleperforms a second query while the stream is still idle,executable_graph_cache::querycan update the samecudaGraphExec_tbeforestackable_ctx_impl::launch_oncecallscudaGraphLaunch. The first caller can then launch the second caller’s graph. Make cache update and launch one atomic operation, or return a lease that blocks updates until launch.
🧹 Nitpick comments (1)
cudax/include/cuda/experimental/__stf/graph/graph_ctx.cuh (1)
506-506: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuesuggestion: Declare these unmodified local bindings
const.
cudax/include/cuda/experimental/__stf/graph/graph_ctx.cuh#L506-L506: useconst auto query_result.cudax/test/stf/graph/get_cache_stats.cu#L22-L22: useconst cudaStream_t stream.cudax/test/stf/graph/get_cache_stats.cu#L40-L40: useauto* const st.cudax/test/stf/graph/graph_cache_policy.cu#L22-L22: useconst cudaStream_t stream.cudax/test/stf/graph/graph_cache_policy.cu#L46-L46: useauto* const st.As per coding guidelines: “All variables that are not modified must be declared
const.”Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0fb4208f-1cb2-4074-88e1-bb4543acac2c
📒 Files selected for processing (8)
cudax/include/cuda/experimental/__stf/graph/graph_ctx.cuhcudax/include/cuda/experimental/__stf/internal/async_resources_handle.cuhcudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuhcudax/include/cuda/experimental/__stf/stackable/stackable_ctx_impl.cuhcudax/test/stf/CMakeLists.txtcudax/test/stf/graph/get_cache_stats.cucudax/test/stf/graph/graph_cache_policy.cucudax/test/stf/stackable/executable_graph_cache_streams.cu
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
The previous pin pointed at an unrelated snapshot of the fork that lacks the two fixes the composed centrality experiment requires. Pin the tip of lab/cugraph-stf-sibling-fixes instead, which combines NVIDIA/cccl#11041 (stream-affine executable graph cache) and NVIDIA/cccl#11036 (preserve nested read-only imports) on upstream main, so a fresh checkout of this branch reproduces the measured composed-centrality behavior. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A cudaStream_t handle value can be recycled after cudaStreamDestroy, so comparing stored handles against the caller could falsely match an entry bound to a dead stream. Entries now bind to the driver-assigned stream id (cudaStreamGetId, unique for the process lifetime); the handle is kept only for idleness probes. Stream probes go through a non-throwing tri-state helper that checks capture status first: cudaStreamQuery on a capturing stream would invalidate that capture, a cross-thread hazard when reclaim probes a stream another context is capturing on. Destroyed-handle errors read as 'unavailable': such entries are skipped by reclaim (destroying the executable is not provably safe while the dead stream's last launch may drain) and can never match a lookup again. Also make the entry constructor advance the LRU counter so entries inserted between refreshes do not share one timestamp. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
andralex
left a comment
andralex
left a comment
There was a problem hiding this comment.
Review by Grégoire (posting from Andrei's account). The design direction is correct and the PR does what it says: binding executables to their launch stream removes the false serialization, the busy-check before update is the right guard, and the test plan has the right shape (same-stream hit, cross-stream miss, busy miss, real overlap). Two defects and two smaller items found; the defects are fixed in 739706b, pushed to this branch.
Fixed in 739706b:
-
Raw-handle identity was unsound (hot path, not just reclaim): a cudaStream_t handle value can be recycled after cudaStreamDestroy, so
e.stream != streamcould falsely MATCH an entry bound to a dead stream against an unrelated new stream on any lookup. Entries now bind to the driver-assigned stream id (cudaStreamGetId, unique for the process lifetime, CUDA >= 12.0); the raw handle is kept only for idleness probes. -
Stream probes could throw or corrupt a capture:
stream_is_idle'scuda_try(status)turned a destroyed stored handle into an exception thrown mid-reclaim, andcudaStreamQueryon a stream that another thread is capturing on INVALIDATES that capture (STF does legacy capture, so this was a real cross-thread hazard). Probes now go through a non-throwing tri-state helper that checkscudaStreamIsCapturingfirst (capture-legal) and reads errors asunavailable. Anunavailableentry is skipped by reclaim: destroying the executable is not provably safe while the dead stream's last launch may still drain. -
LRU tie: the entry constructor set
last_use = cache->indexwithout advancing, so entries inserted between refreshes shared a timestamp and evicted in arbitrary order. Constructor now increments too.
Answers to the review-thread questions: the weak paths were two, not one — identity aliasing was in the hot lookup (every query after any stream churn), while unsafe probes of stored handles were reclaim-only (in query() the only stream probed is the caller's, alive by definition). And no, the cache never learns a stream died, nor can it: CUDA has no liveness query (a recycled handle actively lies), no destruction callback, and no is-this-id-alive API. Entries bound to destroyed streams are permanent zombies — unmatched forever under id binding, skipped by reclaim. That is benign when cache-bound streams outlive the cache, which pool streams from async_resources_handle do by construction; the remaining design decision is inline below.
| /* This will lookup in the cache (if any) and update an existing entry, or | ||
| * instantiate a graph if none is found. */ | ||
| auto query_result = async_resources().cached_graphs_query(nnodes, nedges, *g); | ||
| auto query_result = async_resources().cached_graphs_query(nnodes, nedges, *g, state.submitted_stream); |
There was a problem hiding this comment.
The one open design decision after 739706b: this path binds cache entries to state.submitted_stream, i.e. the USER's stream, whose lifetime STF does not control. With id binding a recycled handle can no longer false-match, but entries bound to destroyed user streams become permanent zombies (never matched, never reclaimed — reclaim cannot prove the dead stream's last launch drained). Options: (a) accept bounded zombie growth (one entry per destroyed-stream x topology, ~10KB/node estimate) and document it; (b) don't insert into the cache on this path — user-stream submits instantiate uncached, losing reuse but keeping the cache zombie-free; (c) have graph_ctx wrap user streams into pool-owned proxies at creation so the ownership invariant covers everything. The stackable path (pick_stream()) is pool-owned and safe by construction either way. My lean is (a) now with a comment, (c) as the eventual clean state.
| continue; | ||
| } | ||
|
|
||
| if (reserved::try_updating_executable_graph(*e.exec_g, g)) |
There was a problem hiding this comment.
Note for a comment rather than a change: the idle-check-then-update sequence is protected by the cache mutex against other QUERIES, but nothing stops a second host thread from LAUNCHING onto this same stream between the check and the update. The design is sound under the invariant "one support stream is submitted to by one thread at a time", which holds today for pool streams and single-ctx user streams, but it is implicit. One sentence here would make it an invariant instead of luck.
Related, pre-existing (the function itself is outside this diff so noting it here at its call site): try_updating_executable_graph's blind cudaGetLastError() also swallows any unrelated earlier pending async error, reporting it as "update failed" (a silent miss) instead of surfacing it. Cheap hardening: check the return value of cudaGraphExecUpdate itself and only clear-and-classify when that call is what failed.
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f662496f-829d-4191-bc05-a25e50e852ed
📒 Files selected for processing (1)
cudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuh
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
| { | ||
| reclaimed += it->second.footprint; | ||
| total_cache_footprint[dev_id] -= it->second.footprint; | ||
| if (query_stream_state(it->second.stream) == stream_state::idle |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
important: Do not probe entry::stream after its owner can call cudaStreamDestroy. A cached entry can retain this raw handle after destruction, and query_stream_state then calls cudaStreamIsCapturing through it during reclamation. CUDA defines use of a stream handle after cudaStreamDestroy as undefined, so the intended unavailable path cannot safely prevent a crash or other invalid behavior. Track launch completion with cache-owned events, or enforce that every bound stream outlives the cache. (docs.nvidia.com)
As per path instructions, focus on lifetime/resource ownership.
Source: Path instructions
2 participants
Summary
cudaGraphExecUpdate, and skip busy entries during LRU reclamationProblem
The cache keyed executable graphs only by device and topology. Concurrent top-level graph scopes with compatible topology could therefore update and launch the same
cudaGraphExec_ton different support streams; CUDA serializes those launches and silently removes legal overlap.Test plan
clang-format --dry-run --Werrorandgit diff --checkKnown trade-offs
Busy entries are conservatively unavailable for update and reclamation, including when unrelated work occupies the bound stream. The cache may temporarily exceed its configured size when every candidate is busy; the implementation never synchronizes to enforce the limit.