Skip to content

Commit f80b7de

Browse files
Andy-Jostclaude
Andy-Jost
and
claude
committed
[doc-only] cuda.core: annotate host callback fn and user_data
Annotate fn and user_data on host_launch and on the callback methods of GraphBuilder, GraphNode, and GraphDefinition. Drop cuda_bindings from the script dependencies of the host_launch example, which does not import it. Follow-up to review comments on #3258. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent 230f0eb commit f80b7de

9 files changed

Lines changed: 31 additions & 10 deletions

File tree

‎cuda_core/cuda/core/_launcher.pyi‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
# This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/_launcher.pyx
22

3+
from collections.abc import Callable
4+
35
from cuda.core._launch_config import LaunchConfig
46
from cuda.core._module import Kernel
57
from cuda.core._stream import Stream
8+
from cuda.core._utils.pycompat import BufferProtocol
69
from cuda.core.graph import GraphBuilder
710
from cuda.core.typing import IsStreamType
811

@@ -27,7 +30,7 @@ def launch(stream: Stream | GraphBuilder | IsStreamType, config: LaunchConfig, k
2730
launching kernel.
2831
2932
"""
30-
def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=None) -> None:
33+
def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> None:
3134
"""Enqueue a host function on a stream in stream order.
3235
3336
The function runs on a driver-owned thread once all work enqueued on the

‎cuda_core/cuda/core/_launcher.pyx‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@ from math import prod
2222
from typing import TYPE_CHECKING
2323

2424
if TYPE_CHECKING:
25+
from collections.abc import Callable
26+
27+
from cuda.core._utils.pycompat import BufferProtocol
2528
from cuda.core.graph import GraphBuilder
2629
from cuda.core.typing import IsStreamType
2730

@@ -70,7 +73,7 @@ def launch(
7073
HANDLE_RETURN(cydriver.cuLaunchKernelEx(&drv_cfg, func_handle, args_ptr, NULL))
7174

7275

73-
def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn, *, user_data=None) -> None:
76+
def host_launch(stream: Stream | GraphBuilder | IsStreamType, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> None:
7477
"""Enqueue a host function on a stream in stream order.
7578

7679
The function runs on a driver-owned thread once all work enqueued on the

‎cuda_core/cuda/core/graph/_graph_builder.pyi‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
# This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/graph/_graph_builder.pyx
22

3+
from collections.abc import Callable
34
from dataclasses import dataclass
45
from typing import overload
56

67
from cuda.core._stream import Stream
78
from cuda.core._utils.cuda_utils import driver
9+
from cuda.core._utils.pycompat import BufferProtocol
810
from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition
911
from cuda.core.graph._subclasses import (ChildGraphNode, EventRecordNode,
1012
EventWaitNode,
@@ -408,7 +410,7 @@ class GraphBuilder:
408410
child : :obj:`~graph.GraphBuilder`
409411
The child graph builder. Must have finished building.
410412
"""
411-
def callback(self, fn, *, user_data=None) -> None:
413+
def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> None:
412414
"""Add a host callback to the graph during stream capture.
413415
414416
The callback runs on the host CPU when the graph reaches this point

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,11 @@ from cuda.core._utils.cuda_utils import (
6262
)
6363

6464
if TYPE_CHECKING:
65+
from collections.abc import Callable
6566
from typing import overload
6667

68+
from cuda.core._utils.pycompat import BufferProtocol
69+
6770
from cuda.core.graph._graph_definition import GraphDefinition
6871
else:
6972
# Cython applies a method decorator by rebinding the type's dict entry, so
@@ -947,7 +950,7 @@ cdef class GraphBuilder:
947950
)
948951
)
949952
950-
def callback(self, fn, *, user_data=None) -> None:
953+
def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> None:
951954
"""Add a host callback to the graph during stream capture.
952955

953956
The callback runs on the host CPU when the graph reaches this point

‎cuda_core/cuda/core/graph/_graph_definition.pyi‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
# This file was generated by stubgen-pyx v0.2.23 from cuda_core/cuda/core/graph/_graph_definition.pyx
22

33
"""GraphDefinition: explicit CUDA graph definition."""
4+
from collections.abc import Callable
5+
46
from cuda.core._device import Device
57
from cuda.core._event import Event
68
from cuda.core._launch_config import LaunchConfig
79
from cuda.core._memory._buffer import Buffer
810
from cuda.core._module import Kernel
911
from cuda.core._utils.cuda_utils import driver
12+
from cuda.core._utils.pycompat import BufferProtocol
1013
from cuda.core.graph._graph_builder import (Graph, GraphCompleteOptions,
1114
GraphDebugPrintOptions)
1215
from cuda.core.graph._graph_node import GraphNode
@@ -127,7 +130,7 @@ class GraphDefinition:
127130
128131
See :meth:`GraphNode.wait` for full documentation.
129132
"""
130-
def callback(self, fn, *, user_data=None) -> HostCallbackNode:
133+
def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> HostCallbackNode:
131134
"""Add an entry-point host callback node (no dependencies).
132135
133136
See :meth:`GraphNode.callback` for full documentation.

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,14 @@ from cuda.core._utils.cuda_utils import driver
3131
from cuda.core.typing import GraphMemoryType
3232

3333
if TYPE_CHECKING:
34+
from collections.abc import Callable
35+
3436
from cuda.core._device import Device
3537
from cuda.core._event import Event
3638
from cuda.core._launch_config import LaunchConfig
3739
from cuda.core._memory._buffer import Buffer
3840
from cuda.core._module import Kernel
41+
from cuda.core._utils.pycompat import BufferProtocol
3942
from cuda.core.graph._graph_builder import (
4043
Graph,
4144
GraphCompleteOptions,
@@ -250,7 +253,7 @@ cdef class GraphDefinition:
250253
"""
251254
return self._entry.wait(event)
252255

253-
def callback(self, fn, *, user_data=None) -> HostCallbackNode:
256+
def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> HostCallbackNode:
254257
"""Add an entry-point host callback node (no dependencies).
255258

256259
See :meth:`GraphNode.callback` for full documentation.

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,15 @@
22

33
"""GraphNode base class — factory, properties, and builder methods."""
44
import weakref
5-
from collections.abc import Iterable
5+
from collections.abc import Callable, Iterable
66

77
from cuda.core._device import Device
88
from cuda.core._event import Event
99
from cuda.core._launch_config import LaunchConfig
1010
from cuda.core._memory._buffer import Buffer
1111
from cuda.core._module import Kernel
1212
from cuda.core._utils.cuda_utils import driver
13+
from cuda.core._utils.pycompat import BufferProtocol
1314
from cuda.core.graph._adjacency_set_proxy import AdjacencySetProxy
1415
from cuda.core.graph._graph_definition import GraphCondition, GraphDefinition
1516
from cuda.core.graph._subclasses import (AllocNode, ChildGraphNode, EmptyNode,
@@ -297,7 +298,7 @@ class GraphNode:
297298
EventWaitNode
298299
A new EventWaitNode representing the event wait operation.
299300
"""
300-
def callback(self, fn, *, user_data=None) -> object:
301+
def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None=None) -> object:
301302
"""Add a host callback node depending on this node.
302303
303304
The callback runs on the host CPU when the graph reaches this node.

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,10 @@ from cuda.core._utils.cuda_utils import driver
7979
from cuda.core.typing import GraphMemoryType
8080

8181
if TYPE_CHECKING:
82+
from collections.abc import Callable
83+
8284
from cuda.core._device import Device
85+
from cuda.core._utils.pycompat import BufferProtocol
8386

8487
__all__ = ['GraphNode']
8588

@@ -495,7 +498,7 @@ cdef class GraphNode:
495498
"""
496499
return GN_wait_event(self, event)
497500
498-
def callback(self, fn, *, user_data=None) -> object:
501+
def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> object:
499502
"""Add a host callback node depending on this node.
500503

501504
The callback runs on the host CPU when the graph reaches this node.

‎cuda_core/examples/host_launch.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
# ################################################################################
1111

1212
# /// script
13-
# dependencies = ["cuda_bindings", "cuda_core"]
13+
# dependencies = ["cuda_core"]
1414
# ///
1515

1616
import ctypes

0 commit comments

Comments
 (0)