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

@@ -48,6 +48,7 @@ class VirtualMemoryResourceOptions:
         Alignment for the VA reservation. If `None`, use the queried granularity.
     peers: Iterable[int]
         Extra device IDs that should be granted access in addition to ``device``.
+        The options store them as a tuple.
     self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
         Access flags for the owning device.
     peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
@@ -74,6 +75,7 @@ class VirtualMemoryResourceOptions:
     _t = driver.CUmemAllocationType
     _allocation_type = {VirtualMemoryAllocationType.PINNED: _t.CU_MEM_ALLOCATION_TYPE_PINNED}
 
+    def __post_init__(self) -> None: ...
     @staticmethod
     def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int: ...
     @staticmethod
@@ -197,9 +199,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 +374,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 (
@@ -117,6 +117,7 @@ class VirtualMemoryResourceOptions:
         Alignment for the VA reservation. If `None`, use the queried granularity.
     peers: Iterable[int]
         Extra device IDs that should be granted access in addition to ``device``.
+        The options store them as a tuple.
     self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
         Access flags for the owning device.
     peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
@@ -165,6 +166,11 @@ class VirtualMemoryResourceOptions:
     if BUILD_CUDA_MAJOR >= 13:
         _allocation_type[VirtualMemoryAllocationType.MANAGED] = _t.CU_MEM_ALLOCATION_TYPE_MANAGED
 
+    def __post_init__(self) -> None:
+        # Allocation and pickling each read the peers, so a one-shot iterable
+        # such as a generator must not reach them.
+        self.peers = tuple(int(peer) for peer in self.peers)
+
     @staticmethod
     def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int:
         flags = VirtualMemoryResourceOptions._access_flags.get(spec)
@@ -427,9 +433,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 +468,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 +1083,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>`__)
+
 - Fixed a ``TypeError`` in :meth:`Program.compile` when retrying after PCH
   heap exhaustion.