From 25381d7d9149f49b36ea5d85101a8ca76fec0d51 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 15 Oct 2026 14:54:55 -0700 Subject: [PATCH 1/2] 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 --- cuda_core/cuda/core/_device.pyi | 13 ++ cuda_core/cuda/core/_device.pyx | 23 +++ .../core/_memory/_device_memory_resource.pyi | 5 +- .../core/_memory/_device_memory_resource.pyx | 16 +- .../core/_memory/_virtual_memory_resource.pyi | 12 +- .../core/_memory/_virtual_memory_resource.pyx | 29 ++- cuda_core/docs/source/release/1.3.0-notes.rst | 9 + .../tests/memory_ipc/test_device_uuid.py | 195 ++++++++++++++++++ 8 files changed, 289 insertions(+), 13 deletions(-) create mode 100644 cuda_core/tests/memory_ipc/test_device_uuid.py diff --git a/cuda_core/cuda/core/_device.pyi b/cuda_core/cuda/core/_device.pyi index 8c2b273a6cd..d0974890acd 100644 --- a/cuda_core/cuda/core/_device.pyi +++ b/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. + """ diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx index 8611cd3d819..5f03c65d238 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/cuda_core/cuda/core/_device.pyx @@ -1633,6 +1633,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 diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyi b/cuda_core/cuda/core/_memory/_device_memory_resource.pyi index 337d03baf4a..efb3948e885 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyi +++ b/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: ... diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index 8a0fed73eee..5e37724d589 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -126,7 +126,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 @@ -341,9 +343,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, alloc_handle: 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) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 3f630f09834..c0403a5f6eb 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/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.""" diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 9a4bb083922..d4a744449df 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/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"" +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. 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 f292bb209a6..8f14057fdcb 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/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 `__) + - :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. diff --git a/cuda_core/tests/memory_ipc/test_device_uuid.py b/cuda_core/tests/memory_ipc/test_device_uuid.py new file mode 100644 index 00000000000..e7bba11dfaa --- /dev/null +++ b/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() From 160a21c343069025c0d09beedbaf76e38de47668 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 15 Oct 2026 14:57:10 -0700 Subject: [PATCH 2/2] 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 --- cuda_core/cuda/core/_memory/_device_memory_resource.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index 13a3930ccdc..f078c3ac7d3 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -348,7 +348,7 @@ def _deep_reduce_device_memory_resource(mr) -> tuple[object, ...]: def _reconstruct_device_memory_resource( - cls: type, device_uuid: str, alloc_handle: IPCAllocationHandle + 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)