Skip to content

Commit 2e6aa4e

Browse files
0z5aclaudeAndy-Jost
authored
Add host-callback contract tests from the stream-callback survey (#2925)
* 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 <Dezhen.lu@student.uni-tuebingen.de> * Address review: restore the hook on failure, add markers, release buffers Install the unraisable hook with monkeypatch so it is restored even if stream or buffer setup raises. Mark the three tests agent_authored, trim the module docstring to the contracts under test, and release every buffer the tests construct through context managers. Signed-off-by: 0z5a <Dezhen.lu@student.uni-tuebingen.de> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Signed-off-by: 0z5a <Dezhen.lu@student.uni-tuebingen.de> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Andy Jost <ajost@nvidia.com>
1 parent fd5f89d commit 2e6aa4e

1 file changed

Lines changed: 143 additions & 0 deletions

File tree

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

0 commit comments

Comments
 (0)