diff --git a/cuda_core/cuda/core/_device.pyi b/cuda_core/cuda/core/_device.pyi index d0974890acd..776c8b0b561 100644 --- a/cuda_core/cuda/core/_device.pyi +++ b/cuda_core/cuda/core/_device.pyi @@ -479,6 +479,11 @@ class Device: Newly returned :obj:`~_device.Device` objects are thread-local singletons for a specified device. + A pickled :obj:`~_device.Device` names its device by UUID. Unpickling + returns the device with that UUID in the receiving process, whatever its + ordinal is there, and raises ``RuntimeError`` if that device is not + visible. The pickled form is for processes on the same machine. + Note ---- Will not initialize the GPU. diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx index 59721c8d1c7..29b515bbe17 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/cuda_core/cuda/core/_device.pyx @@ -983,6 +983,11 @@ class Device: Newly returned :obj:`~_device.Device` objects are thread-local singletons for a specified device. + A pickled :obj:`~_device.Device` names its device by UUID. Unpickling + returns the device with that UUID in the receiving process, whatever its + ordinal is there, and raises ``RuntimeError`` if that device is not + visible. The pickled form is for processes on the same machine. + Note ---- Will not initialize the GPU. @@ -1242,7 +1247,8 @@ class Device: return self._device_id == other._device_id def __reduce__(self) -> tuple[object, ...]: - return Device, (self.device_id,) + # Device ordinals are per process; the receiver maps the UUID to its own device. + return _device_from_uuid, (self.uuid,) def set_current(self, Context ctx: Context | None = None) -> Context | None: """Set device to be used for GPU executions. diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index b522f71d4af..6de513a4dc2 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -113,6 +113,14 @@ Fixes and enhancements one of the devices is not visible in the receiving process. (`#3097 `__) +- A pickled :class:`Device` now names its device by UUID, and unpickling + returns the device with that UUID in the receiving process. Before, it + named the device by ordinal, which names a different device when + ``CUDA_VISIBLE_DEVICES`` differs between the processes. Unpickling raises + ``RuntimeError`` when the device is not visible in the receiving process. + The pickled form is for processes on the same machine. + (`#3196 `__) + - Fixed a ``TypeError`` in :meth:`Program.compile` when retrying after PCH heap exhaustion. diff --git a/cuda_core/tests/memory_ipc/test_device_uuid.py b/cuda_core/tests/memory_ipc/test_device_uuid.py index eeed9257816..97499a64c97 100644 --- a/cuda_core/tests/memory_ipc/test_device_uuid.py +++ b/cuda_core/tests/memory_ipc/test_device_uuid.py @@ -1,12 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Resources sent to another process keep their devices when the ordinals differ (issue #3097). +"""Devices and resources sent to another process keep their devices when the ordinals differ. The child process sees the parent's two devices in reversed order, so each -ordinal names the other device there. The child must still rebuild the -resource for the device that owns the memory, and the peers of a -``VirtualMemoryResource`` must name the same devices as in the parent. +ordinal names the other device there. A ``Device`` must arrive as the device +with the same UUID (issue #3196), the child must rebuild a resource for the +device that owns the memory, and the peers of a ``VirtualMemoryResource`` must +name the same devices as in the parent (issue #3097). """ import contextlib @@ -78,7 +79,9 @@ def _child_main(inbox, outbox): def _receive(obj): - """Describe the received buffer in terms of device UUIDs, which do not depend on the ordinals.""" + """Describe the received object in terms of device UUIDs, which do not depend on the ordinals.""" + if isinstance(obj, Device): + return {"device_id": obj.device_id, "device_uuid": obj.uuid} if not isinstance(obj, Buffer): return {"received": type(obj).__name__} mr = obj.memory_resource @@ -119,21 +122,44 @@ def _check_report(report, owner): assert report["data"] == bytes([PATTERN]) * NBYTES +@pytest.mark.agent_authored(model="claude-opus-5-5") +@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global") +def test_device(device_x2): + """A Device arrives as the device with the same UUID, not the device with the same ordinal.""" + dev0, dev1 = device_x2 + report = _send_to_child(dev0, visible=(dev1, dev0)) + assert "error" not in report, f"the child could not receive the device: {report}" + assert report["device_id"] != dev0.device_id, "CUDA_VISIBLE_DEVICES did not reorder the devices in the child" + assert report["device_uuid"] == dev0.uuid + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global") +def test_device_not_visible(device_x2): + """A child that cannot see the device refuses it and names the missing UUID.""" + dev0, dev1 = device_x2 + report = _send_to_child(dev0, visible=(dev1,)) + assert report["error"] == "RuntimeError", report + assert f"no CUDA device visible to this process has UUID {dev0.uuid}" in report["message"] + + @pytest.mark.agent_authored(model="claude-opus-5-5") @pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global") def test_device_memory_resource_buffer(ipc_mempool_device_x2): """A pool-backed buffer comes back on the device that owns it, not on the device with the same ordinal.""" dev0, dev1 = ipc_mempool_device_x2 mr = DeviceMemoryResource(dev0, options=DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)) - buf = mr.allocate(NBYTES, stream=dev0.default_stream) - _fill(dev0, buf) - - report = _send_to_child(buf, visible=(dev1, dev0)) - _check_report(report, dev0) - - buf.close() - dev0.sync() - mr.close() + buf = None + try: + buf = mr.allocate(NBYTES, stream=dev0.default_stream) + _fill(dev0, buf) + report = _send_to_child(buf, visible=(dev1, dev0)) + _check_report(report, dev0) + finally: + if buf is not None: + buf.close() + dev0.sync() + mr.close() @pytest.mark.agent_authored(model="claude-opus-5-5") @@ -143,14 +169,14 @@ def test_virtual_memory_resource_buffer(vmm_ipc_device_x2): dev0, dev1 = vmm_ipc_device_x2 mr = _vmm_resource(dev0) buf = mr.allocate(NBYTES) - _fill(dev0, buf) - - report = _send_to_child(buf, visible=(dev1, dev0)) - _check_report(report, dev0) - assert report["peer_uuids"] == [] - - buf.close() - dev0.sync() + try: + _fill(dev0, buf) + report = _send_to_child(buf, visible=(dev1, dev0)) + _check_report(report, dev0) + assert report["peer_uuids"] == [] + finally: + buf.close() + dev0.sync() @pytest.mark.agent_authored(model="claude-opus-5-5") @@ -162,15 +188,15 @@ def test_virtual_memory_resource_peers(vmm_ipc_device_x2): pytest.skip("Test requires peer access from the second device to the first") mr = _vmm_resource(dev0, peers=(dev1.device_id,)) buf = mr.allocate(NBYTES) - _fill(dev0, buf) - - report = _send_to_child(buf, visible=(dev1, dev0)) - _check_report(report, dev0) - assert report["peer_uuids"] == [dev1.uuid] - assert report["peer_access"] == [3] - - buf.close() - dev0.sync() + try: + _fill(dev0, buf) + report = _send_to_child(buf, visible=(dev1, dev0)) + _check_report(report, dev0) + assert report["peer_uuids"] == [dev1.uuid] + assert report["peer_access"] == [3] + finally: + buf.close() + dev0.sync() @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): mr = _vmm_resource(dev0, peers=(peer for peer in (dev1.device_id,))) first = mr.allocate(NBYTES) buf = mr.allocate(NBYTES) - _fill(dev0, buf) - assert _access_flags(first, dev1.device_id) == 3 - assert _access_flags(buf, dev1.device_id) == 3 - - report = _send_to_child(buf, visible=(dev1, dev0)) - _check_report(report, dev0) - assert report["peer_uuids"] == [dev1.uuid] - assert report["peer_access"] == [3] - - first.close() - buf.close() - dev0.sync() + try: + _fill(dev0, buf) + assert _access_flags(first, dev1.device_id) == 3 + assert _access_flags(buf, dev1.device_id) == 3 + report = _send_to_child(buf, visible=(dev1, dev0)) + _check_report(report, dev0) + assert report["peer_uuids"] == [dev1.uuid] + assert report["peer_access"] == [3] + finally: + first.close() + buf.close() + dev0.sync() @pytest.mark.agent_authored(model="claude-opus-5-5") diff --git a/cuda_core/tests/test_device.py b/cuda_core/tests/test_device.py index d12b86e0190..44f74620fc2 100644 --- a/cuda_core/tests/test_device.py +++ b/cuda_core/tests/test_device.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import contextlib +import pickle import subprocess import sys from concurrent.futures import ThreadPoolExecutor @@ -552,6 +553,13 @@ def test_device_inequality_different_id_hash(init_cuda): pytest.skip("Test requires at least 2 CUDA devices") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_device_pickle_round_trip(init_cuda): + """Unpickling a Device in the same process and thread returns the same singleton.""" + for device in Device.get_all_devices(): + assert pickle.loads(pickle.dumps(device)) is device # noqa: S301 + + def test_device_dict_key(init_cuda): """Devices should be usable as dictionary keys.""" dev0 = Device(0)