Skip to content

Commit d5f37e0

Browse files
Andy-Jostclaude
andauthored
fix(cuda.core): pickle memory-resource devices by UUID for IPC (#3126)
* 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]> --------- Co-authored-by: Claude Fable 5.1 <[email protected]>
1 parent fd33836 commit d5f37e0

9 files changed

Lines changed: 328 additions & 13 deletions

File tree

‎cuda_core/cuda/core/_device.pyi‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -846,3 +846,16 @@ class Device:
846846
847847
.. versionadded:: 1.1.0
848848
"""
849+
850+
def _device_from_uuid(device_uuid: str) -> Device:
851+
"""Return the visible device whose :attr:`Device.uuid` is ``device_uuid``.
852+
853+
Device ordinals are per process (``CUDA_VISIBLE_DEVICES`` can renumber
854+
devices), so objects that cross a process boundary name their devices by
855+
UUID, and the receiving process maps each UUID back with this function.
856+
857+
Raises
858+
------
859+
RuntimeError
860+
If no device visible to this process has the UUID.
861+
"""

‎cuda_core/cuda/core/_device.pyx‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1628,6 +1628,29 @@ class Device:
16281628
resource, self._context, self._device_id)
16291629

16301630

1631+
def _device_from_uuid(device_uuid: str) -> Device:
1632+
"""Return the visible device whose :attr:`Device.uuid` is ``device_uuid``.
1633+
1634+
Device ordinals are per process (``CUDA_VISIBLE_DEVICES`` can renumber
1635+
devices), so objects that cross a process boundary name their devices by
1636+
UUID, and the receiving process maps each UUID back with this function.
1637+
1638+
Raises
1639+
------
1640+
RuntimeError
1641+
If no device visible to this process has the UUID.
1642+
"""
1643+
devices = Device.get_all_devices()
1644+
for device in devices:
1645+
if device.uuid == device_uuid:
1646+
return device
1647+
visible = ", ".join(f"{device.device_id} ({device.uuid})" for device in devices) or "none"
1648+
raise RuntimeError(
1649+
f"no CUDA device visible to this process has UUID {device_uuid}; the visible devices are: "
1650+
f"{visible}. Make the device visible to this process, for example through CUDA_VISIBLE_DEVICES."
1651+
)
1652+
1653+
16311654
cdef inline int Device_ensure_cuda_initialized() except? -1:
16321655
"""Initialize CUDA driver and check version compatibility (once per process)."""
16331656
global _is_cuInit

‎cuda_core/cuda/core/_memory/_device_memory_resource.pyi‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,9 @@ class DeviceMemoryResource(_MemPool):
101101
102102
For greater efficiency when transferring many buffers, one may also send
103103
MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`,
104-
an MMR is created and registered in the receiving process. Subsequently,
104+
an MMR is created and registered in the receiving process, for the device
105+
with the same UUID as the MR's device, whatever its ordinal is there; if
106+
that device is not visible, unpickling raises ``RuntimeError``. Subsequently,
105107
buffers may be serialized and transferred using ordinary :mod:`pickle`
106108
methods. The reconstruction procedure uses the registry to find the
107109
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:
212214
Access permissions: "rw" for read-write, "r" for read-only, "" for no access.
213215
"""
214216
def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]: ...
217+
def _reconstruct_device_memory_resource(cls: type, device_uuid: str, alloc_handle: IPCAllocationHandle) -> DeviceMemoryResource: ...

‎cuda_core/cuda/core/_memory/_device_memory_resource.pyx‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,9 @@ cdef class DeviceMemoryResource(_MemPool):
125125
126126
For greater efficiency when transferring many buffers, one may also send
127127
MRs and buffers separately. When an MR is sent via :mod:`multiprocessing`,
128-
an MMR is created and registered in the receiving process. Subsequently,
128+
an MMR is created and registered in the receiving process, for the device
129+
with the same UUID as the MR's device, whatever its ordinal is there; if
130+
that device is not visible, unpickling raises ``RuntimeError``. Subsequently,
129131
buffers may be serialized and transferred using ordinary :mod:`pickle`
130132
methods. The reconstruction procedure uses the registry to find the
131133
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):
339341
def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]:
340342
check_multiprocessing_start_method()
341343
from .._device import Device
342-
device = Device(mr.device_id)
344+
# Device ordinals are per process; the receiver maps the UUID to its own device.
345+
device_uuid = Device(mr.device_id).uuid
343346
alloc_handle = mr.allocation_handle
344-
return mr.from_allocation_handle, (device, alloc_handle)
347+
return _reconstruct_device_memory_resource, (type(mr), device_uuid, alloc_handle)
348+
349+
350+
def _reconstruct_device_memory_resource(
351+
cls: type, device_uuid: str, IPCAllocationHandle alloc_handle not None: IPCAllocationHandle
352+
) -> DeviceMemoryResource:
353+
from .._device import _device_from_uuid
354+
return cls.from_allocation_handle(_device_from_uuid(device_uuid), alloc_handle)
345355

