cuda.core: fix pool setup and builder teardown under stream capture (#2838)

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

DeviceMemoryResource(device) with no options raises the release threshold
of the driver's pool with cuMemPoolGetAttribute and cuMemPoolSetAttribute.
The driver refuses both as potentially unsafe calls while the calling
thread is inside a global or thread-local capture, and it invalidates the
capture. Device.memory_resource constructs the resource lazily, so a first
allocation could invalidate a capture in progress. Make the two calls in
relaxed capture mode and restore the thread's previous mode afterwards. A
failure to restore the mode is attached to the propagating error as a note.

Ending an invalidated capture made the builder destroy a graph the driver
had already destroyed. cuStreamEndCapture returns a NULL graph for an
invalidated (or unjoined) capture and releases the capture graph itself,
but the builder kept the owning handle it took from cuStreamGetCaptureInfo
and its deleter called cuGraphDestroy again: a use-after-free that
segfaulted at close() or garbage collection. Add invalidate_root_graph_state
to retire the hierarchy when the driver discards the root graph, and route
end_building(), close() and __dealloc__ through one GB_end_capture helper
that settles graph ownership from the end-capture result. end_building()
now ends an invalidated capture and raises the driver error; the builder
then holds no graph (new CAPTURE_INVALIDATED state), and complete(),
debug_dot_print(), graph_definition, embed() and Graph.update() say so.
close() closes the builder before raising. end_building() on a forked
builder is rejected with RuntimeError instead of invalidating the capture.

Fixes #2834.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* cuda.core: scope the invalidated-capture recovery to top-level builders

The driver also discards the body graph of a conditional node when the
body capture ends invalidated, and the parent graph keeps referring to
it, which cuda.core cannot repair. Say so in end_building(), the release
note, and the GB_end_capture comment. Scope the DeviceMemoryResource
note to the default-pool constructor, since cuMemPoolCreate is still
refused under capture. Drop the end_building() call on a forked builder
from the skip path of test_graph_conditional_on_forked_builder.

Issue #2834

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* cuda.core: close unjoined forks when GraphBuilder.join fails

join() waits on each forked builder's stream and then closes it. When a
wait raised partway through, the forks the loop had not reached stayed
open with capturing streams, and destroying them later during garbage
collection crashed the interpreter.

Close every fork the loop did not reach before the error propagates.
The capture cannot complete without the work captured on those forks,
so end_building() then raises the driver's unjoined-capture error and
the builder closes cleanly, which #2834 made possible.

Issue #2776

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

* cuda.core: fold the join cleanup into the capture-teardown release note

PR #2881 is folded into this PR: its join() cleanup relies on the
invalidated-capture handling here, and the #2776 crash is the same
double destroy of a graph the driver already discarded. Merge the two
release-note entries and correct the test docstring, which attributed
the crash to the fork rather than to the builder's teardown.

Issue #2834, #2776

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* cuda.core: point the conditional-body limitation at issue #2918

An invalidated capture of a conditional body leaves the parent graph
invalid, and ending the parent capture afterwards crashes inside the
driver (#2918, a CUDA driver bug reproduced with the
driver API alone). Say so in the end_building docstring and the release
note instead of explaining the mechanism.

Issue #2834

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* cuda.core: keep GraphBuilder.join's cleanup going when a close fails

join() closes the builders it did not join before its error propagates.
That sweep ran in a finally block and called close() on each builder. A
builder from another capture makes the root's wait fail, and the driver
invalidates that builder's capture as well, so its close() raised too.
The raise stopped the sweep, stranded the later builders with capturing
streams, and replaced the merge error with the close error.

Run the sweep in an except block, close each builder through GB_close,
which returns the driver status instead of raising, and attach a failed
close to the propagating error as a note (error handling policy). The
original error is re-raised unchanged.

Review feedback on #2838.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* test(cuda.core): cover a failing close inside GraphBuilder.join's cleanup

The existing test provokes the join failure with a forked builder, whose
close() cannot raise, so the sweep's own failure path went untested. Join
two separate primaries into a capture: the driver refuses the
cross-capture wait and invalidates the other capture, so closing that
builder fails. The test checks that every builder is closed, that the
merge error propagates, and that the failed close is attached as a note
(or reported as a CUDAWarning on Python 3.10).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* cuda.core: factor retire_graph out of the graph invalidation sweeps

invalidate_child_graph_state and invalidate_root_graph_state repeated the
per-box retirement: detach node handles, drop the registry entry and
attachments, move the box to the graveyard. Both now call retire_graph
and differ only in which boxes they select. No behavior change.

Review feedback on #2838.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: Andy-Jost

cuda_core/cuda/core/_cpp/rt/api.hpp

@@ -357,6 +357,12 @@ void invalidate_child_graph_state(
     const GraphHandle& h_parent,
     CUgraphNode owner_node) noexcept;
 
+// Invalidate cuda.core state for a root graph that CUDA destroyed itself, such
+// as the graph of an invalidated capture ended by cuStreamEndCapture. The
+// owning handle then no longer calls cuGraphDestroy. No-op unless h_root is
+// the live root of its hierarchy.
+void invalidate_root_graph_state(const GraphHandle& h_root) noexcept;
+
 // ============================================================================
 // Graph exec handle functions
 // ============================================================================

cuda_core/cuda/core/_cpp/rt/graph.cpp

@@ -679,6 +679,26 @@ static const GraphNodeBox* get_box(const GraphNodeHandle& h) {
     );
 }
 
