fix(cuda.core): pickle memory-resource devices by UUID for IPC

cuda_core/cuda/core/_device.pyi

@@ -846,3 +846,16 @@ class Device:
 
         .. versionadded:: 1.1.0
         """
+
+def _device_from_uuid(device_uuid: str) -> Device:
+    """Return the visible device whose :attr:`Device.uuid` is ``device_uuid``.
+
+    Device ordinals are per process (``CUDA_VISIBLE_DEVICES`` can renumber
+    devices), so objects that cross a process boundary name their devices by
+    UUID, and the receiving process maps each UUID back with this function.
+
+    Raises
+    ------
+    RuntimeError
+        If no device visible to this process has the UUID.
+    """

cuda_core/cuda/core/_device.pyx

@@ -1628,6 +1628,29 @@ class Device:
             resource, self._context, self._device_id)
 
 
+def _device_from_uuid(device_uuid: str) -> Device:
+    """Return the visible device whose :attr:`Device.uuid` is ``device_uuid``.
+
+    Device ordinals are per process (``CUDA_VISIBLE_DEVICES`` can renumber
+    devices), so objects that cross a process boundary name their devices by
+    UUID, and the receiving process maps each UUID back with this function.
+
+    Raises
+    ------
+    RuntimeError
+        If no device visible to this process has the UUID.
+    """
+    devices = Device.get_all_devices()
+    for device in devices:
+        if device.uuid == device_uuid:
+            return device
+    visible = ", ".join(f"{device.device_id} ({device.uuid})" for device in devices) or "none"
+    raise RuntimeError(
+        f"no CUDA device visible to this process has UUID {device_uuid}; the visible devices are: "
+        f"{visible}. Make the device visible to this process, for example through CUDA_VISIBLE_DEVICES."
+    )
+
+
 cdef inline int Device_ensure_cuda_initialized() except? -1:
     """Initialize CUDA driver and check version compatibility (once per process)."""
     global _is_cuInit

cuda_core/cuda/core/_memory/_device_memory_resource.pyi

@@ -101,7 +101,9 @@ class DeviceMemoryResource(_MemPool):
 
     For greater efficiency when transferring many buffers, one may also send
     MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`,
-    an MMR is created and registered in the receiving process. Subsequently,
+    an MMR is created and registered in the receiving process, for the device
+    with the same UUID as the MR's device, whatever its ordinal is there; if
+    that device is not visible, unpickling raises ``RuntimeError``. Subsequently,
     buffers may be serialized and transferred using ordinary :mod:`pickle`
     methods.  The reconstruction procedure uses the registry to find the
     associated MMR. Unpickling a :class:`Buffer` performs an IPC import from
@@ -212,3 +214,4 @@ def DMR_mempool_get_access(dmr: DeviceMemoryResource, device_id: int) -> str:
         Access permissions: "rw" for read-write, "r" for read-only, "" for no access.
     """
 def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]: ...
+def _reconstruct_device_memory_resource(cls: type, device_uuid: str, alloc_handle: IPCAllocationHandle) -> DeviceMemoryResource: ...

cuda_core/cuda/core/_memory/_device_memory_resource.pyx

