cuda.core: Cythonize GraphBuilder and Graph with handle-layer cleanup (#2008)

* cuda.core: convert GraphBuilder to cdef class with explicit state machine

Refactor GraphBuilder from a Python class using _MembersNeededForFinalize
to a cdef class with explicit _BuilderKind (PRIMARY/FORKED/CONDITIONAL_BODY)
and _CaptureState (NOT_STARTED/CAPTURING/ENDED) tracking. Cleanup moves
into __dealloc__/close, and the builder now uses GraphHandle/StreamHandle
from _resource_handles instead of holding raw driver objects. Drop the
is_stream_owner flag now that StreamHandle owns the lifetime.

End-capture paths in __dealloc__ and close guard on _h_stream so cleanup
is safe even if _init* fails before completing assignment.

Made-with: Cursor

* cuda.core: convert Graph to cdef class with GraphExecHandle

Add a GraphExecHandle to the resource-handle layer (parallel to
GraphHandle) wrapping CUgraphExec with RAII cleanup via
cuGraphExecDestroy on shared_ptr release. Convert Graph from a Python
class using _MembersNeededForFinalize to a cdef class holding a typed
_h_graph_exec attribute, dropping the weakref.finalize machinery.
update/upload/launch move to nogil cydriver paths consistent with the
GraphBuilder rewrite.

Also drop quoted forward-reference annotations on create_graph_builder
and _instantiate_graph/complete now that GraphBuilder is cimported in
_device.pyx and _stream.pyx and Cython accepts the in-module forward
reference to Graph. Clears the related "Strings should no longer be
used for type declarations" warnings.

Made-with: Cursor

* fix(cuda.core): drop unused handle cimports flagged by cython-lint

The cdef-class member declarations live in the .pxd, so the .pyx does
not need to re-cimport GraphExecHandle, GraphHandle, or StreamHandle.

Made-with: Cursor

* fix(cuda.core): break _stream/_device <-> graph._graph_builder import cycle

cimport-ing GraphBuilder at the top of _stream.pyx and _device.pyx made
Cython emit a Python-level import of cuda.core.graph._graph_builder
during _stream module init. That triggered the chain
graph -> _graph_node -> _kernel_arg_handler -> _memory._buffer
-> _device, which then re-entered the still-initializing _stream module
via "from cuda.core._stream import IsStreamT", failing with
ImportError: cannot import name IsStreamT.

Restore the original lazy "import GraphBuilder" inside
create_graph_builder (Stream and Device) and Stream_accept. The return
annotations stay as bare names; "from __future__ import annotations" in
both files defers their evaluation, so they need not resolve at
function-definition time.

Made-with: Cursor

* fix(cuda.core): expose GraphBuilder._init as a Python-callable factory

The previous import-cycle fix changed _stream/_device.create_graph_builder
to a lazy Python "import GraphBuilder" instead of a module-level cimport.
With _init declared as @staticmethod cdef, Python attribute lookup
cannot find it, so every test that builds a graph failed with
"AttributeError: type object 'GraphBuilder' has no attribute '_init'"
at _device.pyx:1376 / _stream.pyx:376.

Convert _init from @staticmethod cdef to @staticmethod def (matches the
Stream._init pattern) and drop the cdef declaration from the .pxd.
_init runs once per builder creation, so the loss of cdef-level
dispatch is irrelevant. Graph._init stays cdef; it is only called
intra-module.

Made-with: Cursor

* fix(cuda.core): pass non-NULL captureStatus to cuStreamGetCaptureInfo

Every graph-builder test failed with CUDA_ERROR_INVALID_VALUE on the
new ``GraphBuilder.begin_building`` path. The driver rejects
``cuStreamGetCaptureInfo`` when ``captureStatus_out`` is NULL, but the
new ``_get_capture_info`` helper accepted a NULL status pointer and
``begin_building`` was calling it that way (it just wanted the freshly
captured graph handle and assumed the status was implied by the
preceding ``cuStreamBeginCapture``).

Pass a stack-local ``CUstreamCaptureStatus`` and document the helper's
requirement that ``status`` be non-NULL. ``graph`` is still allowed to
be NULL (``is_building`` calls it that way and the driver accepts it).

Co-authored-by: Cursor <cursoragent@cursor.com>

* test: verify Graph.close() is idempotent (Glasswing V18.1)

Add coverage that repeated close() does not double-destroy the graph
exec handle. Addresses NVBugs 6268912 / cuda-python-private#370 via the
handle-layer refactor in this PR.

* cuda.core: harden GraphBuilder state machine and expand graph tests

Add CLOSED state with GB_check_open guards, fix begin_building ordering so
capture cleanup runs on partial failure, preserve FORKED primary-graph refs for
conditional nodes, and defer Graph._init until instantiate succeeds. Add tests
for Leo's review gaps including closed-builder errors and forked conditionals.

* fix(cuda.core): resolve merge fallout in resource handle stubs

Drop an unnecessary cimport from _resource_handles.pyx that broke stubgen/mypy
after merging main, regenerate the affected .pyi files, and fix debug_dot_print
path encoding for Cython 3.2.

* fix(cuda.core): pass non-NULL phGraph to cuStreamEndCapture

Use a stack-local CUgraph output for end_building() and
GB_end_capture_if_needed(), and release the GIL inside the helper so
driver calls are nogil regardless of caller (__dealloc__ or close()).

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Ralf W. Grosse-Kunstleve <rwgkio@gmail.com>

Author: 3 people

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -69,6 +69,7 @@ decltype(&cuLibraryGetKernel) p_cuLibraryGetKernel = nullptr;
 
 // Graph
 decltype(&cuGraphDestroy) p_cuGraphDestroy = nullptr;
+decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy = nullptr;
 
 // Linker
 decltype(&cuLinkDestroy) p_cuLinkDestroy = nullptr;
@@ -1136,6 +1137,28 @@ GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent)
     return GraphHandle(box, &box->resource);
 }
 