+// Retire one box of a graph that CUDA has destroyed: detach its node handles,
+// drop its registry entry and attachments, and move it to the graveyard.
+// `graph` must be an iterator into hierarchy.graphs; the splice invalidates it.
+static void retire_graph(
+        GraphHierarchy& hierarchy,
+        std::list<GraphBox>::iterator graph) noexcept {
+    for (auto& entry : graph->node_handles.drain()) {
+        if (GraphNodeHandle h_node = entry.second.lock()) {
+            get_box(h_node)->resource = nullptr;
+        }
+    }
+    if (graph->resource) {
+        graph_registry.unregister_handle(graph->resource);
+        graph->resource = nullptr;
+    }
+    graph->attachments.clear();
+    hierarchy.graveyard.splice(
+        hierarchy.graveyard.end(), hierarchy.graphs, graph);
+}
+
 // graphs is ordered parent-before-child. Nulling a selected box marks its
 // later descendants, whose parent pointers remain valid after list splicing.
 // This permits one allocation-free sweep of the hierarchy.
@@ -701,21 +721,28 @@ void invalidate_child_graph_state(
                              graph->owner_node == owner_node;
         bool is_descendant = graph->parent &&
                              !graph->parent->resource;
-        if (!is_owned_root && !is_descendant) {
-            continue;
+        if (is_owned_root || is_descendant) {
+            retire_graph(hierarchy, graph);
         }
+    }
+}
 
-        // Empty node_handles and invalidate each one.
-        for (auto& entry : graph->node_handles.drain()) {
-            if (GraphNodeHandle h_node = entry.second.lock()) {
-                get_box(h_node)->resource = nullptr;
-            }
-        }
-        graph_registry.unregister_handle(graph->resource);
-        graph->resource = nullptr;
-        graph->attachments.clear();
-        hierarchy.graveyard.splice(
-            hierarchy.graveyard.end(), hierarchy.graphs, graph);
+// CUDA destroyed the root graph and, with it, every child. Retire every box
+// so that no registry entry resolves to the dead graphs and the hierarchy's
+// deleter finds no root to destroy.
+void invalidate_root_graph_state(const GraphHandle& h_root) noexcept {
+    if (!h_root) {
+        return;
+    }
+
+    GraphBox* root = get_box(h_root);
+    if (!root->resource || root->parent) {
+        return;
+    }
+    GraphHierarchy& hierarchy = *root->hierarchy;
+    for (auto it = hierarchy.graphs.begin();
+         it != hierarchy.graphs.end();) {
+        retire_graph(hierarchy, it++);
     }
 }
 

cuda_core/cuda/core/_memory/_memory_pool.pyx

@@ -22,6 +22,7 @@ from cuda.core._rt cimport (
     create_mempool_handle,
     deviceptr_alloc_from_pool,
     get_last_error,
+    attach_rollback_failure,
     as_cu,
     as_py,
 )
@@ -311,24 +312,48 @@ cdef int MP_raise_release_threshold(_MemPool self) except? -1:
     By default the release threshold is 0, meaning memory is returned to
     the OS as soon as there are no active suballocations.  Setting it to
     ULLONG_MAX avoids repeated OS round-trips.
