test(cuda.core): cover graph attachment ownership boundaries

Exercise subgraph deletion, in-flight launches, failed transactions, and independent executable references so attachment lifetime guarantees remain regression-tested.

Author: Andy-Jost

cuda_core/cuda/core/graph/_graph_builder.pyi

@@ -490,4 +490,7 @@ class Graph:
 __all__ = ['Graph', 'GraphBuilder', 'GraphCompleteOptions', 'GraphDebugPrintOptions']
 
 def _instantiate_graph(h_graph, options: GraphCompleteOptions | None=None) -> Graph:
-    ...
+    ...
+
+def _capture_callback_with_tail_failure_for_testing(gb: GraphBuilder, fn, *, user_data=None):
+    """Exercise anonymous attachment retention after capture commits a node."""

cuda_core/cuda/core/graph/_graph_builder.pyx

@@ -852,39 +852,54 @@ cdef class GraphBuilder:
             pointer (caller manages lifetime). If bytes-like, the data is
             copied and its lifetime is tied to the graph.
         """
-        GB_check_open(self)
-        cdef Stream stream = self._stream
-        cdef cydriver.CUstream c_stream = as_cu(stream._h_stream)
-        cdef cydriver.CUstreamCaptureStatus capture_status
+        GB_callback(self, fn, user_data, False)
 
-        with nogil:
-            _get_capture_info(c_stream, &capture_status, NULL)
 
-        if capture_status != cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE:
-            raise RuntimeError("Cannot add callback when graph is not being built")
+cdef inline void GB_callback(
+        GraphBuilder gb, object fn, object user_data,
+        bint fail_tail_discovery_for_testing) except *:
+    GB_check_open(gb)
+    cdef Stream stream = gb._stream
+    cdef cydriver.CUstream c_stream = as_cu(stream._h_stream)
+    cdef cydriver.CUstreamCaptureStatus capture_status
 
-        cdef cydriver.CUhostFn c_fn
-        cdef void* c_user_data = NULL
-        cdef OpaqueHandle fn_owner, data_owner
-        cdef GraphAttachmentsHandle attachments
-        _resolve_host_callback(fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner)
-        attachments = _prepare_host_callback_owners(
-            self._h_graph, fn_owner, data_owner)
+    with nogil:
+        _get_capture_info(c_stream, &capture_status, NULL)
 
-        with nogil:
-            HANDLE_RETURN(cydriver.cuLaunchHostFunc(c_stream, c_fn, c_user_data))
+    if capture_status != cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE:
+        raise RuntimeError("Cannot add callback when graph is not being built")
 
-        # Capturing the host function added a node to the graph; it is now the
-        # stream's sole capture dependency. If tail discovery unexpectedly
-        # fails, keep the owners anonymously graph-retained: rolling them back
-        # would leave the already-captured node with dangling pointers.
-        cdef cydriver.CUgraphNode host_node
-        try:
-            host_node = _capture_tail_node(c_stream)
-        except:
-            graph_abandon_attachments(attachments)
-            raise
-        HANDLE_RETURN(graph_commit_attachments(attachments, host_node))
+    cdef cydriver.CUhostFn c_fn
+    cdef void* c_user_data = NULL
+    cdef OpaqueHandle fn_owner, data_owner
+    cdef GraphAttachmentsHandle attachments
+    _resolve_host_callback(
+        fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner)
+    attachments = _prepare_host_callback_owners(
+        gb._h_graph, fn_owner, data_owner)
+
+    with nogil:
+        HANDLE_RETURN(cydriver.cuLaunchHostFunc(c_stream, c_fn, c_user_data))
+
+    # Capturing the host function added a node to the graph; it is now the
+    # stream's sole capture dependency. If tail discovery unexpectedly fails,
+    # keep the owners anonymously graph-retained: rolling them back would
+    # leave the already-captured node with dangling pointers.
+    cdef cydriver.CUgraphNode host_node
+    try:
+        if fail_tail_discovery_for_testing:
+            raise RuntimeError("forced capture tail discovery failure")
+        host_node = _capture_tail_node(c_stream)
+    except:
+        graph_abandon_attachments(attachments)
+        raise
+    HANDLE_RETURN(graph_commit_attachments(attachments, host_node))
+
+
+def _capture_callback_with_tail_failure_for_testing(
+        GraphBuilder gb, fn, *, user_data=None):
+    """Exercise anonymous attachment retention after capture commits a node."""
+    GB_callback(gb, fn, user_data, True)
 
 
 cdef inline int GB_check_open(GraphBuilder gb) except -1:

cuda_core/cuda/core/graph/_graph_node.pyi

@@ -405,4 +405,7 @@ class GraphNode:
             A new SwitchNode with branches accessible via ``.branches``.
         """
 __all__ = ['GraphNode']
-_node_registry: weakref.WeakValueDictionary[int, GraphNode] = weakref.WeakValueDictionary()
+_node_registry: weakref.WeakValueDictionary[int, GraphNode] = weakref.WeakValueDictionary()
+
+def _fail_callback_node_creation_after_prepare_for_testing(graph_def: GraphDefinition, fn):
+    """Exercise prepared-attachment rollback without relying on driver errors."""

cuda_core/cuda/core/graph/_graph_node.pyx

@@ -485,7 +485,7 @@ cdef class GraphNode:
         HostCallbackNode
             A new HostCallbackNode representing the callback.
         """
-        return GN_callback(self, fn, user_data)
+        return GN_callback(self, fn, user_data, False)
 
     def if_then(self, condition: GraphCondition) -> IfNode:
         """Add an if-conditional node depending on this node.