@@ -125,7 +125,9 @@ cdef class DeviceMemoryResource(_MemPool):
 
     For greater efficiency when transferring many buffers, one may also send
     MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`,
-    an MMR is created and registered in the receiving process. Subsequently,
+    an MMR is created and registered in the receiving process, for the device
+    with the same UUID as the MR's device, whatever its ordinal is there; if
+    that device is not visible, unpickling raises ``RuntimeError``. Subsequently,
     buffers may be serialized and transferred using ordinary :mod:`pickle`
     methods.  The reconstruction procedure uses the registry to find the
     associated MMR. Unpickling a :class:`Buffer` performs an IPC import from
@@ -339,9 +341,17 @@ cpdef str DMR_mempool_get_access(DeviceMemoryResource dmr, int device_id):
 def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]:
     check_multiprocessing_start_method()
     from .._device import Device
-    device = Device(mr.device_id)
+    # Device ordinals are per process; the receiver maps the UUID to its own device.
+    device_uuid = Device(mr.device_id).uuid
     alloc_handle = mr.allocation_handle
-    return mr.from_allocation_handle, (device, alloc_handle)
+    return _reconstruct_device_memory_resource, (type(mr), device_uuid, alloc_handle)
+
+
+def _reconstruct_device_memory_resource(
+    cls: type, device_uuid: str, IPCAllocationHandle alloc_handle not None: IPCAllocationHandle
+) -> DeviceMemoryResource:
+    from .._device import _device_from_uuid
+    return cls.from_allocation_handle(_device_from_uuid(device_uuid), alloc_handle)
 
 
 multiprocessing.reduction.register(DeviceMemoryResource, _deep_reduce_device_memory_resource)

cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi

@@ -197,9 +197,12 @@ class VirtualMemoryResource(MemoryResource):
     access; :meth:`Buffer.from_ipc_descriptor` imports a descriptor with a
     ``VirtualMemoryResource`` of the receiving process for the device that
     owns the memory; and a buffer sent through ``multiprocessing`` does both.
-    Plain ``pickle`` cannot carry the file descriptors; a process with its
-    own file descriptor passing sends the descriptor's ``fds``, ``sizes``,
-    and ``size`` and rebuilds it with
+    A pickled resource names its device and its ``peers`` by UUID, and the
+    receiving process rebuilds it for the devices with those UUIDs, whatever
+    their ordinals are there; it raises ``RuntimeError`` if one of the devices
+    is not visible. Plain ``pickle`` cannot carry the file descriptors; a
+    process with its own file descriptor passing sends the descriptor's
+    ``fds``, ``sizes``, and ``size`` and rebuilds it with
     ``VirtualMemoryIPCBufferDescriptor.from_fds()``. The importing resource
     maps the memory for its device and the devices in its ``peers`` option,
     with its own access options. A descriptor pins the physical memory while
@@ -369,3 +372,6 @@ class VirtualMemoryResource(MemoryResource):
         Returns:
             str: A string describing the object
         """
+
+def _reconstruct_virtual_memory_resource(cls: type, owner_uuid: str | None, peer_uuids: tuple[str, ...], config: VirtualMemoryResourceOptions) -> VirtualMemoryResource:
+    """Rebuild a pickled resource with the receiving process's device ordinals."""

cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx

@@ -52,11 +52,11 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, check_or_create_options
 
 import os
 import platform
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from multiprocessing import context as _mp_context
 from typing import TYPE_CHECKING, Iterable
 
-from cuda.core._device import Device
+from cuda.core._device import Device, _device_from_uuid
 from cuda.core._utils.cuda_utils import driver
 from cuda.core._utils.version import BUILD_CUDA_MAJOR
 from cuda.core.typing import (
@@ -427,9 +427,12 @@ cdef class VirtualMemoryResource(MemoryResource):
     access; :meth:`Buffer.from_ipc_descriptor` imports a descriptor with a
     ``VirtualMemoryResource`` of the receiving process for the device that
     owns the memory; and a buffer sent through ``multiprocessing`` does both.
-    Plain ``pickle`` cannot carry the file descriptors; a process with its
-    own file descriptor passing sends the descriptor's ``fds``, ``sizes``,
-    and ``size`` and rebuilds it with
+    A pickled resource names its device and its ``peers`` by UUID, and the
+    receiving process rebuilds it for the devices with those UUIDs, whatever
+    their ordinals are there; it raises ``RuntimeError`` if one of the devices
+    is not visible. Plain ``pickle`` cannot carry the file descriptors; a
+    process with its own file descriptor passing sends the descriptor's
+    ``fds``, ``sizes``, and ``size`` and rebuilds it with
     ``VirtualMemoryIPCBufferDescriptor.from_fds()``. The importing resource
     maps the memory for its device and the devices in its ``peers`` option,
     with its own access options. A descriptor pins the physical memory while
@@ -459,7 +462,12 @@ cdef class VirtualMemoryResource(MemoryResource):
         # The resource holds no driver object that other processes must share;
         # the receiving process rebuilds an equivalent one from the device and
         # the options. This is what lets a Buffer pickle as (resource, descriptor).
-        return type(self), (self.device.device_id if self.device is not None else 0, self.config)
+        # Device ordinals are per process, so the owning device and the peers
+        # travel by UUID and the receiver maps them to its own devices.
+        owner_uuid = self.device.uuid if self.device is not None else None
+        peer_uuids = tuple(Device(peer).uuid for peer in self.config.peers)
+        return _reconstruct_virtual_memory_resource, (
+            type(self), owner_uuid, peer_uuids, replace(self.config, peers=()))
 
     cdef int _check_config(self, object cfg) except -1:
         """Reject options the driver would reject, before any driver call.
@@ -1069,6 +1077,15 @@ cdef class VirtualMemoryResource(MemoryResource):
         return f"<VirtualMemoryResource device={self.device}>"
 
 
+def _reconstruct_virtual_memory_resource(
+    cls: type, owner_uuid: str | None, peer_uuids: tuple[str, ...], config: VirtualMemoryResourceOptions
+) -> VirtualMemoryResource:
+    """Rebuild a pickled resource with the receiving process's device ordinals."""
+    device = _device_from_uuid(owner_uuid) if owner_uuid is not None else 0
+    peers = tuple(_device_from_uuid(peer_uuid).device_id for peer_uuid in peer_uuids)
+    return cls(device, replace(config, peers=peers))
+
+
 cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf):
     """Export the allocations that back ``buf`` as shareable handles, in address order.
 

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

@@ -104,6 +104,15 @@ New features
 Fixes and enhancements
 ----------------------
 
+- A :class:`~_memory.DeviceMemoryResource` or a :class:`VirtualMemoryResource`
+  sent through ``multiprocessing`` now names its device, and the
+  ``VirtualMemoryResource`` ``peers``, by UUID. The receiving process maps
+  each UUID to its own device ordinal. Before, the receiver used the sender's
+  ordinals, which name different devices when ``CUDA_VISIBLE_DEVICES``
+  differs between the processes. Unpickling now raises ``RuntimeError`` when
+  one of the devices is not visible in the receiving process.
+  (`#3097 <https://github.com/NVIDIA/cuda-python/issues/3097>`__)
+
 - :attr:`system.Device.uuid_without_prefix` preserves UUIDs that NVML already
   reports without a prefix, such as on Orin. This also allows
   :meth:`system.Device.to_cuda_device` to match those devices correctly.

