Skip to content

Commit 0bdf9e9

Browse files
Andy-Jostclaude
Andy-Jost
and
claude
authored
fix(cuda.core): pickle Device by UUID (#3198)
* fix(cuda.core): pickle memory-resource devices by UUID for IPC A DeviceMemoryResource sent through multiprocessing, and a pickled VirtualMemoryResource, named their device by ordinal, and the VirtualMemoryResource also sent its peers as ordinals. Ordinals are per process, so a receiver with a different CUDA_VISIBLE_DEVICES rebuilt the resource on the wrong GPU or granted peer access to the wrong GPU. Send the owner and the peers by UUID and map them back in the receiver with a new private helper. Raise a clear RuntimeError when no visible device has the UUID. The tests reverse the device order in the child through CUDA_VISIBLE_DEVICES and check the imported buffer's device UUID. Closes #3097 Co-Authored-By: Claude Fable 5.1 <[email protected]> * fix(cuda.core): mark the IPC reconstructor's handle parameter not None The check-cython-not-none hook from main requires it on a def parameter annotated with a cdef class. Co-Authored-By: Claude Fable 5.1 <[email protected]> * fix(cuda.core): store VirtualMemoryResourceOptions.peers as a tuple A generator in `peers` was consumed by the first allocation. Later allocations granted no peer access, and a pickled resource sent no peer UUIDs, so the receiver mapped the memory without peer access. `__post_init__` now converts `peers` to a tuple, which also covers options given as a dict, through `replace()`, or per call to `modify_allocation()`. Co-Authored-By: Claude Opus 5.5 <[email protected]> * fix(cuda.core): pickle Device by UUID Device.__reduce__ sent the ordinal, so a process with a different CUDA_VISIBLE_DEVICES unpickled a different GPU. It now sends the UUID, and the receiver maps it back with _device_from_uuid(), which raises RuntimeError when the device is not visible. Closes #3196. Co-Authored-By: Claude Opus 5.5 <[email protected]> * fix(cuda.core): say the Device pickle is for one machine; close test resources on every path - The Device docstring and the release note say the pickled form is for processes on the same machine (asked on #3196). - The buffer tests in test_device_uuid.py release their buffers and pools in finally blocks, so a failed assertion cannot leave an owned pool open for the next test's leak check (CodeRabbit on #3198). - The module docstring cites #3196 for the Device requirement. Co-Authored-By: Claude Fable 5.1 <[email protected]> --------- Co-authored-by: Claude Fable 5.1 <[email protected]>
1 parent 2da9c65 commit 0bdf9e9

5 files changed

Lines changed: 97 additions & 44 deletions

File tree

‎cuda_core/cuda/core/_device.pyi‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -479,6 +479,11 @@ class Device:
479479
Newly returned :obj:`~_device.Device` objects are thread-local singletons
480480
for a specified device.
481481
482+
A pickled :obj:`~_device.Device` names its device by UUID. Unpickling
483+
returns the device with that UUID in the receiving process, whatever its
484+
ordinal is there, and raises ``RuntimeError`` if that device is not
485+
visible. The pickled form is for processes on the same machine.
486+
482487
Note
483488
----
484489
Will not initialize the GPU.

‎cuda_core/cuda/core/_device.pyx‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -983,6 +983,11 @@ class Device:
983983
Newly returned :obj:`~_device.Device` objects are thread-local singletons
984984
for a specified device.
985985

986+
A pickled :obj:`~_device.Device` names its device by UUID. Unpickling
987+
returns the device with that UUID in the receiving process, whatever its
988+
ordinal is there, and raises ``RuntimeError`` if that device is not
989+
visible. The pickled form is for processes on the same machine.
990+
986991
Note
987992
----
988993
Will not initialize the GPU.
@@ -1242,7 +1247,8 @@ class Device:
12421247
return self._device_id == other._device_id
12431248

12441249
def __reduce__(self) -> tuple[object, ...]:
1245-
return Device, (self.device_id,)
1250+
# Device ordinals are per process; the receiver maps the UUID to its own device.
1251+
return _device_from_uuid, (self.uuid,)
12461252

12471253
def set_current(self, Context ctx: Context | None = None) -> Context | None:
12481254
"""Set device to be used for GPU executions.

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,14 @@ Fixes and enhancements
113113
one of the devices is not visible in the receiving process.
114114
(`#3097 <https://github.com/NVIDIA/cuda-python/issues/3097>`__)
115115

116+
- A pickled :class:`Device` now names its device by UUID, and unpickling
117+
returns the device with that UUID in the receiving process. Before, it
118+
named the device by ordinal, which names a different device when
119+
``CUDA_VISIBLE_DEVICES`` differs between the processes. Unpickling raises
120+
``RuntimeError`` when the device is not visible in the receiving process.
121+
The pickled form is for processes on the same machine.
122+
(`#3196 <https://github.com/NVIDIA/cuda-python/issues/3196>`__)
123+
116124
- Fixed a ``TypeError`` in :meth:`Program.compile` when retrying after PCH
117125
heap exhaustion.
118126

‎cuda_core/tests/memory_ipc/test_device_uuid.py‎

Lines changed: 69 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4-
"""Resources sent to another process keep their devices when the ordinals differ (issue #3097).
4+
"""Devices and resources sent to another process keep their devices when the ordinals differ.
55
66
The child process sees the parent's two devices in reversed order, so each
7-
ordinal names the other device there. The child must still rebuild the
8-
resource for the device that owns the memory, and the peers of a
9-
``VirtualMemoryResource`` must name the same devices as in the parent.
7+
ordinal names the other device there. A ``Device`` must arrive as the device
8+
with the same UUID (issue #3196), the child must rebuild a resource for the
9+
device that owns the memory, and the peers of a ``VirtualMemoryResource`` must
10+
name the same devices as in the parent (issue #3097).
1011
"""
1112

1213
import contextlib
@@ -78,7 +79,9 @@ def _child_main(inbox, outbox):
7879

7980

8081
def _receive(obj):
81-
"""Describe the received buffer in terms of device UUIDs, which do not depend on the ordinals."""
82+
"""Describe the received object in terms of device UUIDs, which do not depend on the ordinals."""
83+
if isinstance(obj, Device):
84+
return {"device_id": obj.device_id, "device_uuid": obj.uuid}
8285
if not isinstance(obj, Buffer):
8386
return {"received": type(obj).__name__}
8487
mr = obj.memory_resource
@@ -119,21 +122,44 @@ def _check_report(report, owner):
119122
assert report["data"] == bytes([PATTERN]) * NBYTES
120123

121124

125+
@pytest.mark.agent_authored(model="claude-opus-5-5")
126+
@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global")
127+
def test_device(device_x2):
128+
"""A Device arrives as the device with the same UUID, not the device with the same ordinal."""
129+
dev0, dev1 = device_x2
130+
report = _send_to_child(dev0, visible=(dev1, dev0))
131+
assert "error" not in report, f"the child could not receive the device: {report}"
132+
assert report["device_id"] != dev0.device_id, "CUDA_VISIBLE_DEVICES did not reorder the devices in the child"
133+
assert report["device_uuid"] == dev0.uuid
134+
135+
136+
@pytest.mark.agent_authored(model="claude-opus-5-5")
137+
@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global")
138+
def test_device_not_visible(device_x2):
139+
"""A child that cannot see the device refuses it and names the missing UUID."""
140+
dev0, dev1 = device_x2
141+
report = _send_to_child(dev0, visible=(dev1,))
142+
assert report["error"] == "RuntimeError", report
143+
assert f"no CUDA device visible to this process has UUID {dev0.uuid}" in report["message"]
144+
145+
122146
@pytest.mark.agent_authored(model="claude-opus-5-5")
123147
@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global")
124148
def test_device_memory_resource_buffer(ipc_mempool_device_x2):
125149
"""A pool-backed buffer comes back on the device that owns it, not on the device with the same ordinal."""
126150
dev0, dev1 = ipc_mempool_device_x2
127151
mr = DeviceMemoryResource(dev0, options=DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True))
128-
buf = mr.allocate(NBYTES, stream=dev0.default_stream)
129-
_fill(dev0, buf)
130-
131-
report = _send_to_child(buf, visible=(dev1, dev0))
132-
_check_report(report, dev0)
133-
134-
buf.close()
135-
dev0.sync()
136-
mr.close()
152+
buf = None
153+
try:
154+
buf = mr.allocate(NBYTES, stream=dev0.default_stream)
155+
_fill(dev0, buf)
156+
report = _send_to_child(buf, visible=(dev1, dev0))
157+
_check_report(report, dev0)
158+
finally:
159+
if buf is not None:
160+
buf.close()
161+
dev0.sync()
162+
mr.close()
137163

138164

139165
@pytest.mark.agent_authored(model="claude-opus-5-5")
@@ -143,14 +169,14 @@ def test_virtual_memory_resource_buffer(vmm_ipc_device_x2):
143169
dev0, dev1 = vmm_ipc_device_x2
144170
mr = _vmm_resource(dev0)
145171
buf = mr.allocate(NBYTES)
146-
_fill(dev0, buf)
147-
148-
report = _send_to_child(buf, visible=(dev1, dev0))
149-
_check_report(report, dev0)
150-
assert report["peer_uuids"] == []
151-
152-
buf.close()
153-
dev0.sync()
172+
try:
173+
_fill(dev0, buf)
174+
report = _send_to_child(buf, visible=(dev1, dev0))
175+
_check_report(report, dev0)
176+
assert report["peer_uuids"] == []
177+
finally:
178+
buf.close()
179+
dev0.sync()
154180

155181

156182
@pytest.mark.agent_authored(model="claude-opus-5-5")
@@ -162,15 +188,15 @@ def test_virtual_memory_resource_peers(vmm_ipc_device_x2):
162188
pytest.skip("Test requires peer access from the second device to the first")
163189
mr = _vmm_resource(dev0, peers=(dev1.device_id,))
164190
buf = mr.allocate(NBYTES)
165-
_fill(dev0, buf)
166-
167-
report = _send_to_child(buf, visible=(dev1, dev0))
168-
_check_report(report, dev0)
169-
assert report["peer_uuids"] == [dev1.uuid]
170-
assert report["peer_access"] == [3]
171-
172-
buf.close()
173-
dev0.sync()
191+
try:
192+
_fill(dev0, buf)
193+
report = _send_to_child(buf, visible=(dev1, dev0))
194+
_check_report(report, dev0)
195+
assert report["peer_uuids"] == [dev1.uuid]
196+
assert report["peer_access"] == [3]
197+
finally:
198+
buf.close()
199+
dev0.sync()
174200

175201

176202
@pytest.mark.agent_authored(model="claude-opus-5-5")
@@ -183,18 +209,18 @@ def test_virtual_memory_resource_peers_from_generator(vmm_ipc_device_x2):
183209
mr = _vmm_resource(dev0, peers=(peer for peer in (dev1.device_id,)))
184210
first = mr.allocate(NBYTES)
185211
buf = mr.allocate(NBYTES)
186-
_fill(dev0, buf)
187-
assert _access_flags(first, dev1.device_id) == 3
188-
assert _access_flags(buf, dev1.device_id) == 3
189-
190-
report = _send_to_child(buf, visible=(dev1, dev0))
191-
_check_report(report, dev0)
192-
assert report["peer_uuids"] == [dev1.uuid]
193-
assert report["peer_access"] == [3]
194-
195-
first.close()
196-
buf.close()
197-
dev0.sync()
212+
try:
213+
_fill(dev0, buf)
214+
assert _access_flags(first, dev1.device_id) == 3
215+
assert _access_flags(buf, dev1.device_id) == 3
216+
report = _send_to_child(buf, visible=(dev1, dev0))
217+
_check_report(report, dev0)
218+
assert report["peer_uuids"] == [dev1.uuid]
219+
assert report["peer_access"] == [3]
220+
finally:
221+
first.close()
222+
buf.close()
223+
dev0.sync()
198224

199225

200226
@pytest.mark.agent_authored(model="claude-opus-5-5")

‎cuda_core/tests/test_device.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# SPDX-License-Identifier: Apache-2.0
33

44
import contextlib
5+
import pickle
56
import subprocess
67
import sys
78
from concurrent.futures import ThreadPoolExecutor
@@ -552,6 +553,13 @@ def test_device_inequality_different_id_hash(init_cuda):
552553
pytest.skip("Test requires at least 2 CUDA devices")
553554

554555

556+
@pytest.mark.agent_authored(model="claude-opus-5-5")
557+
def test_device_pickle_round_trip(init_cuda):
558+
"""Unpickling a Device in the same process and thread returns the same singleton."""
559+
for device in Device.get_all_devices():
560+
assert pickle.loads(pickle.dumps(device)) is device # noqa: S301
561+
562+
555563
def test_device_dict_key(init_cuda):
556564
"""Devices should be usable as dictionary keys."""
557565
dev0 = Device(0)

0 commit comments

Comments
 (0)