Add CPU callbacks for stream capture (GraphBuilder.callback)

Andy-Jost · Andy-Jost · commit edbc361d2d46 · 2026-03-24T13:46:17.000-07:00
Implements #1328: host callbacks during stream capture via cuLaunchHostFunc, mirroring the existing GraphDef.callback API. Extracts shared callback infrastructure (_attach_user_object, _attach_host_callback_to_graph, trampoline/destructor) into a new _graph/_utils.pyx module to avoid circular imports between _graph_builder and _graphdef. Made-with: Cursor
diff --git a/cuda_core/cuda/core/_graph/_graph_builder.pyx b/cuda_core/cuda/core/_graph/_graph_builder.pyx
@@ -5,7 +5,12 @@
 import weakref
 from dataclasses import dataclass
 
+from cuda.bindings cimport cydriver
+
+from cuda.core._graph._utils cimport _attach_host_callback_to_graph
+from cuda.core._resource_handles cimport as_cu
 from cuda.core._stream cimport Stream
+from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
 from cuda.core._utils.cuda_utils import (
     driver,
     get_binding_version,
@@ -682,6 +687,57 @@ class GraphBuilder:
             )
         )
 
+    def callback(self, fn, *, user_data=None):
+        """Add a host callback to the graph during stream capture.
+
+        The callback runs on the host CPU when the graph reaches this point
+        in execution. Two modes are supported:
+
+        - **Python callable**: Pass any callable. The GIL is acquired
+          automatically. The callable must take no arguments; use closures
+          or ``functools.partial`` to bind state.
+        - **ctypes function pointer**: Pass a ``ctypes.CFUNCTYPE`` instance.
+          The function receives a single ``void*`` argument (the
+          ``user_data``). The caller must keep the ctypes wrapper alive
+          for the lifetime of the graph.
+
+        .. warning::
+
+            Callbacks must not call CUDA API functions. Doing so may
+            deadlock or corrupt driver state.
+
+        Parameters
+        ----------
+        fn : callable or ctypes function pointer
+            The callback function.
+        user_data : int or bytes-like, optional
+            Only for ctypes function pointers. If ``int``, passed as a raw
+            pointer (caller manages lifetime). If bytes-like, the data is
+            copied and its lifetime is tied to the graph.
+        """
+        cdef Stream stream = <Stream>self._mnff.stream
+        cdef cydriver.CUstream c_stream = as_cu(stream._h_stream)
+        cdef cydriver.CUstreamCaptureStatus capture_status
+        cdef cydriver.CUgraph c_graph = NULL
+
+        with nogil:
+            IF CUDA_CORE_BUILD_MAJOR >= 13:
+                HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
+                    c_stream, &capture_status, NULL, &c_graph, NULL, NULL, NULL))
+            ELSE:
+                HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
+                    c_stream, &capture_status, NULL, &c_graph, NULL, NULL))
+
+        if capture_status != cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE:
+            raise RuntimeError("Cannot add callback when graph is not being built")
+
+        cdef cydriver.CUhostFn c_fn
+        cdef void* c_user_data = NULL
+        _attach_host_callback_to_graph(c_graph, fn, user_data, &c_fn, &c_user_data)
+
+        with nogil:
+            HANDLE_RETURN(cydriver.cuLaunchHostFunc(c_stream, c_fn, c_user_data))
+
 
 class Graph:
     """Represents an executable graph.
diff --git a/cuda_core/cuda/core/_graph/_graphdef.pyx b/cuda_core/cuda/core/_graph/_graphdef.pyx
@@ -30,8 +30,6 @@ GraphNode hierarchy:
 
 from __future__ import annotations
 
-from cpython.ref cimport Py_INCREF
-
 from libc.stddef cimport size_t
 from libc.stdint cimport uintptr_t
 from libc.stdlib cimport malloc, free
@@ -102,16 +100,11 @@ cdef bint _check_node_get_params():
     return _has_cuGraphNodeGetParams
 
 
-cdef extern from "Python.h":
-    void _py_decref "Py_DECREF" (void*)
-
-
-cdef void _py_host_trampoline(void* data) noexcept with gil:
-    (<object>data)()
-
-
-cdef void _py_host_destructor(void* data) noexcept with gil:
-    _py_decref(data)
+from cuda.core._graph._utils cimport (
+    _attach_host_callback_to_graph,
+    _attach_user_object,
+    _is_py_host_trampoline,
+)
 
 
 cdef void _destroy_event_handle_copy(void* ptr) noexcept nogil:
@@ -124,30 +117,6 @@ cdef void _destroy_kernel_handle_copy(void* ptr) noexcept nogil:
     del p
 
 
-cdef void _attach_user_object(
-        cydriver.CUgraph graph, void* ptr,
-        cydriver.CUhostFn destroy) except *:
-    """Create a CUDA user object and transfer ownership to the graph.
-
-    On success the graph owns the resource (via MOVE semantics).
-    On failure the destroy callback is invoked to clean up ptr,
-    then a CUDAError is raised — callers need no try/except.
-    """
-    cdef cydriver.CUuserObject user_obj = NULL
-    cdef cydriver.CUresult ret
-    with nogil:
-        ret = cydriver.cuUserObjectCreate(
-            &user_obj, ptr, destroy, 1,
-            cydriver.CU_USER_OBJECT_NO_DESTRUCTOR_SYNC)
-        if ret == cydriver.CUDA_SUCCESS:
-            ret = cydriver.cuGraphRetainUserObject(
-                graph, user_obj, 1, cydriver.CU_GRAPH_USER_OBJECT_MOVE)
-            if ret != cydriver.CUDA_SUCCESS:
-                cydriver.cuUserObjectRelease(user_obj, 1)
-    if ret != cydriver.CUDA_SUCCESS:
-        if user_obj == NULL:
-            destroy(ptr)
-        HANDLE_RETURN(ret)
 
 
 cdef class Condition:
@@ -1270,56 +1239,20 @@ cdef class GraphNode:
         cdef cydriver.CUgraphNode pred_node = as_cu(self._h_node)
         cdef cydriver.CUgraphNode* deps = NULL
         cdef size_t num_deps = 0
-        cdef void* c_user_data = NULL
-        cdef object callable_obj = None
-        cdef void* fn_pyobj = NULL
 
         if pred_node != NULL:
             deps = &pred_node
             num_deps = 1
 
-        if isinstance(fn, ct._CFuncPtr):
-            Py_INCREF(fn)
-            fn_pyobj = <void*>fn
-            _attach_user_object(
-                as_cu(h_graph), fn_pyobj,
-                <cydriver.CUhostFn>_py_host_destructor)
-            node_params.fn = <cydriver.CUhostFn><uintptr_t>ct.cast(
-                fn, ct.c_void_p).value
-
-            if user_data is not None:
-                if isinstance(user_data, int):
-                    c_user_data = <void*><uintptr_t>user_data
-                else:
-                    buf = bytes(user_data)
-                    c_user_data = malloc(len(buf))
-                    if c_user_data == NULL:
-                        raise MemoryError(
-                            "failed to allocate user_data buffer")
-                    c_memcpy(c_user_data, <const char*>buf, len(buf))
-                    _attach_user_object(
-                        as_cu(h_graph), c_user_data,
-                        <cydriver.CUhostFn>free)
-
-            node_params.userData = c_user_data
-        else:
-            if user_data is not None:
-                raise ValueError(
-                    "user_data is only supported with ctypes "
-                    "function pointers")
-            callable_obj = fn
-            Py_INCREF(fn)
-            fn_pyobj = <void*>fn
-            node_params.fn = <cydriver.CUhostFn>_py_host_trampoline
-            node_params.userData = fn_pyobj
-            _attach_user_object(
-                as_cu(h_graph), fn_pyobj,
-                <cydriver.CUhostFn>_py_host_destructor)
+        _attach_host_callback_to_graph(
+            as_cu(h_graph), fn, user_data,
+            &node_params.fn, &node_params.userData)
 
         with nogil:
             HANDLE_RETURN(cydriver.cuGraphAddHostNode(
                 &new_node, as_cu(h_graph), deps, num_deps, &node_params))
 
+        cdef object callable_obj = fn if not isinstance(fn, ct._CFuncPtr) else None
         self._succ_cache = None
         return HostCallbackNode._create_with_params(
             create_graph_node_handle(new_node, h_graph), callable_obj,
@@ -1947,7 +1880,7 @@ cdef class HostCallbackNode(GraphNode):
             HANDLE_RETURN(cydriver.cuGraphHostNodeGetParams(node, &params))
 
         cdef object callable_obj = None
-        if params.fn == <cydriver.CUhostFn>_py_host_trampoline:
+        if _is_py_host_trampoline(params.fn):
             callable_obj = <object>params.userData
 
         return HostCallbackNode._create_with_params(
diff --git a/cuda_core/cuda/core/_graph/_utils.pxd b/cuda_core/cuda/core/_graph/_utils.pxd
@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cuda.bindings cimport cydriver
+
+
+cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil
+
+cdef void _attach_user_object(
+    cydriver.CUgraph graph, void* ptr,
+    cydriver.CUhostFn destroy) except *
+
+cdef void _attach_host_callback_to_graph(
+    cydriver.CUgraph graph, object fn, object user_data,
+    cydriver.CUhostFn* out_fn, void** out_user_data) except *
diff --git a/cuda_core/cuda/core/_graph/_utils.pyx b/cuda_core/cuda/core/_graph/_utils.pyx
@@ -0,0 +1,106 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from cpython.ref cimport Py_INCREF
+
+from libc.stdint cimport uintptr_t
+from libc.stdlib cimport malloc, free
+from libc.string cimport memcpy as c_memcpy
+
+from cuda.bindings cimport cydriver
+
+from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
+
+
+cdef extern from "Python.h":
+    void _py_decref "Py_DECREF" (void*)
+
+
+cdef void _py_host_trampoline(void* data) noexcept with gil:
+    (<object>data)()
+
+
+cdef void _py_host_destructor(void* data) noexcept with gil:
+    _py_decref(data)
+
+
+cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil:
+    return fn == <cydriver.CUhostFn>_py_host_trampoline
+
+
+cdef void _attach_user_object(
+        cydriver.CUgraph graph, void* ptr,
+        cydriver.CUhostFn destroy) except *:
+    """Create a CUDA user object and transfer ownership to the graph.
+
+    On success the graph owns the resource (via MOVE semantics).
+    On failure the destroy callback is invoked to clean up ptr,
+    then a CUDAError is raised — callers need no try/except.
+    """
+    cdef cydriver.CUuserObject user_obj = NULL
+    cdef cydriver.CUresult ret
+    with nogil:
+        ret = cydriver.cuUserObjectCreate(
+            &user_obj, ptr, destroy, 1,
+            cydriver.CU_USER_OBJECT_NO_DESTRUCTOR_SYNC)
+        if ret == cydriver.CUDA_SUCCESS:
+            ret = cydriver.cuGraphRetainUserObject(
+                graph, user_obj, 1, cydriver.CU_GRAPH_USER_OBJECT_MOVE)
+            if ret != cydriver.CUDA_SUCCESS:
+                cydriver.cuUserObjectRelease(user_obj, 1)
+    if ret != cydriver.CUDA_SUCCESS:
+        if user_obj == NULL:
+            destroy(ptr)
+        HANDLE_RETURN(ret)
+
+
+cdef void _attach_host_callback_to_graph(
+        cydriver.CUgraph graph, object fn, object user_data,
+        cydriver.CUhostFn* out_fn, void** out_user_data) except *:
+    """Resolve a Python callable or ctypes CFuncPtr into a C callback pair.
+
+    Handles Py_INCREF, user-object attachment for lifetime management,
+    and user_data copying.  On return, *out_fn and *out_user_data are
+    ready to pass to cuGraphAddHostNode or cuLaunchHostFunc.
+    """
+    import ctypes as ct
+
+    cdef void* fn_pyobj = NULL
+
+    if isinstance(fn, ct._CFuncPtr):
+        Py_INCREF(fn)
+        fn_pyobj = <void*>fn
+        _attach_user_object(
+            graph, fn_pyobj,
+            <cydriver.CUhostFn>_py_host_destructor)
+        out_fn[0] = <cydriver.CUhostFn><uintptr_t>ct.cast(
+            fn, ct.c_void_p).value
+
+        if user_data is not None:
+            if isinstance(user_data, int):
+                out_user_data[0] = <void*><uintptr_t>user_data
+            else:
+                buf = bytes(user_data)
+                out_user_data[0] = malloc(len(buf))
+                if out_user_data[0] == NULL:
+                    raise MemoryError(
+                        "failed to allocate user_data buffer")
+                c_memcpy(out_user_data[0], <const char*>buf, len(buf))
+                _attach_user_object(
+                    graph, out_user_data[0],
+                    <cydriver.CUhostFn>free)
+        else:
+            out_user_data[0] = NULL
+    else:
+        if user_data is not None:
+            raise ValueError(
+                "user_data is only supported with ctypes "
+                "function pointers")
+        Py_INCREF(fn)
+        fn_pyobj = <void*>fn
+        out_fn[0] = <cydriver.CUhostFn>_py_host_trampoline
+        out_user_data[0] = fn_pyobj
+        _attach_user_object(
+            graph, fn_pyobj,
+            <cydriver.CUhostFn>_py_host_destructor)
diff --git a/cuda_core/tests/graph/test_basic.py b/cuda_core/tests/graph/test_basic.py
@@ -163,3 +163,45 @@ def test_graph_capture_errors(init_cuda):
     with pytest.raises(RuntimeError, match="^Graph has not finished building."):
         gb.complete()
     gb.end_building().complete()
+
+
+def test_graph_capture_callback_python(init_cuda):
+    results = []
+
+    def my_callback():
+        results.append(42)
+
+    launch_stream = Device().create_stream()
+    gb = launch_stream.create_graph_builder().begin_building()
+
+    with pytest.raises(ValueError, match="user_data is only supported"):
+        gb.callback(my_callback, user_data=b"hello")
+
+    gb.callback(my_callback)
+    graph = gb.end_building().complete()
+
+    graph.launch(launch_stream)
+    launch_stream.sync()
+
+    assert results == [42]
+
+
+def test_graph_capture_callback_ctypes(init_cuda):
+    import ctypes
+
+    CALLBACK = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
+    result = [0]
+
+    @CALLBACK
+    def read_byte(data):
+        result[0] = ctypes.cast(data, ctypes.POINTER(ctypes.c_uint8))[0]
+
+    launch_stream = Device().create_stream()
+    gb = launch_stream.create_graph_builder().begin_building()
+    gb.callback(read_byte, user_data=bytes([0xAB]))
+    graph = gb.end_building().complete()
+
+    graph.launch(launch_stream)
+    launch_stream.sync()
+
+    assert result[0] == 0xAB