+
+    Pool attribute reads and writes are potentially unsafe calls while the
+    calling thread is inside a global or thread-local stream capture: the
+    driver refuses them and invalidates the capture. Capture mode is a
+    per-thread property, so the thread is switched to relaxed mode around the
+    two calls and its previous mode is restored afterwards. This has no
+    observable effect when the thread is not capturing. The attribute write
+    executes immediately rather than being recorded into a graph, which is the
+    intent for a process-wide pool setting.
     """
     MP_check_open(self)
     cdef cydriver.cuuint64_t current_threshold
     cdef cydriver.cuuint64_t max_threshold = ULLONG_MAX
+    cdef cydriver.CUstreamCaptureMode mode = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED
+    cdef cydriver.CUresult err
+    cdef cydriver.CUresult restore_err
     with nogil:
-        HANDLE_RETURN(
-            cydriver.cuMemPoolGetAttribute(
-                as_cu(self._h_pool),
-                cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
-                &current_threshold
-            )
+        # On return, mode holds the thread's previous capture mode.
+        HANDLE_RETURN(cydriver.cuThreadExchangeStreamCaptureMode(&mode))
+        err = cydriver.cuMemPoolGetAttribute(
+            as_cu(self._h_pool),
+            cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
+            &current_threshold
         )
-        if current_threshold == 0:
-            HANDLE_RETURN(cydriver.cuMemPoolSetAttribute(
+        if err == cydriver.CUDA_SUCCESS and current_threshold == 0:
+            err = cydriver.cuMemPoolSetAttribute(
                 as_cu(self._h_pool),
                 cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
                 &max_threshold
-            ))
+            )
+        restore_err = cydriver.cuThreadExchangeStreamCaptureMode(&mode)
+    try:
+        HANDLE_RETURN(err)
+    except:
+        if restore_err != cydriver.CUDA_SUCCESS:
+            # The attribute error propagates with the restore failure attached
+            # as a note (error handling policy).
+            attach_rollback_failure(
+                b"cuThreadExchangeStreamCaptureMode", restore_err,
+                b"failed to restore the thread's stream capture mode")
+        raise
+    HANDLE_RETURN(restore_err)
     return 0
 
 

cuda_core/cuda/core/_rt.pxd

@@ -305,6 +305,7 @@ cdef cydriver.CUresult graph_commit_child_graph_update(
     PreparedChildGraphUpdate& prepared, GraphHandle* out_child) except+
 cdef void invalidate_child_graph_state(
     const GraphHandle& h_parent, cydriver.CUgraphNode owner_node) noexcept
+cdef void invalidate_root_graph_state(const GraphHandle& h_root) noexcept
 
 # Graph exec handles
 cdef GraphExecHandle create_graph_exec_handle(

cuda_core/cuda/core/_rt.pyx

@@ -212,6 +212,8 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt":
         PreparedChildGraphUpdate& prepared, GraphHandle* out_child) except+
     void invalidate_child_graph_state "cuda_core::rt::invalidate_child_graph_state" (
         const GraphHandle& h_parent, cydriver.CUgraphNode owner_node) noexcept
+    void invalidate_root_graph_state "cuda_core::rt::invalidate_root_graph_state" (
+        const GraphHandle& h_root) noexcept
 
     # Graph exec handles
     GraphExecHandle create_graph_exec_handle "cuda_core::rt::create_graph_exec_handle" (

cuda_core/cuda/core/graph/_graph_builder.pyi

@@ -136,7 +136,11 @@ class GraphBuilder:
     @staticmethod
     def _init(stream: Stream): ...
     def close(self):
-        """Destroy the graph builder."""
+        """Destroy the graph builder.
+
+        A builder that is still building ends its capture first. The builder
+        is closed even when that fails; the driver error is raised afterwards.
+        """
     @property
     def is_closed(self) -> bool:
         """Whether this graph builder has been closed."""
@@ -212,7 +216,22 @@ class GraphBuilder:
     def is_building(self) -> bool:
         """Returns True if the graph builder is currently building."""
     def end_building(self) -> GraphBuilder:
-        """Ends the building process."""
+        """Ends the building process.
+
+        Raises
+        ------
+        RuntimeError
+            If the builder is not building. A forked builder is ended by
+            :meth:`join`, not by this method.
+        CUDAError
+            If the capture was invalidated, for example by a CUDA call that is
+            not permitted while capturing. The capture is ended and the
+            builder holds no graph afterwards, so :meth:`complete` and
+            :attr:`graph_definition` are unavailable. This applies to a
+            top-level builder. An invalidated capture of a conditional body
+            leaves the parent builder's graph invalid as well; ending it may
+            crash the process (see issue #2918).
+        """
     def complete(self, options: GraphCompleteOptions | None=None) -> Graph:
         """Completes the graph builder and returns the built :obj:`~graph.Graph` object.
 
@@ -261,6 +280,9 @@ class GraphBuilder:
         """Joins multiple graph builders into a single graph builder.
 
         The returned builder inherits work dependencies from the provided builders.
+        If joining fails partway, the builders that were not joined are closed
+        before the error propagates, so none is left capturing. A driver error
+        from one of those closes is attached to the propagating error as a note.
 
         Parameters
         ----------