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 d2119100c7b..59721c8d1c7 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/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 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 ed481334233..f078c3ac7d3 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/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) 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 827677939c1..eef1c09b644 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()