346356

347357
multiprocessing.reduction.register(DeviceMemoryResource, _deep_reduce_device_memory_resource)

‎cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ class VirtualMemoryResourceOptions:
4848
Alignment for the VA reservation. If `None`, use the queried granularity.
4949
peers: Iterable[int]
5050
Extra device IDs that should be granted access in addition to ``device``.
51+
The options store them as a tuple.
5152
self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
5253
Access flags for the owning device.
5354
peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
@@ -74,6 +75,7 @@ class VirtualMemoryResourceOptions:
7475
_t = driver.CUmemAllocationType
7576
_allocation_type = {VirtualMemoryAllocationType.PINNED: _t.CU_MEM_ALLOCATION_TYPE_PINNED}
7677

78+
def __post_init__(self) -> None: ...
7779
@staticmethod
7880
def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int: ...
7981
@staticmethod
@@ -197,9 +199,12 @@ class VirtualMemoryResource(MemoryResource):
197199
access; :meth:`Buffer.from_ipc_descriptor` imports a descriptor with a
198200
``VirtualMemoryResource`` of the receiving process for the device that
199201
owns the memory; and a buffer sent through ``multiprocessing`` does both.
200-
Plain ``pickle`` cannot carry the file descriptors; a process with its
201-
own file descriptor passing sends the descriptor's ``fds``, ``sizes``,
202-
and ``size`` and rebuilds it with
202+
A pickled resource names its device and its ``peers`` by UUID, and the
203+
receiving process rebuilds it for the devices with those UUIDs, whatever
204+
their ordinals are there; it raises ``RuntimeError`` if one of the devices
205+
is not visible. Plain ``pickle`` cannot carry the file descriptors; a
206+
process with its own file descriptor passing sends the descriptor's
207+
``fds``, ``sizes``, and ``size`` and rebuilds it with
203208
``VirtualMemoryIPCBufferDescriptor.from_fds()``. The importing resource
204209
maps the memory for its device and the devices in its ``peers`` option,
205210
with its own access options. A descriptor pins the physical memory while
@@ -369,3 +374,6 @@ class VirtualMemoryResource(MemoryResource):
369374
Returns:
370375
str: A string describing the object
371376
"""
377+
378+
def _reconstruct_virtual_memory_resource(cls: type, owner_uuid: str | None, peer_uuids: tuple[str, ...], config: VirtualMemoryResourceOptions) -> VirtualMemoryResource:
379+
"""Rebuild a pickled resource with the receiving process's device ordinals."""

‎cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,11 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, check_or_create_options
5252

5353
import os
5454
import platform
55-
from dataclasses import dataclass, field
55+
from dataclasses import dataclass, field, replace
5656
from multiprocessing import context as _mp_context
5757
from typing import TYPE_CHECKING, Iterable
5858

