Skip to content

Commit d77759e

Browse files
committed
Add host-callback contract tests from the stream-callback survey
Three behaviours of the existing host-callback helper were not covered: * a Python exception raised inside a callback is reported through sys.unraisablehook (the trampoline is noexcept with gil), the graph launch and the stream sync still return, and the node after the callback still runs; * user_data given as an int is the caller\x27s own pointer, delivered to a ctypes callback unchanged; * two devices running their own graphs each keep their own callback parameters and results. No API is added and no implementation code changes: the ordinary-stream callback entry point does not exist yet, so the stream-side ordering and lifetime matrix for #311 stays blocked on that dependency (reported in the survey that accompanies this change). Signed-off-by: 0z5a <[email protected]>
1 parent 8b6e9f5 commit d77759e

1 file changed

Lines changed: 149 additions & 0 deletions

File tree

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
"""Contract tests for the host-callback helper, from the #311 regression survey.
6+
7+
These add no API. They pin down behaviour of the existing helper that the suite
8+
did not cover:
9+
10+
* how a Python exception raised inside a callback is reported (the trampoline is
11+
``noexcept with gil``, so nothing may cross the C ABI),
12+
* that a caller-owned raw ``user_data`` pointer reaches a ctypes callback
13+
unchanged,
14+
* that two devices running their own graphs keep callback parameters and results
15+
separate.
16+
17+
The ordinary-stream callback entry point does not exist yet, so the stream side
18+
of the #311 matrix stays blocked on that dependency and is not asserted here.
19+
"""
20+
21+
import ctypes
22+
23+
import pytest
24+
25+
from cuda.core import Device, LegacyPinnedMemoryResource
26+
from cuda.core.graph import GraphDefinition
27+
28+
N = 16
29+
30+
31+
def _as_bytes(buffer):
32+
return bytes((ctypes.c_uint8 * N).from_address(int(buffer.handle)))
33+
34+
35+
@pytest.mark.thread_unsafe(reason="replaces the process-global sys.unraisablehook")
36+
def test_python_callback_exception_is_reported_as_unraisable(init_cuda):
37+
"""A raising callback is reported, not propagated, and the graph still finishes."""
38+
import sys
39+
40+
records = []
41+
previous_hook = sys.unraisablehook
42+
sys.unraisablehook = records.append
43+
ran = []
44+
dev = Device(0)
45+
dev.set_current()
46+
stream = dev.create_stream()
47+
device_buffer = dev.allocate(N, stream=stream)
48+
host_buffer = LegacyPinnedMemoryResource().allocate(N)
49+
50+
def throws():
51+
ran.append("callback")
52+
raise RuntimeError("callback raised")
53+
54+
try:
55+
graph_def = GraphDefinition()
56+
graph_def.memset(device_buffer, 0x33, N).callback(throws).memcpy(host_buffer, device_buffer, N)
57+
graph = graph_def.instantiate()
58+
graph.upload(stream)
59+
graph.launch(stream)
60+
stream.sync()
61+
62+
later = []
63+
after_def = GraphDefinition()
64+
after_def.callback(lambda: later.append("later"))
65+
after = after_def.instantiate()
66+
after.upload(stream)
67+
after.launch(stream)
68+
stream.sync()
69+
finally:
70+
sys.unraisablehook = previous_hook
71+
72+
assert ran == ["callback"]
73+
assert [record.exc_type for record in records] == [RuntimeError]
74+
assert "callback raised" in str(records[0].exc_value)
75+
# The node after the raising callback still ran: the launch and the sync
76+
# returned normally instead of turning the exception into a CUDA error.
77+
assert _as_bytes(host_buffer)[:4] == b"\x33" * 4
78+
assert later == ["later"]
79+
80+
81+
def test_ctypes_callback_receives_caller_owned_user_data_pointer(init_cuda):
82+
"""``user_data`` given as an int is the caller's pointer, passed through as-is."""
83+
dev = Device(0)
84+
dev.set_current()
85+
stream = dev.create_stream()
86+
payload = LegacyPinnedMemoryResource().allocate(4)
87+
bytes_view = (ctypes.c_uint8 * 4).from_address(int(payload.handle))
88+
bytes_view[:] = [0x5A, 0x00, 0x00, 0x00]
89+
raw_pointer = int(payload.handle)
90+
91+
callback_type = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
92+
seen = []
93+
94+
@callback_type
95+
def read_byte(data):
96+
seen.append((data, ctypes.cast(data, ctypes.POINTER(ctypes.c_uint8))[0]))
97+
98+
graph_def = GraphDefinition()
99+
graph_def.callback(read_byte, user_data=raw_pointer)
100+
graph = graph_def.instantiate()
101+
graph.upload(stream)
102+
graph.launch(stream)
103+
stream.sync()
104+
105+
assert seen == [(raw_pointer, 0x5A)]
106+
# The caller keeps owning the storage; the graph copied nothing.
107+
bytes_view[:] = [0x00, 0x00, 0x00, 0x00]
108+
109+
110+
def test_two_devices_keep_callback_parameters_separate(device_x2):
111+
"""Two devices, two graphs each: no callback observes another one's value."""
112+
per_stream = 2
113+
slot_count = len(device_x2) * per_stream
114+
# One pinned allocation, cut into one distinct slot per (device, stream), so
115+
# the callbacks cannot alias each other's storage; the owner lives here.
116+
host_slots = LegacyPinnedMemoryResource().allocate(N * slot_count)
117+
host_base = int(host_slots.handle)
118+
runs = []
119+
streams = []
120+
device_buffers = []
121+
expected = []
122+
for device_index, dev in enumerate(device_x2):
123+
dev.set_current()
124+
for stream_index in range(per_stream):
125+
value = 0xA1 + device_index * 0x10 + stream_index
126+
slot = device_index * per_stream + stream_index
127+
expected.append(((device_index, stream_index), value))
128+
device_stream = dev.create_stream()
129+
device_buffer = dev.allocate(N, stream=device_stream)
130+
host_pointer = host_base + slot * N
131+
132+
def snapshot(label=(device_index, stream_index), host_pointer=host_pointer):
133+
runs.append((label, ctypes.c_uint8.from_address(host_pointer).value))
134+
135+
graph_def = GraphDefinition()
136+
graph_def.memset(device_buffer, value, N).memcpy(
137+
host_pointer, device_buffer, N, dst_owner=host_slots
138+
).callback(snapshot)
139+
graph = graph_def.instantiate()
140+
graph.upload(device_stream)
141+
graph.launch(device_stream)
142+
streams.append(device_stream)
143+
device_buffers.append(device_buffer)
144+
145+
for device_stream in streams:
146+
device_stream.sync()
147+
148+
assert sorted(runs) == sorted(expected)
149+
assert len(runs) == slot_count

0 commit comments

Comments
 (0)