+// ============================================================================
+// Graph Exec Handles
+// ============================================================================
+
+namespace {
+struct GraphExecBox {
+    CUgraphExec resource;
+};
+}  // namespace
+
+GraphExecHandle create_graph_exec_handle(CUgraphExec graph_exec) {
+    auto box = std::shared_ptr<const GraphExecBox>(
+        new GraphExecBox{graph_exec},
+        [](const GraphExecBox* b) {
+            GILReleaseGuard gil;
+            p_cuGraphExecDestroy(b->resource);
+            delete b;
+        }
+    );
+    return GraphExecHandle(box, &box->resource);
+}
+
 namespace {
 struct GraphNodeBox {
     mutable CUgraphNode resource;

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -108,6 +108,7 @@ extern decltype(&cuLibraryGetKernel) p_cuLibraryGetKernel;
 
 // Graph
 extern decltype(&cuGraphDestroy) p_cuGraphDestroy;
+extern decltype(&cuGraphExecDestroy) p_cuGraphExecDestroy;
 
 // Linker
 extern decltype(&cuLinkDestroy) p_cuLinkDestroy;
@@ -183,6 +184,7 @@ using MemoryPoolHandle = std::shared_ptr<const CUmemoryPool>;
 using LibraryHandle = std::shared_ptr<const CUlibrary>;
 using KernelHandle = std::shared_ptr<const CUkernel>;
 using GraphHandle = std::shared_ptr<const CUgraph>;
+using GraphExecHandle = std::shared_ptr<const CUgraphExec>;
 using GraphNodeHandle = std::shared_ptr<const CUgraphNode>;
 using GraphicsResourceHandle = std::shared_ptr<const CUgraphicsResource>;
 using NvrtcProgramHandle = std::shared_ptr<const nvrtcProgram>;
@@ -464,6 +466,14 @@ GraphHandle create_graph_handle(CUgraph graph);
 // but h_parent will be prevented from destruction while this handle exists.
 GraphHandle create_graph_handle_ref(CUgraph graph, const GraphHandle& h_parent);
 
+// ============================================================================
+// Graph exec handle functions
+// ============================================================================
+
+// Wrap an externally-created CUgraphExec with RAII cleanup.
+// When the last reference is released, cuGraphExecDestroy is called automatically.
+GraphExecHandle create_graph_exec_handle(CUgraphExec graph_exec);
+
 // ============================================================================
 // Graph node handle functions
 // ============================================================================
@@ -649,6 +659,10 @@ inline CUgraph as_cu(const GraphHandle& h) noexcept {
     return h ? *h : nullptr;
 }
 
+inline CUgraphExec as_cu(const GraphExecHandle& h) noexcept {
+    return h ? *h : nullptr;
+}
+
 inline CUgraphNode as_cu(const GraphNodeHandle& h) noexcept {
     return h ? *h : nullptr;
 }
@@ -729,6 +743,10 @@ inline std::intptr_t as_intptr(const GraphHandle& h) noexcept {
     return reinterpret_cast<std::intptr_t>(as_cu(h));
 }
 
+inline std::intptr_t as_intptr(const GraphExecHandle& h) noexcept {
+    return reinterpret_cast<std::intptr_t>(as_cu(h));
+}
+
 inline std::intptr_t as_intptr(const GraphNodeHandle& h) noexcept {
     return reinterpret_cast<std::intptr_t>(as_cu(h));
 }
@@ -855,6 +873,10 @@ inline PyObject* as_py(const GraphHandle& h) noexcept {
     return detail::make_py("cuda.bindings.driver", "CUgraph", as_intptr(h));
 }
 
+inline PyObject* as_py(const GraphExecHandle& h) noexcept {
+    return detail::make_py("cuda.bindings.driver", "CUgraphExec", as_intptr(h));
+}
+
 inline PyObject* as_py(const GraphNodeHandle& h) noexcept {
     if (!as_intptr(h)) {
         Py_RETURN_NONE;

cuda_core/cuda/core/_device.pyx

@@ -1454,7 +1454,7 @@ class Device:
         from cuda.core.graph._graph_builder import GraphBuilder
 
         self._check_context_initialized()
-        return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)
+        return GraphBuilder._init(self.create_stream())
 
 
 cdef inline int Device_ensure_cuda_initialized() except? -1:

cuda_core/cuda/core/_resource_handles.pxd

@@ -28,6 +28,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     ctypedef shared_ptr[const cydriver.CUlibrary] LibraryHandle
     ctypedef shared_ptr[const cydriver.CUkernel] KernelHandle
     ctypedef shared_ptr[const cydriver.CUgraph] GraphHandle
+    ctypedef shared_ptr[const cydriver.CUgraphExec] GraphExecHandle
     ctypedef shared_ptr[const cydriver.CUgraphNode] GraphNodeHandle
     ctypedef shared_ptr[const cydriver.CUgraphicsResource] GraphicsResourceHandle
     ctypedef shared_ptr[const cynvrtc.nvrtcProgram] NvrtcProgramHandle
@@ -66,6 +67,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     cydriver.CUlibrary as_cu(LibraryHandle h) noexcept nogil
     cydriver.CUkernel as_cu(KernelHandle h) noexcept nogil
     cydriver.CUgraph as_cu(GraphHandle h) noexcept nogil
+    cydriver.CUgraphExec as_cu(GraphExecHandle h) noexcept nogil
     cydriver.CUgraphNode as_cu(GraphNodeHandle h) noexcept nogil
     cydriver.CUgraphicsResource as_cu(GraphicsResourceHandle h) noexcept nogil
     cynvrtc.nvrtcProgram as_cu(NvrtcProgramHandle h) noexcept nogil
@@ -87,6 +89,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     intptr_t as_intptr(LibraryHandle h) noexcept nogil
     intptr_t as_intptr(KernelHandle h) noexcept nogil
     intptr_t as_intptr(GraphHandle h) noexcept nogil
+    intptr_t as_intptr(GraphExecHandle h) noexcept nogil
     intptr_t as_intptr(GraphNodeHandle h) noexcept nogil
     intptr_t as_intptr(GraphicsResourceHandle h) noexcept nogil
     intptr_t as_intptr(NvrtcProgramHandle h) noexcept nogil
@@ -109,6 +112,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     object as_py(LibraryHandle h)
     object as_py(KernelHandle h)
     object as_py(GraphHandle h)
+    object as_py(GraphExecHandle h)
     object as_py(GraphNodeHandle h)
     object as_py(GraphicsResourceHandle h)
     object as_py(NvrtcProgramHandle h)
@@ -219,6 +223,9 @@ cdef LibraryHandle get_kernel_library(const KernelHandle& h) noexcept nogil
 cdef GraphHandle create_graph_handle(cydriver.CUgraph graph) except+ nogil
 cdef GraphHandle create_graph_handle_ref(cydriver.CUgraph graph, const GraphHandle& h_parent) except+ nogil
 
+# Graph exec handles
+cdef GraphExecHandle create_graph_exec_handle(cydriver.CUgraphExec graph_exec) except+ nogil
+
 # Graph node handles
 cdef GraphNodeHandle create_graph_node_handle(cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
 cdef GraphHandle graph_node_get_graph(const GraphNodeHandle& h) noexcept nogil

cuda_core/cuda/core/_resource_handles.pyi

@@ -13,6 +13,7 @@ DevicePtrHandle = shared_ptr
 LibraryHandle = shared_ptr
 KernelHandle = shared_ptr
 GraphHandle = shared_ptr
+GraphExecHandle = shared_ptr
 GraphNodeHandle = shared_ptr
 GraphicsResourceHandle = shared_ptr
 NvrtcProgramHandle = shared_ptr

cuda_core/cuda/core/_resource_handles.pyx

@@ -151,6 +151,10 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     GraphHandle create_graph_handle_ref "cuda_core::create_graph_handle_ref" (
         cydriver.CUgraph graph, const GraphHandle& h_parent) except+ nogil
 
+    # Graph exec handles
+    GraphExecHandle create_graph_exec_handle "cuda_core::create_graph_exec_handle" (
+        cydriver.CUgraphExec graph_exec) except+ nogil
+
     # Graph node handles
     GraphNodeHandle create_graph_node_handle "cuda_core::create_graph_node_handle" (
         cydriver.CUgraphNode node, const GraphHandle& h_graph) except+ nogil
@@ -299,6 +303,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
 
     # Graph
     void* p_cuGraphDestroy "reinterpret_cast<void*&>(cuda_core::p_cuGraphDestroy)"
+    void* p_cuGraphExecDestroy "reinterpret_cast<void*&>(cuda_core::p_cuGraphExecDestroy)"
 
     # Linker
     void* p_cuLinkDestroy "reinterpret_cast<void*&>(cuda_core::p_cuLinkDestroy)"
@@ -358,7 +363,7 @@ cdef void _init_driver_fn_pointers() noexcept:
     global p_cuMemFreeAsync, p_cuMemFree, p_cuMemFreeHost
     global p_cuMemPoolImportPointer
     global p_cuLibraryLoadFromFile, p_cuLibraryLoadData, p_cuLibraryUnload, p_cuLibraryGetKernel
-    global p_cuGraphDestroy
+    global p_cuGraphDestroy, p_cuGraphExecDestroy
     global p_cuLinkDestroy
     global p_cuGraphicsUnmapResources, p_cuGraphicsUnregisterResource
     global p_cuDevSmResourceSplit
@@ -418,6 +423,7 @@ cdef void _init_driver_fn_pointers() noexcept:
 
     # Graph
     p_cuGraphDestroy = _get_driver_fn("cuGraphDestroy")
+    p_cuGraphExecDestroy = _get_driver_fn("cuGraphExecDestroy")
 
     # Linker
     p_cuLinkDestroy = _get_driver_fn("cuLinkDestroy")

cuda_core/cuda/core/_stream.pyx

@@ -414,7 +414,7 @@ cdef class Stream:
         """
         from cuda.core.graph._graph_builder import GraphBuilder
 
-        return GraphBuilder._init(stream=self, is_stream_owner=False)
+        return GraphBuilder._init(self)
 
 
 LEGACY_DEFAULT_STREAM: Stream = Stream._legacy_default()

cuda_core/cuda/core/graph/_graph_builder.pxd

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cuda.bindings cimport cydriver
+
+from cuda.core._resource_handles cimport GraphExecHandle, GraphHandle, StreamHandle
+from cuda.core._stream cimport Stream
+
+
+cdef class GraphBuilder:
+    cdef:
+        GraphHandle _h_graph
+        StreamHandle _h_stream
+        int _kind
+        int _state
+        Stream _stream  # cached to avoid reconstruction from _h_stream handle
+        object __weakref__
+
+
+cdef class Graph:
+    cdef:
+        GraphExecHandle _h_graph_exec
+        object __weakref__
+
+    @staticmethod
+    cdef Graph _init(cydriver.CUgraphExec graph_exec)

cuda_core/cuda/core/graph/_graph_builder.pyi

@@ -8,6 +8,8 @@ from cuda.core._stream import Stream
 from cuda.core._utils.cuda_utils import driver
 from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition
 
+_BuilderKind = int
+_CaptureState = int
 
 @dataclass
 class GraphDebugPrintOptions:
@@ -106,23 +108,19 @@ class GraphBuilder:
 
     """
 
-    class _MembersNeededForFinalize:
-        __slots__ = ('conditional_graph', 'graph', 'is_join_required', 'is_stream_owner', 'stream')
-
-        def __init__(self, graph_builder_obj: GraphBuilder, stream_obj: Stream | None, is_stream_owner: bool, conditional_graph, is_join_required: bool) -> None:
-            ...
-
-        def close(self) -> None:
-            ...
-    __slots__ = ('__weakref__', '_building_ended', '_mnff')
+    def __init__(self):
+        ...
 
-    def __init__(self) -> None:
+    def __dealloc__(self):
         ...
 
-    @classmethod
-    def _init(cls, stream: Stream | None, is_stream_owner: bool, conditional_graph: object=None, is_join_required: bool=False) -> GraphBuilder:
+    @staticmethod
+    def _init(stream: Stream):
         ...
 
+    def close(self):
+        """Destroy the graph builder."""
+
     @property
     def stream(self) -> Stream:
         """Returns the stream associated with the graph builder."""
@@ -155,7 +153,7 @@ class GraphBuilder:
     def end_building(self) -> GraphBuilder:
         """Ends the building process."""
 
-    def complete(self, options: GraphCompleteOptions | None=None) -> 'Graph':
+    def complete(self, options: GraphCompleteOptions | None=None) -> Graph:
         """Completes the graph builder and returns the built :obj:`~graph.Graph` object.
 
         Parameters
@@ -245,9 +243,6 @@ class GraphBuilder:
             A condition variable for controlling conditional execution.
         """
 
-    def _cond_with_params(self, node_params: object) -> tuple[GraphBuilder, ...]:
-        ...
-
     def if_then(self, condition: GraphCondition) -> GraphBuilder:
         """Adds an if condition branch and returns a new graph builder for it.
 
@@ -335,15 +330,7 @@ class GraphBuilder:
 
         """
 
-    def close(self) -> None:
-        """Destroy the graph builder.
-
-        Closes the associated stream if we own it. Borrowed stream
-        object will instead have their references released.
-
-        """
-
-    def embed(self, child: GraphBuilder) -> None:
+    def embed(self, child: GraphBuilder):
         """Embed a previously-built :obj:`~graph.GraphBuilder` as a child node.
 
         Parameters
@@ -392,21 +379,7 @@ class Graph:
 
     """
 
-    class _MembersNeededForFinalize:
-        __slots__ = 'graph'
-
-        def __init__(self, graph_obj: Graph, graph: driver.CUgraphExec) -> None:
-            ...
-
-        def close(self) -> None:
-            ...
-    __slots__ = ('__weakref__', '_mnff')
-
-    def __init__(self) -> None:
-        ...
-
-    @classmethod
-    def _init(cls, graph: driver.CUgraphExec) -> Graph:
+    def __init__(self):
         ...
 
     def close(self) -> None:
@@ -457,5 +430,5 @@ class Graph:
         """
 __all__ = ['Graph', 'GraphBuilder', 'GraphCompleteOptions', 'GraphDebugPrintOptions']
 
-def _instantiate_graph(h_graph, options: GraphCompleteOptions | None=None) -> 'Graph':
+def _instantiate_graph(h_graph, options: GraphCompleteOptions | None=None) -> Graph:
     ...