fix(cuda.core): pickle Device by UUID

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.

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.

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 <https://github.com/NVIDIA/cuda-python/issues/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 <https://github.com/NVIDIA/cuda-python/issues/3196>`__)
+
 - Fixed a ``TypeError`` in :meth:`Program.compile` when retrying after PCH
   heap exhaustion.
 

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")

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)