@@ -1084,7 +1084,9 @@ cdef inline EventWaitNode GN_wait_event(GraphNode self, Event ev):
         create_graph_node_handle(new_node, h_graph), ev._h_event))
 
 
-cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_data):
+cdef inline HostCallbackNode GN_callback(
+        GraphNode self, object fn, object user_data,
+        bint fail_after_prepare_for_testing):
     cdef cydriver.CUDA_HOST_NODE_PARAMS node_params
     cdef cydriver.CUgraphNode new_node = NULL
     cdef GraphHandle h_graph = graph_node_get_graph(self._h_node)
@@ -1104,6 +1106,9 @@ cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_
     attachments = _prepare_host_callback_owners(
         h_graph, fn_owner, data_owner)
 
+    if fail_after_prepare_for_testing:
+        raise RuntimeError("forced failure after attachment preparation")
+
     with nogil:
         HANDLE_RETURN(cydriver.cuGraphAddHostNode(
             &new_node, as_cu(h_graph), deps, num_deps, &node_params))
@@ -1113,3 +1118,12 @@ cdef inline HostCallbackNode GN_callback(GraphNode self, object fn, object user_
     return _registered(HostCallbackNode._create_with_params(
         create_graph_node_handle(new_node, h_graph), callable_obj,
         node_params.fn, node_params.userData))
+
+
+def _fail_callback_node_creation_after_prepare_for_testing(
+        GraphDefinition graph_def, fn):
+    """Exercise prepared-attachment rollback without relying on driver errors."""
+    cdef GraphNode entry = GraphNode.__new__(GraphNode)
+    entry._h_node = create_graph_node_handle(
+        <cydriver.CUgraphNode>NULL, graph_def._h_graph)
+    return GN_callback(entry, fn, None, True)

cuda_core/tests/graph/test_graph_builder.py

@@ -4,6 +4,8 @@
 """GraphBuilder stream capture tests."""
 
 import gc
+import time
+import weakref
 
 import numpy as np
 import pytest
@@ -13,6 +15,18 @@
 
 from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, launch
 from cuda.core.graph import GraphBuilder, GraphDefinition
+from cuda.core.graph._graph_builder import (
+    _capture_callback_with_tail_failure_for_testing,
+)
+
+
+def _wait_until(predicate, timeout=5.0):
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() >= deadline:
+            raise AssertionError(f"condition not satisfied within {timeout}s")
+        gc.collect()
+        time.sleep(0.02)
 
 
 def test_graph_is_building(init_cuda):
@@ -341,6 +355,34 @@ def read_byte(data):
     assert result[0] == 0xAB
 
 
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_capture_tail_discovery_failure_preserves_callback_attachment(init_cuda):
+    """A committed captured node keeps its owner when discovery fails."""
+    called = [False]
+
+    def callback():
+        called[0] = True
+
+    callback_weak = weakref.ref(callback)
+    stream = Device().create_stream()
+    gb = stream.create_graph_builder().begin_building()
+
+    with pytest.raises(RuntimeError, match="forced capture tail discovery failure"):
+        _capture_callback_with_tail_failure_for_testing(gb, callback)
+
+    del callback
+    gc.collect()
+    assert callback_weak() is not None
+
+    graph = gb.end_building().complete()
+    graph.launch(stream)
+    stream.sync()
+    assert called[0]
+
+    del graph, gb
+    _wait_until(lambda: callback_weak() is None)
+
+
 @pytest.mark.skipif(tuple(int(i) for i in np.__version__.split(".")[:2]) < (2, 1), reason="need numpy 2.1.0+")
 def test_graph_child_graph(init_cuda):
     mod = compile_common_kernels()

cuda_core/tests/graph/test_graph_definition_errors.py

@@ -4,6 +4,9 @@
 """Tests for GraphDefinition input validation, error handling, and edge cases."""
 
 import ctypes
+import gc
+import time
+import weakref
 
 import pytest
 from helpers.graph_kernels import compile_common_kernels
@@ -17,6 +20,9 @@
     GraphCondition,
     GraphDefinition,
 )
+from cuda.core.graph._graph_node import (
+    _fail_callback_node_creation_after_prepare_for_testing,
+)
 
 SIZEOF_INT = ctypes.sizeof(ctypes.c_int)
 
@@ -26,6 +32,15 @@ def _skip_if_no_mempool():
         pytest.skip("Device does not support mempool operations")
 
 
+def _wait_until(predicate, timeout=5.0):
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() >= deadline:
+            raise AssertionError(f"condition not satisfied within {timeout}s")
+        gc.collect()
+        time.sleep(0.02)
+
+
 # =============================================================================
 # Type validation — wrong types for conditional node methods
 # =============================================================================
@@ -98,6 +113,24 @@ def test_condition_from_different_graph(init_cuda):
         g2.if_then(condition)
 
 
+@pytest.mark.agent_authored(model="gpt-5.6")
+def test_failed_node_creation_releases_prepared_attachments(init_cuda):
+    """Failure after preparation rolls back ownership without adding a node."""
+
+    def callback():
+        pass
+
+    callback_weak = weakref.ref(callback)
+    graph_def = GraphDefinition()
+
+    with pytest.raises(RuntimeError, match="forced failure after attachment preparation"):
+        _fail_callback_node_creation_after_prepare_for_testing(graph_def, callback)
+
+    del callback
+    _wait_until(lambda: callback_weak() is None)
+    assert graph_def.nodes() == set()
+
+
 # =============================================================================
 # Edge cases — valid but unusual usage patterns
 # =============================================================================