59-
from cuda.core._device import Device
59+
from cuda.core._device import Device, _device_from_uuid
6060
from cuda.core._utils.cuda_utils import driver
6161
from cuda.core._utils.version import BUILD_CUDA_MAJOR
6262
from cuda.core.typing import (
@@ -117,6 +117,7 @@ class VirtualMemoryResourceOptions:
117117
Alignment for the VA reservation. If `None`, use the queried granularity.
118118
peers: Iterable[int]
119119
Extra device IDs that should be granted access in addition to ``device``.
120+
The options store them as a tuple.
120121
self_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
121122
Access flags for the owning device.
122123
peer_access: :obj:`~_memory.VirtualMemoryAccessType` | None | str
@@ -165,6 +166,11 @@ class VirtualMemoryResourceOptions:
165166
if BUILD_CUDA_MAJOR >= 13:
166167
_allocation_type[VirtualMemoryAllocationType.MANAGED] = _t.CU_MEM_ALLOCATION_TYPE_MANAGED
167168

169+
def __post_init__(self) -> None:
170+
# Allocation and pickling each read the peers, so a one-shot iterable
171+
# such as a generator must not reach them.
172+
self.peers = tuple(int(peer) for peer in self.peers)
173+
168174
@staticmethod
169175
def _access_to_flags(spec: VirtualMemoryAccessType | None) -> int:
170176
flags = VirtualMemoryResourceOptions._access_flags.get(spec)
@@ -427,9 +433,12 @@ cdef class VirtualMemoryResource(MemoryResource):
427433
access; :meth:`Buffer.from_ipc_descriptor` imports a descriptor with a
428434
``VirtualMemoryResource`` of the receiving process for the device that
429435
owns the memory; and a buffer sent through ``multiprocessing`` does both.
430-
Plain ``pickle`` cannot carry the file descriptors; a process with its
431-
own file descriptor passing sends the descriptor's ``fds``, ``sizes``,
432-
and ``size`` and rebuilds it with
436+
A pickled resource names its device and its ``peers`` by UUID, and the
437+
receiving process rebuilds it for the devices with those UUIDs, whatever
438+
their ordinals are there; it raises ``RuntimeError`` if one of the devices
439+
is not visible. Plain ``pickle`` cannot carry the file descriptors; a
440+
process with its own file descriptor passing sends the descriptor's
441+
``fds``, ``sizes``, and ``size`` and rebuilds it with
433442
``VirtualMemoryIPCBufferDescriptor.from_fds()``. The importing resource
434443
maps the memory for its device and the devices in its ``peers`` option,
435444
with its own access options. A descriptor pins the physical memory while
@@ -459,7 +468,12 @@ cdef class VirtualMemoryResource(MemoryResource):
459468
# The resource holds no driver object that other processes must share;
460469
# the receiving process rebuilds an equivalent one from the device and
461470
# the options. This is what lets a Buffer pickle as (resource, descriptor).
462-
return type(self), (self.device.device_id if self.device is not None else 0, self.config)
471+
# Device ordinals are per process, so the owning device and the peers
472+
# travel by UUID and the receiver maps them to its own devices.
473+
owner_uuid = self.device.uuid if self.device is not None else None
474+
peer_uuids = tuple(Device(peer).uuid for peer in self.config.peers)
475+
return _reconstruct_virtual_memory_resource, (
476+
type(self), owner_uuid, peer_uuids, replace(self.config, peers=()))
463477
464478
cdef int _check_config(self, object cfg) except -1:
465479
"""Reject options the driver would reject, before any driver call.
@@ -1069,6 +1083,15 @@ cdef class VirtualMemoryResource(MemoryResource):
10691083
return f"<VirtualMemoryResource device={self.device}>"
10701084
10711085
1086+
def _reconstruct_virtual_memory_resource(
1087+
cls: type, owner_uuid: str | None, peer_uuids: tuple[str, ...], config: VirtualMemoryResourceOptions
1088+
) -> VirtualMemoryResource:
1089+
"""Rebuild a pickled resource with the receiving process's device ordinals."""
1090+
device = _device_from_uuid(owner_uuid) if owner_uuid is not None else 0
1091+
peers = tuple(_device_from_uuid(peer_uuid).device_id for peer_uuid in peer_uuids)
1092+
return cls(device, replace(config, peers=peers))
1093+
1094+
10721095
cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf):
10731096
"""Export the allocations that back ``buf`` as shareable handles, in address order.
10741097

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,15 @@ New features
104104
Fixes and enhancements
105105
----------------------
106106

107+
- A :class:`~_memory.DeviceMemoryResource` or a :class:`VirtualMemoryResource`
108+
sent through ``multiprocessing`` now names its device, and the
109+
``VirtualMemoryResource`` ``peers``, by UUID. The receiving process maps
110+
each UUID to its own device ordinal. Before, the receiver used the sender's
111+
ordinals, which name different devices when ``CUDA_VISIBLE_DEVICES``
112+
differs between the processes. Unpickling now raises ``RuntimeError`` when
113+
one of the devices is not visible in the receiving process.
114+
(`#3097 <https://github.com/NVIDIA/cuda-python/issues/3097>`__)
115+
107116
- Fixed a ``TypeError`` in :meth:`Program.compile` when retrying after PCH
108117
heap exhaustion.
109118

0 commit comments

Comments
 (0)