Skip to content

Commit 35d1722

Browse files
cuda.core: keep kernel-argument objects alive in graph kernel nodes (#2041)
* cuda.core: keep kernel-argument objects alive in graph kernel nodes `GraphDefinition.launch()` did not extend the lifetime of the Python kernel-argument objects to the lifetime of the graph. The `ParamHolder` built in `GN_launch` held the only references to those objects and was destroyed when `GN_launch` returned. The driver only stores the raw pointer values in the kernel node, so a `Buffer` reachable only through the call could be GC'd before the graph ran, leaving the graph with a stale device pointer. Attach the `kernel_args` tuple to the graph as a CUDA user object, mirroring the existing handling of `KernelHandle` and `EventHandle`. This reuses the `_py_host_destructor` path already used by the host callback machinery. Closes #2039 Co-authored-by: Cursor <[email protected]> * test: assert kernel-arg Buffer is freed when graph is released Addresses review feedback (PR #2041): the existing test only proved the graph kept the Buffer alive, not that the user-object machinery actually releases it once the graph is destroyed. Without the symmetric check, a working attachment is indistinguishable from a permanent leak. Co-authored-by: Cursor <[email protected]> * test: poll for kernel-arg Buffer release to handle async cleanup The freeing assertion at the end of test_kernel_args_buffer_lifetime failed on free-threaded Python (py3.14t) because cuGraphExecDestroy releases its user-object references via an asynchronous DPC, and free- threaded CPython's deferred ref counting can need an extra GC pass to settle. Poll the weakref with a bounded timeout and per-iteration GC instead of asserting eagerly. Co-authored-by: Cursor <[email protected]> --------- Co-authored-by: Cursor <[email protected]>
1 parent 57a538c commit 35d1722

3 files changed

Lines changed: 117 additions & 1 deletion

File tree

‎cuda_core/cuda/core/graph/_graph_node.pyx‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
from __future__ import annotations
88

9+
from cpython.ref cimport Py_INCREF
10+
911
from libc.stddef cimport size_t
1012
from libc.stdint cimport uintptr_t
1113
from libc.string cimport memset as c_memset
@@ -54,6 +56,7 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
5456
from cuda.core.graph._utils cimport (
5557
_attach_host_callback_to_graph,
5658
_attach_user_object,
59+
_py_host_destructor,
5760
)
5861

5962
import weakref
@@ -617,6 +620,12 @@ cdef inline KernelNode GN_launch(GraphNode self, LaunchConfig conf, Kernel ker,
617620
_attach_user_object(as_cu(h_graph), <void*>new KernelHandle(ker._h_kernel),
618621
<cydriver.CUhostFn>_destroy_kernel_handle_copy)
619622

623+
cdef object kernel_args = ker_args.kernel_args
624+
if kernel_args is not None:
625+
Py_INCREF(kernel_args)
626+
_attach_user_object(as_cu(h_graph), <void*>kernel_args,
627+
<cydriver.CUhostFn>_py_host_destructor)
628+
620629
return _registered(KernelNode._create_with_params(
621630
create_graph_node_handle(new_node, h_graph),
622631
conf.grid, conf.block, conf.shmem_size,

‎cuda_core/cuda/core/graph/_utils.pxd‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ from cuda.bindings cimport cydriver
77

88
cdef bint _is_py_host_trampoline(cydriver.CUhostFn fn) noexcept nogil
99

10+
cdef void _py_host_destructor(void* data) noexcept with gil
11+
1012
cdef void _attach_user_object(
1113
cydriver.CUgraph graph, void* ptr,
1214
cydriver.CUhostFn destroy) except *

‎cuda_core/tests/graph/test_graph_definition_lifetime.py‎

Lines changed: 106 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,34 @@
33

44
"""Tests for GraphDefinition resource lifetime management and RAII correctness."""
55

6+
import ctypes
67
import gc
8+
import time
9+
import weakref
710

811
import pytest
912
from helpers.graph_kernels import compile_common_kernels
1013
from helpers.misc import try_create_condition
1114

12-
from cuda.core import Device, EventOptions, Kernel, LaunchConfig
15+
16+
def _wait_until(predicate, timeout=2.0, interval=0.01):
17+
"""Poll predicate() until True or timeout, driving gc each iteration.
18+
19+
Used for assertions about resource cleanup that may be delayed by CUDA's
20+
asynchronous user-object destructor pump (DPC) or, on free-threaded
21+
Python, by deferred reference-count processing. A bounded poll keeps the
22+
test correct without depending on undocumented driver timing guarantees.
23+
"""
24+
deadline = time.monotonic() + timeout
25+
while time.monotonic() < deadline:
26+
gc.collect()
27+
if predicate():
28+
return
29+
time.sleep(interval)
30+
raise AssertionError(f"condition not satisfied within {timeout}s")
31+
32+
33+
from cuda.core import Device, DeviceMemoryResource, EventOptions, Kernel, LaunchConfig
1334
from cuda.core.graph import (
1435
ChildGraphNode,
1536
ConditionalNode,
@@ -485,3 +506,87 @@ def test_kernel_node_reconstruction_preserves_validity(init_cuda):
485506
stream = Device().create_stream()
486507
graph.launch(stream)
487508
stream.sync()
509+
510+
511+
# =============================================================================
512+
# Kernel argument lifetime — kernel nodes should keep argument objects alive
513+
# =============================================================================
514+
515+
516+
def test_kernel_args_buffer_lifetime(init_cuda):
517+
"""Buffer passed as a kernel arg is kept alive by the graph, the kernel
518+
executes against its memory after the original Python ref drops, and the
519+
Buffer is released once the graph is destroyed.
520+
521+
Without the user-object attachment, the ParamHolder is destroyed when the
522+
kernel node is added, the Buffer is GC'd, and the graph is left with a
523+
stale device pointer.
524+
525+
The final freeing assertion uses a bounded poll because CUgraphExec
526+
releases its user-object references via an asynchronous DPC, and on
527+
free-threaded Python the resulting Py_DECREF chain may need an extra
528+
GC pass to settle.
529+
"""
530+
from cuda.core._utils.cuda_utils import driver, handle_return
531+
532+
_skip_if_no_mempool()
533+
dev = Device()
534+
mr = DeviceMemoryResource(dev)
535+
add_one = compile_common_kernels().get_kernel("add_one")
536+
buf = mr.allocate(ctypes.sizeof(ctypes.c_int), stream=dev.default_stream)
537+
buf.fill(0, stream=dev.default_stream)
538+
dev.default_stream.sync()
539+
buf_weak = weakref.ref(buf)
540+
dptr = int(buf.handle)
541+
542+
g = GraphDefinition()
543+
g.launch(LaunchConfig(grid=1, block=1), add_one, buf)
544+
545+
del buf
546+
gc.collect()
547+
assert buf_weak() is not None # graph kept the Buffer alive
548+
549+
stream = dev.create_stream()
550+
g.instantiate().launch(stream)
551+
stream.sync()
552+
553+
out = (ctypes.c_int * 1)(0)
554+
handle_return(driver.cuMemcpyDtoH(out, dptr, ctypes.sizeof(ctypes.c_int)))
555+
assert out[0] == 1
556+
557+
del g
558+
_wait_until(lambda: buf_weak() is None)
559+
560+
561+
def test_kernel_args_survive_graph_clone(init_cuda):
562+
"""Cloned graph keeps Buffer alive via CUDA user objects.
563+
564+
A graph clone does not inherit Python-level references, so only user
565+
objects (which propagate through cuGraphClone) can keep the args alive.
566+
"""
567+
from cuda.core._utils.cuda_utils import driver, handle_return
568+
569+
_skip_if_no_mempool()
570+
dev = Device()
571+
mr = DeviceMemoryResource(dev)
572+
add_one = compile_common_kernels().get_kernel("add_one")
573+
buf = mr.allocate(ctypes.sizeof(ctypes.c_int), stream=dev.default_stream)
574+
buf.fill(0, stream=dev.default_stream)
575+
dev.default_stream.sync()
576+
dptr = int(buf.handle)
577+
578+
g = GraphDefinition()
579+
g.launch(LaunchConfig(grid=1, block=1), add_one, buf)
580+
cloned_cu_graph = handle_return(driver.cuGraphClone(driver.CUgraph(g.handle)))
581+
582+
del buf, g
583+
gc.collect()
584+
585+
graph_exec = handle_return(driver.cuGraphInstantiate(cloned_cu_graph, 0))
586+
stream = dev.create_stream()
587+
handle_return(driver.cuGraphLaunch(graph_exec, driver.CUstream(int(stream.handle))))
588+
stream.sync()
589+
590+
out = (ctypes.c_int * 1)(0)
591+
handle_return(driver.cuMemcpyDtoH(out, dptr, ctypes.sizeof(ctypes.c_int)))
592+
assert out[0] == 1

0 commit comments

Comments
 (0)