cuda.core: fix pool setup and builder teardown under stream capture

[Andy-Jost]

Summary

Fixes #2834, both parts, and #2776. DeviceMemoryResource(device) can be constructed for the device's default pool, and Device.memory_resource first touched, while the calling thread is inside a global or thread_local stream capture, without failing or invalidating the capture. Ending, closing, or collecting a top-level GraphBuilder whose capture was invalidated, or left unjoined by a join() that failed partway, no longer segfaults.

Both minimal reproducers from #2834 were run against this branch on a local GPU host: every scenario of the first now leaves the capture active, and the second survives 200 invalidate-and-collect rounds in each capture mode.

Changes

• _memory_pool.pyx: MP_raise_release_threshold makes its cuMemPoolGetAttribute/cuMemPoolSetAttribute calls in relaxed capture mode via cuThreadExchangeStreamCaptureMode and restores the thread's previous mode. Capture mode is per thread, so this needs no stream and has no effect when the thread is not capturing. A failed restore is attached to the propagating error as a note, per the error handling policy. This mirrors the libcu++ fix in [libcu++] Make resolving a default memory pool legal under stream capture cccl#11360.
• _cpp/rt/graph.cpp, api.hpp, _rt.pxd, _rt.pyx: new invalidate_root_graph_state(h_root) retires a hierarchy whose root graph CUDA destroyed itself, so the owning handle's deleter no longer calls cuGraphDestroy on it.
• _graph_builder.pyx: GB_end_capture replaces GB_end_capture_if_needed and settles graph ownership from the cuStreamEndCapture result. A NULL graph means the driver discarded the capture graph; the builder drops its handle and enters the new CAPTURE_INVALIDATED state. end_building() ends an invalidated capture and raises the driver error, close() closes the builder and then raises, and __dealloc__ reports as before. complete(), debug_dot_print(), graph_definition, embed() and Graph.update() reject a graph-less builder with a clear message. end_building() on a forked builder raises RuntimeError instead of ending capture on the wrong stream, which invalidated the whole capture.
• _graph_builder.pyx: join() closes every fork the loop did not reach before the error propagates, so no fork is left capturing on its private stream. The capture then ends with the driver's unjoined-work error from end_building(), which the invalidated-capture handling above turns into a clean error instead of a second cuGraphDestroy (cuda.core: GraphBuilder.join leaves forked builders in a state that segfaults at garbage collection if it raises midway #2776, folded in from cuda.core: close unjoined forks when GraphBuilder.join fails #2881).
• Tests: construction and lazy first touch under each capture mode keep the capture valid; the thread's capture mode is restored afterwards; an invalidated capture can be ended, closed, and collected in every mode, and collection reports a CUDAWarning; forked end_building() is rejected; a join() that fails partway closes the unjoined forks and end_building() reports the unjoined work.
• Release notes for 1.3.0.

Related Work

• Issues cuda.core: DeviceMemoryResource setup is refused under non-relaxed stream capture, and tearing down the invalidated GraphBuilder segfaults #2834 and cuda.core: GraphBuilder.join leaves forked builders in a state that segfaults at garbage collection if it raises midway #2776. cuda.core: close unjoined forks when GraphBuilder.join fails #2881 carried the join() cleanup as a stacked PR and is folded in here, since its safe failure mode depends on the invalidated-capture handling in this PR.
• [libcu++] Make resolving a default memory pool legal under stream capture cccl#11360 (same root cause in libcu++).
• Not covered: the body builder of a conditional node. When a body capture ends invalidated, the driver discards the body graph that the parent graph still refers to, so the parent graph is invalid as well and cuda.core cannot repair it. The docstrings and release note say so. Creating a new pool with options (cuMemPoolCreate) during a global or thread_local capture is still refused by the driver.

🤖 Generated with Claude Code

[copy-pr-bot]

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.

[Andy-Jost]

/ok to test

[github-actions]

Doc Preview CI
🚀 View preview at https://nvidia.github.io/cuda-python/pr-preview/pr-2838/
https://nvidia.github.io/cuda-python/pr-preview/pr-2838/cuda-core/
https://nvidia.github.io/cuda-python/pr-preview/pr-2838/cuda-bindings/
https://nvidia.github.io/cuda-python/pr-preview/pr-2838/cuda-pathfinder/
Preview will be ready when the GitHub Pages deployment is complete.

[Andy-Jost]

/ok to test

[Andy-Jost]

/ok to test

[brandon-b-miller]

Hmm I think we have to pass a primary whose end capture is actually refused to fully exercize things here. Provoking the RuntimeError from something we created via split makes it so that GB_end_capture simply short circuits.

import gc
from cuda.core import Device
from cuda.core.graph import GraphBuilder
dev = Device()
dev.set_current()
def bad_primary():
    """A primary whose close() fails: it has an unjoined fork, so the driver
    refuses cuStreamEndCapture."""
    p = dev.create_graph_builder().begin_building()
    p, _unjoined = p.split(2)
    return p
root = dev.create_graph_builder().begin_building()
root, fork = root.split(2)
bad1, bad2 = bad_primary(), bad_primary()
try:
    GraphBuilder.join(root, bad1, bad2, fork)
except Exception as e:
    print(f"join raised: {type(e).__name__}: {str(e).split(':')[0]}")
# join() documents that on failure "the builders that were not joined are
# closed before the error propagates, so none is left capturing".
left_open = [n for n, b in (("bad1", bad1), ("bad2", bad2), ("fork", fork)) if not b.is_closed]
print(f"left open: {left_open}")
del bad1, bad2, fork
gc.collect()

this gives me, on this branch,

join raised: CUDAError: CUDA_ERROR_STREAM_CAPTURE_INVALIDATED
left open: ['bad2', 'fork']

So the first truly failing close stranded the ones later on in the close list.

[Andy-Jost]

Thanks for catching that. I will upload a fix.

[brandon-b-miller]

This looks like parts of it can be factored out from here and invalidate_child_graph_state below, would be a bit more DRY

[brandon-b-miller]

retire_graph(hierarchy, graph) ?

[Andy-Jost]

Will do

[brandon-b-miller]

This doesn't cut across

cuda-python/cuda_core/docs/source/error_handling.rst

Line 136 in a04b5a3

``cuda.core`` never terminates the process. A CUDA error is raised, or reported

right?

[Andy-Jost]

The crash is a SEGV due to a driver bug with no known workaround. Issue #2918 has more details.

[brandon-b-miller]

Couple Q's otherwise LGTM.

[Andy-Jost]

64b5925 includes fixes for the things @brandon-b-miller pointed out