From 25381d7d9149f49b36ea5d85101a8ca76fec0d51 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 15 Oct 2026 14:54:55 -0700 Subject: [PATCH 1/5] 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/5] 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) From a55b31ad5f74e198251dc6fe18a2a4a40f0b814c Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 22 Oct 2026 16:20:10 -0700 Subject: [PATCH 3/5] 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 --- .../core/_memory/_virtual_memory_resource.pyi | 2 ++ .../core/_memory/_virtual_memory_resource.pyx | 6 +++++ .../tests/memory_ipc/test_device_uuid.py | 24 +++++++++++++++++++ cuda_core/tests/test_memory.py | 7 ++++++ 4 files changed, 39 insertions(+) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index c0403a5f6eb..7dc7bbecce9 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/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 diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index eef1c09b644..bc783855db1 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -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) diff --git a/cuda_core/tests/memory_ipc/test_device_uuid.py b/cuda_core/tests/memory_ipc/test_device_uuid.py index e7bba11dfaa..eeed9257816 100644 --- a/cuda_core/tests/memory_ipc/test_device_uuid.py +++ b/cuda_core/tests/memory_ipc/test_device_uuid.py @@ -173,6 +173,30 @@ def test_virtual_memory_resource_peers(vmm_ipc_device_x2): 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_from_generator(vmm_ipc_device_x2): + """Peers given as a generator reach every allocation and the pickled resource, not only the first allocation.""" + 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=(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() + + @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"]) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 6076857279b..db492288fb5 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -3116,6 +3116,13 @@ def test_vmm_options_handle_type_win32_raises(): VirtualMemoryResourceOptions._handle_type_to_driver("win32") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_vmm_options_store_peers_as_tuple(): + """The options store a one-shot peers iterable as a tuple, so every reader sees all the peers.""" + options = VirtualMemoryResourceOptions(peers=(peer for peer in (1, 2))) + assert options.peers == (1, 2) + + @pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.parametrize("location_type", ["host", "host_numa", "host_numa_current"]) def test_vmm_host_location_types_report_host_accessible(init_cuda, location_type): From c26672ff1abf7c9d8d865b69124a69b49c8a6b1e Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 22 Oct 2026 16:21:20 -0700 Subject: [PATCH 4/5] fix(cuda.core): pickle Device by UUID Device.__reduce__ sent the ordinal, so a process with a different CUDA_VISIBLE_DEVICES unpickled a different GPU. It now sends the UUID, and the receiver maps it back with _device_from_uuid(), which raises RuntimeError when the device is not visible. Closes #3196. Co-Authored-By: Claude Opus 5.5 --- cuda_core/cuda/core/_device.pyi | 5 +++ cuda_core/cuda/core/_device.pyx | 8 ++++- cuda_core/docs/source/release/1.3.0-notes.rst | 7 ++++ .../tests/memory_ipc/test_device_uuid.py | 34 ++++++++++++++++--- cuda_core/tests/test_device.py | 8 +++++ 5 files changed, 56 insertions(+), 6 deletions(-) diff --git a/cuda_core/cuda/core/_device.pyi b/cuda_core/cuda/core/_device.pyi index d0974890acd..dd009d2c097 100644 --- a/cuda_core/cuda/core/_device.pyi +++ b/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. + Note ---- Will not initialize the GPU. diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx index 59721c8d1c7..0eed71279f2 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/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. + 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. 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 8f14057fdcb..bd48a5ab3e2 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -113,6 +113,13 @@ Fixes and enhancements one of the devices is not visible in the receiving process. (`#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. + (`#3196 `__) + - :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 index e7bba11dfaa..9fb8521ad05 100644 --- a/cuda_core/tests/memory_ipc/test_device_uuid.py +++ b/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, 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,6 +122,27 @@ 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): diff --git a/cuda_core/tests/test_device.py b/cuda_core/tests/test_device.py index d12b86e0190..44f74620fc2 100644 --- a/cuda_core/tests/test_device.py +++ b/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) From 2197a4da2cb8df976a13dca6ab4a630331123b84 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Fri, 23 Oct 2026 12:05:29 -0700 Subject: [PATCH 5/5] fix(cuda.core): say the Device pickle is for one machine; close test resources on every path - The Device docstring and the release note say the pickled form is for processes on the same machine (asked on #3196). - The buffer tests in test_device_uuid.py release their buffers and pools in finally blocks, so a failed assertion cannot leave an owned pool open for the next test's leak check (CodeRabbit on #3198). - The module docstring cites #3196 for the Device requirement. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_device.pyi | 2 +- cuda_core/cuda/core/_device.pyx | 2 +- cuda_core/docs/source/release/1.3.0-notes.rst | 1 + .../tests/memory_ipc/test_device_uuid.py | 84 ++++++++++--------- 4 files changed, 46 insertions(+), 43 deletions(-) diff --git a/cuda_core/cuda/core/_device.pyi b/cuda_core/cuda/core/_device.pyi index dd009d2c097..776c8b0b561 100644 --- a/cuda_core/cuda/core/_device.pyi +++ b/cuda_core/cuda/core/_device.pyi @@ -482,7 +482,7 @@ class 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. + visible. The pickled form is for processes on the same machine. Note ---- diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx index 0eed71279f2..29b515bbe17 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/cuda_core/cuda/core/_device.pyx @@ -986,7 +986,7 @@ class 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. + visible. The pickled form is for processes on the same machine. Note ---- 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 127caedc3d8..6de513a4dc2 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -118,6 +118,7 @@ Fixes and enhancements 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 `__) - Fixed a ``TypeError`` in :meth:`Program.compile` when retrying after PCH diff --git a/cuda_core/tests/memory_ipc/test_device_uuid.py b/cuda_core/tests/memory_ipc/test_device_uuid.py index f3ae7be8bd7..97499a64c97 100644 --- a/cuda_core/tests/memory_ipc/test_device_uuid.py +++ b/cuda_core/tests/memory_ipc/test_device_uuid.py @@ -5,9 +5,9 @@ The child process sees the parent's two devices in reversed order, so each ordinal names the other device there. A ``Device`` must arrive as the device -with the same UUID, 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). +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 @@ -149,15 +149,17 @@ 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") @@ -167,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") @@ -186,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") @@ -207,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")