diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index ce1f9dab90b..fbc7b3d2ae4 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -204,7 +204,7 @@ what makes that safe: the allocation is released exactly once, when its last map A `VirtualMemoryBuffer` is shared the way pool-backed buffers are, through `Buffer.ipc_descriptor` and `Buffer.from_ipc_descriptor`, when the resource's -`handle_type` can travel between processes (`posix_fd` on Linux) and the memory +`handle_type` can travel between processes (`posix_fd` or `fabric` on Linux) and the memory lives on a device (`is_ipc_enabled` reports both; host-located memory is not shared yet). The exporter calls `cuMemExportToShareableHandle` once per mapping that covers the buffer's size, in address order, and the @@ -298,16 +298,43 @@ A `modify_allocation` config must keep the resource's `handle_type`: every chunk of a buffer must be exportable the same way. Fabric handles. `CU_MEM_HANDLE_TYPE_FABRIC` exports a `CUmemFabricHandle`, 64 -bytes of plain data that travel through `pickle`, MPI, or a socket, need no -close, and cannot hold a reference to the allocation while in flight. The -importer needs access to the exporter's IMEX channel rather than a matching -device index, and the memory may live on another node. The file-descriptor -rules above (ownership, pinning, `pickle` raising, the same-device check) do -not hold for it, so fabric sharing gets a descriptor class of its own instead -of a widened `VirtualMemoryFdBufferDescriptor`, whose `handles` stay -`IPCAllocationHandle`s. `from_fds()` is POSIX-only by construction for the -same reason and takes no handle type; a fabric descriptor is plain data and -needs no out-of-band constructor. Tracked on the VMM epic (#2980). +opaque bytes cached by the driver at `cuMemCreate`. Repeated exports of one +allocation return the same bytes once the handle struct is zeroed before the +export (the driver fills its fields and leaves the padding, which otherwise +made two exports of one buffer compare unequal on B300); export creates no OS +object and takes no reference. `VirtualMemoryFabricBufferDescriptor` carries one `bytes` per +allocation, its size, the buffer size, and the owner's device UUID. It +compares and hashes by (sizes, handles, size); the UUID is diagnostic metadata. +`from_handles()` validates externally transported handles and sizes. Plain +`pickle`, `multiprocessing`, MPI, or a socket can carry the descriptor, with +no reduction registration or spawning-Popen retention. The fd descriptor +keeps its ownership and identity semantics and its POSIX-only `from_fds()`. + +Nothing pins the allocation between export and import. The exporter must keep +the buffer alive until every importer has imported it; a descriptor parked +in a queue holds no driver reference. Import takes references of its own, +so the imported buffer is independent of the exporter after import. A stale +or forged handle fails with a driver error; an import failure carries a note +about stale handles, forgery, and differing IMEX channels. Imported memory +cannot be re-exported; the same provenance check covers both transports. + +Both concrete descriptors derive from `_VirtualMemoryDescriptor`, whose +`handle_ptr` supplies an fd cast to `void*` or fills a caller-provided fabric +scratch struct and returns its address. The import loop checks the handle +type before driver calls and uses the same location check, reservation, +mappings, and range for both transports. An imported fabric allocation +reports DEVICE and the owner's device ordinal as numbered by the importer. +A location mismatch includes the owner's UUID and its local ordinal, or says +that the UUID is not visible in this process. + +Construction checks `Device.properties.handle_type_fabric_supported` once +for a device-located fabric resource. It requires an NVLink fabric and an +accessible IMEX channel (`/dev/nvidia-caps-imex-channels/channel*`); otherwise +construction raises a useful `RuntimeError` before an opaque allocation +failure. `is_ipc_enabled` remains a configuration query without driver calls. +The types are available in both supported bindings majors; no version gate +is needed. Sharing is supported on one node; multi-node IMEX sharing is out +of scope and untested. Tracked on the VMM epic (#2980). ## Why `modify_allocation` returns a new buffer diff --git a/cuda_core/cuda/core/_memory/_ipc.pxd b/cuda_core/cuda/core/_memory/_ipc.pxd index 75d41c51193..947657fba54 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pxd +++ b/cuda_core/cuda/core/_memory/_ipc.pxd @@ -55,10 +55,20 @@ cdef class IPCAllocationHandle: # Descriptor of a VirtualMemoryResource buffer: one exported handle per # physical allocation that backs the buffer, in address order. -cdef class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): +cdef class _VirtualMemoryDescriptor(IPCBufferDescriptor): cdef readonly int _handle_type cdef readonly tuple _chunk_sizes + + cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL + + +cdef class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor): + cdef readonly tuple _handles + + +cdef class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor): cdef readonly tuple _handles + cdef object _owner_uuid # Buffer IPC Implementation diff --git a/cuda_core/cuda/core/_memory/_ipc.pyi b/cuda_core/cuda/core/_memory/_ipc.pyi index ffc3f39f602..ce59f53d108 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyi +++ b/cuda_core/cuda/core/_memory/_ipc.pyi @@ -58,7 +58,19 @@ class IPCAllocationHandle: @property def uuid(self) -> uuid.UUID: ... -class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): +class _VirtualMemoryDescriptor(IPCBufferDescriptor): + """Common allocation layout and driver-handle interface for VMM descriptors.""" + _handle_type: int + _chunk_sizes: tuple[Any, ...] + + @property + def handle_type(self) -> VirtualMemoryHandleType: + """The export handle type; the importing resource must use the same type.""" + @property + def sizes(self) -> tuple[int, ...]: + """Size of each physical allocation, in address order; the sum is at least ``size``.""" + +class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor): """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. The descriptor holds one exported POSIX file descriptor per physical @@ -66,8 +78,7 @@ class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): :class:`IPCAllocationHandle` that the descriptor owns and closes when it is released), with each allocation's size (``sizes``) and the buffer's size (``size``, which may be smaller than their sum). ``handle_type`` is - ``"posix_fd"``, the only handle type cuda.core shares today; the rules - below describe file descriptors. + ``"posix_fd"``; the rules below describe file descriptors. A descriptor pins the physical memory while it exists, in every process that holds one, and nothing else does: the exporting buffer keeps no @@ -92,8 +103,6 @@ class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): which checks what it can and fails instead of mapping a size the driver rejects. """ - _handle_type: int - _chunk_sizes: tuple[Any, ...] _handles: tuple[Any, ...] @staticmethod @@ -133,12 +142,6 @@ class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): If an ``int`` is not an open file descriptor. """ @property - def handle_type(self) -> VirtualMemoryHandleType: - """The handle type the exporter used; the importing resource must be configured with the same.""" - @property - def sizes(self) -> tuple[int, ...]: - """Size of each physical allocation, in address order; the sum is at least ``size``.""" - @property def handles(self) -> tuple[IPCAllocationHandle, ...]: """The exported handles, one :class:`IPCAllocationHandle` per allocation, in address order. @@ -155,5 +158,60 @@ class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): """ def __reduce__(self) -> tuple[object, ...]: ... +class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor): + """Plain-data descriptor of a :class:`VirtualMemoryBuffer` exported with fabric handles. + + Each physical allocation contributes one opaque 64-byte handle, in address + order. Plain ``pickle``, ``multiprocessing``, or a transport of your own + can carry these bytes. Use :meth:`from_handles` to rebuild a descriptor + received outside cuda.core, then :meth:`Buffer.from_ipc_descriptor` to + import it. ``handle_type`` is always ``"fabric"``. + + Nothing is pinned: unlike a file descriptor, a fabric descriptor parked + in a queue does not retain the memory. The exporter must keep the buffer + alive until every importer has imported it. After import, the imported + allocation holds its own driver references and is independent of the + exporter. Importing a stale or forged handle fails with a driver error. + + Sharing is supported on one node; multi-node sharing is untested. Fabric + handles require an NVLink fabric and an accessible IMEX channel. Treat + handles and sizes received from another process as untrusted and import + only through :meth:`Buffer.from_ipc_descriptor`. + """ + _handles: tuple[Any, ...] + + @staticmethod + def _from_exports(sizes: tuple[int, ...], handles: tuple[bytes, ...], size: int, owner_uuid: str | None) -> VirtualMemoryFabricBufferDescriptor: ... + @classmethod + def from_handles(cls, handles: Iterable[bytes], sizes: Iterable[int], *, size: int | None=None, owner_uuid: str | None=None) -> VirtualMemoryFabricBufferDescriptor: + """Build a descriptor from fabric handles received through a transport of your own. + + Parameters + ---------- + handles : Iterable[bytes] + One opaque 64-byte handle per physical allocation, in address order. + sizes : Iterable[int] + One positive allocation size per handle, in the same order. + size : int, optional + Buffer size, between zero and the sum of ``sizes``; the sum by default. + owner_uuid : str, optional + Exporting device's UUID, used to identify the owner in error messages. + + Raises + ------ + ValueError + If counts differ, a handle is not 64 bytes, an allocation size is + not a positive integer, or the buffer size is outside the allocation sizes. + """ + @property + def handles(self) -> tuple[bytes, ...]: + """One opaque 64-byte fabric handle per allocation, in address order.""" + @property + def owner_uuid(self) -> str | None: + """The exporting device's UUID, if supplied by the exporter.""" + def __reduce__(self) -> tuple[object, ...]: ... + def __eq__(self, other) -> bool: ... + def __hash__(self) -> int: ... + def _reduce_allocation_handle(alloc_handle: IPCAllocationHandle) -> tuple[object, ...]: ... def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None) -> IPCAllocationHandle: ... diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index 5091d39f0fc..90fb4c87925 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -7,6 +7,8 @@ from __future__ import annotations cimport cpython from libc.stddef cimport size_t +from libc.stdint cimport uintptr_t +from libc.string cimport memcpy from cuda.bindings cimport cydriver from cuda.core._memory._buffer cimport Buffer, Buffer_check_open, Buffer_from_deviceptr_handle from cuda.core._memory._memory_pool cimport _MemPool, MP_check_open @@ -177,7 +179,29 @@ def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handle) -cdef class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): +cdef class _VirtualMemoryDescriptor(IPCBufferDescriptor): + """Common allocation layout and driver-handle interface for VMM descriptors.""" + + @property + def handle_type(self) -> VirtualMemoryHandleType: + """The export handle type; the importing resource must use the same type.""" + from cuda.core._memory._virtual_memory_resource import VirtualMemoryResourceOptions + + for spec, value in VirtualMemoryResourceOptions._handle_types.items(): + if spec is not None and int(value) == self._handle_type: + return spec + raise ValueError(f"the descriptor carries an unknown handle type ({self._handle_type})") + + @property + def sizes(self) -> tuple[int, ...]: + """Size of each physical allocation, in address order; the sum is at least ``size``.""" + return self._chunk_sizes + + cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL: + raise NotImplementedError("a concrete VMM descriptor is required") + + +cdef class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor): """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. The descriptor holds one exported POSIX file descriptor per physical @@ -185,8 +209,7 @@ cdef class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): :class:`IPCAllocationHandle` that the descriptor owns and closes when it is released), with each allocation's size (``sizes``) and the buffer's size (``size``, which may be smaller than their sum). ``handle_type`` is - ``"posix_fd"``, the only handle type cuda.core shares today; the rules - below describe file descriptors. + ``"posix_fd"``; the rules below describe file descriptors. A descriptor pins the physical memory while it exists, in every process that holds one, and nothing else does: the exporting buffer keeps no @@ -294,20 +317,10 @@ cdef class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): handles.append(handle) return VirtualMemoryFdBufferDescriptor._from_exports(driver_type, sizes_t, tuple(handles), size) - @property - def handle_type(self) -> VirtualMemoryHandleType: - """The handle type the exporter used; the importing resource must be configured with the same.""" - from cuda.core._memory._virtual_memory_resource import VirtualMemoryResourceOptions - - for spec, value in VirtualMemoryResourceOptions._handle_types.items(): - if spec is not None and int(value) == self._handle_type: - return spec - raise ValueError(f"the descriptor carries an unknown handle type ({self._handle_type})") - - @property - def sizes(self) -> tuple[int, ...]: - """Size of each physical allocation, in address order; the sum is at least ``size``.""" - return self._chunk_sizes + cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL: + # fd zero is valid, so NULL is a valid result unless an exception is set. + cdef int fd = int(self._handles[i]) + return fd @property def handles(self) -> tuple[IPCAllocationHandle, ...]: @@ -334,6 +347,114 @@ cdef class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor): self._handle_type, self._chunk_sizes, self._handles, self._size) +cdef class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor): + """Plain-data descriptor of a :class:`VirtualMemoryBuffer` exported with fabric handles. + + Each physical allocation contributes one opaque 64-byte handle, in address + order. Plain ``pickle``, ``multiprocessing``, or a transport of your own + can carry these bytes. Use :meth:`from_handles` to rebuild a descriptor + received outside cuda.core, then :meth:`Buffer.from_ipc_descriptor` to + import it. ``handle_type`` is always ``"fabric"``. + + Nothing is pinned: unlike a file descriptor, a fabric descriptor parked + in a queue does not retain the memory. The exporter must keep the buffer + alive until every importer has imported it. After import, the imported + allocation holds its own driver references and is independent of the + exporter. Importing a stale or forged handle fails with a driver error. + + Sharing is supported on one node; multi-node sharing is untested. Fabric + handles require an NVLink fabric and an accessible IMEX channel. Treat + handles and sizes received from another process as untrusted and import + only through :meth:`Buffer.from_ipc_descriptor`. + """ + + @staticmethod + def _from_exports( + sizes: tuple[int, ...], handles: tuple[bytes, ...], size: int, owner_uuid: str | None + ) -> VirtualMemoryFabricBufferDescriptor: + return VirtualMemoryFabricBufferDescriptor.from_handles( + handles, sizes, size=size, owner_uuid=owner_uuid) + + @classmethod + def from_handles( + cls, handles: Iterable[bytes], sizes: Iterable[int], *, + size: int | None = None, owner_uuid: str | None = None, + ) -> VirtualMemoryFabricBufferDescriptor: + """Build a descriptor from fabric handles received through a transport of your own. + + Parameters + ---------- + handles : Iterable[bytes] + One opaque 64-byte handle per physical allocation, in address order. + sizes : Iterable[int] + One positive allocation size per handle, in the same order. + size : int, optional + Buffer size, between zero and the sum of ``sizes``; the sum by default. + owner_uuid : str, optional + Exporting device's UUID, used to identify the owner in error messages. + + Raises + ------ + ValueError + If counts differ, a handle is not 64 bytes, an allocation size is + not a positive integer, or the buffer size is outside the allocation sizes. + """ + cdef tuple handles_t = tuple(handles) + try: + sizes_t = tuple(operator.index(n) for n in sizes) + size = None if size is None else operator.index(size) + except TypeError as exc: + raise ValueError("allocation sizes and size must be integers") from exc + if len(handles_t) != len(sizes_t): + raise ValueError(f"{len(handles_t)} handles for {len(sizes_t)} sizes; one size per handle is required") + for handle in handles_t: + if not isinstance(handle, bytes) or len(handle) != 64: + raise ValueError("each fabric handle must be 64 bytes") + for n in sizes_t: + if n <= 0: + raise ValueError(f"allocation sizes must be positive, got {n}") + total = sum(sizes_t) + size = total if size is None else size + if size < 0 or size > total: + raise ValueError(f"size {size} is outside 0..{total}, the sum of the allocation sizes") + cdef VirtualMemoryFabricBufferDescriptor self = VirtualMemoryFabricBufferDescriptor.__new__(cls) + self._payload = b"" + self._handle_type = cydriver.CU_MEM_HANDLE_TYPE_FABRIC + self._chunk_sizes = sizes_t + self._handles = handles_t + self._size = size + self._owner_uuid = owner_uuid + return self + + @property + def handles(self) -> tuple[bytes, ...]: + """One opaque 64-byte fabric handle per allocation, in address order.""" + return self._handles + + @property + def owner_uuid(self) -> str | None: + """The exporting device's UUID, if supplied by the exporter.""" + return self._owner_uuid + + cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL: + cdef bytes handle = self._handles[i] + memcpy(scratch.data, handle, 64) + return scratch + + def __reduce__(self) -> tuple[object, ...]: + return VirtualMemoryFabricBufferDescriptor._from_exports, ( + self._chunk_sizes, self._handles, self._size, self._owner_uuid) + + def __eq__(self, other) -> bool: + if not isinstance(other, VirtualMemoryFabricBufferDescriptor): + return NotImplemented + cdef VirtualMemoryFabricBufferDescriptor rhs = other + return (self._chunk_sizes, self._handles, self._size) == (rhs._chunk_sizes, rhs._handles, rhs._size) + + def __hash__(self) -> int: + return hash((self._chunk_sizes, self._handles, self._size)) + + # Buffer IPC Implementation # ------------------------- cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self): @@ -359,7 +480,7 @@ cdef Buffer Buffer_from_ipc_descriptor( cls, object mr, IPCBufferDescriptor ipc_descriptor, stream ): """Import a buffer that was exported from another process.""" - if isinstance(ipc_descriptor, VirtualMemoryFdBufferDescriptor): + if isinstance(ipc_descriptor, _VirtualMemoryDescriptor): # Imported here rather than cimported: _virtual_memory_resource # cimports this module. from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 4716c141b24..60ab2619cea 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -6,7 +6,9 @@ from typing import Iterable from cuda.core._device import Device from cuda.core._memory._buffer import Buffer, MemoryResource -from cuda.core._memory._ipc import VirtualMemoryFdBufferDescriptor +from cuda.core._memory._ipc import (VirtualMemoryFabricBufferDescriptor, + VirtualMemoryFdBufferDescriptor, + _VirtualMemoryDescriptor) from cuda.core._stream import Stream from cuda.core._utils.cuda_utils import driver from cuda.core.graph import GraphBuilder @@ -18,7 +20,7 @@ from cuda.core.typing import (DevicePointerType, VirtualMemoryAccessType, __all__ = ['VirtualMemoryBuffer', 'VirtualMemoryResource', 'VirtualMemoryResourceOptions'] _HOST_LOCATION_TYPES = frozenset({VirtualMemoryLocationType.HOST, VirtualMemoryLocationType.HOST_NUMA, VirtualMemoryLocationType.HOST_NUMA_CURRENT}) -_IPC_HANDLE_TYPES = frozenset({VirtualMemoryHandleType.POSIX_FD}) if platform.system() == 'Linux' else frozenset() +_IPC_HANDLE_TYPES = frozenset({VirtualMemoryHandleType.POSIX_FD, VirtualMemoryHandleType.FABRIC}) if platform.system() == 'Linux' else frozenset() @dataclass class VirtualMemoryResourceOptions: @@ -36,10 +38,20 @@ class VirtualMemoryResourceOptions: handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str Export handle type for the physical allocation, and the switch for sharing buffers with other processes (see - :attr:`VirtualMemoryResource.is_ipc_enabled`). Use ``"posix_fd"`` on - Linux if you plan to import/export the allocation. Use `None` if you - don't need an exportable handle. ``location_type="host"`` requires - `None`. + :attr:`VirtualMemoryResource.is_ipc_enabled`). Use ``"posix_fd"`` or + ``"fabric"`` on Linux if you plan to import/export the allocation. + Fabric handles work with plain ``pickle`` and pin nothing: keep the + exporting buffer alive until every importer has imported it. Sharing + is supported on one node; multi-node sharing is untested. Fabric needs + an NVLink fabric and an accessible IMEX channel + (``/dev/nvidia-caps-imex-channels/channel*``); construction raises + ``RuntimeError`` if the device reports no fabric handle support. A + device can report support while the allocation is refused, because + the attribute does not check for an NVLink fabric and older drivers + do not check for an IMEX channel; the allocation then fails with a + driver error that says so. + Use `None` if you don't need an exportable handle. + ``location_type="host"`` requires `None`. gpu_direct_rdma: bool Hint that the allocation should be GDR-capable (if supported). granularity: :obj:`~_memory.VirtualMemoryGranularityType` | str @@ -107,10 +119,15 @@ class VirtualMemoryBuffer(Buffer): freed when the last buffer that maps it closes. """ @property - def ipc_descriptor(self) -> VirtualMemoryFdBufferDescriptor: + def ipc_descriptor(self) -> VirtualMemoryFdBufferDescriptor | VirtualMemoryFabricBufferDescriptor: """A new descriptor for sharing this buffer with other processes. - Every access exports the physical allocations that back the buffer + The descriptor type follows the resource's ``handle_type``. Fabric + descriptors carry plain bytes and pin nothing: keep the exporting + buffer alive until every importer has imported it. Plain ``pickle`` + works for fabric descriptors and buffers. + + For ``"posix_fd"``, every access exports the physical allocations again and returns a new descriptor that owns one file descriptor per allocation; the file descriptors close when the descriptor is released. Keep the descriptor in a variable while you use its @@ -184,8 +201,9 @@ class VirtualMemoryResource(MemoryResource): A shared buffer's descriptor comes from the exporting process and is not trusted: :meth:`Buffer.from_ipc_descriptor` checks what it can, and the driver rejects a size that does not match the exported allocation. Unpickling a shared buffer performs a live import, - so unpickle buffers only from a trusted principal. A descriptor pins the physical memory - while it exists; release it when the receiver has imported it. + so unpickle buffers only from a trusted principal. A POSIX fd descriptor pins the physical + memory while it exists; release it when the receiver has imported it. Fabric descriptors + pin nothing; keep the exporting buffer alive until every importer has imported it. Notes ----- @@ -200,7 +218,7 @@ class VirtualMemoryResource(MemoryResource): :meth:`Buffer.set_deallocation_stream`. Buffers can be shared with other processes when ``config.handle_type`` - is ``"posix_fd"`` (Linux), the only handle type cuda.core shares today; + is ``"posix_fd"`` or ``"fabric"`` on Linux; ``handle_type`` is the switch and :attr:`is_ipc_enabled` the query. The rules that follow describe POSIX file descriptors. :attr:`Buffer.ipc_descriptor` exports the physical allocations that back @@ -217,8 +235,22 @@ class VirtualMemoryResource(MemoryResource): ``VirtualMemoryFdBufferDescriptor.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 - it exists, in every process that holds one, and nothing else does. A - buffer that contains imported memory (an import, or anything + it exists, in every process that holds one, and nothing else does. + Fabric descriptors carry one opaque 64-byte handle per allocation and work + with plain ``pickle``; they pin nothing, so the exporter must keep the + buffer alive until every importer has imported it. Rebuild handles sent + through a transport of your own with + ``VirtualMemoryFabricBufferDescriptor.from_handles()``. Fabric sharing is + supported on one node; multi-node sharing is untested. It requires a device + on an NVLink fabric and an accessible IMEX channel + (``/dev/nvidia-caps-imex-channels/channel*``); construction raises + ``RuntimeError`` when the device reports no fabric handle support. A device + can report support while the allocation is refused, because the attribute + does not check for an NVLink fabric and older drivers do not check for an + IMEX channel; the allocation then fails with a driver error that says so. + Fabric-handle sharing was validated on NVIDIA internal hardware and is not + covered by the project's CI in this release. + A buffer that contains imported memory (an import, or anything :meth:`modify_allocation` derived from one) cannot be exported again; forward the descriptor it came from, or copy into a buffer you own. """ @@ -347,15 +379,15 @@ class VirtualMemoryResource(MemoryResource): def is_ipc_enabled(self) -> bool: """Whether buffers of this resource can be shared with other processes. - True when ``config.handle_type`` is ``"posix_fd"`` on Linux and the - memory lives on a device; ``handle_type`` is the only switch, there is + True when ``config.handle_type`` is ``"posix_fd"`` or ``"fabric"`` on + Linux and the memory lives on a device; ``handle_type`` is the only switch, there is no separate option. With ``handle_type=None`` the allocations cannot - be exported, Win32 KMT handles are not transported by cuda.core, and - fabric handles are not supported yet. Host-located virtual memory - (``location_type`` ``"host"``, ``"host_numa"``, or + be exported, and Win32 KMT handles are not transported by cuda.core. + This is a configuration query and makes no driver call. Host-located + virtual memory (``location_type`` ``"host"``, ``"host_numa"``, or ``"host_numa_current"``) is not shared yet. """ - def _import_ipc_buffer(self, desc: VirtualMemoryFdBufferDescriptor, stream) -> VirtualMemoryBuffer: + def _import_ipc_buffer(self, desc: _VirtualMemoryDescriptor, stream) -> VirtualMemoryBuffer: """Import a buffer that another process exported; see :meth:`Buffer.from_ipc_descriptor`. Every exported allocation is imported with this resource's handle diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 69a04ccc4aa..fbbd091a691 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -4,10 +4,11 @@ from __future__ import annotations -from libc.stdint cimport uintptr_t from libc.string cimport memset from libcpp.vector cimport vector +cimport cpython + from cuda.bindings cimport cydriver from cuda.core._memory._buffer cimport ( Buffer, @@ -18,6 +19,8 @@ from cuda.core._memory._buffer cimport ( from cuda.core._memory._ipc cimport ( IPCAllocationHandle, IPCDataForBuffer, + _VirtualMemoryDescriptor, + VirtualMemoryFabricBufferDescriptor, VirtualMemoryFdBufferDescriptor, ) from cuda.core._rt cimport ( @@ -84,9 +87,12 @@ _HOST_LOCATION_TYPES = frozenset( ) # Handle types whose exported handles cuda.core can carry to another process. -# POSIX file descriptors travel through multiprocessing on Linux. Win32 KMT -# handles have no transport here, and fabric handles are not supported yet. -_IPC_HANDLE_TYPES = frozenset({VirtualMemoryHandleType.POSIX_FD}) if platform.system() == "Linux" else frozenset() +# POSIX file descriptors and plain fabric handles are transported on Linux. +# Win32 KMT handles have no transport here. +_IPC_HANDLE_TYPES = ( + frozenset({VirtualMemoryHandleType.POSIX_FD, VirtualMemoryHandleType.FABRIC}) + if platform.system() == "Linux" else frozenset() +) @dataclass @@ -105,10 +111,20 @@ class VirtualMemoryResourceOptions: handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str Export handle type for the physical allocation, and the switch for sharing buffers with other processes (see - :attr:`VirtualMemoryResource.is_ipc_enabled`). Use ``"posix_fd"`` on - Linux if you plan to import/export the allocation. Use `None` if you - don't need an exportable handle. ``location_type="host"`` requires - `None`. + :attr:`VirtualMemoryResource.is_ipc_enabled`). Use ``"posix_fd"`` or + ``"fabric"`` on Linux if you plan to import/export the allocation. + Fabric handles work with plain ``pickle`` and pin nothing: keep the + exporting buffer alive until every importer has imported it. Sharing + is supported on one node; multi-node sharing is untested. Fabric needs + an NVLink fabric and an accessible IMEX channel + (``/dev/nvidia-caps-imex-channels/channel*``); construction raises + ``RuntimeError`` if the device reports no fabric handle support. A + device can report support while the allocation is refused, because + the attribute does not check for an NVLink fabric and older drivers + do not check for an IMEX channel; the allocation then fails with a + driver error that says so. + Use `None` if you don't need an exportable handle. + ``location_type="host"`` requires `None`. gpu_direct_rdma: bool Hint that the allocation should be GDR-capable (if supported). granularity: :obj:`~_memory.VirtualMemoryGranularityType` | str @@ -268,6 +284,18 @@ cdef int _raise_last_error_noting(str detail) except -1: return 0 +cdef int _raise_allocation_error(const cydriver.CUmemAllocationProp* prop) except -1: + """Raise the status of a failed physical allocation; a fabric resource gets the IMEX note.""" + if prop.requestedHandleTypes == cydriver.CU_MEM_HANDLE_TYPE_FABRIC: + _raise_last_error_noting( + "while creating a physical allocation with a fabric handle; fabric handles need an NVLink fabric " + "and an accessible IMEX channel, and a device can report support while the allocation is refused " + "(the attribute does not check for an NVLink fabric, and older drivers do not check for an IMEX " + "channel)" + ) + return _raise_last_error() + + cdef int _check_noting(cydriver.CUresult status, str detail) except -1: """``HANDLE_RETURN(status)``, with ``detail`` attached to the raised error as a note.""" if status == cydriver.CUDA_SUCCESS: @@ -305,10 +333,15 @@ cdef class VirtualMemoryBuffer(Buffer): """ @property - def ipc_descriptor(self) -> VirtualMemoryFdBufferDescriptor: + def ipc_descriptor(self) -> VirtualMemoryFdBufferDescriptor | VirtualMemoryFabricBufferDescriptor: """A new descriptor for sharing this buffer with other processes. - Every access exports the physical allocations that back the buffer + The descriptor type follows the resource's ``handle_type``. Fabric + descriptors carry plain bytes and pin nothing: keep the exporting + buffer alive until every importer has imported it. Plain ``pickle`` + works for fabric descriptors and buffers. + + For ``"posix_fd"``, every access exports the physical allocations again and returns a new descriptor that owns one file descriptor per allocation; the file descriptors close when the descriptor is released. Keep the descriptor in a variable while you use its @@ -346,6 +379,8 @@ cdef class VirtualMemoryBuffer(Buffer): # Buffers from a trusted principal. Must not serialize the parent's stream. Buffer_check_open(self) desc = _export_ipc_descriptor(self) + if isinstance(desc, VirtualMemoryFabricBufferDescriptor): + return Buffer._reduce_helper, (self.memory_resource, desc) popen = _mp_context.get_spawning_popen() if popen is not None: # A Queue, Pipe, or Pool duplicates the file descriptors while this @@ -418,8 +453,9 @@ cdef class VirtualMemoryResource(MemoryResource): A shared buffer's descriptor comes from the exporting process and is not trusted: :meth:`Buffer.from_ipc_descriptor` checks what it can, and the driver rejects a size that does not match the exported allocation. Unpickling a shared buffer performs a live import, - so unpickle buffers only from a trusted principal. A descriptor pins the physical memory - while it exists; release it when the receiver has imported it. + so unpickle buffers only from a trusted principal. A POSIX fd descriptor pins the physical + memory while it exists; release it when the receiver has imported it. Fabric descriptors + pin nothing; keep the exporting buffer alive until every importer has imported it. Notes ----- @@ -434,7 +470,7 @@ cdef class VirtualMemoryResource(MemoryResource): :meth:`Buffer.set_deallocation_stream`. Buffers can be shared with other processes when ``config.handle_type`` - is ``"posix_fd"`` (Linux), the only handle type cuda.core shares today; + is ``"posix_fd"`` or ``"fabric"`` on Linux; ``handle_type`` is the switch and :attr:`is_ipc_enabled` the query. The rules that follow describe POSIX file descriptors. :attr:`Buffer.ipc_descriptor` exports the physical allocations that back @@ -451,8 +487,22 @@ cdef class VirtualMemoryResource(MemoryResource): ``VirtualMemoryFdBufferDescriptor.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 - it exists, in every process that holds one, and nothing else does. A - buffer that contains imported memory (an import, or anything + it exists, in every process that holds one, and nothing else does. + Fabric descriptors carry one opaque 64-byte handle per allocation and work + with plain ``pickle``; they pin nothing, so the exporter must keep the + buffer alive until every importer has imported it. Rebuild handles sent + through a transport of your own with + ``VirtualMemoryFabricBufferDescriptor.from_handles()``. Fabric sharing is + supported on one node; multi-node sharing is untested. It requires a device + on an NVLink fabric and an accessible IMEX channel + (``/dev/nvidia-caps-imex-channels/channel*``); construction raises + ``RuntimeError`` when the device reports no fabric handle support. A device + can report support while the allocation is refused, because the attribute + does not check for an NVLink fabric and older drivers do not check for an + IMEX channel; the allocation then fails with a driver error that says so. + Fabric-handle sharing was validated on NVIDIA internal hardware and is not + covered by the project's CI in this release. + A buffer that contains imported memory (an import, or anything :meth:`modify_allocation` derived from one) cannot be exported again; forward the descriptor it came from, or copy into a buffer you own. """ @@ -471,6 +521,13 @@ cdef class VirtualMemoryResource(MemoryResource): self.device = None if self.device is not None and not self.device.properties.virtual_memory_management_supported: raise RuntimeError("VirtualMemoryResource requires CUDA VMM API support") + if (self.config.handle_type == VirtualMemoryHandleType.FABRIC + and self.device is not None and not self.device.properties.handle_type_fabric_supported): + raise RuntimeError( + "fabric handles need a device on an NVLink fabric and an accessible IMEX channel " + "(/dev/nvidia-caps-imex-channels/channel*); " + f"device {self.device.device_id} reports no fabric handle support" + ) self._check_config(self.config) def __reduce__(self) -> tuple[object, ...]: @@ -683,7 +740,7 @@ cdef class VirtualMemoryResource(MemoryResource): with nogil: h_alloc = create_mem_allocation_handle(aligned, prop, descs.data(), descs.size()) if not h_alloc: - _raise_last_error() + _raise_allocation_error(&prop) with nogil: h_res = create_va_reservation_handle(aligned, addr_align, hint) if not h_res: @@ -836,7 +893,7 @@ cdef class VirtualMemoryResource(MemoryResource): with nogil: h_alloc = create_mem_allocation_handle(add, prop, descs.data(), descs.size()) if not h_alloc: - _raise_last_error() + _raise_allocation_error(&prop) with nogil: h_map = create_va_mapping_handle(base + total, h_alloc, h_res) if not h_map: @@ -869,7 +926,7 @@ cdef class VirtualMemoryResource(MemoryResource): with nogil: h_alloc = create_mem_allocation_handle(add, prop, descs.data(), descs.size()) if not h_alloc: - _raise_last_error() + _raise_allocation_error(&prop) with nogil: m2 = create_va_mapping_handle(base_new + offset, h_alloc, h_res_new) if not m2: @@ -938,12 +995,12 @@ cdef class VirtualMemoryResource(MemoryResource): def is_ipc_enabled(self) -> bool: """Whether buffers of this resource can be shared with other processes. - True when ``config.handle_type`` is ``"posix_fd"`` on Linux and the - memory lives on a device; ``handle_type`` is the only switch, there is + True when ``config.handle_type`` is ``"posix_fd"`` or ``"fabric"`` on + Linux and the memory lives on a device; ``handle_type`` is the only switch, there is no separate option. With ``handle_type=None`` the allocations cannot - be exported, Win32 KMT handles are not transported by cuda.core, and - fabric handles are not supported yet. Host-located virtual memory - (``location_type`` ``"host"``, ``"host_numa"``, or + be exported, and Win32 KMT handles are not transported by cuda.core. + This is a configuration query and makes no driver call. Host-located + virtual memory (``location_type`` ``"host"``, ``"host_numa"``, or ``"host_numa_current"``) is not shared yet. """ return ( @@ -951,7 +1008,7 @@ cdef class VirtualMemoryResource(MemoryResource): and self.config.location_type not in _HOST_LOCATION_TYPES ) - def _import_ipc_buffer(self, VirtualMemoryFdBufferDescriptor desc not None, stream) -> VirtualMemoryBuffer: + def _import_ipc_buffer(self, _VirtualMemoryDescriptor desc not None, stream) -> VirtualMemoryBuffer: """Import a buffer that another process exported; see :meth:`Buffer.from_ipc_descriptor`. Every exported allocation is imported with this resource's handle @@ -980,7 +1037,8 @@ cdef class VirtualMemoryResource(MemoryResource): cdef size_t gran, gran_min = 0, addr_align, chunk, total = 0, offset = 0 cdef Py_ssize_t n, i cdef cydriver.CUdeviceptr hint - cdef int handle_type, fd + cdef int handle_type + cdef cydriver.CUmemFabricHandle fh cdef void* os_handle if not self.is_ipc_enabled: @@ -992,7 +1050,7 @@ cdef class VirtualMemoryResource(MemoryResource): f"handle_type={str(cfg.handle_type)!r}; import it with a resource configured like the exporter's" ) n = len(desc._chunk_sizes) - if n != len(desc._handles): + if n != len(desc.handles): raise ValueError("malformed VirtualMemoryResource buffer descriptor") if stream is not None: s = Stream_accept(stream) @@ -1044,19 +1102,25 @@ cdef class VirtualMemoryResource(MemoryResource): # buffer therefore does not keep the descriptor; its file descriptors # live only as long as the caller holds the descriptor object. for i in range(n): - fd = int(desc._handles[i]) # raises if the handle was closed chunk = desc._chunk_sizes[i] - os_handle = fd + os_handle = desc.handle_ptr(i, &fh) with nogil: h_alloc = import_mem_allocation_handle( os_handle, handle_type, chunk, descs.data(), descs.size()) if not h_alloc: - _raise_last_error_noting( - f"while importing chunk {i + 1} of {n} of a VirtualMemoryResource buffer descriptor; either " - "the shared handle does not refer to a CUDA allocation exported with " - f"handle_type={str(cfg.handle_type)!r}, or the device that owns the memory is not usable in " - "this process (for example, CUDA_VISIBLE_DEVICES hides it)" - ) + if handle_type == cydriver.CU_MEM_HANDLE_TYPE_FABRIC: + note = ( + f"while importing chunk {i + 1} of {n} of a fabric buffer descriptor; the handle may be " + "stale (the exporting buffer was closed), forged, or exported on a different IMEX channel" + ) + else: + note = ( + f"while importing chunk {i + 1} of {n} of a VirtualMemoryResource buffer descriptor; either " + "the shared handle does not refer to a CUDA allocation exported with " + f"handle_type={str(cfg.handle_type)!r}, or the device that owns the memory is not usable in " + "this process (for example, CUDA_VISIBLE_DEVICES hides it)" + ) + _raise_last_error_noting(note) # The exporter chose where the memory lives. This resource must map # memory of that kind on that device, or its access descriptors # would describe the wrong location. @@ -1066,8 +1130,23 @@ cdef class VirtualMemoryResource(MemoryResource): prop.location.type == cydriver.CU_MEM_LOCATION_TYPE_DEVICE and imported_prop.location.id != prop.location.id ): + owner_detail = "" + owner_uuid = getattr(desc, "owner_uuid", None) + if owner_uuid is not None: + try: + owner = _device_from_uuid(owner_uuid) + except RuntimeError: + owner_detail = ( + f"memory belongs to the device with UUID {owner_uuid}, " + "which is not visible in this process; " + ) + else: + owner_detail = ( + f"memory belongs to the device with UUID {owner_uuid}, " + f"which is device {owner.device_id} in this process; " + ) raise ValueError( - f"chunk {i + 1} of the descriptor is memory on {_describe_location(imported_prop.location)}, " + f"{owner_detail}chunk {i + 1} of the descriptor is memory on {_describe_location(imported_prop.location)}, " f"but this resource maps memory on {_describe_location(prop.location)}; import the descriptor " "with a VirtualMemoryResource for the device that owns the memory, and name other devices " "in its peers option" @@ -1127,7 +1206,7 @@ def _reconstruct_virtual_memory_resource( return cls(device, replace(config, peers=peers)) -cdef VirtualMemoryFdBufferDescriptor _export_ipc_descriptor(Buffer buf): +cdef _VirtualMemoryDescriptor _export_ipc_descriptor(Buffer buf): """Export the allocations that back ``buf`` as shareable handles, in address order. Only the allocations that cover ``buf.size`` are exported; a buffer's size @@ -1147,6 +1226,7 @@ cdef VirtualMemoryFdBufferDescriptor _export_ipc_descriptor(Buffer buf): cdef size_t i, n, chunk, offset = 0 cdef int handle_type cdef int fd = -1 + cdef cydriver.CUmemFabricHandle fh cdef list sizes = [] cdef list handles = [] @@ -1171,6 +1251,16 @@ cdef VirtualMemoryFdBufferDescriptor _export_ipc_descriptor(Buffer buf): for i in range(n): h_alloc = allocs[i] chunk = mem_allocation_size(h_alloc) + if handle_type == cydriver.CU_MEM_HANDLE_TYPE_FABRIC: + # The driver fills its fields of the 64 bytes and leaves the padding; zero it first so that two + # exports of one allocation carry the same bytes and no stack bytes travel in a descriptor. + memset(&fh, 0, sizeof(fh)) + with nogil: + HANDLE_RETURN(cydriver.cuMemExportToShareableHandle( + &fh, as_cu(h_alloc), handle_type, 0)) + handles.append(cpython.PyBytes_FromStringAndSize(fh.data, sizeof(fh.data))) + sizes.append(chunk) + continue with nogil: status = cydriver.cuMemExportToShareableHandle( &fd, as_cu(h_alloc), handle_type, 0) @@ -1189,4 +1279,7 @@ cdef VirtualMemoryFdBufferDescriptor _export_ipc_descriptor(Buffer buf): # The handle owns the fd now; if append fails, the handle closes it once. handles.append(handle) sizes.append(chunk) + if handle_type == cydriver.CU_MEM_HANDLE_TYPE_FABRIC: + return VirtualMemoryFabricBufferDescriptor._from_exports( + tuple(sizes), tuple(handles), buf._size, mr.device.uuid) return VirtualMemoryFdBufferDescriptor._from_exports(handle_type, tuple(sizes), tuple(handles), buf._size) diff --git a/cuda_core/docs/source/api_private.rst b/cuda_core/docs/source/api_private.rst index bc8c511b356..f39af0c5eb1 100644 --- a/cuda_core/docs/source/api_private.rst +++ b/cuda_core/docs/source/api_private.rst @@ -49,6 +49,7 @@ CUDA runtime _memory._ipc.IPCAllocationHandle _memory._ipc.IPCBufferDescriptor _memory._ipc.VirtualMemoryFdBufferDescriptor + _memory._ipc.VirtualMemoryFabricBufferDescriptor _memory._managed_buffer.AccessedBySetProxy 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 5d6521d30a8..f33dd190077 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -84,21 +84,39 @@ New features and proceeds without the synchronization. - :class:`VirtualMemoryResource` buffers can now be shared with other - processes the way pool-backed buffers are. :attr:`Buffer.ipc_descriptor` - exports every physical allocation that backs a :class:`VirtualMemoryBuffer` - as a POSIX file descriptor, :meth:`Buffer.from_ipc_descriptor` imports the - descriptor with a ``VirtualMemoryResource`` of the receiving process for the - device that owns the memory, and a buffer or its descriptor can be sent - through ``multiprocessing`` directly (plain ``pickle`` cannot carry file - descriptors). :attr:`VirtualMemoryResource.is_ipc_enabled` is ``True`` when - the resource's ``handle_type`` is ``"posix_fd"`` on Linux and the memory - lives on a device; host-located virtual memory is not shared yet. A grown buffer - imports as one contiguous range. A descriptor holds one file descriptor per - physical allocation and pins the memory while it exists; nothing else does. - A buffer that contains imported memory cannot be exported again. For a - transport of your own, the descriptor exposes ``fds``, ``sizes``, and - ``size``, and ``VirtualMemoryFdBufferDescriptor.from_fds()`` rebuilds it - in the receiver. Fabric handles are not shared yet. + processes the way pool-backed buffers are, through two transports chosen + by the resource's ``handle_type``. :attr:`Buffer.ipc_descriptor` exports + every physical allocation that backs a :class:`VirtualMemoryBuffer`, + :meth:`Buffer.from_ipc_descriptor` imports the descriptor with a + ``VirtualMemoryResource`` of the receiving process for the device that owns + the memory, and :attr:`VirtualMemoryResource.is_ipc_enabled` is ``True`` + when ``handle_type`` is ``"posix_fd"`` or ``"fabric"`` on Linux and the + memory lives on a device; host-located virtual memory is not shared yet. A + grown buffer imports as one contiguous range, and a buffer that contains + imported memory cannot be exported again. + + With ``"posix_fd"``, each allocation is exported as a POSIX file + descriptor; a buffer or its descriptor travels through ``multiprocessing`` + directly (plain ``pickle`` cannot carry file descriptors). The descriptor + holds one file descriptor per physical allocation and pins the memory while + it exists; nothing else does. For a transport of your own, it exposes + ``fds``, ``sizes``, and ``size``, and + ``VirtualMemoryFdBufferDescriptor.from_fds()`` rebuilds it in the receiver. + + With ``"fabric"``, each allocation is exported as a 64-byte fabric handle; + the descriptor is plain data that ``pickle`` carries, and it pins nothing: + keep the exporting buffer alive until every importer has imported it. + ``VirtualMemoryFabricBufferDescriptor.from_handles()`` rebuilds a descriptor + received through a transport of your own. Sharing is supported on one node; + multi-node sharing is untested. Fabric requires an NVLink fabric and an + accessible IMEX channel (``/dev/nvidia-caps-imex-channels/channel*``); + resource construction raises ``RuntimeError`` when the device reports no + fabric handle support. A device can report support while the allocation is + refused, because the attribute does not check for an NVLink fabric and + older drivers do not check for an IMEX channel; the allocation then fails + with a driver error that says so. Fabric-handle sharing was + validated on NVIDIA internal hardware and is not covered by the project's + CI in this release. (`#2980 `__) Fixes and enhancements diff --git a/cuda_core/tests/conftest.py b/cuda_core/tests/conftest.py index e62f89b8251..35ee9f2adce 100644 --- a/cuda_core/tests/conftest.py +++ b/cuda_core/tests/conftest.py @@ -425,6 +425,70 @@ def vmm_ipc_device_x2(device_x2): yield devices +_fabric_allocation_error = {} # device ordinal -> "" when a fabric allocation works there, else the status + + +def _require_vmm_fabric_device(device): + """Return a VMM device with fabric handles and an accessible IMEX channel, or skip.""" + from cuda_python_test_helpers import IS_WINDOWS, IS_WSL + + if IS_WINDOWS or IS_WSL: + pytest.skip("Sharing fabric handles requires Linux") + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + if not device.properties.handle_type_fabric_supported: + pytest.skip("Fabric handles need an NVLink fabric and an accessible IMEX channel") + # The attribute does not check for an NVLink fabric, and older drivers do not check for an IMEX + # channel either, so it is true on the CI runners (no fabric, no channel), where the allocation + # itself is refused. Probe once per device. + if device.device_id not in _fabric_allocation_error: + from cuda.core import Device, VirtualMemoryResource, VirtualMemoryResourceOptions + from cuda.core._utils.cuda_utils import CUDAError + + current = Device() + device.set_current() + try: + mr = VirtualMemoryResource(device, VirtualMemoryResourceOptions(handle_type="fabric")) + mr.allocate(1).close() + _fabric_allocation_error[device.device_id] = "" + except CUDAError as exc: + # Only a refused physical allocation means "no fabric here" (NOT_PERMITTED on the CI runners, + # INVALID_VALUE on an H200 without an IMEX channel). VirtualMemoryResource attaches its note about + # the fabric and IMEX requirements at that stage alone (on Python 3.10 it lands in the message); + # a failure of the reservation, the mapping, or the stream setup is a real error and propagates. + if "fabric handle" not in " ".join((str(exc), *getattr(exc, "__notes__", ()))): + raise + _fabric_allocation_error[device.device_id] = str(exc).split(":", 1)[0] + finally: + current.set_current() + if _fabric_allocation_error[device.device_id]: + pytest.skip( + "The device reports fabric handle support, but a fabric allocation is refused here: " + + _fabric_allocation_error[device.device_id] + ) + return device + + +@pytest.fixture +def vmm_fabric_device(init_cuda): + """A fabric-capable VMM device, tracking spawned children for teardown.""" + from helpers.child_processes import track_child_processes + + device = _require_vmm_fabric_device(init_cuda) + with track_child_processes(): + yield device + + +@pytest.fixture +def vmm_fabric_device_x2(device_x2): + """Two fabric-capable VMM devices, tracking spawned children for teardown.""" + from helpers.child_processes import track_child_processes + + devices = tuple(_require_vmm_fabric_device(device) for device in device_x2) + with track_child_processes(): + yield devices + + def _make_ipc_memory_resource(kind, device): """Create an IPC-enabled memory resource of the given kind, or skip when the device cannot.""" if kind == "device": diff --git a/cuda_core/tests/memory_ipc/test_vmm_fabric.py b/cuda_core/tests/memory_ipc/test_vmm_fabric.py new file mode 100644 index 00000000000..af497bea769 --- /dev/null +++ b/cuda_core/tests/memory_ipc/test_vmm_fabric.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-node VMM sharing with plain fabric handles (issue #2980).""" + +import ctypes +import gc +import multiprocessing as mp +import os +import pickle +from contextlib import ExitStack, closing + +import pytest +from helpers.child_processes import child_timeout_sec, kill_subprocesses, warmup_multiprocessing + +from cuda.bindings import driver +from cuda.core import Buffer, Device, VirtualMemoryResource, VirtualMemoryResourceOptions +from cuda.core._device import _device_from_uuid +from cuda.core._memory._ipc import VirtualMemoryFabricBufferDescriptor +from cuda.core._utils.cuda_utils import CUDAError, handle_return +from cuda_python_test_helpers import IS_WINDOWS, IS_WSL + +CHILD_TIMEOUT_SEC = child_timeout_sec() +pytestmark = pytest.mark.parallel_threads_limit(4) + + +def _resource(device, **options): + options.setdefault("handle_type", "fabric") + return VirtualMemoryResource(device, VirtualMemoryResourceOptions(**options)) + + +def _fill(buf, value): + handle_return(driver.cuMemsetD8(int(buf.handle), value, buf.size)) + handle_return(driver.cuCtxSynchronize()) + + +def _read(buf): + host = (ctypes.c_ubyte * 16)() + handle_return(driver.cuMemcpyDtoH(ctypes.addressof(host), int(buf.handle), 16)) + return bytes(host) + + +def _notes(exc): + # Python 3.10 appends the note to the message instead of using PEP 678. + return " ".join((str(exc), *getattr(exc, "__notes__", ()))) + + +def _check_import_error(exc): + # The driver's status for a bad handle varies with the IMEX setup; the note is cuda.core's and does not. + for detail in ("stale", "exporting buffer was closed", "forged", "different IMEX channel"): + assert detail in _notes(exc) + + +def _exchange(owner_uuid, payload): + device = _device_from_uuid(owner_uuid) + device.set_current() + buf = ( + Buffer.from_ipc_descriptor(_resource(device), payload, stream=device.default_stream) + if isinstance(payload, VirtualMemoryFabricBufferDescriptor) + else payload + ) + with buf: + seen = _read(buf) + _fill(buf, 0x62) + return seen + + +def _child_exchange(owner_uuid, mode, incoming, outgoing): + if mode == "queue": + payload = incoming.get(timeout=CHILD_TIMEOUT_SEC) + elif mode == "pipe": + assert incoming.poll(CHILD_TIMEOUT_SEC), "no payload from parent" + payload = incoming.recv() + else: + payload = incoming + outgoing.send(_exchange(owner_uuid, payload)) + outgoing.close() + + +def _round_trip(mode, owner_uuid, payload): + if mode == "pool": + with mp.Pool(1) as pool: + return pool.apply_async(_exchange, (owner_uuid, payload)).get(timeout=CHILD_TIMEOUT_SEC) + with ExitStack() as stack: + receive, send = mp.Pipe(duplex=False) + stack.enter_context(receive) + stack.enter_context(send) + if mode == "queue": + incoming = mp.Queue() + stack.callback(incoming.join_thread) + stack.callback(incoming.close) + child_input = incoming + elif mode == "pipe": + child_input, incoming = mp.Pipe(duplex=False) + stack.enter_context(child_input) + stack.enter_context(incoming) + else: + child_input = payload + process = stack.enter_context( + closing(mp.Process(target=_child_exchange, args=(owner_uuid, mode, child_input, send))) + ) + try: + process.start() + send.close() + if mode == "queue": + incoming.put(payload) + elif mode == "pipe": + child_input.close() + incoming.send(payload) + assert receive.poll(CHILD_TIMEOUT_SEC), "child did not respond" + result = receive.recv() + process.join(timeout=CHILD_TIMEOUT_SEC) + finally: + survivors = kill_subprocesses(process) + assert not survivors, "child did not exit within timeout" + assert process.exitcode == 0 + return result + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_export_shape(vmm_fabric_device): + device = vmm_fabric_device + mr = _resource(device) + with mr.allocate(1) as first, mr.modify_allocation(first, 2 * first.size) as grown: + desc = grown.ipc_descriptor + assert desc.handle_type == "fabric" + assert desc.owner_uuid == device.uuid + assert desc.sizes == (first.size, first.size) + assert desc.size == grown.size == sum(desc.sizes) + assert all(isinstance(handle, bytes) and len(handle) == 64 for handle in desc.handles) + assert len(desc.handles) == 2 + assert desc.handles[0] != desc.handles[1] + assert ( + desc == grown.ipc_descriptor + ) # the driver caches the packet per allocation; the export zeroes its padding + assert pickle.loads(pickle.dumps(desc)) == desc # noqa: S301 our own plain data + prefix = VirtualMemoryFabricBufferDescriptor.from_handles(desc.handles, desc.sizes, size=1) + with Buffer.from_ipc_descriptor(mr, prefix, stream=device.default_stream) as alias: + assert alias.size == 1 < sum(prefix.sizes) + assert not any(hasattr(desc, name) for name in ("fds", "close", "is_closed")) + with mr.allocate(0) as empty: + desc = empty.ipc_descriptor + assert (desc.handles, desc.sizes, desc.size) == ((), (), 0) + with Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) as alias: + assert alias.is_mapped and alias.size == 0 + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_descriptor_pickle_and_validation(): + descriptor = VirtualMemoryFabricBufferDescriptor + with pytest.raises(RuntimeError, match="cannot be instantiated directly"): + descriptor() + desc = descriptor.from_handles([bytes(64)], [4096], size=2048, owner_uuid="owner") + rebuilt = pickle.loads(pickle.dumps(desc)) # noqa: S301 our own plain data + assert rebuilt == desc and rebuilt is not desc + assert hash(rebuilt) == hash(desc) + assert rebuilt.owner_uuid == "owner" + assert desc == descriptor.from_handles(desc.handles, desc.sizes, size=desc.size, owner_uuid="other") + assert desc != descriptor.from_handles(desc.handles, desc.sizes) + assert desc != descriptor.from_handles([b"x" * 64], desc.sizes, size=desc.size) + assert desc != descriptor.from_handles(desc.handles, [8192], size=desc.size) + assert desc != object() + assert descriptor.from_handles(desc.handles, desc.sizes).size == 4096 + for handles, sizes, size in ( + ([bytes(63)], [4096], None), + ([bytes(65)], [4096], None), + ([bytearray(64)], [4096], None), + ([bytes(64)], [], None), + ([bytes(64)], [0], None), + ([bytes(64)], [-1], None), + ([bytes(64)], [1.5], None), + ([bytes(64)], [4096], 4097), + ([bytes(64)], [4096], -1), + ([bytes(64)], [4096], 1.5), + ): + with pytest.raises(ValueError): + descriptor.from_handles(handles, sizes, size=size) + + +@pytest.mark.parametrize("mode", ["queue", "pipe", "argument", "pool"]) +@pytest.mark.parametrize("send_buffer", [False, True], ids=["descriptor", "buffer"]) +@pytest.mark.thread_unsafe(reason="counts process-global file descriptors") +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_cross_process(vmm_fabric_device, mode, send_buffer): + warmup_multiprocessing() + device = vmm_fabric_device + with _resource(device).allocate(1) as buf: + _fill(buf, 0x51) + gc.collect() + before = len(os.listdir("/proc/self/fd")) + payload = buf if send_buffer else buf.ipc_descriptor + assert len(os.listdir("/proc/self/fd")) == before + assert _round_trip(mode, device.uuid, payload) == b"\x51" * 16 + assert _read(buf) == b"\x62" * 16 + gc.collect() + assert len(os.listdir("/proc/self/fd")) == before + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_second_gpu_and_wrong_owner(vmm_fabric_device_x2): + owner, peer = vmm_fabric_device_x2 + if not peer.can_access_peer(owner): + pytest.skip("Second device cannot access the owner's memory") + mr = _resource(owner, peers=[peer.device_id]) + with mr.allocate(1) as buf: + _fill(buf, 0x53) + desc = buf.ipc_descriptor + with pytest.raises(ValueError) as info: + Buffer.from_ipc_descriptor(_resource(peer), desc, stream=peer.default_stream) + assert f"memory belongs to the device with UUID {owner.uuid}" in str(info.value) + assert f"which is device {owner.device_id} in this process" in str(info.value) + hidden_owner = VirtualMemoryFabricBufferDescriptor.from_handles( + desc.handles, desc.sizes, size=desc.size, owner_uuid="missing-owner" + ) + with pytest.raises(ValueError, match="UUID missing-owner, which is not visible in this process"): + Buffer.from_ipc_descriptor(_resource(peer), hidden_owner, stream=peer.default_stream) + with Buffer.from_ipc_descriptor(mr, desc, stream=owner.default_stream) as alias: + # set_current replaces the top of the context stack, so switch back the same way. + peer.set_current() + try: + assert _read(alias) == b"\x53" * 16 + _fill(alias, 0x54) + finally: + owner.set_current() + assert _read(buf) == b"\x54" * 16 + buf.close() + assert _read(alias) == b"\x54" * 16 # the imported allocation owns a reference + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_handle_type_mismatch(vmm_fabric_device): + device = vmm_fabric_device + fabric, fd = _resource(device), _resource(device, handle_type="posix_fd") + # Empty descriptors would need no import call; even these must reject a mismatch. + for exporting, importing in ((fabric, fd), (fd, fabric)): + with exporting.allocate(0) as buf, pytest.raises(ValueError, match="configured like the exporter's"): + Buffer.from_ipc_descriptor(importing, buf.ipc_descriptor, stream=device.default_stream) + + +# Fixed patterns: a random first byte of 0x01 would read as the driver's own handle type and take another error path. +@pytest.mark.parametrize("handle", [bytes(64), b"\xff" * 64], ids=["zero", "ones"]) +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_garbage(vmm_fabric_device, handle): + device = vmm_fabric_device + mr = _resource(device) + with mr.allocate(1) as buf: + desc = VirtualMemoryFabricBufferDescriptor.from_handles([handle], [buf.size]) + with pytest.raises(CUDAError) as info: + Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) + _check_import_error(info.value) + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_stale_descriptor(vmm_fabric_device): + device = vmm_fabric_device + mr = _resource(device) + with mr.allocate(1) as buf: + desc = buf.ipc_descriptor + with pytest.raises(CUDAError) as info: + Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) + _check_import_error(info.value) + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_reexport_imported_derivatives(vmm_fabric_device): + device = vmm_fabric_device + mr = _resource(device) + with ExitStack() as stack: + buf = stack.enter_context(mr.allocate(1)) + alias = stack.enter_context(Buffer.from_ipc_descriptor(mr, buf.ipc_descriptor, stream=device.default_stream)) + same = stack.enter_context(mr.modify_allocation(alias, alias.size)) + smaller_request = stack.enter_context(mr.modify_allocation(alias, 1)) + grown = stack.enter_context(mr.modify_allocation(alias, 2 * alias.size)) + for imported in (alias, same, smaller_request, grown): + with pytest.raises(RuntimeError, match="imported from another process"): + _ = imported.ipc_descriptor + with pytest.raises(RuntimeError, match="imported from another process"): + pickle.dumps(imported) + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_configuration_and_buffer_pickle(vmm_fabric_device): + device = vmm_fabric_device + mr = _resource(device) + assert mr.is_ipc_enabled + for location in ("host", "host_numa", "host_numa_current"): + host = _resource(device, location_type=location, handle_type=None if location == "host" else "fabric") + assert not host.is_ipc_enabled + with mr.allocate(1) as buf: + _fill(buf, 0x55) + with pickle.loads(pickle.dumps(buf)) as alias: # noqa: S301 our own live buffer + assert alias.is_mapped + assert _read(alias) == b"\x55" * 16 + _fill(alias, 0x56) + assert _read(buf) == b"\x56" * 16 + + +def _child_stale_queue(queue, ready, outgoing): + assert ready.wait(CHILD_TIMEOUT_SEC), "parent did not release the export" + try: + buf = queue.get(timeout=CHILD_TIMEOUT_SEC) + except CUDAError as exc: + _check_import_error(exc) + outgoing.send("stale") + else: + buf.close() + raise AssertionError("queued fabric buffer pinned the allocation") + assert queue.get(timeout=CHILD_TIMEOUT_SEC) == "next item" + outgoing.close() + + +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_queued_buffer_does_not_pin(vmm_fabric_device): + device = vmm_fabric_device + mr = _resource(device) + with ExitStack() as stack: + buf = stack.enter_context(mr.allocate(1)) + unaffected = stack.enter_context(mr.allocate(1)) + _fill(unaffected, 0x57) + queue = mp.Queue() + stack.callback(queue.join_thread) + stack.callback(queue.close) + ready = mp.Event() + receive, send = mp.Pipe(duplex=False) + stack.enter_context(receive) + stack.enter_context(send) + process = stack.enter_context(closing(mp.Process(target=_child_stale_queue, args=(queue, ready, send)))) + try: + process.start() + send.close() + queue.put(buf) + queue.put("next item") + # Finish pickling before close, but hold the child before unpickling. + queue.close() + queue.join_thread() + buf.close() + ready.set() + assert receive.poll(CHILD_TIMEOUT_SEC), "child did not report the stale import" + assert receive.recv() == "stale" + process.join(timeout=CHILD_TIMEOUT_SEC) + finally: + survivors = kill_subprocesses(process) + assert not survivors, "child did not exit within timeout" + assert process.exitcode == 0 + assert _read(unaffected) == b"\x57" * 16 + with mr.allocate(1) as fresh: + _fill(fresh, 0x58) + assert _read(fresh) == b"\x58" * 16 + + +@pytest.mark.skipif(IS_WINDOWS or IS_WSL, reason="Fabric IPC requires Linux") +@pytest.mark.agent_authored(model="gpt-6.1-sol") +def test_unsupported_fabric_construction(init_cuda): + unsupported = [ + device + for device in Device.get_all_devices() + if device.properties.virtual_memory_management_supported and not device.properties.handle_type_fabric_supported + ] + if not unsupported: + pytest.skip("No visible VMM device without fabric handle support") + for device in unsupported: + with pytest.raises(RuntimeError) as info: + _resource(device) + message = str(info.value) + assert "NVLink fabric" in message + assert "accessible IMEX channel (/dev/nvidia-caps-imex-channels/channel*)" in message + assert f"device {device.device_id} reports no fabric handle support" in message diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index 728040fb147..16b85242fec 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -178,7 +178,6 @@ def _on_queue_feeder_error(e, _obj): [ pytest.param({"handle_type": "posix_fd"}, True, id="posix_fd"), pytest.param({"handle_type": None}, False, id="none"), - pytest.param({"handle_type": "fabric"}, False, id="fabric"), pytest.param({"handle_type": "win32_kmt"}, False, id="win32_kmt"), pytest.param({"location_type": "host", "handle_type": None}, False, id="host"), pytest.param( @@ -187,7 +186,7 @@ def _on_queue_feeder_error(e, _obj): ], ) def test_is_ipc_enabled(vmm_ipc_device, options, expected): - """Only POSIX file descriptors on device memory have a transport; host-located memory is not shared yet.""" + """POSIX file descriptors on device memory have a transport; host-located memory is not shared yet.""" mr = VirtualMemoryResource(vmm_ipc_device, config=VirtualMemoryResourceOptions(**options)) assert mr.is_ipc_enabled is expected