cuda_core/tests/memory_ipc/test_device_uuid.py

@@ -0,0 +1,195 @@
+# 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).
+
+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.
+"""
+
+import contextlib
+import ctypes
+import multiprocessing as mp
+import os
+
+import pytest
+from helpers.child_processes import child_timeout_sec, kill_subprocesses
+from helpers.constants import POOL_SIZE
+
+from cuda.bindings import driver
+from cuda.core import (
+    Buffer,
+    Device,
+    DeviceMemoryResource,
+    DeviceMemoryResourceOptions,
+    VirtualMemoryResource,
+    VirtualMemoryResourceOptions,
+)
+from cuda.core._utils.cuda_utils import handle_return
+
+CHILD_TIMEOUT_SEC = child_timeout_sec()
+NBYTES = 64
+PATTERN = 0x5A
+
+
+@contextlib.contextmanager
+def _visible_devices(*devices):
+    """Make only ``devices``, in this order, visible to processes started in the block.
+
+    The spawn start method copies ``os.environ`` when ``Process.start()`` runs.
+    The devices are named by UUID, so the order does not depend on how this
+    process numbers them.
+    """
+    saved = os.environ.get("CUDA_VISIBLE_DEVICES")
+    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(f"GPU-{device.uuid}" for device in devices)
+    try:
+        yield
+    finally:
+        if saved is None:
+            del os.environ["CUDA_VISIBLE_DEVICES"]
+        else:
+            os.environ["CUDA_VISIBLE_DEVICES"] = saved
+
+
+def _send_to_child(obj, visible):
+    """Send ``obj`` to a child that sees only ``visible``, and return the child's report."""
+    inbox, outbox = mp.Queue(), mp.Queue()
+    process = mp.Process(target=_child_main, args=(inbox, outbox))
+    with _visible_devices(*visible):
+        process.start()
+    inbox.put(obj)
+    try:
+        return outbox.get(timeout=CHILD_TIMEOUT_SEC)
+    finally:
+        process.join(timeout=CHILD_TIMEOUT_SEC)
+        survivors = kill_subprocesses(process)
+        assert not survivors, "child did not exit within timeout"
+        assert process.exitcode == 0, f"child exited with {process.exitcode}"
+
+
+def _child_main(inbox, outbox):
+    try:
+        report = _receive(inbox.get(timeout=CHILD_TIMEOUT_SEC))
+    except Exception as exc:  # report every failure to the parent, which asserts on it
+        report = {"error": type(exc).__name__, "message": str(exc)}
+    outbox.put(report)
+
+
+def _receive(obj):
+    """Describe the received buffer in terms of device UUIDs, which do not depend on the ordinals."""
+    if not isinstance(obj, Buffer):
+        return {"received": type(obj).__name__}
+    mr = obj.memory_resource
+    report = {"device_id": obj.device_id, "device_uuid": Device(obj.device_id).uuid}
+    host = (ctypes.c_ubyte * NBYTES)()
+    handle_return(driver.cuMemcpyDtoH(ctypes.addressof(host), int(obj.handle), NBYTES))
+    report["data"] = bytes(host)
+    if isinstance(mr, VirtualMemoryResource):
+        report["peer_uuids"] = [Device(peer).uuid for peer in mr.config.peers]
+        report["peer_access"] = [_access_flags(obj, peer) for peer in mr.config.peers]
+    obj.close()
+    return report
+
+
+def _access_flags(buf, device_id):
+    """The access a device has to ``buf``: 0 for none, 3 for read-write."""
+    location = driver.CUmemLocation()
+    location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
+    location.id = device_id
+    return int(handle_return(driver.cuMemGetAccess(location, int(buf.handle))))
+
+
+def _fill(device, buf):
+    device.set_current()
+    handle_return(driver.cuMemsetD8(int(buf.handle), PATTERN, NBYTES))
+    device.sync()
+
+
+def _vmm_resource(device, peers=()):
+    return VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(handle_type="posix_fd", peers=peers))
+
+
+def _check_report(report, owner):
+    assert "error" not in report, f"the child could not receive the buffer: {report}"
+    # The child numbers the devices the other way round; otherwise this test proves nothing.
+    assert report["device_id"] != owner.device_id, "CUDA_VISIBLE_DEVICES did not reorder the devices in the child"
+    assert report["device_uuid"] == owner.uuid
+    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_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()
+
+
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global")
+def test_virtual_memory_resource_buffer(vmm_ipc_device_x2):
+    """A VirtualMemoryResource buffer comes back on the device that owns it."""
+    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()
+
+
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global")
+def test_virtual_memory_resource_peers(vmm_ipc_device_x2):
+    """The peers of a VirtualMemoryResource name the same devices in the child, and those devices get access."""
+    dev0, dev1 = vmm_ipc_device_x2
+    if not dev1.can_access_peer(dev0):
+        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()
+
+
+@pytest.mark.agent_authored(model="claude-opus-5-5")
+@pytest.mark.thread_unsafe(reason="sets CUDA_VISIBLE_DEVICES, which is process-global")
+@pytest.mark.parametrize("kind", ["device", "vmm"])
+def test_owning_device_not_visible(request, kind):
+    """A child that cannot see the owning device refuses the resource and names the missing UUID."""
+    if kind == "device":
+        dev0, dev1 = request.getfixturevalue("ipc_mempool_device_x2")
+        mr = DeviceMemoryResource(dev0, options=DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True))
+    else:
+        dev0, dev1 = request.getfixturevalue("vmm_ipc_device_x2")
+        mr = _vmm_resource(dev0)
+
+    # The child receives the resource itself: it fails before it needs any buffer's handles.
+    report = _send_to_child(mr, visible=(dev1,))
+    assert report["error"] == "RuntimeError", report
+    assert f"no CUDA device visible to this process has UUID {dev0.uuid}" in report["message"]
+    assert dev1.uuid in report["message"]
+
+    if kind == "device":
+        mr.close()