From 2231df7df59ada757c78a3450b33f5f252aff6a3 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 5 Oct 2026 17:03:51 -0700 Subject: [PATCH 01/11] cuda.core: share VirtualMemoryResource buffers between processes Buffers of a VirtualMemoryResource whose handle_type is "posix_fd" can now be exported and imported like pool-backed buffers. Buffer.ipc_descriptor exports every physical allocation that backs a VirtualMemoryBuffer as a file descriptor, Buffer.from_ipc_descriptor imports the descriptor with a VirtualMemoryResource of the receiving process, and a buffer or its descriptor can be sent through multiprocessing directly. A grown buffer imports as one contiguous range with the same byte layout. The import reuses MemAllocationHandle: import_mem_allocation_handle wraps cuMemImportFromShareableHandle in the same box and cuMemRelease deleter as create_mem_allocation_handle, with the importer's access descriptors. VirtualMemoryIPCBufferDescriptor carries one IPCAllocationHandle and one size per allocation and owns the file descriptors; multiprocessing duplicates them into the receiving process. VirtualMemoryResource pickles as (device, options), which lets a buffer pickle as (resource, descriptor). Fabric handles are not shared yet: no available test system reports fabric support or exposes an IMEX channel, so that path could not be verified. Part of #2980. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 29 ++ cuda_core/cuda/core/_cpp/rt/api.hpp | 12 + cuda_core/cuda/core/_cpp/rt/driver_api.hpp | 1 + .../cuda/core/_cpp/rt/virtual_memory.cpp | 39 ++- cuda_core/cuda/core/_memory/_buffer.pyi | 14 +- cuda_core/cuda/core/_memory/_buffer.pyx | 12 +- cuda_core/cuda/core/_memory/_ipc.pxd | 10 +- cuda_core/cuda/core/_memory/_ipc.pyi | 28 +- cuda_core/cuda/core/_memory/_ipc.pyx | 65 ++++- .../core/_memory/_virtual_memory_resource.pyi | 42 ++- .../core/_memory/_virtual_memory_resource.pyx | 188 +++++++++++- cuda_core/cuda/core/_rt.pxd | 3 + cuda_core/cuda/core/_rt.pyx | 3 + cuda_core/docs/source/release/1.3.0-notes.rst | 11 + cuda_core/tests/memory_ipc/test_vmm_ipc.py | 267 ++++++++++++++++++ cuda_core/tests/test_memory.py | 3 +- 16 files changed, 699 insertions(+), 28 deletions(-) create mode 100644 cuda_core/tests/memory_ipc/test_vmm_ipc.py diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index c2001b5da19..846af404ea2 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -199,6 +199,35 @@ what makes that safe: the allocation is released exactly once, when its last map its own `cuMemCreate` reference. It is not called for buffers from `allocate()`, whose ranges free themselves. A subclass override of `deallocate()` therefore does not run for them. +## Sharing across processes + +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; +`is_ipc_enabled` reports it). The exporter calls `cuMemExportToShareableHandle` +once per mapping that covers the buffer's size, in address order, and the +descriptor carries the handles with each allocation's size. Each file +descriptor is owned by an `IPCAllocationHandle` and closed when the descriptor +goes; `multiprocessing` duplicates it into the receiving process. The sizes +travel because the OS handle does not expose them to the importer; a wrong size +fails the importer's `cuMemMap`. + +The importer is a `VirtualMemoryResource` of the receiving process with the +same `handle_type`. `import_mem_allocation_handle` calls +`cuMemImportFromShareableHandle` and wraps the result in the same box and +deleter as `create_mem_allocation_handle`: the driver gives an imported handle +one reference and frees the memory when all references are released and no +mapping remains, so `cuMemRelease` is the right teardown for both. The access +descriptors stored in the box are the importer's, built from its options for +its device; the exporter's access does not travel. The import then reserves +`sum(sizes)` with the importer's alignment, maps each allocation in order, and +builds a range and a `VirtualMemoryBuffer` exactly as `allocate()` does, so a +grown buffer imports as one contiguous range with the same byte layout, and the +imported buffer frees itself through the same deleters. A descriptor imported +in the exporting process yields an alias of the exporter's memory at a new +address. A `VirtualMemoryResource` pickles as (device, options), which is what +lets a `Buffer` pickle as (resource, descriptor). + ## Why `modify_allocation` returns a new buffer The input buffer stays open and aliases the result. The chunks the input already mapped are diff --git a/cuda_core/cuda/core/_cpp/rt/api.hpp b/cuda_core/cuda/core/_cpp/rt/api.hpp index d77e57c86e5..fb9ffcd85f3 100644 --- a/cuda_core/cuda/core/_cpp/rt/api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/api.hpp @@ -262,6 +262,18 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat // Size of the allocation; the only size cuMemMap accepts for it. size_t mem_allocation_size(const MemAllocationHandle& h) noexcept; +// Import a physical allocation that another process exported with +// cuMemExportToShareableHandle. `os_handle` is the shareable handle in the +// form cuMemImportFromShareableHandle expects for `handle_type` (a POSIX file +// descriptor is its value cast to void*). `size` is the allocation's size, +// which the OS handle does not expose to the importer; the caller takes it +// from the exporter, and a wrong size fails the later cuMemMap. The result is +// a MemAllocationHandle like one from create_mem_allocation_handle: the +// access descriptors are applied to every mapping, and the last reference +// calls cuMemRelease. Returns empty handle on error (caller must check). +MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type, + size_t size, const CUmemAccessDesc* descs, size_t count); + // Reserve an address range via cuMemAddressReserve. Pass alignment 0 for the // driver default. When the last reference is released, cuMemAddressFree is // called with the exact reserved pair. Returns empty handle on error. diff --git a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp index 9efc1186c2c..9237e0ae723 100644 --- a/cuda_core/cuda/core/_cpp/rt/driver_api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/driver_api.hpp @@ -135,6 +135,7 @@ namespace cuda_core::rt { X(cuMemMap, 10020) \ X(cuMemUnmap, 10020) \ X(cuMemSetAccess, 10020) \ + X(cuMemImportFromShareableHandle, 10020) \ /* cuda-bindings requests 7000 (PTDS) or 2000 (legacy) */ \ X(cuStreamSynchronize, 7000) \ X(cuStreamIsCapturing, 10000) \ diff --git a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp index 8a946180794..505111a9743 100644 --- a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp @@ -52,6 +52,22 @@ const Box* box_of(const Handle& h) noexcept { reinterpret_cast(h.get()) - offsetof(Box, resource)); } +// Wrap an allocation handle this process holds one reference to, whether +// cuMemCreate or cuMemImportFromShareableHandle produced it. The last +// reference releases it; the memory is freed once no mapping remains. +MemAllocationHandle wrap_mem_allocation(CUmemGenericAllocationHandle handle, size_t size, + std::vector&& access) { + auto box = std::shared_ptr( + new MemAllocationBox{{handle}, size, std::move(access)}, + [](const MemAllocationBox* b) { + GILReleaseGuard gil; + pw_cuMemRelease(b->resource.raw); + delete b; + } + ); + return MemAllocationHandle(box, &box->resource); +} + } // namespace // ============================================================================ @@ -68,15 +84,20 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemCreate, &handle, size, &prop, 0))) { return {}; } - auto box = std::shared_ptr( - new MemAllocationBox{{handle}, size, std::move(access)}, - [](const MemAllocationBox* b) { - GILReleaseGuard gil; - pw_cuMemRelease(b->resource.raw); - delete b; - } - ); - return MemAllocationHandle(box, &box->resource); + return wrap_mem_allocation(handle, size, std::move(access)); +} + +MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type, + size_t size, const CUmemAccessDesc* descs, size_t count) { + // Copy the descriptors before the driver call so a failed copy leaves + // nothing to undo. + std::vector access(descs, descs + count); + GILReleaseGuard gil; + CUmemGenericAllocationHandle handle = 0; + if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemImportFromShareableHandle, &handle, os_handle, handle_type))) { + return {}; + } + return wrap_mem_allocation(handle, size, std::move(access)); } size_t mem_allocation_size(const MemAllocationHandle& h) noexcept { diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index c2a03808968..531900ee8d8 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -2,10 +2,9 @@ from typing import TypedDict +from _typeshed import Incomplete from cuda.core._memory._copy_enums import CopyOptions -from cuda.core._memory._device_memory_resource import DeviceMemoryResource from cuda.core._memory._ipc import IPCBufferDescriptor -from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource from cuda.core._stream import Stream from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph import GraphBuilder @@ -82,13 +81,18 @@ class Buffer: :class:`Buffer` is closed or garbage collected. """ @classmethod - def from_ipc_descriptor(cls, mr: DeviceMemoryResource | PinnedMemoryResource, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer: + def from_ipc_descriptor(cls, mr: Incomplete, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer: """Import a buffer that was exported from another process. Parameters ---------- - mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` - The IPC-enabled memory resource matching the exporting process. + mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` | :obj:`~_memory.VirtualMemoryResource` + The IPC-enabled memory resource matching the exporting process. A + descriptor exported from a :class:`VirtualMemoryBuffer` is imported + with a :class:`VirtualMemoryResource` of this process whose + ``handle_type`` matches the exporter's; the import maps the shared + physical memory for that resource's device with that resource's + access options. ipc_descriptor : :obj:`~_memory.IPCBufferDescriptor` The descriptor exported from another process. stream : :obj:`~_stream.Stream` diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index b1ee356c245..291f1b31a30 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -388,15 +388,21 @@ cdef class Buffer: @classmethod def from_ipc_descriptor( - cls, mr: DeviceMemoryResource | PinnedMemoryResource, ipc_descriptor: IPCBufferDescriptor, + cls, mr: DeviceMemoryResource | PinnedMemoryResource | VirtualMemoryResource, + ipc_descriptor: IPCBufferDescriptor, *, stream: Stream ) -> Buffer: """Import a buffer that was exported from another process. Parameters ---------- - mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` - The IPC-enabled memory resource matching the exporting process. + mr : :obj:`~_memory.DeviceMemoryResource` | :obj:`~_memory.PinnedMemoryResource` | :obj:`~_memory.VirtualMemoryResource` + The IPC-enabled memory resource matching the exporting process. A + descriptor exported from a :class:`VirtualMemoryBuffer` is imported + with a :class:`VirtualMemoryResource` of this process whose + ``handle_type`` matches the exporter's; the import maps the shared + physical memory for that resource's device with that resource's + access options. ipc_descriptor : :obj:`~_memory.IPCBufferDescriptor` The descriptor exported from another process. stream : :obj:`~_stream.Stream` diff --git a/cuda_core/cuda/core/_memory/_ipc.pxd b/cuda_core/cuda/core/_memory/_ipc.pxd index 5321d48fe13..7d2feb57271 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pxd +++ b/cuda_core/cuda/core/_memory/_ipc.pxd @@ -53,10 +53,18 @@ cdef class IPCAllocationHandle: cpdef close(self) +# Descriptor of a VirtualMemoryResource buffer: one exported handle per +# physical allocation that backs the buffer, in address order. +cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): + cdef readonly int _handle_type + cdef readonly tuple _chunk_sizes + cdef readonly tuple _handles + + # Buffer IPC Implementation # ------------------------- cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer) -cdef Buffer Buffer_from_ipc_descriptor(cls, _MemPool, IPCBufferDescriptor, stream) +cdef Buffer Buffer_from_ipc_descriptor(cls, object, IPCBufferDescriptor, stream) # _MemPool IPC Implementation diff --git a/cuda_core/cuda/core/_memory/_ipc.pyi b/cuda_core/cuda/core/_memory/_ipc.pyi index ef8f4110f3f..f34075a293f 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyi +++ b/cuda_core/cuda/core/_memory/_ipc.pyi @@ -1,6 +1,7 @@ # This file was generated by stubgen-pyx v0.2.22 from cuda_core/cuda/core/_memory/_ipc.pyx import uuid +from typing import Any __all__ = [] @@ -39,7 +40,8 @@ class IPCBufferDescriptor: def size(self) -> int: ... class IPCAllocationHandle: - """Shareable handle to an IPC-enabled device memory pool.""" + """Shareable OS handle to an IPC-enabled memory pool or to a physical allocation + of a :class:`VirtualMemoryResource` buffer.""" def __init__(self, *arg, **kwargs) -> None: ... @classmethod def _init(cls, handle: int, uuid: uuid.UUID | None) -> IPCAllocationHandle: ... @@ -54,5 +56,29 @@ class IPCAllocationHandle: @property def uuid(self) -> uuid.UUID: ... +class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): + """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. + + The descriptor holds one exported handle per physical allocation that + backs the buffer, in address order, with each allocation's size. With + POSIX file descriptors the handles are :class:`IPCAllocationHandle` + objects: the descriptor owns them and closes them when it is released, and + ``multiprocessing`` duplicates them into the receiving process. Plain + ``pickle`` cannot carry file descriptors. + + Note + ---- + The sizes are controlled by the exporting peer. Receivers must treat them + as untrusted and import only through :meth:`Buffer.from_ipc_descriptor`, + which fails instead of mapping a size the driver rejects. + """ + _handle_type: int + _chunk_sizes: tuple[Any, ...] + _handles: tuple[Any, ...] + + @staticmethod + def _from_exports(handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple[Any, ...], size: int) -> VirtualMemoryIPCBufferDescriptor: ... + def __reduce__(self) -> tuple[object, ...]: ... + 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 7923c21fc00..ab1c18bd674 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -116,7 +116,8 @@ cdef inline int IPCAllocationHandle_check_open(IPCAllocationHandle self) except cdef class IPCAllocationHandle: - """Shareable handle to an IPC-enabled device memory pool.""" + """Shareable OS handle to an IPC-enabled memory pool or to a physical allocation + of a :class:`VirtualMemoryResource` buffer.""" def __init__(self, *arg, **kwargs) -> None: raise RuntimeError("IPCAllocationHandle objects cannot be instantiated directly. Please use MemoryResource APIs.") @@ -169,12 +170,52 @@ def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handle) +cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): + """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. + + The descriptor holds one exported handle per physical allocation that + backs the buffer, in address order, with each allocation's size. With + POSIX file descriptors the handles are :class:`IPCAllocationHandle` + objects: the descriptor owns them and closes them when it is released, and + ``multiprocessing`` duplicates them into the receiving process. Plain + ``pickle`` cannot carry file descriptors. + + Note + ---- + The sizes are controlled by the exporting peer. Receivers must treat them + as untrusted and import only through :meth:`Buffer.from_ipc_descriptor`, + which fails instead of mapping a size the driver rejects. + """ + + @staticmethod + def _from_exports( + handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple, size: int + ) -> VirtualMemoryIPCBufferDescriptor: + cdef VirtualMemoryIPCBufferDescriptor self = VirtualMemoryIPCBufferDescriptor.__new__( + VirtualMemoryIPCBufferDescriptor) + self._payload = b"" + self._size = size + self._handle_type = handle_type + self._chunk_sizes = tuple(chunk_sizes) + self._handles = tuple(handles) + return self + + def __reduce__(self) -> tuple[object, ...]: + return VirtualMemoryIPCBufferDescriptor._from_exports, ( + self._handle_type, self._chunk_sizes, self._handles, self._size) + + # Buffer IPC Implementation # ------------------------- cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self): Buffer_check_open(self) if not self.memory_resource.is_ipc_enabled: raise RuntimeError("Memory resource is not IPC-enabled") + if not isinstance(self.memory_resource, _MemPool): + # VirtualMemoryBuffer overrides ipc_descriptor; a plain Buffer wrapped + # with Buffer.from_handle(mr=) lands here. + raise TypeError( + f"a Buffer that {type(self.memory_resource).__name__} did not allocate cannot be shared") cdef cydriver.CUmemPoolPtrExportData data with nogil: HANDLE_RETURN( @@ -186,10 +227,24 @@ cdef IPCBufferDescriptor Buffer_get_ipc_descriptor(Buffer self): return IPCBufferDescriptor._init(data_b, self.size) cdef Buffer Buffer_from_ipc_descriptor( - cls, _MemPool mr, IPCBufferDescriptor ipc_descriptor, stream + cls, object mr, IPCBufferDescriptor ipc_descriptor, stream ): """Import a buffer that was exported from another process.""" - MP_check_open(mr) + if isinstance(ipc_descriptor, VirtualMemoryIPCBufferDescriptor): + # Imported here rather than cimported: _virtual_memory_resource + # cimports this module. + from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource + if not isinstance(mr, VirtualMemoryResource): + raise TypeError( + "this descriptor was exported from a VirtualMemoryResource buffer and must be " + f"imported with a VirtualMemoryResource, not {type(mr).__name__}") + return mr._import_ipc_buffer(ipc_descriptor, stream) + if not isinstance(mr, _MemPool): + raise TypeError( + "this descriptor was exported from a memory pool buffer and must be imported with a " + f"DeviceMemoryResource or PinnedMemoryResource, not {type(mr).__name__}") + cdef _MemPool pool = <_MemPool>mr + MP_check_open(pool) if not mr.is_ipc_enabled: raise RuntimeError("Memory resource is not IPC-enabled") cdef size_t payload_size = len(ipc_descriptor._payload) @@ -201,7 +256,7 @@ cdef Buffer Buffer_from_ipc_descriptor( ) cdef Stream s = Stream_accept(stream) cdef DevicePtrHandle h_ptr = deviceptr_import_ipc( - mr._h_pool, + pool._h_pool, ipc_descriptor.payload_ptr(), s._h_stream ) @@ -220,7 +275,7 @@ cdef Buffer Buffer_from_ipc_descriptor( f"IPC buffer descriptor size ({claimed_size}) exceeds " f"mapped allocation extent ({mapped_size} bytes)" ) - return Buffer_from_deviceptr_handle(h_ptr, claimed_size, mr, ipc_descriptor) + return Buffer_from_deviceptr_handle(h_ptr, claimed_size, pool, ipc_descriptor) # _MemPool IPC Implementation diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 2deaca761e0..2546e82baed 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -1,10 +1,13 @@ # This file was generated by stubgen-pyx v0.2.22 from cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +import platform from dataclasses import dataclass, field from typing import Iterable from cuda.core._device import Device from cuda.core._memory._buffer import Buffer, MemoryResource +from cuda.core._memory._ipc import (IPCBufferDescriptor, + VirtualMemoryIPCBufferDescriptor) from cuda.core._stream import Stream from cuda.core._utils.cuda_utils import driver from cuda.core.graph import GraphBuilder @@ -16,6 +19,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() @dataclass class VirtualMemoryResourceOptions: @@ -90,6 +94,18 @@ class VirtualMemoryBuffer(Buffer): was grown from: the two share their physical memory, and that memory is freed when the last buffer that maps it closes. """ + @property + def ipc_descriptor(self) -> IPCBufferDescriptor: + """Descriptor for sharing this buffer with other processes. + + The descriptor exports every physical allocation that backs the + buffer as a shareable handle and is cached on the buffer. It requires + a resource whose ``handle_type`` can be shared; see + :attr:`VirtualMemoryResource.is_ipc_enabled`. Import it with + :meth:`Buffer.from_ipc_descriptor` and a + :class:`VirtualMemoryResource` of the receiving process, or send the + buffer itself through ``multiprocessing``. + """ def close(self, stream: Stream | GraphBuilder | None=None) -> None: """Release this buffer's share of its address range. @@ -141,11 +157,20 @@ class VirtualMemoryResource(MemoryResource): garbage collector waits at that point instead. To control when the wait happens, close the buffer explicitly or record an idle stream with :meth:`Buffer.set_deallocation_stream`. + + Buffers can be shared with other processes when ``config.handle_type`` + is ``"posix_fd"`` (Linux). :attr:`Buffer.ipc_descriptor` exports the + physical allocations that back a buffer, :meth:`Buffer.from_ipc_descriptor` + imports them with a ``VirtualMemoryResource`` of the receiving process, + and a buffer sent through ``multiprocessing`` does both. The importing + resource maps the shared memory for its own device with its own access + options; see :attr:`is_ipc_enabled`. """ device: object config: object def __init__(self, device_id: Device | int, config: VirtualMemoryResourceOptions | None=None) -> None: ... + def __reduce__(self) -> tuple[object, ...]: ... def allocate(self, size: int, *, stream: Stream | GraphBuilder | None=None) -> VirtualMemoryBuffer: """ Allocate a buffer of the given size using CUDA virtual memory. @@ -261,7 +286,22 @@ class VirtualMemoryResource(MemoryResource): """ @property def is_ipc_enabled(self) -> bool: - """Return False. Buffers of this resource cannot be shared through IPC descriptors.""" + """Whether buffers of this resource can be shared with other processes. + + True when ``config.handle_type`` is ``"posix_fd"`` on Linux. 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. + """ + def _import_ipc_buffer(self, desc: VirtualMemoryIPCBufferDescriptor, 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 + type, then mapped in order into one new address reservation, with the + access descriptors this resource's options produce for its device. + The returned buffer owns the imported handles like a buffer from + :meth:`allocate` owns created ones, and reports ``is_mapped``. + """ @property def device_id(self) -> int: """ diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 66b0c4fc0db..100812ddc9f 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -4,6 +4,8 @@ from __future__ import annotations +cimport cython +from libc.stdint cimport uintptr_t from libc.string cimport memset from libcpp.vector cimport vector @@ -14,6 +16,12 @@ from cuda.core._memory._buffer cimport ( Buffer_from_deviceptr_handle, MemoryResource, ) +from cuda.core._memory._ipc cimport ( + IPCAllocationHandle, + IPCBufferDescriptor, + IPCDataForBuffer, + VirtualMemoryIPCBufferDescriptor, +) from cuda.core._rt cimport ( ContextHandle, DevicePtrHandle, @@ -32,6 +40,7 @@ from cuda.core._rt cimport ( deviceptr_create_vmm, get_last_error, get_primary_context, + import_mem_allocation_handle, mem_allocation_size, set_deallocation_stream, va_mapping_allocation, @@ -42,6 +51,8 @@ from cuda.core._rt cimport ( from cuda.core._stream cimport Stream, Stream_accept, Stream_handle_is_default_token, Stream_is_default_token from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, check_or_create_options +import os +import platform from dataclasses import dataclass, field from typing import TYPE_CHECKING, Iterable @@ -72,6 +83,11 @@ _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() + @dataclass class VirtualMemoryResourceOptions: @@ -225,6 +241,24 @@ cdef class VirtualMemoryBuffer(Buffer): freed when the last buffer that maps it closes. """ + @property + @cython.critical_section + def ipc_descriptor(self) -> IPCBufferDescriptor: + """Descriptor for sharing this buffer with other processes. + + The descriptor exports every physical allocation that backs the + buffer as a shareable handle and is cached on the buffer. It requires + a resource whose ``handle_type`` can be shared; see + :attr:`VirtualMemoryResource.is_ipc_enabled`. Import it with + :meth:`Buffer.from_ipc_descriptor` and a + :class:`VirtualMemoryResource` of the receiving process, or send the + buffer itself through ``multiprocessing``. + """ + Buffer_check_open(self) + if self._ipc_data is None: + self._ipc_data = IPCDataForBuffer(_export_ipc_descriptor(self), False) + return self._ipc_data.ipc_descriptor + def close(self, stream: Stream | GraphBuilder | None = None) -> None: """Release this buffer's share of its address range. @@ -295,6 +329,14 @@ cdef class VirtualMemoryResource(MemoryResource): garbage collector waits at that point instead. To control when the wait happens, close the buffer explicitly or record an idle stream with :meth:`Buffer.set_deallocation_stream`. + + Buffers can be shared with other processes when ``config.handle_type`` + is ``"posix_fd"`` (Linux). :attr:`Buffer.ipc_descriptor` exports the + physical allocations that back a buffer, :meth:`Buffer.from_ipc_descriptor` + imports them with a ``VirtualMemoryResource`` of the receiving process, + and a buffer sent through ``multiprocessing`` does both. The importing + resource maps the shared memory for its own device with its own access + options; see :attr:`is_ipc_enabled`. """ cdef: @@ -313,6 +355,12 @@ cdef class VirtualMemoryResource(MemoryResource): raise RuntimeError("VirtualMemoryResource requires CUDA VMM API support") self._check_config(self.config) + def __reduce__(self) -> tuple[object, ...]: + # 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 VirtualMemoryResource, (self.device.device_id if self.device is not None else 0, self.config) + cdef int _check_config(self, object cfg) except -1: """Reject options the driver would reject, before any driver call. @@ -729,8 +777,106 @@ cdef class VirtualMemoryResource(MemoryResource): @property def is_ipc_enabled(self) -> bool: - """Return False. Buffers of this resource cannot be shared through IPC descriptors.""" - return False + """Whether buffers of this resource can be shared with other processes. + + True when ``config.handle_type`` is ``"posix_fd"`` on Linux. 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. + """ + return self.config.handle_type in _IPC_HANDLE_TYPES + + def _import_ipc_buffer(self, VirtualMemoryIPCBufferDescriptor 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 + type, then mapped in order into one new address reservation, with the + access descriptors this resource's options produce for its device. + The returned buffer owns the imported handles like a buffer from + :meth:`allocate` owns created ones, and reports ``is_mapped``. + """ + cdef Stream s = None + cdef object cfg = self.config + cdef cydriver.CUmemAllocationProp prop + cdef vector[cydriver.CUmemAccessDesc] descs + cdef vector[MemAllocationHandle] allocs + cdef vector[VaMappingHandle] mappings + cdef MemAllocationHandle h_alloc + cdef VaReservationHandle h_res + cdef VaMappingHandle h_map + cdef VmmRangeHandle rng + cdef DevicePtrHandle h_ptr + cdef size_t gran, addr_align, chunk, i, n, total = 0, offset = 0 + cdef cydriver.CUdeviceptr hint + cdef int handle_type, fd + cdef void* os_handle + + if not self.is_ipc_enabled: + raise RuntimeError("Memory resource is not IPC-enabled") + handle_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(cfg.handle_type)) + if desc._handle_type != handle_type: + raise ValueError( + "the descriptor's handle type does not match this resource's " + 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): + raise ValueError("malformed VirtualMemoryResource buffer descriptor") + if stream is not None: + s = Stream_accept(stream) + if n == 0: + # An exported empty buffer: an empty range, as allocate(0) returns. + h_ptr = deviceptr_create_vmm(0, create_vmm_range(mappings)) + if s is not None and not Stream_is_default_token(s): + HANDLE_RETURN(set_deallocation_stream(h_ptr, s._h_stream)) + return Buffer_from_deviceptr_handle(h_ptr, 0, self, desc, VirtualMemoryBuffer) + + self._fill_prop(cfg, &prop) + self._fill_access(cfg, &prop, descs) + gran = self._granularity(cfg, &prop) + for chunk_obj in desc._chunk_sizes: + chunk = chunk_obj + if chunk == 0 or chunk % gran != 0: + raise ValueError( + f"descriptor chunk size {chunk} is not a positive multiple of the {gran}-byte granularity" + ) + if chunk > -1 - total: + raise OverflowError("the descriptor's chunk sizes do not fit in size_t") + total += chunk + if desc._size > total: + raise ValueError(f"descriptor size ({desc._size}) exceeds the exported allocations ({total} bytes)") + addr_align = cfg.addr_align or gran + hint = cfg.addr_hint or 0 + + # Every handle below is a local: if a later step fails, the locals die + # in reverse order and release everything imported so far. + for i in range(n): + fd = int(desc._handles[i]) # raises if the handle was closed + chunk = desc._chunk_sizes[i] + os_handle = fd + with nogil: + h_alloc = import_mem_allocation_handle( + os_handle, handle_type, chunk, descs.data(), descs.size()) + if not h_alloc: + _raise_last_error() + allocs.push_back(h_alloc) + with nogil: + h_res = create_va_reservation_handle(total, addr_align, hint) + if not h_res: + _raise_last_error() + for i in range(n): + with nogil: + h_map = create_va_mapping_handle(as_cu(h_res) + offset, allocs[i], h_res) + if not h_map: + _raise_last_error() + mappings.push_back(h_map) + offset += mem_allocation_size(allocs[i]) + rng = create_vmm_range(mappings) + h_ptr = deviceptr_create_vmm(as_cu(h_res), rng) + if not h_ptr: + _raise_last_error() + self._record_deallocation_stream(h_ptr, s) + return Buffer_from_deviceptr_handle(h_ptr, desc._size, self, desc, VirtualMemoryBuffer) @property def device_id(self) -> int: @@ -750,3 +896,41 @@ cdef class VirtualMemoryResource(MemoryResource): str: A string describing the object """ return f"" + + +cdef VirtualMemoryIPCBufferDescriptor _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 + is a prefix of its range. Each file descriptor is owned by an + :class:`IPCAllocationHandle` in the descriptor and closed with it. + """ + cdef VirtualMemoryResource mr = buf._memory_resource + cdef vector[VaMappingHandle] mappings + cdef MemAllocationHandle h_alloc + cdef size_t i, chunk, offset = 0 + cdef int handle_type + cdef int fd = -1 + cdef list sizes = [] + cdef list handles = [] + + if not mr.is_ipc_enabled: + raise RuntimeError("Memory resource is not IPC-enabled") + handle_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(mr.config.handle_type)) + mappings = vmm_range_mappings(vmm_range(buf._h_ptr)) + for i in range(mappings.size()): + if offset >= buf._size: + break + h_alloc = va_mapping_allocation(mappings[i]) + chunk = mem_allocation_size(h_alloc) + with nogil: + HANDLE_RETURN(cydriver.cuMemExportToShareableHandle( + &fd, as_cu(h_alloc), handle_type, 0)) + try: + handles.append(IPCAllocationHandle._init(fd, None)) + except: # noqa: E722 rollback-then-raise: the fd is not owned yet + os.close(fd) + raise + sizes.append(chunk) + offset += chunk + return VirtualMemoryIPCBufferDescriptor._from_exports(handle_type, tuple(sizes), tuple(handles), buf._size) diff --git a/cuda_core/cuda/core/_rt.pxd b/cuda_core/cuda/core/_rt.pxd index 30709187f38..c99e6c3ce18 100644 --- a/cuda_core/cuda/core/_rt.pxd +++ b/cuda_core/cuda/core/_rt.pxd @@ -299,6 +299,9 @@ cdef MemAllocationHandle create_mem_allocation_handle( size_t size, const cydriver.CUmemAllocationProp& prop, const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil cdef size_t mem_allocation_size(const MemAllocationHandle& h) noexcept nogil +cdef MemAllocationHandle import_mem_allocation_handle( + void* os_handle, cydriver.CUmemAllocationHandleType handle_type, size_t size, + const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil cdef VaReservationHandle create_va_reservation_handle( size_t size, size_t alignment, cydriver.CUdeviceptr hint) except+ nogil cdef size_t va_reservation_size(const VaReservationHandle& h) noexcept nogil diff --git a/cuda_core/cuda/core/_rt.pyx b/cuda_core/cuda/core/_rt.pyx index 8eef0d5877c..b948a6bdafa 100644 --- a/cuda_core/cuda/core/_rt.pyx +++ b/cuda_core/cuda/core/_rt.pyx @@ -180,6 +180,9 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil size_t mem_allocation_size "cuda_core::rt::mem_allocation_size" ( const MemAllocationHandle& h) noexcept nogil + MemAllocationHandle import_mem_allocation_handle "cuda_core::rt::import_mem_allocation_handle" ( + void* os_handle, cydriver.CUmemAllocationHandleType handle_type, size_t size, + const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil VaReservationHandle create_va_reservation_handle "cuda_core::rt::create_va_reservation_handle" ( size_t size, size_t alignment, cydriver.CUdeviceptr hint) except+ nogil size_t va_reservation_size "cuda_core::rt::va_reservation_size" ( 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 81321e71006..8e457363087 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -78,6 +78,17 @@ New features disturb a capture in its context is reported as a :class:`CUDAWarning` 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, and a + buffer or its descriptor can be sent through ``multiprocessing`` directly. + :attr:`VirtualMemoryResource.is_ipc_enabled` is ``True`` when the resource's + ``handle_type`` is ``"posix_fd"`` on Linux. A grown buffer imports as one + contiguous range. Fabric handles are not shared yet. + (`#2980 `__) + Fixes and enhancements ---------------------- diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py new file mode 100644 index 00000000000..a8b102e2db1 --- /dev/null +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Sharing VirtualMemoryResource buffers between processes (issue #2980).""" + +import ctypes +import gc +import multiprocessing as mp +import os +import pickle + +import pytest +from helpers import IS_WINDOWS, IS_WSL +from helpers.child_processes import child_timeout_sec, kill_subprocesses, track_child_processes + +from cuda.bindings import driver +from cuda.core import Buffer, Device, MemoryResource, VirtualMemoryResource, VirtualMemoryResourceOptions +from cuda.core._memory._ipc import IPCBufferDescriptor +from cuda.core._memory._virtual_memory_resource import VirtualMemoryBuffer +from cuda.core._utils.cuda_utils import handle_return + +CHILD_TIMEOUT_SEC = child_timeout_sec() + +# These tests spawn processes and pass file descriptors, which fails for very many threads. +pytestmark = pytest.mark.parallel_threads_limit(4) + + +@pytest.fixture +def vmm_ipc_device(init_cuda): + """A device that can export VMM allocations as POSIX file descriptors, or skip.""" + device = init_cuda + props = device.properties + if not props.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + if IS_WINDOWS or IS_WSL or not props.handle_type_posix_file_descriptor_supported: + pytest.skip("Sharing virtual memory needs POSIX file descriptor handles") + with track_child_processes(): + yield device + + +def _resource(device, **options): + options.setdefault("handle_type", "posix_fd") + return VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(**options)) + + +def _fill(buf, value, *, offset=0, size=None): + """Write ``value`` to ``size`` bytes of ``buf`` at ``offset`` and wait for the write.""" + size = buf.size - offset if size is None else size + handle_return(driver.cuMemsetD8(int(buf.handle) + offset, value, size)) + handle_return(driver.cuCtxSynchronize()) + + +def _read(buf, offset, size): + """Copy ``size`` bytes of ``buf`` at ``offset`` to the host.""" + host = (ctypes.c_ubyte * size)() + handle_return(driver.cuMemcpyDtoH(ctypes.addressof(host), int(buf.handle) + offset, size)) + return bytes(host) + + +def _is_mapped(ptr): + """Whether the driver still has a mapping at ``ptr``.""" + status, handle = driver.cuMemRetainAllocationHandle(ptr) + if status == driver.CUresult.CUDA_SUCCESS: + handle_return(driver.cuMemRelease(handle)) + return True + return False + + +class _SharedStub(MemoryResource): + """Reports IPC support without being a pool, to reach the base-class export path.""" + + def __init__(self, device): + self._device_id = device.device_id + + def allocate(self, size, *, stream=None): + raise NotImplementedError + + def deallocate(self, ptr, size, *, stream=None): + pass + + @property + def is_device_accessible(self): + return True + + @property + def is_host_accessible(self): + return False + + @property + def device_id(self): + return self._device_id + + @property + def is_ipc_enabled(self): + return True + + +def _run_child(target, *args): + """Run ``target(*args)`` in a spawned child and fail if it does not exit cleanly.""" + process = mp.Process(target=target, args=args) + process.start() + 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}" + + +class TestVmmIpcDescriptor: + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_descriptor_round_trip(self, vmm_ipc_device): + """A child imports the descriptor, reads the parent's data, writes its own, and releases.""" + device = vmm_ipc_device + mr = _resource(device) + assert mr.is_ipc_enabled + buf = mr.allocate(1) + _fill(buf, 0xAB) + desc = buf.ipc_descriptor + assert isinstance(desc, IPCBufferDescriptor) + assert desc.size == buf.size + assert buf.ipc_descriptor is desc # cached + assert not buf.is_mapped + + queue = mp.Queue() + _run_child(self.child_main, device.device_id, desc, queue) + reads = queue.get(timeout=CHILD_TIMEOUT_SEC) + assert reads["is_mapped"] is True + assert reads["size"] == buf.size + assert reads["first_bytes"] == b"\xab" * 16 + assert reads["released"] is True + assert _read(buf, 0, 16) == b"\x5a" * 16 + buf.close() + + @staticmethod + def child_main(device_id, desc, queue): + device = Device(device_id) + device.set_current() + mr = _resource(device) + imported = Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) + assert isinstance(imported, VirtualMemoryBuffer) + result = { + "is_mapped": imported.is_mapped, + "size": imported.size, + "first_bytes": _read(imported, 0, 16), + } + _fill(imported, 0x5A) + ptr = int(imported.handle) + imported.close() + result["released"] = not _is_mapped(ptr) + queue.put(result) + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_buffer_pickles_through_queue(self, vmm_ipc_device): + """A buffer sent directly pickles as (resource, descriptor) and imports on arrival.""" + device = vmm_ipc_device + mr = _resource(device) + buf = mr.allocate(1) + _fill(buf, 0x11) + queue = mp.Queue() + queue.put(buf) + _run_child(self.child_buffer, device.device_id, queue) + assert _read(buf, 0, 16) == b"\x22" * 16 + buf.close() + + @staticmethod + def child_buffer(device_id, queue): + Device(device_id).set_current() + imported = queue.get(timeout=CHILD_TIMEOUT_SEC) + assert isinstance(imported, VirtualMemoryBuffer) + assert imported.is_mapped + assert isinstance(imported.memory_resource, VirtualMemoryResource) + assert imported.memory_resource.device_id == device_id + assert _read(imported, 0, 16) == b"\x11" * 16 + _fill(imported, 0x22) + imported.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_grown_buffer_imports_as_one_range(self, vmm_ipc_device): + """A buffer with two physical allocations exports both and imports contiguously.""" + device = vmm_ipc_device + mr = _resource(device) + first = mr.allocate(1) + gran = first.size + grown = mr.modify_allocation(first, 2 * gran) + assert grown.size == 2 * gran + _fill(grown, 0x11, offset=0, size=gran) + _fill(grown, 0x22, offset=gran, size=gran) + desc = grown.ipc_descriptor + assert desc.size == 2 * gran + assert len(desc._chunk_sizes) == 2 + + queue = mp.Queue() + _run_child(self.child_grown, device.device_id, desc, gran, queue) + assert queue.get(timeout=CHILD_TIMEOUT_SEC) == {"low": b"\x11" * 16, "high": b"\x22" * 16} + first.close() + grown.close() + + @staticmethod + def child_grown(device_id, desc, gran, queue): + device = Device(device_id) + device.set_current() + imported = Buffer.from_ipc_descriptor(_resource(device), desc, stream=device.default_stream) + assert imported.size == 2 * gran + queue.put({"low": _read(imported, 0, 16), "high": _read(imported, gran, 16)}) + imported.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_import_in_same_process_aliases_memory(self, vmm_ipc_device): + """Importing a descriptor where it was exported maps the same memory at a new address.""" + device = vmm_ipc_device + mr = _resource(device) + buf = mr.allocate(1) + alias = Buffer.from_ipc_descriptor(mr, buf.ipc_descriptor, stream=device.default_stream) + assert alias.is_mapped + assert int(alias.handle) != int(buf.handle) + _fill(buf, 0x33) + assert _read(alias, 0, 16) == b"\x33" * 16 + _fill(alias, 0x44) + assert _read(buf, 0, 16) == b"\x44" * 16 + ptr_buf, ptr_alias = int(buf.handle), int(alias.handle) + buf.close() + assert _is_mapped(ptr_alias) # the import keeps the memory alive + alias.close() + assert not _is_mapped(ptr_alias) + assert not _is_mapped(ptr_buf) + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_descriptor_owns_its_file_descriptors(self, vmm_ipc_device): + """The exported file descriptors close when the descriptor is released.""" + device = vmm_ipc_device + mr = _resource(device) + buf = mr.allocate(1) + desc = buf.ipc_descriptor + fd = int(desc._handles[0]) + os.fstat(fd) # open + with pytest.raises(TypeError): + pickle.dumps(desc) # file descriptors need multiprocessing + buf.close() # drops the cached descriptor + del desc + gc.collect() + with pytest.raises(OSError): + os.fstat(fd) + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_errors(self, vmm_ipc_device): + device = vmm_ipc_device + stream = device.default_stream + shared = _resource(device) + private = _resource(device, handle_type=None) + assert not private.is_ipc_enabled + with private.allocate(1) as buf, pytest.raises(RuntimeError, match="not IPC-enabled"): + _ = buf.ipc_descriptor + with shared.allocate(1) as buf: + desc = buf.ipc_descriptor + with pytest.raises(RuntimeError, match="not IPC-enabled"): + Buffer.from_ipc_descriptor(private, desc, stream=stream) + with pytest.raises(TypeError, match="VirtualMemoryResource"): + Buffer.from_ipc_descriptor(device.memory_resource, desc, stream=stream) + pool_desc = IPCBufferDescriptor._init(b"\x00" * 64, 64) + with pytest.raises(TypeError, match="memory pool"): + Buffer.from_ipc_descriptor(shared, pool_desc, stream=stream) + # The base-class export path serves pool buffers only: a Buffer that a + # non-pool resource did not allocate is refused before any driver call. + with device.memory_resource.allocate(64, stream=stream) as backing: + wrapped = Buffer.from_handle(int(backing.handle), 64, mr=_SharedStub(device)) + with pytest.raises(TypeError, match="did not allocate"): + _ = wrapped.ipc_descriptor + wrapped.close() diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index f33e17ddccf..c705c859eae 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1390,7 +1390,8 @@ def test_vmm_allocator_basic_allocation(use_device_object, handle_type): assert buffer.size >= 4096 # May be aligned up assert buffer.device_id == device.device_id assert buffer.memory_resource == vmm_mr - assert vmm_mr.is_ipc_enabled is False + # POSIX file descriptors can be shared; Win32 KMT handles have no transport. + assert vmm_mr.is_ipc_enabled is (handle_type == "posix_fd") # Test deallocation buffer.close() From 682db0f4b1d4bdb3f0905ce9dff9f85481dc37fe Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 5 Oct 2026 17:14:26 -0700 Subject: [PATCH 02/11] cuda.core: compare the descriptor's chunk count with a signed length The CI wheel build compiles with -Werror and -Wsign-compare, and the VMM import path compared a size_t count with the Py_ssize_t that len() returns. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 100812ddc9f..0919372250f 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -806,7 +806,8 @@ cdef class VirtualMemoryResource(MemoryResource): cdef VaMappingHandle h_map cdef VmmRangeHandle rng cdef DevicePtrHandle h_ptr - cdef size_t gran, addr_align, chunk, i, n, total = 0, offset = 0 + cdef size_t gran, addr_align, chunk, total = 0, offset = 0 + cdef Py_ssize_t n, i cdef cydriver.CUdeviceptr hint cdef int handle_type, fd cdef void* os_handle @@ -866,11 +867,11 @@ cdef class VirtualMemoryResource(MemoryResource): _raise_last_error() for i in range(n): with nogil: - h_map = create_va_mapping_handle(as_cu(h_res) + offset, allocs[i], h_res) + h_map = create_va_mapping_handle(as_cu(h_res) + offset, allocs[i], h_res) if not h_map: _raise_last_error() mappings.push_back(h_map) - offset += mem_allocation_size(allocs[i]) + offset += mem_allocation_size(allocs[i]) rng = create_vmm_range(mappings) h_ptr = deviceptr_create_vmm(as_cu(h_res), rng) if not h_ptr: From 4ed2c03aba693abe1bacd06063f7ec9511ace01f Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Mon, 5 Oct 2026 17:42:27 -0700 Subject: [PATCH 03/11] cuda.core: mark two VMM IPC tests thread-unsafe Under pytest-run-parallel the same test body runs in several threads. The alias test checks mapping state by address, and a freed address can be reused by another thread's allocation; the descriptor test checks a closed file descriptor number, which another thread can reuse. Co-Authored-By: Claude Fable 5.1 --- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index a8b102e2db1..1061c2198ad 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -204,6 +204,7 @@ def child_grown(device_id, desc, gran, queue): imported.close() @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="checks mapping state by address; another thread can reuse a freed address") def test_import_in_same_process_aliases_memory(self, vmm_ipc_device): """Importing a descriptor where it was exported maps the same memory at a new address.""" device = vmm_ipc_device @@ -224,6 +225,9 @@ def test_import_in_same_process_aliases_memory(self, vmm_ipc_device): assert not _is_mapped(ptr_buf) @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe( + reason="file descriptor numbers are process-global; another thread can reuse a closed one" + ) def test_descriptor_owns_its_file_descriptors(self, vmm_ipc_device): """The exported file descriptors close when the descriptor is released.""" device = vmm_ipc_device From ef9b7129f4b58841301982a114c44c76f1d7d23b Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Tue, 6 Oct 2026 09:52:20 -0700 Subject: [PATCH 04/11] cuda.core: validate and harden VirtualMemoryResource IPC imports Address the review of the VirtualMemoryResource IPC path. Import: check every imported allocation's location against the resource's before mapping (the same device, or the same host location type) and raise a ValueError that names both; reject a descriptor with a size but no allocations; attach notes to driver failures that say which chunk failed and why a peer-controlled size fails at the map. The imported buffer no longer keeps the descriptor: the driver copies the shared allocation into a handle of its own and keeps no reference to the file descriptor, so the descriptors live only as long as the caller holds the descriptor object. Export: an imported buffer, and a buffer grown from one, cannot be exported again because the driver exports only allocations created with the requested handle type; say so instead of surfacing INVALID_VALUE. modify_allocation() rejects a config whose handle_type differs from the resource's, so every chunk of a buffer is exportable the same way. __reduce__ preserves subclasses. Tests: parametrize the memory_ipc suite over VirtualMemoryResource (VirtualMR), with the pool-only scenarios guarded or skipped; PatternGen works on the first `size` bytes of a larger buffer; the fd-leak harness asserts on the success path; new VMM-only tests for many-chunk ranges, a second device with and without peers, descriptor validation and rollback, re-export, empty buffers, and a DLPack consumer. Docs: file descriptor and memory lifetime, multiprocessing-only transport, trust boundary, device rule, re-export limit; VirtualMemoryIPCBufferDescriptor in api_private; the from_ipc_descriptor stub has its mr type again. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 44 +- cuda_core/cuda/core/_memory/_buffer.pyi | 6 +- cuda_core/cuda/core/_memory/_buffer.pyx | 1 + cuda_core/cuda/core/_memory/_ipc.pyi | 12 +- cuda_core/cuda/core/_memory/_ipc.pyx | 12 +- .../core/_memory/_virtual_memory_resource.pyi | 53 +- .../core/_memory/_virtual_memory_resource.pyx | 197 ++++++- cuda_core/docs/source/api_private.rst | 1 + cuda_core/docs/source/release/1.3.0-notes.rst | 13 +- cuda_core/tests/conftest.py | 100 +++- cuda_core/tests/helpers/buffers.py | 24 +- cuda_core/tests/memory_ipc/test_errors.py | 75 ++- cuda_core/tests/memory_ipc/test_event_ipc.py | 8 +- cuda_core/tests/memory_ipc/test_leaks.py | 33 +- cuda_core/tests/memory_ipc/test_memory_ipc.py | 50 +- .../tests/memory_ipc/test_send_buffers.py | 17 +- cuda_core/tests/memory_ipc/test_serialize.py | 50 +- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 480 ++++++++++++++---- cuda_core/tests/memory_ipc/test_workerpool.py | 26 +- 19 files changed, 954 insertions(+), 248 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index 846af404ea2..2c8e2fc4ab7 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -213,20 +213,40 @@ travel because the OS handle does not expose them to the importer; a wrong size fails the importer's `cuMemMap`. The importer is a `VirtualMemoryResource` of the receiving process with the -same `handle_type`. `import_mem_allocation_handle` calls -`cuMemImportFromShareableHandle` and wraps the result in the same box and +same `handle_type`, for the device that owns the memory. `import_mem_allocation_handle` +calls `cuMemImportFromShareableHandle` and wraps the result in the same box and deleter as `create_mem_allocation_handle`: the driver gives an imported handle one reference and frees the memory when all references are released and no -mapping remains, so `cuMemRelease` is the right teardown for both. The access -descriptors stored in the box are the importer's, built from its options for -its device; the exporter's access does not travel. The import then reserves -`sum(sizes)` with the importer's alignment, maps each allocation in order, and -builds a range and a `VirtualMemoryBuffer` exactly as `allocate()` does, so a -grown buffer imports as one contiguous range with the same byte layout, and the -imported buffer frees itself through the same deleters. A descriptor imported -in the exporting process yields an alias of the exporter's memory at a new -address. A `VirtualMemoryResource` pickles as (device, options), which is what -lets a `Buffer` pickle as (resource, descriptor). +mapping remains, so `cuMemRelease` is the right teardown for both. After each +import, `cuMemGetAllocationPropertiesFromHandle` must report the resource's +location (the same device, or the same host location type); a mismatch is a +`ValueError` before anything is mapped, because the access descriptors would +describe the wrong location. The access descriptors stored in the box are the +importer's, built from its options for its device and peers; the exporter's +access does not travel. The import then reserves `sum(sizes)` with the +importer's alignment, maps each allocation in order, and builds a range and a +`VirtualMemoryBuffer` exactly as `allocate()` does, so a grown buffer imports +as one contiguous range with the same byte layout, and the imported buffer +frees itself through the same deleters. A descriptor imported in the exporting +process yields an alias of the exporter's memory at a new address. A +`VirtualMemoryResource` pickles as (device, options), which is what lets a +`Buffer` pickle as (resource, descriptor). + +File descriptors and memory. The driver copies each shared allocation into a +handle of its own during the import and keeps no reference to the file +descriptor, so the imported buffer does not keep the descriptor; it records +only that it was imported (`is_mapped`). The descriptor's file descriptors live +as long as the descriptor object, on the exporter (where the buffer caches it) +and in every process that unpickled a copy, and each open descriptor keeps the +physical memory allocated. A descriptor on a `Queue` pins the memory in the +sender until the receiver unpickles it. + +Re-export. The driver exports only allocations created with the requested +handle type, and an imported allocation carries none, so an imported buffer, +and any buffer grown from one, cannot be exported again; `ipc_descriptor` says +so instead of surfacing the driver's INVALID_VALUE. A `modify_allocation` +config must keep the resource's `handle_type` for the same reason: every chunk +of a buffer must be exportable the same way. ## Why `modify_allocation` returns a new buffer diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 531900ee8d8..3cc88e6b2ca 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -2,9 +2,11 @@ from typing import TypedDict -from _typeshed import Incomplete from cuda.core._memory._copy_enums import CopyOptions +from cuda.core._memory._device_memory_resource import DeviceMemoryResource from cuda.core._memory._ipc import IPCBufferDescriptor +from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource +from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource from cuda.core._stream import Stream from cuda.core._utils.pycompat import BufferProtocol from cuda.core.graph import GraphBuilder @@ -81,7 +83,7 @@ class Buffer: :class:`Buffer` is closed or garbage collected. """ @classmethod - def from_ipc_descriptor(cls, mr: Incomplete, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer: + def from_ipc_descriptor(cls, mr: DeviceMemoryResource | PinnedMemoryResource | VirtualMemoryResource, ipc_descriptor: IPCBufferDescriptor, *, stream: Stream) -> Buffer: """Import a buffer that was exported from another process. Parameters diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 291f1b31a30..c63d3e6eb42 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -44,6 +44,7 @@ from cuda.core._dlpack import classify_dl_device, make_py_capsule from cuda.core._device import Device if TYPE_CHECKING: + from cuda.core._memory._virtual_memory_resource import VirtualMemoryResource from cuda.core.graph import GraphBuilder diff --git a/cuda_core/cuda/core/_memory/_ipc.pyi b/cuda_core/cuda/core/_memory/_ipc.pyi index f34075a293f..23934ed9903 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyi +++ b/cuda_core/cuda/core/_memory/_ipc.pyi @@ -64,13 +64,21 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): POSIX file descriptors the handles are :class:`IPCAllocationHandle` objects: the descriptor owns them and closes them when it is released, and ``multiprocessing`` duplicates them into the receiving process. Plain - ``pickle`` cannot carry file descriptors. + ``pickle`` cannot carry file descriptors; send the descriptor, or the + buffer, through ``multiprocessing`` (a ``Queue``, a ``Pipe``, ``Process`` + arguments, or a ``Pool``). + + A live descriptor costs one file descriptor per allocation and keeps the + physical memory allocated, in every process that holds a copy, until it is + released. A descriptor placed on a ``Queue`` or sent to a ``Pool`` pins the + memory in the sender until the receiver has unpickled it. Note ---- The sizes are controlled by the exporting peer. Receivers must treat them as untrusted and import only through :meth:`Buffer.from_ipc_descriptor`, - which fails instead of mapping a size the driver rejects. + which checks what it can and fails instead of mapping a size the driver + rejects. """ _handle_type: int _chunk_sizes: tuple[Any, ...] diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index ab1c18bd674..87206a02409 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -178,13 +178,21 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): POSIX file descriptors the handles are :class:`IPCAllocationHandle` objects: the descriptor owns them and closes them when it is released, and ``multiprocessing`` duplicates them into the receiving process. Plain - ``pickle`` cannot carry file descriptors. + ``pickle`` cannot carry file descriptors; send the descriptor, or the + buffer, through ``multiprocessing`` (a ``Queue``, a ``Pipe``, ``Process`` + arguments, or a ``Pool``). + + A live descriptor costs one file descriptor per allocation and keeps the + physical memory allocated, in every process that holds a copy, until it is + released. A descriptor placed on a ``Queue`` or sent to a ``Pool`` pins the + memory in the sender until the receiver has unpickled it. Note ---- The sizes are controlled by the exporting peer. Receivers must treat them as untrusted and import only through :meth:`Buffer.from_ipc_descriptor`, - which fails instead of mapping a size the driver rejects. + which checks what it can and fails instead of mapping a size the driver + rejects. """ @staticmethod diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 2546e82baed..727bf15c512 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -99,12 +99,17 @@ class VirtualMemoryBuffer(Buffer): """Descriptor for sharing this buffer with other processes. The descriptor exports every physical allocation that backs the - buffer as a shareable handle and is cached on the buffer. It requires - a resource whose ``handle_type`` can be shared; see + buffer as a shareable handle and is cached on the buffer until the + buffer closes. Each exported allocation takes one file descriptor, + and a live descriptor keeps the physical memory allocated in every + process that holds a copy. It requires a resource whose + ``handle_type`` can be shared; see :attr:`VirtualMemoryResource.is_ipc_enabled`. Import it with :meth:`Buffer.from_ipc_descriptor` and a :class:`VirtualMemoryResource` of the receiving process, or send the - buffer itself through ``multiprocessing``. + buffer itself through ``multiprocessing``. A buffer imported from + another process cannot be exported again; share the descriptor it + was imported from. """ def close(self, stream: Stream | GraphBuilder | None=None) -> None: """Release this buffer's share of its address range. @@ -146,6 +151,11 @@ class VirtualMemoryResource(MemoryResource): how CUDA Virtual Memory Management works before using this. Other MemoryResource subclasses in cuda.core should already meet the common needs. + 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. + Notes ----- Every buffer this resource returns is a :class:`VirtualMemoryBuffer` that @@ -160,11 +170,15 @@ class VirtualMemoryResource(MemoryResource): Buffers can be shared with other processes when ``config.handle_type`` is ``"posix_fd"`` (Linux). :attr:`Buffer.ipc_descriptor` exports the - physical allocations that back a buffer, :meth:`Buffer.from_ipc_descriptor` - imports them with a ``VirtualMemoryResource`` of the receiving process, - and a buffer sent through ``multiprocessing`` does both. The importing - resource maps the shared memory for its own device with its own access - options; see :attr:`is_ipc_enabled`. + physical allocations that back a buffer, one file descriptor each; + :meth:`Buffer.from_ipc_descriptor` imports them 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. The importing resource + maps the memory for its device and the devices in its ``peers`` option, + with its own access options. A live descriptor keeps the memory allocated + in every process that holds a copy. An imported buffer, and a buffer grown + from one, cannot be exported again. See :attr:`is_ipc_enabled`. """ device: object config: object @@ -210,7 +224,9 @@ class VirtualMemoryResource(MemoryResource): freed when the last of the two closes. When the driver can extend the address range in place, the returned buffer has the same pointer; otherwise it has a new one and the existing contents are reachable - through both. Closing the returned buffer never closes ``buf``. + through both. Closing the returned buffer never closes ``buf``. A + buffer grown from an imported buffer maps imported memory and cannot + be exported again. Concurrent calls on buffers that alias one another are safe. @@ -224,7 +240,8 @@ class VirtualMemoryResource(MemoryResource): Configuration for the new physical memory chunk only. Existing chunks keep the access they were created with, and the resource's own configuration is unchanged. It must name the resource's - ``location_type`` and passes the same checks as the constructor. + ``location_type`` and ``handle_type`` and passes the same checks + as the constructor. When ``buf`` already covers ``new_size`` there is no new chunk, so ``config`` has no effect. This method never changes the access of memory that is already mapped. @@ -241,8 +258,8 @@ class VirtualMemoryResource(MemoryResource): TypeError If ``buf`` did not come from this resource. ValueError - If ``config`` names a different location than the resource, or - the constructor would reject it. + If ``config`` names a different location or handle type than the + resource, or the constructor would reject it. RuntimeError If ``buf`` is closed, or ``config`` requests GPUDirect RDMA on a device without support. @@ -297,10 +314,14 @@ class VirtualMemoryResource(MemoryResource): """Import a buffer that another process exported; see :meth:`Buffer.from_ipc_descriptor`. Every exported allocation is imported with this resource's handle - type, then mapped in order into one new address reservation, with the - access descriptors this resource's options produce for its device. - The returned buffer owns the imported handles like a buffer from - :meth:`allocate` owns created ones, and reports ``is_mapped``. + type and checked to live where this resource allocates (its device, + or its host location). The allocations are then mapped in order into + one new address reservation with the access descriptors this + resource's options produce for its device and peers, and the + deallocation stream is recorded as :meth:`allocate` does. The + returned buffer owns the imported handles like a buffer from + :meth:`allocate` owns created ones, reports ``is_mapped``, and cannot + be exported again. """ @property def device_id(self) -> int: diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 0919372250f..acfc3bb39fd 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -230,6 +230,52 @@ cdef int _raise_last_error() except -1: ) +cdef inline int _attach_note(exc, str detail) except -1: + # PEP 678 notes (Python 3.11+) keep the detail separable from the message; + # older interpreters get it appended to the message instead. + add_note = getattr(exc, "add_note", None) + if add_note is not None: + add_note(detail) + else: + exc.args = (f"{exc.args[0]} ({detail})", *exc.args[1:]) + return 0 + + +cdef int _raise_last_error_noting(str detail) except -1: + """Raise the status a handle factory recorded, with ``detail`` attached as a note.""" + try: + _raise_last_error() + except Exception as exc: + _attach_note(exc, detail) + raise + return 0 + + +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: + return 0 + try: + HANDLE_RETURN(status) + except Exception as exc: + _attach_note(exc, detail) + raise + return 0 + + +cdef str _describe_location(cydriver.CUmemLocation loc): + """Name a driver memory location for an error message.""" + if loc.type == cydriver.CU_MEM_LOCATION_TYPE_DEVICE: + return f"device {loc.id}" + if loc.type == cydriver.CU_MEM_LOCATION_TYPE_HOST: + return "host memory" + if loc.type == cydriver.CU_MEM_LOCATION_TYPE_HOST_NUMA: + return f"host NUMA node {loc.id}" + if loc.type == cydriver.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT: + return "the current host NUMA node" + return f"location type {loc.type}" + + cdef class VirtualMemoryBuffer(Buffer): """A :class:`Buffer` returned by :class:`VirtualMemoryResource`. @@ -247,16 +293,26 @@ cdef class VirtualMemoryBuffer(Buffer): """Descriptor for sharing this buffer with other processes. The descriptor exports every physical allocation that backs the - buffer as a shareable handle and is cached on the buffer. It requires - a resource whose ``handle_type`` can be shared; see + buffer as a shareable handle and is cached on the buffer until the + buffer closes. Each exported allocation takes one file descriptor, + and a live descriptor keeps the physical memory allocated in every + process that holds a copy. It requires a resource whose + ``handle_type`` can be shared; see :attr:`VirtualMemoryResource.is_ipc_enabled`. Import it with :meth:`Buffer.from_ipc_descriptor` and a :class:`VirtualMemoryResource` of the receiving process, or send the - buffer itself through ``multiprocessing``. + buffer itself through ``multiprocessing``. A buffer imported from + another process cannot be exported again; share the descriptor it + was imported from. """ Buffer_check_open(self) if self._ipc_data is None: self._ipc_data = IPCDataForBuffer(_export_ipc_descriptor(self), False) + elif self._ipc_data.is_mapped: + raise RuntimeError( + "a VirtualMemoryBuffer imported from another process cannot be exported again: the CUDA driver " + "does not export imported allocations. Share the descriptor it was imported from." + ) return self._ipc_data.ipc_descriptor def close(self, stream: Stream | GraphBuilder | None = None) -> None: @@ -318,6 +374,11 @@ cdef class VirtualMemoryResource(MemoryResource): how CUDA Virtual Memory Management works before using this. Other MemoryResource subclasses in cuda.core should already meet the common needs. + 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. + Notes ----- Every buffer this resource returns is a :class:`VirtualMemoryBuffer` that @@ -332,11 +393,15 @@ cdef class VirtualMemoryResource(MemoryResource): Buffers can be shared with other processes when ``config.handle_type`` is ``"posix_fd"`` (Linux). :attr:`Buffer.ipc_descriptor` exports the - physical allocations that back a buffer, :meth:`Buffer.from_ipc_descriptor` - imports them with a ``VirtualMemoryResource`` of the receiving process, - and a buffer sent through ``multiprocessing`` does both. The importing - resource maps the shared memory for its own device with its own access - options; see :attr:`is_ipc_enabled`. + physical allocations that back a buffer, one file descriptor each; + :meth:`Buffer.from_ipc_descriptor` imports them 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. The importing resource + maps the memory for its device and the devices in its ``peers`` option, + with its own access options. A live descriptor keeps the memory allocated + in every process that holds a copy. An imported buffer, and a buffer grown + from one, cannot be exported again. See :attr:`is_ipc_enabled`. """ cdef: @@ -359,7 +424,7 @@ 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 VirtualMemoryResource, (self.device.device_id if self.device is not None else 0, self.config) + return type(self), (self.device.device_id if self.device is not None else 0, self.config) cdef int _check_config(self, object cfg) except -1: """Reject options the driver would reject, before any driver call. @@ -379,6 +444,13 @@ cdef class VirtualMemoryResource(MemoryResource): 'virtual memory with location_type="host" cannot have an exportable handle type; ' "pass handle_type=None" ) + # The handle type is part of the resource's sharing contract, like the + # location: every chunk of a buffer must be exportable the same way. + if cfg.handle_type != self.config.handle_type: + raise ValueError( + f"config.handle_type {str(cfg.handle_type)!r} does not match the resource's " + f"{str(self.config.handle_type)!r}; the handle type of a buffer's chunks cannot change" + ) if cfg.gpu_direct_rdma and self.device is not None and not self.device.properties.gpu_direct_rdma_supported: raise RuntimeError("GPU Direct RDMA is not supported on this device") return 0 @@ -556,7 +628,9 @@ cdef class VirtualMemoryResource(MemoryResource): freed when the last of the two closes. When the driver can extend the address range in place, the returned buffer has the same pointer; otherwise it has a new one and the existing contents are reachable - through both. Closing the returned buffer never closes ``buf``. + through both. Closing the returned buffer never closes ``buf``. A + buffer grown from an imported buffer maps imported memory and cannot + be exported again. Concurrent calls on buffers that alias one another are safe. @@ -570,7 +644,8 @@ cdef class VirtualMemoryResource(MemoryResource): Configuration for the new physical memory chunk only. Existing chunks keep the access they were created with, and the resource's own configuration is unchanged. It must name the resource's - ``location_type`` and passes the same checks as the constructor. + ``location_type`` and ``handle_type`` and passes the same checks + as the constructor. When ``buf`` already covers ``new_size`` there is no new chunk, so ``config`` has no effect. This method never changes the access of memory that is already mapped. @@ -587,8 +662,8 @@ cdef class VirtualMemoryResource(MemoryResource): TypeError If ``buf`` did not come from this resource. ValueError - If ``config`` names a different location than the resource, or - the constructor would reject it. + If ``config`` names a different location or handle type than the + resource, or the constructor would reject it. RuntimeError If ``buf`` is closed, or ``config`` requests GPUDirect RDMA on a device without support. @@ -790,14 +865,19 @@ cdef class VirtualMemoryResource(MemoryResource): """Import a buffer that another process exported; see :meth:`Buffer.from_ipc_descriptor`. Every exported allocation is imported with this resource's handle - type, then mapped in order into one new address reservation, with the - access descriptors this resource's options produce for its device. - The returned buffer owns the imported handles like a buffer from - :meth:`allocate` owns created ones, and reports ``is_mapped``. + type and checked to live where this resource allocates (its device, + or its host location). The allocations are then mapped in order into + one new address reservation with the access descriptors this + resource's options produce for its device and peers, and the + deallocation stream is recorded as :meth:`allocate` does. The + returned buffer owns the imported handles like a buffer from + :meth:`allocate` owns created ones, reports ``is_mapped``, and cannot + be exported again. """ cdef Stream s = None cdef object cfg = self.config cdef cydriver.CUmemAllocationProp prop + cdef cydriver.CUmemAllocationProp imported_prop cdef vector[cydriver.CUmemAccessDesc] descs cdef vector[MemAllocationHandle] allocs cdef vector[VaMappingHandle] mappings @@ -806,6 +886,7 @@ cdef class VirtualMemoryResource(MemoryResource): cdef VaMappingHandle h_map cdef VmmRangeHandle rng cdef DevicePtrHandle h_ptr + cdef Buffer buf cdef size_t gran, addr_align, chunk, total = 0, offset = 0 cdef Py_ssize_t n, i cdef cydriver.CUdeviceptr hint @@ -826,11 +907,18 @@ cdef class VirtualMemoryResource(MemoryResource): if stream is not None: s = Stream_accept(stream) if n == 0: + if desc._size != 0: + raise ValueError( + f"malformed VirtualMemoryResource buffer descriptor: a size of {desc._size} bytes with no " + "exported allocations" + ) # An exported empty buffer: an empty range, as allocate(0) returns. h_ptr = deviceptr_create_vmm(0, create_vmm_range(mappings)) if s is not None and not Stream_is_default_token(s): HANDLE_RETURN(set_deallocation_stream(h_ptr, s._h_stream)) - return Buffer_from_deviceptr_handle(h_ptr, 0, self, desc, VirtualMemoryBuffer) + buf = Buffer_from_deviceptr_handle(h_ptr, 0, self, None, VirtualMemoryBuffer) + buf._ipc_data = IPCDataForBuffer(None, True) + return buf self._fill_prop(cfg, &prop) self._fill_access(cfg, &prop, descs) @@ -851,6 +939,14 @@ cdef class VirtualMemoryResource(MemoryResource): # Every handle below is a local: if a later step fails, the locals die # in reverse order and release everything imported so far. + # + # The driver copies each shared allocation into a handle of its own + # during cuMemImportFromShareableHandle and keeps no reference to the + # file descriptor afterwards: the shareable handle and the allocation + # handle are independent references to the memory, and the CUDA + # samples and PyTorch close the descriptor right after the import. The + # 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] @@ -859,25 +955,51 @@ cdef class VirtualMemoryResource(MemoryResource): h_alloc = import_mem_allocation_handle( os_handle, handle_type, chunk, descs.data(), descs.size()) if not h_alloc: - _raise_last_error() + _raise_last_error_noting( + f"while importing chunk {i + 1} of {n} of a VirtualMemoryResource buffer descriptor; the " + "shared handle may not refer to a CUDA allocation exported with " + f"handle_type={str(cfg.handle_type)!r}" + ) + # 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. + with nogil: + HANDLE_RETURN(cydriver.cuMemGetAllocationPropertiesFromHandle(&imported_prop, as_cu(h_alloc))) + if imported_prop.location.type != prop.location.type or ( + prop.location.type == cydriver.CU_MEM_LOCATION_TYPE_DEVICE + and imported_prop.location.id != prop.location.id + ): + raise ValueError( + f"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" + ) allocs.push_back(h_alloc) with nogil: h_res = create_va_reservation_handle(total, addr_align, hint) if not h_res: _raise_last_error() for i in range(n): + chunk = mem_allocation_size(allocs[i]) with nogil: h_map = create_va_mapping_handle(as_cu(h_res) + offset, allocs[i], h_res) if not h_map: - _raise_last_error() + _raise_last_error_noting( + f"while mapping chunk {i + 1} of {n} ({chunk} bytes) of a VirtualMemoryResource buffer " + "descriptor; a chunk size that does not match the exported allocation fails here, and the " + "sizes come from the exporting peer" + ) mappings.push_back(h_map) - offset += mem_allocation_size(allocs[i]) + offset += chunk rng = create_vmm_range(mappings) h_ptr = deviceptr_create_vmm(as_cu(h_res), rng) if not h_ptr: _raise_last_error() self._record_deallocation_stream(h_ptr, s) - return Buffer_from_deviceptr_handle(h_ptr, desc._size, self, desc, VirtualMemoryBuffer) + buf = Buffer_from_deviceptr_handle(h_ptr, desc._size, self, None, VirtualMemoryBuffer) + buf._ipc_data = IPCDataForBuffer(None, True) + return buf @property def device_id(self) -> int: @@ -904,12 +1026,16 @@ cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf): Only the allocations that cover ``buf.size`` are exported; a buffer's size is a prefix of its range. Each file descriptor is owned by an - :class:`IPCAllocationHandle` in the descriptor and closed with it. + :class:`IPCAllocationHandle` in the descriptor and closed with it. An + allocation this process imported cannot be exported: the driver exports + only allocations created with the requested handle type. """ cdef VirtualMemoryResource mr = buf._memory_resource cdef vector[VaMappingHandle] mappings cdef MemAllocationHandle h_alloc - cdef size_t i, chunk, offset = 0 + cdef cydriver.CUmemAllocationProp chunk_prop + cdef cydriver.CUresult status + cdef size_t i, n, chunk, offset = 0 cdef int handle_type cdef int fd = -1 cdef list sizes = [] @@ -919,14 +1045,31 @@ cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf): raise RuntimeError("Memory resource is not IPC-enabled") handle_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(mr.config.handle_type)) mappings = vmm_range_mappings(vmm_range(buf._h_ptr)) - for i in range(mappings.size()): + n = mappings.size() + for i in range(n): if offset >= buf._size: break h_alloc = va_mapping_allocation(mappings[i]) chunk = mem_allocation_size(h_alloc) + # The driver exports only allocations created with the requested + # handle type; an imported allocation carries none, so the export would + # fail with INVALID_VALUE. Say what the limitation is instead. with nogil: - HANDLE_RETURN(cydriver.cuMemExportToShareableHandle( - &fd, as_cu(h_alloc), handle_type, 0)) + HANDLE_RETURN(cydriver.cuMemGetAllocationPropertiesFromHandle(&chunk_prop, as_cu(h_alloc))) + if (chunk_prop.requestedHandleTypes & handle_type) == 0: + raise RuntimeError( + f"chunk {i + 1} of this buffer cannot be exported: it was imported from another process, and " + "the CUDA driver does not export imported allocations. Share the descriptor the memory was " + "imported from instead." + ) + with nogil: + status = cydriver.cuMemExportToShareableHandle( + &fd, as_cu(h_alloc), handle_type, 0) + _check_noting( + status, + f"while exporting chunk {i + 1} of {n} of a VirtualMemoryResource buffer; each exported chunk takes " + "one file descriptor, so a process near its file descriptor limit fails here", + ) try: handles.append(IPCAllocationHandle._init(fd, None)) except: # noqa: E722 rollback-then-raise: the fd is not owned yet diff --git a/cuda_core/docs/source/api_private.rst b/cuda_core/docs/source/api_private.rst index 80675799c07..41963145bb2 100644 --- a/cuda_core/docs/source/api_private.rst +++ b/cuda_core/docs/source/api_private.rst @@ -48,6 +48,7 @@ CUDA runtime _device.DeviceProperties _memory._ipc.IPCAllocationHandle _memory._ipc.IPCBufferDescriptor + _memory._ipc.VirtualMemoryIPCBufferDescriptor _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 8e457363087..87a6599d3be 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -82,11 +82,14 @@ New features 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, and a - buffer or its descriptor can be sent through ``multiprocessing`` directly. - :attr:`VirtualMemoryResource.is_ipc_enabled` is ``True`` when the resource's - ``handle_type`` is ``"posix_fd"`` on Linux. A grown buffer imports as one - contiguous range. Fabric handles are not shared yet. + 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. A grown buffer + imports as one contiguous range. A live descriptor holds one file descriptor + per physical allocation and keeps the memory allocated; an imported buffer + cannot be exported again. Fabric handles are not shared yet. (`#2980 `__) Fixes and enhancements diff --git a/cuda_core/tests/conftest.py b/cuda_core/tests/conftest.py index a8678292422..6c99a733d3e 100644 --- a/cuda_core/tests/conftest.py +++ b/cuda_core/tests/conftest.py @@ -41,6 +41,8 @@ ManagedMemoryResourceOptions, PinnedMemoryResource, PinnedMemoryResourceOptions, + VirtualMemoryResource, + VirtualMemoryResourceOptions, _device, ) from cuda.core._utils.cuda_utils import handle_return @@ -271,32 +273,108 @@ def ipc_device(init_cuda): yield device +def _require_vmm_ipc_device(device): + """Return ``device`` if it can export VMM allocations as POSIX file descriptors, otherwise skip.""" + from cuda_python_test_helpers import IS_WINDOWS, IS_WSL + + props = device.properties + if not props.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + if IS_WINDOWS or IS_WSL or not props.handle_type_posix_file_descriptor_supported: + pytest.skip("Sharing virtual memory needs POSIX file descriptor handles") + return device + + +@pytest.fixture +def vmm_ipc_device(init_cuda): + """A device whose VirtualMemoryResource buffers can be shared between processes, or skip. + + Tracks spawned ``multiprocessing.Process`` children like :func:`ipc_device`. + """ + from helpers.child_processes import track_child_processes + + device = _require_vmm_ipc_device(init_cuda) + with track_child_processes(): + yield device + + +@pytest.fixture +def vmm_ipc_device_x2(device_x2): + """Two devices whose VirtualMemoryResource buffers can be shared between processes, or skip.""" + from helpers.child_processes import track_child_processes + + devices = tuple(_require_vmm_ipc_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": + options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) + return DeviceMemoryResource(device, options=options) + if kind == "pinned": + skip_if_pinned_memory_unsupported(device) + options = PinnedMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) + return PinnedMemoryResource(options=options) + assert kind == "vmm", kind + _require_vmm_ipc_device(device) + return VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(handle_type="posix_fd")) + + +def _close_memory_resource(mr): + # A VirtualMemoryResource owns no pool and has no close(); its buffers free themselves. + close = getattr(mr, "close", None) + if close is not None: + close() + + @pytest.fixture( params=[ pytest.param("device", id="DeviceMR"), pytest.param("pinned", id="PinnedMR"), + pytest.param("vmm", id="VirtualMR"), ] ) def ipc_memory_resource(request, ipc_device): - """Provides IPC-enabled memory resource (either Device or Pinned).""" - mr_type = request.param - - if mr_type == "device": - options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) - mr = DeviceMemoryResource(ipc_device, options=options) - else: # pinned - skip_if_pinned_memory_unsupported(ipc_device) - options = PinnedMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) - mr = PinnedMemoryResource(options=options) + """Provides an IPC-enabled memory resource: a device pool, a pinned pool, or a VirtualMemoryResource. + The VMM parameter needs VMM support and POSIX file descriptor handles on + top of what :func:`ipc_device` requires; it skips otherwise. Tests that + use the pool-only half of the IPC protocol (``allocation_handle``, + ``from_allocation_handle``, the registry, ``uuid``, ``is_mapped`` on the + resource) guard or skip for it. + """ + mr = _make_ipc_memory_resource(request.param, ipc_device) assert mr.is_ipc_enabled yield mr - mr.close() + _close_memory_resource(mr) # TODO(seberg): Make sure the `mr` and it's buffers are fully torn down. # May be unnecessary as `mr.close()` is not parallel with other work. ipc_device.sync() +@pytest.fixture( + params=[ + pytest.param("device", id="DeviceMR"), + pytest.param("vmm", id="VirtualMR"), + ] +) +def ipc_memory_resource_factory(request, ipc_device): + """Provides a callable that creates fresh IPC-enabled resources of one kind; all are closed at teardown.""" + created = [] + + def make(): + mr = _make_ipc_memory_resource(request.param, ipc_device) + created.append(mr) + return mr + + yield make + for mr in created: + _close_memory_resource(mr) + ipc_device.sync() + + @pytest.fixture def mempool_device(init_cuda): """Obtains a device suitable for mempool tests, or skips.""" diff --git a/cuda_core/tests/helpers/buffers.py b/cuda_core/tests/helpers/buffers.py index 92a51429ab3..b576185119d 100644 --- a/cuda_core/tests/helpers/buffers.py +++ b/cuda_core/tests/helpers/buffers.py @@ -22,6 +22,7 @@ "compare_equal_buffers", "make_instrumented_memory_resource", "make_scratch_buffer", + "prefix_view", "thread_unsafe_on_windows", ] @@ -247,19 +248,17 @@ def __init__(self, device, size, *, stream): self.pattern_buffers = {} def fill_buffer(self, buffer, seed=None, value=None): - """Fill a device buffer with a sequential test pattern using unified memory.""" - assert buffer.size == self.size + """Fill the first ``size`` bytes of a device buffer with a test pattern using unified memory.""" pattern_buffer = self._get_pattern_buffer(seed, value) - buffer.copy_from(pattern_buffer, stream=self.stream) + prefix_view(buffer, self.size).copy_from(pattern_buffer, stream=self.stream) def verify_buffer(self, buffer, seed=None, value=None): - """Verify the buffer contents against a sequential pattern.""" - assert buffer.size == self.size + """Verify the first ``size`` bytes of the buffer against a sequential pattern.""" scratch_buffer = DummyUnifiedMemoryResource(self.device).allocate(self.size) ptr_test = self._ptr(scratch_buffer) pattern_buffer = self._get_pattern_buffer(seed, value) ptr_expected = self._ptr(pattern_buffer) - scratch_buffer.copy_from(buffer, stream=self.stream) + scratch_buffer.copy_from(prefix_view(buffer, self.size), stream=self.stream) self.stream.sync() assert libc.memcmp(ptr_test, ptr_expected, self.size) == 0 @@ -287,6 +286,19 @@ def _get_pattern_buffer(self, seed, value): return pattern_buffer +def prefix_view(buffer, nbytes): + """Return ``buffer`` itself when it has ``nbytes`` bytes, else a non-owning view of its first ``nbytes``. + + Resources that round sizes up (VirtualMemoryResource rounds to its + granularity) return buffers larger than requested, and the copy helpers + need equal sizes. The view keeps ``buffer`` alive and frees nothing. + """ + if buffer.size == nbytes: + return buffer + assert buffer.size > nbytes, f"buffer of {buffer.size} bytes is smaller than the {nbytes}-byte pattern" + return Buffer.from_handle(int(buffer.handle), nbytes, owner=buffer) + + def make_scratch_buffer(device, value, nbytes): """Create a unified memory buffer with the specified value.""" buffer = DummyUnifiedMemoryResource(device).allocate(nbytes) diff --git a/cuda_core/tests/memory_ipc/test_errors.py b/cuda_core/tests/memory_ipc/test_errors.py index c57e63442d8..b4884c463c9 100644 --- a/cuda_core/tests/memory_ipc/test_errors.py +++ b/cuda_core/tests/memory_ipc/test_errors.py @@ -18,6 +18,7 @@ DeviceMemoryResource, DeviceMemoryResourceOptions, PinnedMemoryResource, + VirtualMemoryResource, ) from cuda.core._memory._ipc import IPCAllocationHandle, IPCBufferDescriptor from cuda.core._utils.cuda_utils import CUDAError @@ -30,6 +31,11 @@ pytestmark = pytest.mark.parallel_threads_limit(4) +def _is_pool(mr): + """Whether ``mr`` has the pool half of the IPC protocol (allocation handle, registry, is_mapped).""" + return isinstance(mr, (DeviceMemoryResource, PinnedMemoryResource)) + + def test_outer_timeout_marker_is_applied(request): """Verify that memory_ipc/conftest.py applies the outer pytest-timeout marker. @@ -57,7 +63,12 @@ def test_ipc_types_cannot_be_constructed_directly(): def test_import_truncated_buffer_descriptor(ipc_device, ipc_memory_resource): """Truncated IPC buffer descriptor payload is rejected before driver import.""" desc = IPCBufferDescriptor._init(b"\x00" * 8, NBYTES) - with pytest.raises(ValueError, match=r"payload is 8 bytes; expected at least 64"): + if _is_pool(ipc_memory_resource): + expected = pytest.raises(ValueError, match=r"payload is 8 bytes; expected at least 64") + else: + # A pool descriptor cannot be imported with a VirtualMemoryResource at all. + expected = pytest.raises(TypeError, match="memory pool") + with expected: Buffer.from_ipc_descriptor(ipc_memory_resource, desc, stream=ipc_device.default_stream) @@ -102,6 +113,8 @@ def test_ipc_allocation_handle_state_tracks_close(): @pytest.mark.agent_authored(model="gpt-5.6") def test_closed_ipc_allocation_handle_rejected_before_registry_hit(ipc_device, ipc_memory_resource): mr = ipc_memory_resource + if not _is_pool(mr): + pytest.skip("allocation handles and the registry are pool-only") handle = IPCAllocationHandle._init(os.dup(int(mr.allocation_handle)), mr.uuid) assert mr.register(mr.uuid) is mr handle.close() @@ -116,7 +129,10 @@ def test_closed_ipc_allocation_handle_rejected_before_registry_hit(ipc_device, i class ChildErrorHarness: """Test harness for checking errors in child processes. Subclasses override - PARENT_ACTION, CHILD_ACTION, and ASSERT (see below for examples).""" + PARENT_ACTION, CHILD_ACTION, and ASSERT (see below for examples). A subclass + that exercises the pool half of the protocol sets POOL_ONLY.""" + + POOL_ONLY = False @pytest.mark.thread_unsafe( reason=( @@ -132,6 +148,8 @@ def test_main(self, ipc_device, ipc_memory_resource): self.device = ipc_device self.mr = ipc_memory_resource self._extra_mrs = [] + if self.POOL_ONLY and not _is_pool(self.mr): + pytest.skip("this scenario needs the pool half of the IPC protocol") try: # Start a child process to generate error info. Target a fresh @@ -178,8 +196,14 @@ class TestImportOversizedBufferDescriptorSize(ChildErrorHarness): def PARENT_ACTION(self, queue): stream = self.device.default_stream self.buffer = self.mr.allocate(NBYTES, stream=stream) - payload, _ = self.buffer.ipc_descriptor.__reduce__()[1] - oversized = IPCBufferDescriptor._init(payload, NBYTES * 100) + desc = self.buffer.ipc_descriptor + if _is_pool(self.mr): + payload, _ = desc.__reduce__()[1] + oversized = IPCBufferDescriptor._init(payload, NBYTES * 100) + else: + # The VMM descriptor carries the exported handles and their sizes; + # forge only the buffer size. + oversized = type(desc)._from_exports(desc._handle_type, desc._chunk_sizes, desc._handles, desc.size * 100) stream.sync() queue.put(oversized) @@ -190,31 +214,43 @@ def CHILD_ACTION(self, queue): def ASSERT(self, exc_type, exc_msg): assert exc_type is ValueError assert "exceeds" in exc_msg - assert "mapped allocation extent" in exc_msg + if _is_pool(self.mr): + assert "mapped allocation extent" in exc_msg + else: + assert "exported allocations" in exc_msg class TestAllocFromImportedMr(ChildErrorHarness): - """Error when attempting to allocate from an import memory resource.""" + """Allocating from a received memory resource: an error for an imported pool, fine for a rebuilt VMM resource.""" def PARENT_ACTION(self, queue): queue.put(self.mr) def CHILD_ACTION(self, queue): mr = queue.get(timeout=CHILD_TIMEOUT_SEC) - mr.allocate(NBYTES, stream=self.device.default_stream) + buffer = mr.allocate(NBYTES, stream=self.device.default_stream) + buffer.close() def ASSERT(self, exc_type, exc_msg): - assert exc_type is TypeError - assert exc_msg == "Cannot allocate from a mapped IPC-enabled memory resource" + if _is_pool(self.mr): + assert exc_type is TypeError + assert exc_msg == "Cannot allocate from a mapped IPC-enabled memory resource" + else: + # A VirtualMemoryResource pickles as (device, options) and the + # receiver's copy is an ordinary resource, not a mapping. + assert exc_type is None class TestImportWrongMR(ChildErrorHarness): - """Error when importing a buffer from the wrong memory resource.""" + """Importing a buffer with a resource other than the exporter's: an error for pools, fine for VMM.""" def PARENT_ACTION(self, queue): - options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) - mr2 = DeviceMemoryResource(self.device, options=options) - self._extra_mrs.append(mr2) + if _is_pool(self.mr): + options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) + mr2 = DeviceMemoryResource(self.device, options=options) + self._extra_mrs.append(mr2) + else: + mr2 = VirtualMemoryResource(self.device, config=self.mr.config) stream = self.device.default_stream self.buffer = mr2.allocate(NBYTES, stream=stream) stream.sync() @@ -222,11 +258,16 @@ def PARENT_ACTION(self, queue): def CHILD_ACTION(self, queue): mr, buffer_desc = queue.get(timeout=CHILD_TIMEOUT_SEC) - Buffer.from_ipc_descriptor(mr, buffer_desc, stream=self.device.default_stream) + Buffer.from_ipc_descriptor(mr, buffer_desc, stream=self.device.default_stream).close() def ASSERT(self, exc_type, exc_msg): - assert exc_type is CUDAError - assert "CUDA_ERROR_INVALID_VALUE" in exc_msg + if _is_pool(self.mr): + assert exc_type is CUDAError + assert "CUDA_ERROR_INVALID_VALUE" in exc_msg + else: + # A VMM import never references the exporting resource; any + # resource with the same handle type and location imports it. + assert exc_type is None class TestImportBuffer(ChildErrorHarness): @@ -255,6 +296,8 @@ class TestDanglingBuffer(ChildErrorHarness): resource. """ + POOL_ONLY = True # plain pickle of a VMM buffer is refused (fds), and VMM has no registry + def PARENT_ACTION(self, queue): options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) mr2 = DeviceMemoryResource(self.device, options=options) diff --git a/cuda_core/tests/memory_ipc/test_event_ipc.py b/cuda_core/tests/memory_ipc/test_event_ipc.py index 7608b9c8794..be2af397124 100644 --- a/cuda_core/tests/memory_ipc/test_event_ipc.py +++ b/cuda_core/tests/memory_ipc/test_event_ipc.py @@ -4,7 +4,7 @@ import multiprocessing as mp import pytest -from helpers.buffers import compare_equal_buffers, make_scratch_buffer +from helpers.buffers import compare_equal_buffers, make_scratch_buffer, prefix_view from helpers.child_processes import child_timeout_sec, kill_subprocesses from helpers.latch import LatchKernel from helpers.logging import TimestampedLogger @@ -56,7 +56,7 @@ def test_main(self, ipc_device, ipc_memory_resource): log("enqueuing latch kernel on stream1") latch.launch(stream1) log("enqueuing copy on stream1") - buffer.copy_from(ones, stream=stream1) + prefix_view(buffer, NBYTES).copy_from(ones, stream=stream1) ipc_event_options = EventOptions(ipc_enabled=True) e = stream1.record(options=ipc_event_options) @@ -78,7 +78,7 @@ def test_main(self, ipc_device, ipc_memory_resource): log("done") # Finish up. - target.copy_from(buffer, stream=stream1) + target.copy_from(prefix_view(buffer, NBYTES), stream=stream1) stream1.sync() assert compare_equal_buffers(target, twos) @@ -95,7 +95,7 @@ def child_main(self, log, q_in, q_out): log(f"got event ({hex(e.handle)})") stream2.wait(e) log("enqueuing copy on stream2") - buffer.copy_from(twos, stream=stream2) + prefix_view(buffer, NBYTES).copy_from(twos, stream=stream2) log("signaling parent") q_out.put(None) log("waiting") diff --git a/cuda_core/tests/memory_ipc/test_leaks.py b/cuda_core/tests/memory_ipc/test_leaks.py index bca7948737c..8bb29922f3c 100644 --- a/cuda_core/tests/memory_ipc/test_leaks.py +++ b/cuda_core/tests/memory_ipc/test_leaks.py @@ -31,6 +31,8 @@ def test_alloc_handle(ipc_memory_resource): """Check for fd leaks in allocation_handle.""" mr = ipc_memory_resource + if not hasattr(mr, "allocation_handle"): + pytest.skip("allocation handles are pool-only") with CheckFDLeaks(): [mr.allocation_handle for _ in range(10)] @@ -86,16 +88,29 @@ def __reduce__(self): raise RuntimeError("Irreducible") +def _get_alloc_handle(mr, _stream): + if not hasattr(mr, "allocation_handle"): + pytest.skip("allocation handles are pool-only") + return mr.allocation_handle + + +def _get_mr(mr, _stream): + return mr + + +def _get_buffer(mr, stream): + return mr.allocate(NBYTES, stream=stream) + + +def _get_buffer_desc(mr, stream): + return mr.allocate(NBYTES, stream=stream).ipc_descriptor + + @pytest.mark.flaky(reruns=2) @skip_if_unrunnable @pytest.mark.parametrize( "getobject", - [ - lambda mr, _stream: mr.allocation_handle, - lambda mr, _stream: mr, - lambda mr, stream: mr.allocate(NBYTES, stream=stream), - lambda mr, stream: mr.allocate(NBYTES, stream=stream).ipc_descriptor, - ], + [_get_alloc_handle, _get_mr, _get_buffer, _get_buffer_desc], ids=["alloc_handle", "mr", "buffer", "buffer_desc"], ) @pytest.mark.parametrize("launcher", [exec_success, exec_launch_failure, exec_reduce_failure]) @@ -128,10 +143,12 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): - if exc_type is not None: + # Only the success path is checked: a failing block has already + # reported its own error, and its cleanup may legitimately be partial. + if exc_type is None: gc.collect() final_fds = self.process.num_fds() - assert final_fds == self.initial_fds + assert final_fds == self.initial_fds, f"file descriptors leaked: {self.initial_fds} -> {final_fds}" return False diff --git a/cuda_core/tests/memory_ipc/test_memory_ipc.py b/cuda_core/tests/memory_ipc/test_memory_ipc.py index 0eace7f154c..836d4b0a271 100644 --- a/cuda_core/tests/memory_ipc/test_memory_ipc.py +++ b/cuda_core/tests/memory_ipc/test_memory_ipc.py @@ -7,7 +7,7 @@ from helpers.buffers import PatternGen from helpers.child_processes import child_timeout_sec, kill_subprocesses -from cuda.core import Buffer, DeviceMemoryResource +from cuda.core import Buffer, DeviceMemoryResource, PinnedMemoryResource, VirtualMemoryResource CHILD_TIMEOUT_SEC = child_timeout_sec() NBYTES = 64 @@ -18,6 +18,23 @@ pytestmark = pytest.mark.parallel_threads_limit(4) +def _is_pool(mr): + """Whether ``mr`` has the pool half of the IPC protocol (allocation handle, registry, is_mapped).""" + return isinstance(mr, (DeviceMemoryResource, PinnedMemoryResource)) + + +def _share_token(mr): + """What a child needs to rebuild ``mr``: a pool's allocation handle, or the VMM resource itself.""" + return mr.allocation_handle if _is_pool(mr) else mr + + +def _import_token(device, token): + """Rebuild a memory resource in a child from :func:`_share_token`'s result.""" + if isinstance(token, VirtualMemoryResource): + return token + return DeviceMemoryResource.from_allocation_handle(device, token) + + class TestIpcMempool: @pytest.mark.flaky(reruns=2) def test_main(self, ipc_device, ipc_memory_resource): @@ -25,7 +42,8 @@ def test_main(self, ipc_device, ipc_memory_resource): # Set up the IPC-enabled memory pool and share it. device = ipc_device mr = ipc_memory_resource - assert not mr.is_mapped + if _is_pool(mr): + assert not mr.is_mapped stream = device.default_stream pgen = PatternGen(device, NBYTES, stream=stream) @@ -56,7 +74,8 @@ def test_main(self, ipc_device, ipc_memory_resource): def child_main(self, device, mr, queue): device.set_current() - assert mr.is_mapped + if _is_pool(mr): + assert mr.is_mapped buffer = queue.get(timeout=CHILD_TIMEOUT_SEC) assert buffer.is_mapped stream = device.default_stream @@ -132,16 +151,16 @@ def test_main(self, ipc_device, ipc_memory_resource): Demonstrate that a memory pool allocation handle can be reused for IPC with multiple processes. Uses buffer descriptors. """ - # Set up the IPC-enabled memory pool and share it using one handle. + # Set up the IPC-enabled memory resource and share it using one token. device = ipc_device mr = ipc_memory_resource - alloc_handle = mr.allocation_handle + token = _share_token(mr) stream = device.default_stream # Start children. q1, q2 = (mp.Queue() for _ in range(2)) - p1 = mp.Process(target=self.child_main, args=(device, alloc_handle, False, q1)) - p2 = mp.Process(target=self.child_main, args=(device, alloc_handle, True, q2)) + p1 = mp.Process(target=self.child_main, args=(device, token, False, q1)) + p2 = mp.Process(target=self.child_main, args=(device, token, True, q2)) p1.start() p2.start() @@ -168,12 +187,12 @@ def test_main(self, ipc_device, ipc_memory_resource): buffer2.close() stream.sync() - def child_main(self, device, alloc_handle, seed, queue): + def child_main(self, device, token, seed, queue): """Fills a shared memory buffer.""" # In this case, the device needs to be set up (passing the mr does it # implicitly in other tests). device.set_current() - mr = DeviceMemoryResource.from_allocation_handle(device, alloc_handle) + mr = _import_token(device, token) buffer_descriptor = queue.get(timeout=CHILD_TIMEOUT_SEC) stream = device.default_stream buffer = Buffer.from_ipc_descriptor(mr, buffer_descriptor, stream=stream) @@ -192,13 +211,13 @@ def test_main(self, ipc_device, ipc_memory_resource): """ device = ipc_device mr = ipc_memory_resource - alloc_handle = mr.allocation_handle + token = _share_token(mr) stream = device.default_stream # Start children. q1, q2 = (mp.Queue() for _ in range(2)) - p1 = mp.Process(target=self.child_main, args=(device, alloc_handle, False, q1)) - p2 = mp.Process(target=self.child_main, args=(device, alloc_handle, True, q2)) + p1 = mp.Process(target=self.child_main, args=(device, token, False, q1)) + p2 = mp.Process(target=self.child_main, args=(device, token, True, q2)) p1.start() p2.start() @@ -225,12 +244,13 @@ def test_main(self, ipc_device, ipc_memory_resource): buffer2.close() stream.sync() - def child_main(self, device, alloc_handle, seed, queue): + def child_main(self, device, token, seed, queue): """Fills a shared memory buffer.""" device.set_current() - # Register the memory resource. - DeviceMemoryResource.from_allocation_handle(device, alloc_handle) + # Register the memory resource (a pool); a VirtualMemoryResource needs + # no registration because a buffer carries it. + _import_token(device, token) # Now get buffers. buffer = queue.get(timeout=CHILD_TIMEOUT_SEC) diff --git a/cuda_core/tests/memory_ipc/test_send_buffers.py b/cuda_core/tests/memory_ipc/test_send_buffers.py index d2a59f1cce6..4a897d8d203 100644 --- a/cuda_core/tests/memory_ipc/test_send_buffers.py +++ b/cuda_core/tests/memory_ipc/test_send_buffers.py @@ -7,9 +7,8 @@ import pytest from helpers.buffers import PatternGen from helpers.child_processes import child_timeout_sec, kill_subprocesses -from helpers.constants import POOL_SIZE -from cuda.core import Device, DeviceMemoryResource, DeviceMemoryResourceOptions +from cuda.core import Device, VirtualMemoryResource CHILD_TIMEOUT_SEC = child_timeout_sec() NBYTES = 64 @@ -23,12 +22,11 @@ class TestIpcSendBuffers: @pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("nmrs", (1, NMRS)) - def test_main(self, ipc_device, nmrs): + def test_main(self, ipc_device, ipc_memory_resource_factory, nmrs): """Test passing buffers sourced from multiple memory resources.""" - # Set up several IPC-enabled memory pools. + # Set up several IPC-enabled memory resources. device = ipc_device - options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) - mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)] + mrs = [ipc_memory_resource_factory() for _ in range(nmrs)] buffers = [] stream = device.default_stream @@ -60,8 +58,6 @@ def test_main(self, ipc_device, nmrs): for buffer in buffers: buffer.close() stream.sync() - for mr in mrs: - mr.close() def child_main(self, device, buffers): device.set_current() @@ -96,6 +92,11 @@ def test_main(self, ipc_device, ipc_memory_resource): # Allocate, fill a buffer. mr = ipc_memory_resource + if isinstance(mr, VirtualMemoryResource): + pytest.skip( + "an imported VirtualMemoryBuffer cannot be exported again: the driver does not export " + "imported allocations, and cuda.core does not forward the original descriptor" + ) stream = device.default_stream pgen = PatternGen(device, NBYTES, stream=stream) buffer = mr.allocate(NBYTES, stream=stream) diff --git a/cuda_core/tests/memory_ipc/test_serialize.py b/cuda_core/tests/memory_ipc/test_serialize.py index f2805c31e63..798cb7ecb41 100644 --- a/cuda_core/tests/memory_ipc/test_serialize.py +++ b/cuda_core/tests/memory_ipc/test_serialize.py @@ -9,7 +9,7 @@ from helpers.buffers import PatternGen from helpers.child_processes import child_timeout_sec, kill_subprocesses -from cuda.core import Buffer, Device, DeviceMemoryResource, PinnedMemoryResource +from cuda.core import Buffer, Device, DeviceMemoryResource, PinnedMemoryResource, VirtualMemoryResource CHILD_TIMEOUT_SEC = child_timeout_sec() NBYTES = 64 @@ -18,6 +18,11 @@ pytestmark = pytest.mark.parallel_threads_limit(4) +def _is_pool(mr): + """Whether ``mr`` has the pool half of the IPC protocol (allocation handle, registry, uuid).""" + return isinstance(mr, (DeviceMemoryResource, PinnedMemoryResource)) + + class TestObjectSerializationDirect: """ Test the low-level interface for sharing memory resources. @@ -37,9 +42,15 @@ def test_main(self, ipc_device, ipc_memory_resource): process = mp.Process(target=self.child_main, args=(child_conn,)) process.start() - # Send a memory resource by allocation handle. - alloc_handle = mr.allocation_handle - mp.reduction.send_handle(parent_conn, alloc_handle.handle, process.pid) + # Send a memory resource: a pool by its allocation handle over the raw + # fd channel, a VirtualMemoryResource as the object itself (it holds no + # driver object; a VMM descriptor carries its fds per allocation). + parent_conn.send(_is_pool(mr)) + if _is_pool(mr): + alloc_handle = mr.allocation_handle + mp.reduction.send_handle(parent_conn, alloc_handle.handle, process.pid) + else: + parent_conn.send(mr) # Send a buffer. buffer1 = mr.allocate(NBYTES, stream=stream) @@ -69,9 +80,13 @@ def child_main(self, conn): device.set_current() # Receive the memory resource. - handle = mp.reduction.recv_handle(conn) - mr = DeviceMemoryResource.from_allocation_handle(device, handle) - os.close(handle) + if conn.recv(): # pool + handle = mp.reduction.recv_handle(conn) + mr = DeviceMemoryResource.from_allocation_handle(device, handle) + os.close(handle) + else: + mr = conn.recv() + assert isinstance(mr, VirtualMemoryResource) # Receive the buffers. buffer1 = conn.recv() # directly @@ -106,7 +121,7 @@ def test_main(self, ipc_device, ipc_memory_resource): # being passed when spawning the child. pipe[0].put(mr) uuid = pipe[1].get(timeout=CHILD_TIMEOUT_SEC) - assert uuid == mr.uuid + assert uuid == getattr(mr, "uuid", None) # Send a buffer. buffer = mr.allocate(NBYTES, stream=stream) @@ -131,11 +146,16 @@ def child_main(self, pipe, _): # Memory resource. mr = pipe[0].get(timeout=CHILD_TIMEOUT_SEC) - pipe[1].put(mr.uuid) + pipe[1].put(getattr(mr, "uuid", None)) # Buffer. buffer = pipe[0].get(timeout=CHILD_TIMEOUT_SEC) - assert buffer.memory_resource.handle == mr.handle + if _is_pool(mr): + assert buffer.memory_resource.handle == mr.handle + else: + # A VMM resource pickles as (device, options); the buffer's rebuilt + # resource is equivalent, not identical. + assert buffer.memory_resource.config == mr.config stream = device.default_stream pgen = PatternGen(device, NBYTES, stream=stream) pgen.fill_buffer(buffer, seed=True) @@ -157,7 +177,7 @@ def test_main(self, ipc_device, ipc_memory_resource): # Define the objects. device = ipc_device mr = ipc_memory_resource - alloc_handle = mr.allocation_handle + alloc_handle = mr.allocation_handle if _is_pool(mr) else None stream = device.default_stream buffer = mr.allocate(NBYTES, stream=stream) buffer_desc = buffer.ipc_descriptor @@ -185,10 +205,16 @@ def child_main(self, alloc_handle, mr1, buffer_desc, buffer): with pytest.raises(TypeError): DeviceMemoryResource.from_allocation_handle(device, alloc_handle) mr2 = PinnedMemoryResource.from_allocation_handle(alloc_handle) - else: + elif isinstance(mr1, DeviceMemoryResource): with pytest.raises(TypeError): PinnedMemoryResource.from_allocation_handle(alloc_handle) mr2 = DeviceMemoryResource.from_allocation_handle(device, alloc_handle) + else: + # A VirtualMemoryResource has no allocation handle; the resource + # argument itself is the rebuilt resource. + assert alloc_handle is None + assert isinstance(mr1, VirtualMemoryResource) + mr2 = mr1 stream = device.default_stream pgen = PatternGen(device, NBYTES, stream=stream) diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index 1061c2198ad..827d6dde2f4 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -1,7 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Sharing VirtualMemoryResource buffers between processes (issue #2980).""" +"""Sharing VirtualMemoryResource buffers between processes (issue #2980). + +The scenarios that pool-backed and VMM buffers have in common run in the rest +of this directory through the ``VirtualMR`` parameter of the +``ipc_memory_resource`` fixture. This module covers what is specific to +virtual memory: ranges made of several physical allocations, which device the +memory belongs to, the validation of a descriptor from an untrusted peer, and +the limits of re-export. +""" import ctypes import gc @@ -10,14 +18,16 @@ import pickle import pytest -from helpers import IS_WINDOWS, IS_WSL -from helpers.child_processes import child_timeout_sec, kill_subprocesses, track_child_processes +from helpers.child_processes import child_timeout_sec, kill_subprocesses +from helpers.contexts import assert_no_cuda_warning from cuda.bindings import driver from cuda.core import Buffer, Device, MemoryResource, VirtualMemoryResource, VirtualMemoryResourceOptions -from cuda.core._memory._ipc import IPCBufferDescriptor +from cuda.core._dlpack import DLDeviceType +from cuda.core._memory._ipc import IPCAllocationHandle, IPCBufferDescriptor, VirtualMemoryIPCBufferDescriptor from cuda.core._memory._virtual_memory_resource import VirtualMemoryBuffer -from cuda.core._utils.cuda_utils import handle_return +from cuda.core._utils.cuda_utils import CUDAError, handle_return +from cuda.core.utils import StridedMemoryView CHILD_TIMEOUT_SEC = child_timeout_sec() @@ -25,19 +35,6 @@ pytestmark = pytest.mark.parallel_threads_limit(4) -@pytest.fixture -def vmm_ipc_device(init_cuda): - """A device that can export VMM allocations as POSIX file descriptors, or skip.""" - device = init_cuda - props = device.properties - if not props.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") - if IS_WINDOWS or IS_WSL or not props.handle_type_posix_file_descriptor_supported: - pytest.skip("Sharing virtual memory needs POSIX file descriptor handles") - with track_child_processes(): - yield device - - def _resource(device, **options): options.setdefault("handle_type", "posix_fd") return VirtualMemoryResource(device, config=VirtualMemoryResourceOptions(**options)) @@ -66,6 +63,43 @@ def _is_mapped(ptr): return False +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 _reserve_at(addr, size): + """Reserve ``size`` bytes at ``addr``. Return the reservation, or None if the driver placed it elsewhere.""" + ptr = handle_return(driver.cuMemAddressReserve(size, 0, addr, 0)) + if int(ptr) != addr: + handle_return(driver.cuMemAddressFree(ptr, size)) + return None + return ptr + + +def _open_fds(): + return len(os.listdir("/proc/self/fd")) + + +def _has_note(exc, text): + """Whether ``text`` is in one of the exception's notes, or in its message on Python 3.10.""" + return any(text in note for note in getattr(exc, "__notes__", ())) or text in str(exc) + + +def _forge(desc, **fields): + """A copy of ``desc`` with some fields replaced, as a hostile or buggy peer could send.""" + fields.setdefault("handle_type", desc._handle_type) + fields.setdefault("chunk_sizes", desc._chunk_sizes) + fields.setdefault("handles", desc._handles) + fields.setdefault("size", desc.size) + return VirtualMemoryIPCBufferDescriptor._from_exports( + fields["handle_type"], fields["chunk_sizes"], fields["handles"], fields["size"] + ) + + class _SharedStub(MemoryResource): """Reports IPC support without being a pool, to reach the base-class export path.""" @@ -95,6 +129,10 @@ def is_ipc_enabled(self): return True +class _CustomResource(VirtualMemoryResource): + """A subclass, to check that pickling preserves the type.""" + + def _run_child(target, *args): """Run ``target(*args)`` in a spawned child and fail if it does not exit cleanly.""" process = mp.Process(target=target, args=args) @@ -105,104 +143,374 @@ def _run_child(target, *args): assert process.exitcode == 0, f"child exited with {process.exitcode}" -class TestVmmIpcDescriptor: +@pytest.mark.agent_authored(model="claude-fable-5-1") +@pytest.mark.parametrize( + ("options", "expected"), + [ + 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({"location_type": "host_numa", "handle_type": "posix_fd"}, True, id="host_numa-posix_fd"), + ], +) +def test_is_ipc_enabled(vmm_ipc_device, options, expected): + """Only POSIX file descriptors have a transport; the location does not matter.""" + mr = VirtualMemoryResource(vmm_ipc_device, config=VirtualMemoryResourceOptions(**options)) + assert mr.is_ipc_enabled is expected + + +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_resource_pickles_as_its_own_type(vmm_ipc_device): + """A VirtualMemoryResource pickles as (device, options) and a subclass comes back as the subclass.""" + mr = _CustomResource(vmm_ipc_device, config=VirtualMemoryResourceOptions(handle_type="posix_fd", peers=())) + copy = pickle.loads(pickle.dumps(mr)) # noqa: S301 our own bytes; the resource carries no handle + assert type(copy) is _CustomResource + assert copy.device_id == mr.device_id + assert copy.config == mr.config + + +class TestManyChunkBuffer: @pytest.mark.agent_authored(model="claude-fable-5-1") - def test_descriptor_round_trip(self, vmm_ipc_device): - """A child imports the descriptor, reads the parent's data, writes its own, and releases.""" + @pytest.mark.thread_unsafe(reason="counts process-global file descriptors and records process-global warnings") + def test_main(self, vmm_ipc_device): + """A range of three physical allocations, the last placed by a relocating grow, exports and imports whole. + + Each allocation costs one file descriptor while the descriptor lives; + the child reads every chunk, writes one back, and tears the import + down without a warning. + """ device = vmm_ipc_device mr = _resource(device) - assert mr.is_ipc_enabled + queue = mp.Queue() # before the descriptor count: the queue's pipe is not part of it + buffers = [] + try: + buf = mr.allocate(1) + buffers.append(buf) + gran = buf.size + buf = mr.modify_allocation(buf, 2 * gran) # extends in place or moves + buffers.append(buf) + # Take the range after the buffer so that this grow must remap the + # existing chunks at a new address. + decoy = _reserve_at(int(buf.handle) + buf.size, gran) + if decoy is None: + pytest.skip("the driver did not grant a reservation right after the buffer") + try: + moved = mr.modify_allocation(buf, 3 * gran) + finally: + handle_return(driver.cuMemAddressFree(decoy, gran)) + buffers.append(moved) + assert int(moved.handle) != int(buf.handle) + assert moved.size == 3 * gran + values = (0x11, 0x22, 0x33) + for i, value in enumerate(values): + _fill(moved, value, offset=i * gran, size=gran) + + fds_before = _open_fds() + desc = moved.ipc_descriptor + assert desc.size == 3 * gran + assert desc._chunk_sizes == (gran, gran, gran) + assert _open_fds() == fds_before + 3 + fds = [int(handle) for handle in desc._handles] + + _run_child(self.child_main, device.device_id, desc, gran, queue) + result = queue.get(timeout=CHILD_TIMEOUT_SEC) + assert result["chunks"] == [bytes([value]) * 16 for value in values] + assert result["released"] is True + assert _read(moved, 2 * gran, 16) == b"\x44" * 16 + finally: + for buffer in buffers: + buffer.close() + # Closing the buffer drops its cached descriptor; the last reference goes here. + del desc + gc.collect() + for fd in fds: + with pytest.raises(OSError): + os.fstat(fd) + + @staticmethod + def child_main(device_id, desc, gran, queue): + device = Device(device_id) + device.set_current() + mr = _resource(device) + with assert_no_cuda_warning(): + imported = Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) + assert isinstance(imported, VirtualMemoryBuffer) + assert imported.is_mapped + assert imported.size == 3 * gran + chunks = [_read(imported, i * gran, 16) for i in range(3)] + _fill(imported, 0x44, offset=2 * gran, size=gran) + ptr = int(imported.handle) + imported.close() + released = not _is_mapped(ptr) + queue.put({"chunks": chunks, "released": released}) + + +class TestImportOnSecondDevice: + """The memory belongs to the exporter's device; the importer maps it for that device and may add peers.""" + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_main(self, vmm_ipc_device_x2): + dev0, dev1 = vmm_ipc_device_x2 + dev1.set_current() + mr = _resource(dev1) buf = mr.allocate(1) - _fill(buf, 0xAB) + _fill(buf, 0x5A) desc = buf.ipc_descriptor - assert isinstance(desc, IPCBufferDescriptor) - assert desc.size == buf.size - assert buf.ipc_descriptor is desc # cached - assert not buf.is_mapped + # In this process: a resource for the other device refuses the memory before mapping anything. + with pytest.raises(ValueError, match=f"memory on device {dev1.device_id}, but this resource maps memory on"): + Buffer.from_ipc_descriptor(_resource(dev0), desc, stream=dev0.default_stream) + + peers_possible = dev1.can_access_peer(dev0) queue = mp.Queue() - _run_child(self.child_main, device.device_id, desc, queue) - reads = queue.get(timeout=CHILD_TIMEOUT_SEC) - assert reads["is_mapped"] is True - assert reads["size"] == buf.size - assert reads["first_bytes"] == b"\xab" * 16 - assert reads["released"] is True - assert _read(buf, 0, 16) == b"\x5a" * 16 + _run_child(self.child_main, dev0.device_id, dev1.device_id, desc, peers_possible, queue) + result = queue.get(timeout=CHILD_TIMEOUT_SEC) + assert f"device {dev0.device_id}" in result["wrong_device"] + assert result["device_id"] == dev1.device_id + assert result["access_without_peers"] == (3, 0) + assert result["data"] == b"\x5a" * 16 + if peers_possible: + assert result["access_with_peers"] == (3, 3) + assert result["peer_read"] == b"\x5a" * 16 + else: + assert "access_with_peers" not in result + buf.close() + dev1.sync() @staticmethod - def child_main(device_id, desc, queue): - device = Device(device_id) - device.set_current() - mr = _resource(device) - imported = Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) - assert isinstance(imported, VirtualMemoryBuffer) - result = { - "is_mapped": imported.is_mapped, - "size": imported.size, - "first_bytes": _read(imported, 0, 16), - } - _fill(imported, 0x5A) - ptr = int(imported.handle) + def child_main(id0, id1, desc, peers_possible, queue): + dev0, dev1 = Device(id0), Device(id1) + dev1.set_current() + result = {} + + # A resource for the wrong device: refused with the device ids named. + with pytest.raises(ValueError) as info: + Buffer.from_ipc_descriptor(_resource(dev0), desc, stream=dev0.default_stream) + result["wrong_device"] = str(info.value) + + # The owning device, without peers: only that device can reach the memory. + imported = Buffer.from_ipc_descriptor(_resource(dev1), desc, stream=dev1.default_stream) + result["device_id"] = imported.device_id + result["access_without_peers"] = (_access_flags(imported, id1), _access_flags(imported, id0)) + result["data"] = _read(imported, 0, 16) imported.close() - result["released"] = not _is_mapped(ptr) + + # With the other device as a peer, it can read the import directly. + if peers_possible: + imported = Buffer.from_ipc_descriptor(_resource(dev1, peers=[id0]), desc, stream=dev1.default_stream) + result["access_with_peers"] = (_access_flags(imported, id1), _access_flags(imported, id0)) + dev0.set_current() + result["peer_read"] = _read(imported, 0, 16) + dev1.set_current() + imported.close() queue.put(result) + +class TestDescriptorValidation: + """A descriptor is data from another process; each malformed field fails before memory is mapped.""" + @pytest.mark.agent_authored(model="claude-fable-5-1") - def test_buffer_pickles_through_queue(self, vmm_ipc_device): - """A buffer sent directly pickles as (resource, descriptor) and imports on arrival.""" + def test_rejected_before_import(self, vmm_ipc_device): device = vmm_ipc_device + stream = device.default_stream mr = _resource(device) - buf = mr.allocate(1) - _fill(buf, 0x11) - queue = mp.Queue() - queue.put(buf) - _run_child(self.child_buffer, device.device_id, queue) - assert _read(buf, 0, 16) == b"\x22" * 16 + with mr.allocate(1) as buf: + gran = buf.size + desc = buf.ipc_descriptor + fabric = int(driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC) + cases = [ + (_forge(desc, handle_type=fabric), ValueError, "handle type does not match"), + (_forge(desc, chunk_sizes=(gran, gran)), ValueError, "malformed"), + (_forge(desc, chunk_sizes=(0,), size=0), ValueError, "positive multiple"), + (_forge(desc, chunk_sizes=(gran + 1,)), ValueError, "positive multiple"), + (_forge(desc, size=2 * gran), ValueError, "exceeds the exported allocations"), + (_forge(desc, chunk_sizes=(), handles=(), size=gran), ValueError, "no exported allocations"), + ] + for forged, exc_type, text in cases: + with pytest.raises(exc_type, match=text): + Buffer.from_ipc_descriptor(mr, forged, stream=stream) + + # A handle that was closed on this side. + closed = IPCAllocationHandle._init(os.dup(int(desc._handles[0])), None) + closed.close() + with pytest.raises(ValueError, match="closed"): + Buffer.from_ipc_descriptor(mr, _forge(desc, handles=(closed,)), stream=stream) + + # A file descriptor that is not a CUDA allocation: the driver refuses it, and the error says where. + read_end, write_end = os.pipe() + os.close(write_end) + not_cuda = IPCAllocationHandle._init(read_end, None) # owns and closes read_end + with pytest.raises(CUDAError) as info: + Buffer.from_ipc_descriptor(mr, _forge(desc, handles=(not_cuda,)), stream=stream) + assert _has_note(info.value, "while importing chunk 1 of 1") + assert _has_note(info.value, "handle_type='posix_fd'") + + # The buffer is untouched and the genuine descriptor still imports. + _fill(buf, 0x66) + alias = Buffer.from_ipc_descriptor(mr, desc, stream=stream) + assert _read(alias, 0, 16) == b"\x66" * 16 + alias.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="checks mapping state by address; another thread can reuse a freed address") + def test_wrong_chunk_size_fails_at_map_and_rolls_back(self, vmm_ipc_device): + """A chunk size that does not match the allocation fails in the driver's map; earlier chunks are undone. + + The importer is given an address hint so that the rollback is + observable: after the failure nothing is mapped at that address. + """ + device = vmm_ipc_device + stream = device.default_stream + mr = _resource(device) + first = mr.allocate(1) + gran = first.size + buf = mr.modify_allocation(first, 2 * gran) + desc = buf.ipc_descriptor + assert desc._chunk_sizes == (gran, gran) + + # An address the importer will use for a three-chunk range. + probe = int(handle_return(driver.cuMemAddressReserve(3 * gran, 0, 0, 0))) + handle_return(driver.cuMemAddressFree(probe, 3 * gran)) + importer = _resource(device, addr_hint=probe) + alias = Buffer.from_ipc_descriptor(importer, desc, stream=stream) + honored = int(alias.handle) == probe + alias.close() + if not honored: + first.close() + buf.close() + pytest.skip("the driver did not honor the address hint") + + forged = _forge(desc, chunk_sizes=(gran, 2 * gran), size=3 * gran) + with pytest.raises(CUDAError) as info: + Buffer.from_ipc_descriptor(importer, forged, stream=stream) + assert _has_note(info.value, f"while mapping chunk 2 of 2 ({2 * gran} bytes)") + assert not _is_mapped(probe) # chunk 1 was mapped there and has been undone + reservation = _reserve_at(probe, 3 * gran) + assert reservation is not None, "the importer's address reservation was not released" + handle_return(driver.cuMemAddressFree(reservation, 3 * gran)) + + # The exporter's memory is intact and the genuine descriptor still imports. + _fill(buf, 0x77, offset=gran, size=gran) + alias = Buffer.from_ipc_descriptor(importer, desc, stream=stream) + assert _read(alias, gran, 16) == b"\x77" * 16 + alias.close() + first.close() buf.close() - @staticmethod - def child_buffer(device_id, queue): - Device(device_id).set_current() - imported = queue.get(timeout=CHILD_TIMEOUT_SEC) - assert isinstance(imported, VirtualMemoryBuffer) + +class TestReexport: + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_imported_buffer_cannot_be_exported(self, vmm_ipc_device): + """The driver exports only allocations it created with the handle type; imports and grows from them fail.""" + device = vmm_ipc_device + stream = device.default_stream + mr = _resource(device) + buf = mr.allocate(1) + gran = buf.size + desc = buf.ipc_descriptor + imported = Buffer.from_ipc_descriptor(mr, desc, stream=stream) assert imported.is_mapped - assert isinstance(imported.memory_resource, VirtualMemoryResource) - assert imported.memory_resource.device_id == device_id - assert _read(imported, 0, 16) == b"\x11" * 16 - _fill(imported, 0x22) - imported.close() + with pytest.raises(RuntimeError, match="cannot be exported again"): + _ = imported.ipc_descriptor + with pytest.raises(RuntimeError, match="cannot be exported again"): + mp.reduction.ForkingPickler.dumps(imported) # a buffer pickles through its descriptor + + # A buffer grown from an import maps the imported chunk, so it cannot be exported either. + grown = mr.modify_allocation(imported, 2 * gran) + assert not grown.is_mapped + with pytest.raises(RuntimeError, match="imported from another process"): + _ = grown.ipc_descriptor + + # The original descriptor still serves new importers. + _fill(buf, 0x3C) + alias = Buffer.from_ipc_descriptor(mr, desc, stream=stream) + assert _read(alias, 0, 16) == b"\x3c" * 16 + for buffer in (alias, grown, imported, buf): + buffer.close() @pytest.mark.agent_authored(model="claude-fable-5-1") - def test_grown_buffer_imports_as_one_range(self, vmm_ipc_device): - """A buffer with two physical allocations exports both and imports contiguously.""" + def test_grow_cannot_change_handle_type(self, vmm_ipc_device): + """Every chunk of a buffer must be exportable the same way.""" device = vmm_ipc_device mr = _resource(device) - first = mr.allocate(1) - gran = first.size - grown = mr.modify_allocation(first, 2 * gran) - assert grown.size == 2 * gran - _fill(grown, 0x11, offset=0, size=gran) - _fill(grown, 0x22, offset=gran, size=gran) - desc = grown.ipc_descriptor - assert desc.size == 2 * gran - assert len(desc._chunk_sizes) == 2 + with mr.allocate(1) as buf, pytest.raises(ValueError, match="handle_type"): + mr.modify_allocation(buf, 2 * buf.size, config=VirtualMemoryResourceOptions(handle_type=None)) + private = _resource(device, handle_type=None) + with private.allocate(1) as buf, pytest.raises(ValueError, match="handle_type"): + private.modify_allocation(buf, 2 * buf.size, config=VirtualMemoryResourceOptions()) + +class TestEmptyBuffer: + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_round_trip(self, vmm_ipc_device): + """An empty buffer exports no allocations and imports as an empty mapped buffer.""" + device = vmm_ipc_device + mr = _resource(device) + empty = mr.allocate(0) + desc = empty.ipc_descriptor + assert desc.size == 0 + assert desc._chunk_sizes == () + assert desc._handles == () + alias = Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) + assert alias.size == 0 + assert alias.is_mapped + alias.close() queue = mp.Queue() - _run_child(self.child_grown, device.device_id, desc, gran, queue) - assert queue.get(timeout=CHILD_TIMEOUT_SEC) == {"low": b"\x11" * 16, "high": b"\x22" * 16} - first.close() - grown.close() + _run_child(self.child_main, device.device_id, desc, queue) + assert queue.get(timeout=CHILD_TIMEOUT_SEC) == {"size": 0, "is_mapped": True} + empty.close() @staticmethod - def child_grown(device_id, desc, gran, queue): + def child_main(device_id, desc, queue): device = Device(device_id) device.set_current() imported = Buffer.from_ipc_descriptor(_resource(device), desc, stream=device.default_stream) - assert imported.size == 2 * gran - queue.put({"low": _read(imported, 0, 16), "high": _read(imported, gran, 16)}) + queue.put({"size": imported.size, "is_mapped": imported.is_mapped}) + imported.close() + + +class TestDLPackConsumer: + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_imported_buffer_feeds_a_consumer(self, vmm_ipc_device): + """An imported buffer, received as a Process argument, is a complete DLPack producer.""" + device = vmm_ipc_device + mr = _resource(device) + buf = mr.allocate(1) + _fill(buf, 0x7E) + queue = mp.Queue() + _run_child(self.child_main, device.device_id, buf, queue) + result = queue.get(timeout=CHILD_TIMEOUT_SEC) + assert result["dl_device"] == (int(DLDeviceType.kDLCUDA), device.device_id) + assert result["shape"] == (buf.size,) + assert result["device_id"] == device.device_id + assert result["is_device_accessible"] is True + assert result["ptr_matches"] is True + assert result["first_bytes"] == b"\x7e" * 16 + buf.close() + + @staticmethod + def child_main(device_id, imported, queue): + Device(device_id).set_current() + assert isinstance(imported, VirtualMemoryBuffer) + assert imported.is_mapped + view = StridedMemoryView.from_dlpack(imported, stream_ptr=-1) + result = { + "dl_device": tuple(imported.__dlpack_device__()), + "shape": tuple(view.shape), + "device_id": view.device_id, + "is_device_accessible": view.is_device_accessible, + "ptr_matches": view.ptr == int(imported.handle), + "first_bytes": _read(imported, 0, 16), + } + del view imported.close() + queue.put(result) + +class TestSameProcess: @pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.thread_unsafe(reason="checks mapping state by address; another thread can reuse a freed address") def test_import_in_same_process_aliases_memory(self, vmm_ipc_device): @@ -229,7 +537,7 @@ def test_import_in_same_process_aliases_memory(self, vmm_ipc_device): reason="file descriptor numbers are process-global; another thread can reuse a closed one" ) def test_descriptor_owns_its_file_descriptors(self, vmm_ipc_device): - """The exported file descriptors close when the descriptor is released.""" + """The exported file descriptors close when the descriptor is released, and plain pickle refuses them.""" device = vmm_ipc_device mr = _resource(device) buf = mr.allocate(1) @@ -238,6 +546,8 @@ def test_descriptor_owns_its_file_descriptors(self, vmm_ipc_device): os.fstat(fd) # open with pytest.raises(TypeError): pickle.dumps(desc) # file descriptors need multiprocessing + with pytest.raises(TypeError): + pickle.dumps(buf) buf.close() # drops the cached descriptor del desc gc.collect() diff --git a/cuda_core/tests/memory_ipc/test_workerpool.py b/cuda_core/tests/memory_ipc/test_workerpool.py index f2a5fb1d50c..da60582ac87 100644 --- a/cuda_core/tests/memory_ipc/test_workerpool.py +++ b/cuda_core/tests/memory_ipc/test_workerpool.py @@ -7,9 +7,8 @@ import pytest from helpers.buffers import PatternGen -from helpers.constants import POOL_SIZE -from cuda.core import Buffer, Device, DeviceMemoryResource, DeviceMemoryResourceOptions +from cuda.core import Buffer, Device, VirtualMemoryResource NBYTES = 64 NWORKERS = 2 @@ -31,10 +30,9 @@ class TestIpcWorkerPool: @pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("nmrs", (1, NMRS)) - def test_main(self, ipc_device, nmrs): + def test_main(self, ipc_device, ipc_memory_resource_factory, nmrs): device = ipc_device - options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) - mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)] + mrs = [ipc_memory_resource_factory() for _ in range(nmrs)] buffers = [] stream = device.default_stream @@ -52,8 +50,6 @@ def test_main(self, ipc_device, nmrs): for buffer in buffers: buffer.close() stream.sync() - for mr in mrs: - mr.close() def process_buffer(self, buffer): device = Device(buffer.memory_resource.device_id) @@ -80,10 +76,9 @@ def init_worker(mrs): @pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("nmrs", (1, NMRS)) - def test_main(self, ipc_device, nmrs): + def test_main(self, ipc_device, ipc_memory_resource_factory, nmrs): device = ipc_device - options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) - mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)] + mrs = [ipc_memory_resource_factory() for _ in range(nmrs)] buffers = [] stream = device.default_stream @@ -104,8 +99,6 @@ def test_main(self, ipc_device, nmrs): for buffer in buffers: buffer.close() stream.sync() - for mr in mrs: - mr.close() def process_buffer(self, mr_idx, buffer_desc): mr = self.mrs[mr_idx] @@ -137,10 +130,11 @@ def init_worker(mrs): @pytest.mark.flaky(reruns=2) @pytest.mark.parametrize("nmrs", (1, NMRS)) - def test_main(self, ipc_device, nmrs): + def test_main(self, ipc_device, ipc_memory_resource_factory, nmrs): device = ipc_device - options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True) - mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)] + mrs = [ipc_memory_resource_factory() for _ in range(nmrs)] + if isinstance(mrs[0], VirtualMemoryResource): + pytest.skip("plain pickle cannot carry the file descriptors of a VirtualMemoryResource buffer") buffers = [] stream = device.default_stream @@ -158,8 +152,6 @@ def test_main(self, ipc_device, nmrs): for buffer in buffers: buffer.close() stream.sync() - for mr in mrs: - mr.close() def process_buffer(self, device, buffer_s): device.set_current() From b804e18cf9af825101bbdba633340e6adb1eee11 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Tue, 6 Oct 2026 15:44:44 -0700 Subject: [PATCH 05/11] cuda.core: no descriptor cache, no re-export, raw-fd transport for VMM IPC Apply the decisions on the VirtualMemoryResource IPC design. No exporter-side cache. VirtualMemoryBuffer.ipc_descriptor exports on every access and returns a new descriptor that owns its file descriptors; they close when the descriptor is released. Pickling a buffer for multiprocessing builds a transient descriptor, so a sender holds no file descriptors between sends. A spawned Process is the one exception: it receives the file descriptors when it is created, after pickling, so the transient descriptor stays on the spawning Popen until the Process object is released. A descriptor pins the physical memory while it exists, and nothing else does. No re-export of imported memory. Every physical allocation records whether cuMemImportFromShareableHandle produced it (mem_allocation_is_imported), and modify_allocation reuses the allocation handles of its input, so the mark reaches every alias, in-place grow, and move. Exporting such a buffer, directly or through pickling, raises a RuntimeError before any driver call that says the buffer contains memory imported from another process, that it cannot be exported again, and that the caller should forward the descriptor it imported from or copy into a buffer it owns. A Queue feeder thread reports the same error and drops the item. Raw-fd transport. VirtualMemoryIPCBufferDescriptor exposes handle_type, sizes, handles, and fds, and from_fds() rebuilds a descriptor in the receiving process from file descriptors received another way (the integers are duplicated; IPCAllocationHandle objects are shared). The descriptor's own pickling stays multiprocessing-only. No ipc_enabled option: handle_type is the switch and is_ipc_enabled the query, and the docs say so. Tests: file descriptor counts return to baseline on both sides, export under a lowered RLIMIT_NOFILE, the re-export errors for imported and derived buffers (direct, Pipe, Process arguments, Queue feeder thread), and raw-fd round trips through os.dup and a Unix socket with SCM_RIGHTS. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 55 ++- cuda_core/cuda/core/_cpp/rt/api.hpp | 7 + .../cuda/core/_cpp/rt/virtual_memory.cpp | 15 +- cuda_core/cuda/core/_memory/_buffer.pyi | 6 +- cuda_core/cuda/core/_memory/_buffer.pyx | 6 +- cuda_core/cuda/core/_memory/_ipc.pyi | 88 +++- cuda_core/cuda/core/_memory/_ipc.pyx | 145 +++++- .../core/_memory/_virtual_memory_resource.pyi | 81 ++-- .../core/_memory/_virtual_memory_resource.pyx | 146 +++--- cuda_core/cuda/core/_rt.pxd | 1 + cuda_core/cuda/core/_rt.pyx | 2 + cuda_core/docs/source/release/1.3.0-notes.rst | 10 +- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 445 ++++++++++++++++-- 13 files changed, 846 insertions(+), 161 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index 2c8e2fc4ab7..b9c0bdbd9b9 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -232,21 +232,46 @@ process yields an alias of the exporter's memory at a new address. A `VirtualMemoryResource` pickles as (device, options), which is what lets a `Buffer` pickle as (resource, descriptor). -File descriptors and memory. The driver copies each shared allocation into a -handle of its own during the import and keeps no reference to the file -descriptor, so the imported buffer does not keep the descriptor; it records -only that it was imported (`is_mapped`). The descriptor's file descriptors live -as long as the descriptor object, on the exporter (where the buffer caches it) -and in every process that unpickled a copy, and each open descriptor keeps the -physical memory allocated. A descriptor on a `Queue` pins the memory in the -sender until the receiver unpickles it. - -Re-export. The driver exports only allocations created with the requested -handle type, and an imported allocation carries none, so an imported buffer, -and any buffer grown from one, cannot be exported again; `ipc_descriptor` says -so instead of surfacing the driver's INVALID_VALUE. A `modify_allocation` -config must keep the resource's `handle_type` for the same reason: every chunk -of a buffer must be exportable the same way. +File descriptors and memory. A descriptor pins the physical memory while it +exists, in every process that holds one, and nothing else does. The exporting +buffer keeps no descriptor: `ipc_descriptor` exports again on every access and +returns a new descriptor that owns one file descriptor per allocation, and +`VirtualMemoryBuffer.__reduce__` builds a transient one while `multiprocessing` +pickles the buffer, so a sender holds no file descriptors between sends. One +transport needs more: a `Queue`, `Pipe`, or `Pool` duplicates the file +descriptors while pickling (the resource sharer), but a spawned `Process` +records their numbers during pickling and passes them to the child only when +it is created (`spawnv_passfds`), so `__reduce__` keeps the transient +descriptor on the spawning `Popen` object, which dies with the `Process` +object. Without that, the numbers are reused by the spawn's own pipes +("bad value(s) in fds_to_keep") or the child imports a stranger's fd. The +driver copies +each shared allocation into a handle of its own during the import and keeps no +reference to the file descriptor, so the imported buffer does not keep the +descriptor either; it records only that it was imported (`is_mapped`). A +descriptor parked in a `Queue` or sent to a `Pool` pins the memory in the +sender until the receiver unpickles it. Besides `multiprocessing`, a process +with its own file descriptor passing sends the descriptor's `fds`, `sizes`, +`handle_type`, and `size` and rebuilds it with +`VirtualMemoryIPCBufferDescriptor.from_fds()`, which duplicates the integers +it is given. + +Re-export. The driver exports only allocations this process created with the +requested handle type; an imported allocation has none, and the driver reports +INVALID_VALUE. cuda.core does not rely on that: `MemAllocationBox.imported` +records provenance per allocation (`mem_allocation_is_imported`), and since +`modify_allocation` reuses the allocation handles of the input, the mark +travels into every alias, in-place grow, and move. `_export_ipc_descriptor` +checks the allocations that cover the buffer before any driver call and raises +a RuntimeError that says the buffer contains memory imported from another +process, that such memory cannot be exported again, and that the caller should +forward the descriptor it imported from or copy into a buffer it owns. The same +exception surfaces from `pickle`/`multiprocessing` (a `Queue` feeder thread +reports it through `Queue._on_queue_feeder_error` and drops the item; the queue +stays usable). PyTorch applies the same contract to received CUDA tensors +("Attempted to send CUDA tensor received from another process"). There is no +retention option. A `modify_allocation` config must keep the resource's +`handle_type`: every chunk of a buffer must be exportable the same way. ## Why `modify_allocation` returns a new buffer diff --git a/cuda_core/cuda/core/_cpp/rt/api.hpp b/cuda_core/cuda/core/_cpp/rt/api.hpp index fb9ffcd85f3..cd9a0c45bb7 100644 --- a/cuda_core/cuda/core/_cpp/rt/api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/api.hpp @@ -262,6 +262,13 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat // Size of the allocation; the only size cuMemMap accepts for it. size_t mem_allocation_size(const MemAllocationHandle& h) noexcept; +// Whether cuMemImportFromShareableHandle produced the allocation. The driver +// exports only allocations this process created with a handle type, so an +// imported allocation cannot be exported again; callers check this before +// cuMemExportToShareableHandle instead of relying on the driver's error. The +// mark travels with the allocation into every range that maps it. +bool mem_allocation_is_imported(const MemAllocationHandle& h) noexcept; + // Import a physical allocation that another process exported with // cuMemExportToShareableHandle. `os_handle` is the shareable handle in the // form cuMemImportFromShareableHandle expects for `handle_type` (a POSIX file diff --git a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp index 505111a9743..39d6219b975 100644 --- a/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp @@ -28,9 +28,10 @@ using namespace detail; namespace { struct MemAllocationBox { - MemAllocationValue resource; // from cuMemCreate + MemAllocationValue resource; // from cuMemCreate or cuMemImportFromShareableHandle size_t size; // the only size cuMemMap accepts for it std::vector access; // applied to every mapping of this allocation + bool imported; // another process created it; it cannot be exported again }; struct VaReservationBox { @@ -56,9 +57,9 @@ const Box* box_of(const Handle& h) noexcept { // cuMemCreate or cuMemImportFromShareableHandle produced it. The last // reference releases it; the memory is freed once no mapping remains. MemAllocationHandle wrap_mem_allocation(CUmemGenericAllocationHandle handle, size_t size, - std::vector&& access) { + std::vector&& access, bool imported) { auto box = std::shared_ptr( - new MemAllocationBox{{handle}, size, std::move(access)}, + new MemAllocationBox{{handle}, size, std::move(access), imported}, [](const MemAllocationBox* b) { GILReleaseGuard gil; pw_cuMemRelease(b->resource.raw); @@ -84,7 +85,7 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemCreate, &handle, size, &prop, 0))) { return {}; } - return wrap_mem_allocation(handle, size, std::move(access)); + return wrap_mem_allocation(handle, size, std::move(access), false); } MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type, @@ -97,13 +98,17 @@ MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocatio if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemImportFromShareableHandle, &handle, os_handle, handle_type))) { return {}; } - return wrap_mem_allocation(handle, size, std::move(access)); + return wrap_mem_allocation(handle, size, std::move(access), true); } size_t mem_allocation_size(const MemAllocationHandle& h) noexcept { return h ? box_of(h)->size : 0; } +bool mem_allocation_is_imported(const MemAllocationHandle& h) noexcept { + return h ? box_of(h)->imported : false; +} + // ============================================================================ // Address reservations // ============================================================================ diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 3cc88e6b2ca..94429295948 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -109,7 +109,11 @@ class Buffer: """ @property def ipc_descriptor(self) -> IPCBufferDescriptor: - """Descriptor for sharing this buffer with other processes.""" + """Descriptor for sharing this buffer with other processes. + + A pool-backed buffer caches its descriptor; a + :class:`VirtualMemoryBuffer` exports a new one on every access. + """ def close(self, stream: Stream | GraphBuilder | None=None) -> None: """Deallocate this buffer asynchronously on the given stream. diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index c63d3e6eb42..db92f077cac 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -421,7 +421,11 @@ cdef class Buffer: @property @cython.critical_section def ipc_descriptor(self) -> IPCBufferDescriptor: - """Descriptor for sharing this buffer with other processes.""" + """Descriptor for sharing this buffer with other processes. + + A pool-backed buffer caches its descriptor; a + :class:`VirtualMemoryBuffer` exports a new one on every access. + """ Buffer_check_open(self) cdef object ipc_data if self._ipc_data is None: diff --git a/cuda_core/cuda/core/_memory/_ipc.pyi b/cuda_core/cuda/core/_memory/_ipc.pyi index 23934ed9903..a85c15363de 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyi +++ b/cuda_core/cuda/core/_memory/_ipc.pyi @@ -1,7 +1,9 @@ # This file was generated by stubgen-pyx v0.2.22 from cuda_core/cuda/core/_memory/_ipc.pyx import uuid -from typing import Any +from typing import Any, Iterable + +from cuda.core.typing import VirtualMemoryHandleType __all__ = [] @@ -60,18 +62,27 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. The descriptor holds one exported handle per physical allocation that - backs the buffer, in address order, with each allocation's size. With - POSIX file descriptors the handles are :class:`IPCAllocationHandle` - objects: the descriptor owns them and closes them when it is released, and - ``multiprocessing`` duplicates them into the receiving process. Plain - ``pickle`` cannot carry file descriptors; send the descriptor, or the - buffer, through ``multiprocessing`` (a ``Queue``, a ``Pipe``, ``Process`` - arguments, or a ``Pool``). - - A live descriptor costs one file descriptor per allocation and keeps the - physical memory allocated, in every process that holds a copy, until it is - released. A descriptor placed on a ``Queue`` or sent to a ``Pool`` pins the - memory in the sender until the receiver has unpickled it. + backs the buffer, in address order (``handles``), with each allocation's + size (``sizes``) and the buffer's size (``size``, which may be smaller + than their sum). With POSIX file descriptors the handles are + :class:`IPCAllocationHandle` objects: the descriptor owns them and closes + them when it is released. + + A descriptor pins the physical memory while it exists, in every process + that holds one, and nothing else does: the exporting buffer keeps no + descriptor, and an imported buffer does not keep the one it came from. A + descriptor parked in a ``multiprocessing`` ``Queue`` or sent to a ``Pool`` + pins the memory in the sender until the receiver has unpickled it; a + buffer passed as a ``Process`` argument pins it until the ``Process`` + object is released, because the child receives the file descriptors when + it is spawned. + + Two transports exist. ``multiprocessing`` (a ``Queue``, a ``Pipe``, + ``Process`` arguments, or a ``Pool``) pickles the descriptor, or a buffer, + and duplicates the file descriptors into the receiving process; plain + ``pickle`` cannot carry them and raises. A process with its own file + descriptor passing sends ``fds``, ``sizes``, ``handle_type``, and ``size`` + itself, and the receiver rebuilds the descriptor with :meth:`from_fds`. Note ---- @@ -86,6 +97,57 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): @staticmethod def _from_exports(handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple[Any, ...], size: int) -> VirtualMemoryIPCBufferDescriptor: ... + @classmethod + def from_fds(cls, fds: Iterable[int | IPCAllocationHandle], sizes: Iterable[int], *, handle_type: VirtualMemoryHandleType | str='posix_fd', size: int | None=None) -> VirtualMemoryIPCBufferDescriptor: + """Build a descriptor from file descriptors received outside ``multiprocessing``. + + The exporter sends ``fds``, ``sizes``, ``handle_type``, and ``size`` + of its descriptor through a transport of its own (a Unix socket with + ``SCM_RIGHTS``, for example); the receiver rebuilds the descriptor + here and imports it with :meth:`Buffer.from_ipc_descriptor`. + + Parameters + ---------- + fds : Iterable[int | IPCAllocationHandle] + One handle per physical allocation, in address order, as the + exporter's ``fds`` listed them. An ``int`` is duplicated: the + caller keeps ownership of the original and closes it. An + :class:`IPCAllocationHandle` is shared and closes with its last + holder. + sizes : Iterable[int] + The exporter's ``sizes``, one per handle. + handle_type : VirtualMemoryHandleType | str, optional + The exporter's ``handle_type``; ``"posix_fd"`` by default. + size : int, optional + The exporter's ``size``, at most the sum of ``sizes``; the sum by + default. + + Raises + ------ + ValueError + If the counts differ, a size is not positive, ``size`` exceeds the + sum of ``sizes``, or ``handle_type`` cannot be shared. + RuntimeError + If an :class:`IPCAllocationHandle` is closed. + OSError + 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 per allocation, in address order; the descriptor owns them.""" + @property + def fds(self) -> tuple[int, ...]: + """The exported file descriptors as integers, for a transport of your own. + + They belong to this descriptor: send or duplicate them while it is + alive. Each one is released when the descriptor is. + """ def __reduce__(self) -> tuple[object, ...]: ... def _reduce_allocation_handle(alloc_handle: IPCAllocationHandle) -> tuple[object, ...]: ... diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index 87206a02409..81251bc595a 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -2,6 +2,8 @@ # # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + cimport cpython from libc.stddef cimport size_t @@ -23,10 +25,15 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN from cuda.core._utils.cuda_utils import check_multiprocessing_start_method import multiprocessing +import operator import os import platform import uuid import weakref +from typing import TYPE_CHECKING, Iterable + +if TYPE_CHECKING: + from cuda.core.typing import VirtualMemoryHandleType __all__ = [] @@ -174,18 +181,27 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. The descriptor holds one exported handle per physical allocation that - backs the buffer, in address order, with each allocation's size. With - POSIX file descriptors the handles are :class:`IPCAllocationHandle` - objects: the descriptor owns them and closes them when it is released, and - ``multiprocessing`` duplicates them into the receiving process. Plain - ``pickle`` cannot carry file descriptors; send the descriptor, or the - buffer, through ``multiprocessing`` (a ``Queue``, a ``Pipe``, ``Process`` - arguments, or a ``Pool``). - - A live descriptor costs one file descriptor per allocation and keeps the - physical memory allocated, in every process that holds a copy, until it is - released. A descriptor placed on a ``Queue`` or sent to a ``Pool`` pins the - memory in the sender until the receiver has unpickled it. + backs the buffer, in address order (``handles``), with each allocation's + size (``sizes``) and the buffer's size (``size``, which may be smaller + than their sum). With POSIX file descriptors the handles are + :class:`IPCAllocationHandle` objects: the descriptor owns them and closes + them when it is released. + + A descriptor pins the physical memory while it exists, in every process + that holds one, and nothing else does: the exporting buffer keeps no + descriptor, and an imported buffer does not keep the one it came from. A + descriptor parked in a ``multiprocessing`` ``Queue`` or sent to a ``Pool`` + pins the memory in the sender until the receiver has unpickled it; a + buffer passed as a ``Process`` argument pins it until the ``Process`` + object is released, because the child receives the file descriptors when + it is spawned. + + Two transports exist. ``multiprocessing`` (a ``Queue``, a ``Pipe``, + ``Process`` arguments, or a ``Pool``) pickles the descriptor, or a buffer, + and duplicates the file descriptors into the receiving process; plain + ``pickle`` cannot carry them and raises. A process with its own file + descriptor passing sends ``fds``, ``sizes``, ``handle_type``, and ``size`` + itself, and the receiver rebuilds the descriptor with :meth:`from_fds`. Note ---- @@ -208,7 +224,112 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): self._handles = tuple(handles) return self + @classmethod + def from_fds( + cls, + fds: Iterable[int | IPCAllocationHandle], + sizes: Iterable[int], + *, + handle_type: VirtualMemoryHandleType | str = "posix_fd", + size: int | None = None, + ) -> VirtualMemoryIPCBufferDescriptor: + """Build a descriptor from file descriptors received outside ``multiprocessing``. + + The exporter sends ``fds``, ``sizes``, ``handle_type``, and ``size`` + of its descriptor through a transport of its own (a Unix socket with + ``SCM_RIGHTS``, for example); the receiver rebuilds the descriptor + here and imports it with :meth:`Buffer.from_ipc_descriptor`. + + Parameters + ---------- + fds : Iterable[int | IPCAllocationHandle] + One handle per physical allocation, in address order, as the + exporter's ``fds`` listed them. An ``int`` is duplicated: the + caller keeps ownership of the original and closes it. An + :class:`IPCAllocationHandle` is shared and closes with its last + holder. + sizes : Iterable[int] + The exporter's ``sizes``, one per handle. + handle_type : VirtualMemoryHandleType | str, optional + The exporter's ``handle_type``; ``"posix_fd"`` by default. + size : int, optional + The exporter's ``size``, at most the sum of ``sizes``; the sum by + default. + + Raises + ------ + ValueError + If the counts differ, a size is not positive, ``size`` exceeds the + sum of ``sizes``, or ``handle_type`` cannot be shared. + RuntimeError + If an :class:`IPCAllocationHandle` is closed. + OSError + If an ``int`` is not an open file descriptor. + """ + # Lazy import: this module is loaded while _virtual_memory_resource imports it. + from cuda.core._memory._virtual_memory_resource import VirtualMemoryResourceOptions + from cuda.core.typing import VirtualMemoryHandleType + + if handle_type is None: + raise ValueError("handle_type=None cannot be shared; the exporter's handle_type is required") + driver_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(VirtualMemoryHandleType(handle_type))) + cdef tuple sizes_t = tuple(operator.index(n) for n in sizes) + cdef list given = list(fds) + if len(given) != len(sizes_t): + raise ValueError(f"{len(given)} handles for {len(sizes_t)} sizes; one size per handle is required") + 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 operator.index(size) + if size < 0 or size > total: + raise ValueError(f"size {size} is outside 0..{total}, the sum of the allocation sizes") + cdef list handles = [] + cdef int dup + for item in given: + if isinstance(item, IPCAllocationHandle): + IPCAllocationHandle_check_open(item) + handles.append(item) + continue + dup = os.dup(operator.index(item)) + try: + handles.append(IPCAllocationHandle._init(dup, None)) + except: # noqa: E722 rollback-then-raise: the duplicate is not owned yet + os.close(dup) + raise + return VirtualMemoryIPCBufferDescriptor._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 + + @property + def handles(self) -> tuple[IPCAllocationHandle, ...]: + """The exported handles, one per allocation, in address order; the descriptor owns them.""" + return self._handles + + @property + def fds(self) -> tuple[int, ...]: + """The exported file descriptors as integers, for a transport of your own. + + They belong to this descriptor: send or duplicate them while it is + alive. Each one is released when the descriptor is. + """ + return tuple(int(handle) for handle in self._handles) + def __reduce__(self) -> tuple[object, ...]: + # multiprocessing duplicates the IPCAllocationHandle fds into the receiver; plain pickle raises. return VirtualMemoryIPCBufferDescriptor._from_exports, ( self._handle_type, self._chunk_sizes, self._handles, self._size) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index 727bf15c512..ade6e744067 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -6,8 +6,7 @@ from typing import Iterable from cuda.core._device import Device from cuda.core._memory._buffer import Buffer, MemoryResource -from cuda.core._memory._ipc import (IPCBufferDescriptor, - VirtualMemoryIPCBufferDescriptor) +from cuda.core._memory._ipc import VirtualMemoryIPCBufferDescriptor from cuda.core._stream import Stream from cuda.core._utils.cuda_utils import driver from cuda.core.graph import GraphBuilder @@ -33,7 +32,9 @@ class VirtualMemoryResourceOptions: location_type: :obj:`~_memory.VirtualMemoryLocationType` | str Controls the location of the allocation. handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str - Export handle type for the physical allocation. Use ``"posix_fd"`` on + 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`. @@ -95,22 +96,37 @@ class VirtualMemoryBuffer(Buffer): freed when the last buffer that maps it closes. """ @property - def ipc_descriptor(self) -> IPCBufferDescriptor: - """Descriptor for sharing this buffer with other processes. - - The descriptor exports every physical allocation that backs the - buffer as a shareable handle and is cached on the buffer until the - buffer closes. Each exported allocation takes one file descriptor, - and a live descriptor keeps the physical memory allocated in every - process that holds a copy. It requires a resource whose - ``handle_type`` can be shared; see - :attr:`VirtualMemoryResource.is_ipc_enabled`. Import it with + def ipc_descriptor(self) -> VirtualMemoryIPCBufferDescriptor: + """A new descriptor for sharing this buffer with other processes. + + Every access exports the physical allocations that back the buffer + again and returns a new descriptor that owns one file descriptor per + allocation; the file descriptors close when the descriptor is + released. A descriptor pins the physical memory while it exists, in + every process that holds one, and nothing else does: the buffer + itself keeps no descriptor, so it pins nothing beyond its own mapping. + Pickling the buffer through ``multiprocessing`` builds a transient + descriptor, so the sender holds no file descriptors between sends; + a buffer passed as a ``Process`` argument is the exception, because + the child receives the file descriptors when it is spawned, after + pickling, so that descriptor lives until the ``Process`` object is + released. Requires a resource whose ``handle_type`` can be shared; see + :attr:`VirtualMemoryResource.is_ipc_enabled`. Import with :meth:`Buffer.from_ipc_descriptor` and a :class:`VirtualMemoryResource` of the receiving process, or send the - buffer itself through ``multiprocessing``. A buffer imported from - another process cannot be exported again; share the descriptor it - was imported from. + buffer itself. + + Raises + ------ + RuntimeError + If the resource is not IPC-enabled, or the buffer contains memory + imported from another process (an import, or a buffer + :meth:`VirtualMemoryResource.modify_allocation` derived from + one): such memory cannot be exported again. Forward the descriptor + it was imported from, or copy the data into a buffer this process + allocated. The check happens before any driver call. """ + def __reduce__(self) -> tuple[object, ...]: ... def close(self, stream: Stream | GraphBuilder | None=None) -> None: """Release this buffer's share of its address range. @@ -154,7 +170,8 @@ 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. + 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. Notes ----- @@ -169,16 +186,23 @@ class VirtualMemoryResource(MemoryResource): :meth:`Buffer.set_deallocation_stream`. Buffers can be shared with other processes when ``config.handle_type`` - is ``"posix_fd"`` (Linux). :attr:`Buffer.ipc_descriptor` exports the - physical allocations that back a buffer, one file descriptor each; - :meth:`Buffer.from_ipc_descriptor` imports them 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. The importing resource - maps the memory for its device and the devices in its ``peers`` option, - with its own access options. A live descriptor keeps the memory allocated - in every process that holds a copy. An imported buffer, and a buffer grown - from one, cannot be exported again. See :attr:`is_ipc_enabled`. + is ``"posix_fd"`` (Linux); ``handle_type`` is the switch and + :attr:`is_ipc_enabled` the query. :attr:`Buffer.ipc_descriptor` exports + the physical allocations that back a buffer, one file descriptor each, + into a new descriptor on every 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``, ``handle_type``, 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 it exists, in every process that holds one, and + nothing else does. 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. """ device: object config: object @@ -305,7 +329,8 @@ 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. With + True when ``config.handle_type`` is ``"posix_fd"`` on Linux; + ``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. diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index acfc3bb39fd..a647b3a7c79 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -4,7 +4,6 @@ from __future__ import annotations -cimport cython from libc.stdint cimport uintptr_t from libc.string cimport memset from libcpp.vector cimport vector @@ -18,7 +17,6 @@ from cuda.core._memory._buffer cimport ( ) from cuda.core._memory._ipc cimport ( IPCAllocationHandle, - IPCBufferDescriptor, IPCDataForBuffer, VirtualMemoryIPCBufferDescriptor, ) @@ -41,6 +39,7 @@ from cuda.core._rt cimport ( get_last_error, get_primary_context, import_mem_allocation_handle, + mem_allocation_is_imported, mem_allocation_size, set_deallocation_stream, va_mapping_allocation, @@ -54,6 +53,7 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, check_or_create_options import os import platform from dataclasses import dataclass, field +from multiprocessing import context as _mp_context from typing import TYPE_CHECKING, Iterable from cuda.core._device import Device @@ -101,7 +101,9 @@ class VirtualMemoryResourceOptions: location_type: :obj:`~_memory.VirtualMemoryLocationType` | str Controls the location of the allocation. handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str - Export handle type for the physical allocation. Use ``"posix_fd"`` on + 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`. @@ -288,32 +290,53 @@ cdef class VirtualMemoryBuffer(Buffer): """ @property - @cython.critical_section - def ipc_descriptor(self) -> IPCBufferDescriptor: - """Descriptor for sharing this buffer with other processes. - - The descriptor exports every physical allocation that backs the - buffer as a shareable handle and is cached on the buffer until the - buffer closes. Each exported allocation takes one file descriptor, - and a live descriptor keeps the physical memory allocated in every - process that holds a copy. It requires a resource whose - ``handle_type`` can be shared; see - :attr:`VirtualMemoryResource.is_ipc_enabled`. Import it with + def ipc_descriptor(self) -> VirtualMemoryIPCBufferDescriptor: + """A new descriptor for sharing this buffer with other processes. + + Every access exports the physical allocations that back the buffer + again and returns a new descriptor that owns one file descriptor per + allocation; the file descriptors close when the descriptor is + released. A descriptor pins the physical memory while it exists, in + every process that holds one, and nothing else does: the buffer + itself keeps no descriptor, so it pins nothing beyond its own mapping. + Pickling the buffer through ``multiprocessing`` builds a transient + descriptor, so the sender holds no file descriptors between sends; + a buffer passed as a ``Process`` argument is the exception, because + the child receives the file descriptors when it is spawned, after + pickling, so that descriptor lives until the ``Process`` object is + released. Requires a resource whose ``handle_type`` can be shared; see + :attr:`VirtualMemoryResource.is_ipc_enabled`. Import with :meth:`Buffer.from_ipc_descriptor` and a :class:`VirtualMemoryResource` of the receiving process, or send the - buffer itself through ``multiprocessing``. A buffer imported from - another process cannot be exported again; share the descriptor it - was imported from. + buffer itself. + + Raises + ------ + RuntimeError + If the resource is not IPC-enabled, or the buffer contains memory + imported from another process (an import, or a buffer + :meth:`VirtualMemoryResource.modify_allocation` derived from + one): such memory cannot be exported again. Forward the descriptor + it was imported from, or copy the data into a buffer this process + allocated. The check happens before any driver call. """ Buffer_check_open(self) - if self._ipc_data is None: - self._ipc_data = IPCDataForBuffer(_export_ipc_descriptor(self), False) - elif self._ipc_data.is_mapped: - raise RuntimeError( - "a VirtualMemoryBuffer imported from another process cannot be exported again: the CUDA driver " - "does not export imported allocations. Share the descriptor it was imported from." - ) - return self._ipc_data.ipc_descriptor + return _export_ipc_descriptor(self) + + def __reduce__(self) -> tuple[object, ...]: + # Unpickling performs a live import from the descriptor; only deserialize + # Buffers from a trusted principal. Must not serialize the parent's stream. + Buffer_check_open(self) + desc = _export_ipc_descriptor(self) + popen = _mp_context.get_spawning_popen() + if popen is not None: + # A Queue, Pipe, or Pool duplicates the file descriptors while this + # transient descriptor is alive. A spawned Process does not: it + # records their numbers now and hands them to the child when it is + # created, after pickling has finished. Keep the descriptor on the + # Popen until then; the Popen dies with the Process object. + vars(popen).setdefault("_cuda_core_descriptors_until_spawned", []).append(desc) + return Buffer._reduce_helper, (self.memory_resource, desc) def close(self, stream: Stream | GraphBuilder | None = None) -> None: """Release this buffer's share of its address range. @@ -377,7 +400,8 @@ 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. + 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. Notes ----- @@ -392,16 +416,23 @@ cdef class VirtualMemoryResource(MemoryResource): :meth:`Buffer.set_deallocation_stream`. Buffers can be shared with other processes when ``config.handle_type`` - is ``"posix_fd"`` (Linux). :attr:`Buffer.ipc_descriptor` exports the - physical allocations that back a buffer, one file descriptor each; - :meth:`Buffer.from_ipc_descriptor` imports them 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. The importing resource - maps the memory for its device and the devices in its ``peers`` option, - with its own access options. A live descriptor keeps the memory allocated - in every process that holds a copy. An imported buffer, and a buffer grown - from one, cannot be exported again. See :attr:`is_ipc_enabled`. + is ``"posix_fd"`` (Linux); ``handle_type`` is the switch and + :attr:`is_ipc_enabled` the query. :attr:`Buffer.ipc_descriptor` exports + the physical allocations that back a buffer, one file descriptor each, + into a new descriptor on every 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``, ``handle_type``, 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 it exists, in every process that holds one, and + nothing else does. 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. """ cdef: @@ -854,7 +885,8 @@ 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. With + True when ``config.handle_type`` is ``"posix_fd"`` on Linux; + ``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. @@ -1026,14 +1058,17 @@ cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf): Only the allocations that cover ``buf.size`` are exported; a buffer's size is a prefix of its range. Each file descriptor is owned by an - :class:`IPCAllocationHandle` in the descriptor and closed with it. An - allocation this process imported cannot be exported: the driver exports - only allocations created with the requested handle type. + :class:`IPCAllocationHandle` in the new descriptor and closed with it; + nothing is cached on the buffer. A buffer that contains memory imported + from another process is refused before any driver call: the driver + exports only allocations this process created with the handle type, and + every allocation carries its provenance (``mem_allocation_is_imported``) + into every range that maps it. """ cdef VirtualMemoryResource mr = buf._memory_resource cdef vector[VaMappingHandle] mappings + cdef vector[MemAllocationHandle] allocs cdef MemAllocationHandle h_alloc - cdef cydriver.CUmemAllocationProp chunk_prop cdef cydriver.CUresult status cdef size_t i, n, chunk, offset = 0 cdef int handle_type @@ -1045,30 +1080,30 @@ cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf): raise RuntimeError("Memory resource is not IPC-enabled") handle_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(mr.config.handle_type)) mappings = vmm_range_mappings(vmm_range(buf._h_ptr)) - n = mappings.size() - for i in range(n): + for i in range(mappings.size()): if offset >= buf._size: break h_alloc = va_mapping_allocation(mappings[i]) - chunk = mem_allocation_size(h_alloc) - # The driver exports only allocations created with the requested - # handle type; an imported allocation carries none, so the export would - # fail with INVALID_VALUE. Say what the limitation is instead. - with nogil: - HANDLE_RETURN(cydriver.cuMemGetAllocationPropertiesFromHandle(&chunk_prop, as_cu(h_alloc))) - if (chunk_prop.requestedHandleTypes & handle_type) == 0: + allocs.push_back(h_alloc) + offset += mem_allocation_size(h_alloc) + n = allocs.size() + for i in range(n): + if mem_allocation_is_imported(allocs[i]): raise RuntimeError( - f"chunk {i + 1} of this buffer cannot be exported: it was imported from another process, and " - "the CUDA driver does not export imported allocations. Share the descriptor the memory was " - "imported from instead." + f"this buffer contains memory imported from another process (allocation {i + 1} of {n}); imported " + "memory cannot be exported again. Forward the descriptor it was imported from, or copy the data " + "into a buffer this process allocated." ) + for i in range(n): + h_alloc = allocs[i] + chunk = mem_allocation_size(h_alloc) with nogil: status = cydriver.cuMemExportToShareableHandle( &fd, as_cu(h_alloc), handle_type, 0) _check_noting( status, - f"while exporting chunk {i + 1} of {n} of a VirtualMemoryResource buffer; each exported chunk takes " - "one file descriptor, so a process near its file descriptor limit fails here", + f"while exporting allocation {i + 1} of {n} of a VirtualMemoryResource buffer; each exported " + "allocation takes one file descriptor, so a process near its file descriptor limit fails here", ) try: handles.append(IPCAllocationHandle._init(fd, None)) @@ -1076,5 +1111,4 @@ cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf): os.close(fd) raise sizes.append(chunk) - offset += chunk return VirtualMemoryIPCBufferDescriptor._from_exports(handle_type, tuple(sizes), tuple(handles), buf._size) diff --git a/cuda_core/cuda/core/_rt.pxd b/cuda_core/cuda/core/_rt.pxd index c99e6c3ce18..2b8014016ba 100644 --- a/cuda_core/cuda/core/_rt.pxd +++ b/cuda_core/cuda/core/_rt.pxd @@ -299,6 +299,7 @@ cdef MemAllocationHandle create_mem_allocation_handle( size_t size, const cydriver.CUmemAllocationProp& prop, const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil cdef size_t mem_allocation_size(const MemAllocationHandle& h) noexcept nogil +cdef bint mem_allocation_is_imported(const MemAllocationHandle& h) noexcept nogil cdef MemAllocationHandle import_mem_allocation_handle( void* os_handle, cydriver.CUmemAllocationHandleType handle_type, size_t size, const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil diff --git a/cuda_core/cuda/core/_rt.pyx b/cuda_core/cuda/core/_rt.pyx index b948a6bdafa..ae44b9ed52c 100644 --- a/cuda_core/cuda/core/_rt.pyx +++ b/cuda_core/cuda/core/_rt.pyx @@ -180,6 +180,8 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil size_t mem_allocation_size "cuda_core::rt::mem_allocation_size" ( const MemAllocationHandle& h) noexcept nogil + bint mem_allocation_is_imported "cuda_core::rt::mem_allocation_is_imported" ( + const MemAllocationHandle& h) noexcept nogil MemAllocationHandle import_mem_allocation_handle "cuda_core::rt::import_mem_allocation_handle" ( void* os_handle, cydriver.CUmemAllocationHandleType handle_type, size_t size, const cydriver.CUmemAccessDesc* descs, size_t count) except+ nogil 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 87a6599d3be..af61df694f4 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -87,9 +87,13 @@ New features 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. A grown buffer - imports as one contiguous range. A live descriptor holds one file descriptor - per physical allocation and keeps the memory allocated; an imported buffer - cannot be exported again. Fabric handles are not shared yet. + 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``, + ``handle_type``, and ``size``, and + ``VirtualMemoryIPCBufferDescriptor.from_fds()`` rebuilds it in the receiver. + Fabric handles are not shared yet. (`#2980 `__) Fixes and enhancements diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index 827d6dde2f4..25b22a870a8 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -13,9 +13,14 @@ import ctypes import gc +import json import multiprocessing as mp +import multiprocessing.queues import os import pickle +import resource +import socket +import threading import pytest from helpers.child_processes import child_timeout_sec, kill_subprocesses @@ -143,6 +148,31 @@ def _run_child(target, *args): assert process.exitcode == 0, f"child exited with {process.exitcode}" +def _noop(*args): + pass + + +def _close_argument(buffer): + buffer.close() + + +_REEXPORT_MESSAGE = "imported from another process" +_FEEDER_ERRORS: list = [] +_FEEDER_EVENT = threading.Event() + + +class _ReportingQueue(multiprocessing.queues.Queue): + """A Queue whose feeder-thread errors are recorded, not only printed to stderr.""" + + def __init__(self): + super().__init__(ctx=mp.get_context()) + + @staticmethod + def _on_queue_feeder_error(e, _obj): + _FEEDER_ERRORS.append((type(e), str(e))) + _FEEDER_EVENT.set() + + @pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.parametrize( ("options", "expected"), @@ -222,7 +252,7 @@ def test_main(self, vmm_ipc_device): finally: for buffer in buffers: buffer.close() - # Closing the buffer drops its cached descriptor; the last reference goes here. + # The descriptor is the only holder of the file descriptors; the buffer caches nothing. del desc gc.collect() for fd in fds: @@ -402,9 +432,27 @@ def test_wrong_chunk_size_fails_at_map_and_rolls_back(self, vmm_ipc_device): class TestReexport: + """Memory imported from another process cannot be exported again; the error is clear and early.""" + + @staticmethod + def _derived(mr, imported, gran): + """Buffers modify_allocation derives from an import: an alias, a grow, and (if the driver allows) a move.""" + derived = { + "alias": mr.modify_allocation(imported, imported.size), + "grow": mr.modify_allocation(imported, 2 * gran), + } + base = derived["grow"] + decoy = _reserve_at(int(base.handle) + base.size, gran) + if decoy is not None: + try: + derived["move"] = mr.modify_allocation(base, base.size + gran) + finally: + handle_return(driver.cuMemAddressFree(decoy, gran)) + assert int(derived["move"].handle) != int(base.handle) + return derived + @pytest.mark.agent_authored(model="claude-fable-5-1") - def test_imported_buffer_cannot_be_exported(self, vmm_ipc_device): - """The driver exports only allocations it created with the handle type; imports and grows from them fail.""" + def test_export_and_pickling_raise_before_any_driver_call(self, vmm_ipc_device): device = vmm_ipc_device stream = device.default_stream mr = _resource(device) @@ -413,36 +461,378 @@ def test_imported_buffer_cannot_be_exported(self, vmm_ipc_device): desc = buf.ipc_descriptor imported = Buffer.from_ipc_descriptor(mr, desc, stream=stream) assert imported.is_mapped - with pytest.raises(RuntimeError, match="cannot be exported again"): - _ = imported.ipc_descriptor - with pytest.raises(RuntimeError, match="cannot be exported again"): - mp.reduction.ForkingPickler.dumps(imported) # a buffer pickles through its descriptor - - # A buffer grown from an import maps the imported chunk, so it cannot be exported either. - grown = mr.modify_allocation(imported, 2 * gran) - assert not grown.is_mapped - with pytest.raises(RuntimeError, match="imported from another process"): - _ = grown.ipc_descriptor - - # The original descriptor still serves new importers. - _fill(buf, 0x3C) - alias = Buffer.from_ipc_descriptor(mr, desc, stream=stream) - assert _read(alias, 0, 16) == b"\x3c" * 16 - for buffer in (alias, grown, imported, buf): - buffer.close() + buffers = {"import": imported, **self._derived(mr, imported, gran)} + try: + for name, buffer in buffers.items(): + with pytest.raises(RuntimeError, match=_REEXPORT_MESSAGE) as info: + _ = buffer.ipc_descriptor + message = str(info.value) + assert "cannot be exported again" in message, name + assert "Forward the descriptor" in message and "copy the data" in message, name + with pytest.raises(RuntimeError, match=_REEXPORT_MESSAGE): + mp.reduction.ForkingPickler.dumps(buffer) + # The transports that pickle in the caller's thread raise to the caller. + parent_conn, child_conn = mp.Pipe() + with parent_conn, child_conn, pytest.raises(RuntimeError, match=_REEXPORT_MESSAGE): + parent_conn.send(imported) + with pytest.raises(RuntimeError, match=_REEXPORT_MESSAGE): + mp.Process(target=_noop, args=(imported,)).start() + # The original descriptor still serves new importers, and the data is intact. + _fill(buf, 0x3C) + alias = Buffer.from_ipc_descriptor(mr, desc, stream=stream) + assert _read(alias, 0, 16) == b"\x3c" * 16 + alias.close() + finally: + for buffer in reversed(list(buffers.values())): + buffer.close() + buf.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="records feeder-thread errors in module state") + def test_queue_feeder_reports_the_error_and_stays_usable(self, vmm_ipc_device): + """A Queue pickles in its feeder thread: the error is reported there, the item is dropped, and the queue works on.""" + device = vmm_ipc_device + stream = device.default_stream + mr = _resource(device) + buf = mr.allocate(1) + imported = Buffer.from_ipc_descriptor(mr, buf.ipc_descriptor, stream=stream) + _FEEDER_ERRORS.clear() + _FEEDER_EVENT.clear() + queue = _ReportingQueue() + try: + queue.put(imported) + assert _FEEDER_EVENT.wait(timeout=CHILD_TIMEOUT_SEC), "the feeder thread did not report the error" + ((exc_type, message),) = _FEEDER_ERRORS + assert exc_type is RuntimeError + assert _REEXPORT_MESSAGE in message + assert "cannot be exported again" in message + # A consumer is not left waiting: the next item arrives. + queue.put("next item") + assert queue.get(timeout=CHILD_TIMEOUT_SEC) == "next item" + finally: + queue.close() + queue.join_thread() + imported.close() + buf.close() @pytest.mark.agent_authored(model="claude-fable-5-1") def test_grow_cannot_change_handle_type(self, vmm_ipc_device): """Every chunk of a buffer must be exportable the same way.""" device = vmm_ipc_device mr = _resource(device) - with mr.allocate(1) as buf, pytest.raises(ValueError, match="handle_type"): + with ( + mr.allocate(1) as buf, + pytest.raises(ValueError, match="handle_type"), + ): mr.modify_allocation(buf, 2 * buf.size, config=VirtualMemoryResourceOptions(handle_type=None)) private = _resource(device, handle_type=None) - with private.allocate(1) as buf, pytest.raises(ValueError, match="handle_type"): + with ( + private.allocate(1) as buf, + pytest.raises(ValueError, match="handle_type"), + ): private.modify_allocation(buf, 2 * buf.size, config=VirtualMemoryResourceOptions()) +class TestFileDescriptorLifetime: + """A descriptor is the only holder of file descriptors, on both sides.""" + + @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="counts process-global file descriptors") + def test_exporter_holds_none_without_a_descriptor(self, vmm_ipc_device): + """Each access exports anew; the file descriptors close with the descriptor, not with the buffer.""" + device = vmm_ipc_device + mr = _resource(device) + first = mr.allocate(1) + gran = first.size + buf = mr.modify_allocation(first, 2 * gran) + fds0 = _open_fds() + desc = buf.ipc_descriptor + assert desc.sizes == (gran, gran) + assert _open_fds() == fds0 + 2 + again = buf.ipc_descriptor + assert again is not desc + assert set(again.fds).isdisjoint(desc.fds) + assert _open_fds() == fds0 + 4 + del again + gc.collect() + assert _open_fds() == fds0 + 2 + numbers = desc.fds + del desc + gc.collect() + assert _open_fds() == fds0, "the buffer alone must hold no file descriptors" + for fd in numbers: + with pytest.raises(OSError): + os.fstat(fd) + buf.close() + first.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="counts process-global file descriptors") + def test_process_argument_releases_with_the_process_object(self, vmm_ipc_device): + """A buffer passed to a spawned Process imports there; its descriptor goes when the Process object does.""" + device = vmm_ipc_device + mr = _resource(device) + buf = mr.allocate(1) + fds0 = _open_fds() + process = mp.Process(target=_close_argument, args=(buf,)) + process.start() + 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}" + del process + gc.collect() + assert _open_fds() == fds0 + buf.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_importer_holds_none_after_the_import(self, vmm_ipc_device): + """The import adds no file descriptors, and releasing the received descriptor closes its own.""" + device = vmm_ipc_device + mr = _resource(device) + first = mr.allocate(1) + gran = first.size + buf = mr.modify_allocation(first, 2 * gran) + _fill(buf, 0x5B) + to_child, from_child = mp.Queue(), mp.Queue() + process = mp.Process(target=self.child_main, args=(device.device_id, to_child, from_child)) + process.start() + to_child.put(buf.ipc_descriptor) # transient: nothing is kept on this side + counts = from_child.get(timeout=CHILD_TIMEOUT_SEC) + 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}" + assert counts["data"] == b"\x5b" * 16 + assert counts["after_import"] == counts["with_descriptor"] + assert counts["after_release"] == counts["with_descriptor"] - 2 + assert counts["after_close"] == counts["after_release"] + buf.close() + first.close() + + @staticmethod + def child_main(device_id, to_child, from_child): + device = Device(device_id) + device.set_current() + mr = _resource(device) + desc = to_child.get(timeout=CHILD_TIMEOUT_SEC) + # A first import lets the driver finish any lazy initialization before counting. + Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream).close() + counts = {"with_descriptor": _open_fds()} + imported = Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) + counts["after_import"] = _open_fds() + del desc + gc.collect() + counts["after_release"] = _open_fds() + counts["data"] = _read(imported, 0, 16) # the import, not the descriptor, holds the memory now + imported.close() + counts["after_close"] = _open_fds() + from_child.put(counts) + + +class TestFileDescriptorLimit: + @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="changes the process-wide file descriptor limit") + def test_export_near_the_limit(self, vmm_ipc_device): + """Under a lowered RLIMIT_NOFILE an export either succeeds or fails with a readable error and no leak.""" + device = vmm_ipc_device + stream = device.default_stream + mr = _resource(device) + first = mr.allocate(1) + gran = first.size + second = mr.modify_allocation(first, 2 * gran) + buf = mr.modify_allocation(second, 3 * gran) + values = (0x61, 0x62, 0x63) + for i, value in enumerate(values): + _fill(buf, value, offset=i * gran, size=gran) + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + fillers = [] + imported = desc = None + try: + # Fill the holes in the descriptor table so that the limit alone decides how many more fit. + highest = max(int(name) for name in os.listdir("/proc/self/fd")) + while True: + fd = os.open(os.devnull, os.O_RDONLY) + if fd > highest: + os.close(fd) + break + fillers.append(fd) + baseline = _open_fds() + # Room for two more descriptors: a three-allocation export cannot complete. + resource.setrlimit(resource.RLIMIT_NOFILE, (highest + 1 + 2, hard)) + failure = None + try: + desc = buf.ipc_descriptor + except CUDAError as exc: + failure = exc + if failure is not None: + assert _has_note(failure, "file descriptor limit") + assert _open_fds() == baseline, "a failed export must close the descriptors it already took" + else: + # Acceptable only if the driver did not need one descriptor per allocation. + assert len(desc.fds) == 3 + desc = None + gc.collect() + # Room for the three descriptors and the listing: export and import both succeed. + resource.setrlimit(resource.RLIMIT_NOFILE, (highest + 1 + 6, hard)) + desc = buf.ipc_descriptor + assert len(desc.fds) == 3 + imported = Buffer.from_ipc_descriptor(mr, desc, stream=stream) + assert _open_fds() == baseline + 3, "the import must not take file descriptors of its own" + finally: + resource.setrlimit(resource.RLIMIT_NOFILE, (soft, hard)) + for fd in fillers: + os.close(fd) + assert [_read(imported, i * gran, 16) for i in range(3)] == [bytes([value]) * 16 for value in values] + del desc + gc.collect() + for buffer in (imported, buf, second, first): + buffer.close() + + +class TestRawFileDescriptorTransport: + """The descriptor's fds, sizes, handle_type, and size travel over any file-descriptor-passing transport.""" + + @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="counts process-global file descriptors") + def test_round_trip_with_os_dup(self, vmm_ipc_device): + device = vmm_ipc_device + stream = device.default_stream + mr = _resource(device) + first = mr.allocate(1) + gran = first.size + buf = mr.modify_allocation(first, 2 * gran) + _fill(buf, 0x71, size=gran) + _fill(buf, 0x72, offset=gran, size=gran) + fds0 = _open_fds() + desc = buf.ipc_descriptor + assert desc.handle_type == "posix_fd" + assert desc.sizes == (gran, gran) + assert desc.size == 2 * gran + assert desc.fds == tuple(int(handle) for handle in desc.handles) + # A transport of our own: duplicate the integers, then let the descriptor go. + wire = { + "fds": [os.dup(fd) for fd in desc.fds], + "sizes": list(desc.sizes), + "handle_type": str(desc.handle_type), + "size": desc.size, + } + del desc + gc.collect() + assert _open_fds() == fds0 + 2 + rebuilt = VirtualMemoryIPCBufferDescriptor.from_fds( + wire["fds"], wire["sizes"], handle_type=wire["handle_type"], size=wire["size"] + ) + assert _open_fds() == fds0 + 4, "from_fds duplicates the integers it is given" + for fd in wire["fds"]: + os.close(fd) # the caller keeps ownership of its own + assert _open_fds() == fds0 + 2 + assert rebuilt.handle_type == "posix_fd" + assert rebuilt.sizes == (gran, gran) + assert rebuilt.size == 2 * gran + imported = Buffer.from_ipc_descriptor(mr, rebuilt, stream=stream) + assert _read(imported, 0, 16) == b"\x71" * 16 + assert _read(imported, gran, 16) == b"\x72" * 16 + del rebuilt + gc.collect() + assert _open_fds() == fds0 + imported.close() + buf.close() + first.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_from_fds_shares_handle_objects_and_rejects_bad_input(self, vmm_ipc_device): + device = vmm_ipc_device + stream = device.default_stream + mr = _resource(device) + buf = mr.allocate(1) + gran = buf.size + desc = buf.ipc_descriptor + from_fds = VirtualMemoryIPCBufferDescriptor.from_fds + + shared = from_fds(desc.handles, desc.sizes) + assert shared.handles[0] is desc.handles[0] # an IPCAllocationHandle is shared, not duplicated + assert shared.size == gran + Buffer.from_ipc_descriptor(mr, shared, stream=stream).close() + + with pytest.raises(ValueError, match="one size per handle"): + from_fds(desc.fds, ()) + with pytest.raises(ValueError, match="positive"): + from_fds(desc.fds, (0,)) + with pytest.raises(ValueError, match="sum of the allocation sizes"): + from_fds(desc.fds, desc.sizes, size=gran + 1) + with pytest.raises(ValueError): + from_fds(desc.fds, desc.sizes, handle_type="win32") + with pytest.raises(ValueError, match="handle_type=None"): + from_fds(desc.fds, desc.sizes, handle_type=None) + with pytest.raises(OSError): + from_fds([1 << 20], desc.sizes) + closed = IPCAllocationHandle._init(os.dup(desc.fds[0]), None) + closed.close() + with pytest.raises(RuntimeError, match="closed"): + from_fds([closed], desc.sizes) + # A handle type other than the resource's is refused at import, before any driver call. + other = from_fds(desc.fds, desc.sizes, handle_type="fabric") + with pytest.raises(ValueError, match="handle type does not match"): + Buffer.from_ipc_descriptor(mr, other, stream=stream) + buf.close() + + @pytest.mark.agent_authored(model="claude-fable-5-1") + def test_round_trip_through_a_unix_socket(self, vmm_ipc_device): + """A child receives the file descriptors with SCM_RIGHTS and rebuilds the descriptor itself.""" + device = vmm_ipc_device + mr = _resource(device) + first = mr.allocate(1) + gran = first.size + buf = mr.modify_allocation(first, 2 * gran) + _fill(buf, 0x73, size=gran) + _fill(buf, 0x74, offset=gran, size=gran) + parent_sock, child_sock = socket.socketpair() + queue = mp.Queue() + process = mp.Process(target=self.child_main, args=(device.device_id, child_sock, queue)) + process.start() + child_sock.close() + desc = buf.ipc_descriptor + header = json.dumps({"sizes": list(desc.sizes), "handle_type": str(desc.handle_type), "size": desc.size}) + socket.send_fds(parent_sock, [header.encode()], list(desc.fds)) + del desc # the kernel duplicated the descriptors into the message + gc.collect() + result = queue.get(timeout=CHILD_TIMEOUT_SEC) + 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}" + parent_sock.close() + assert result == {"low": b"\x73" * 16, "high": b"\x74" * 16, "sizes": [gran, gran], "size": 2 * gran} + assert _read(buf, 0, 16) == b"\x75" * 16 + buf.close() + first.close() + + @staticmethod + def child_main(device_id, sock, queue): + device = Device(device_id) + device.set_current() + with sock: + msg, fds, _flags, _addr = socket.recv_fds(sock, 4096, 16) + header = json.loads(msg) + desc = VirtualMemoryIPCBufferDescriptor.from_fds( + fds, header["sizes"], handle_type=header["handle_type"], size=header["size"] + ) + for fd in fds: + os.close(fd) # from_fds duplicated them; the received ones are ours to close + imported = Buffer.from_ipc_descriptor(_resource(device), desc, stream=device.default_stream) + gran = header["sizes"][0] + result = { + "low": _read(imported, 0, 16), + "high": _read(imported, gran, 16), + "sizes": list(desc.sizes), + "size": desc.size, + } + _fill(imported, 0x75, size=gran) + imported.close() + queue.put(result) + + class TestEmptyBuffer: @pytest.mark.agent_authored(model="claude-fable-5-1") def test_round_trip(self, vmm_ipc_device): @@ -452,8 +842,9 @@ def test_round_trip(self, vmm_ipc_device): empty = mr.allocate(0) desc = empty.ipc_descriptor assert desc.size == 0 - assert desc._chunk_sizes == () - assert desc._handles == () + assert desc.sizes == () + assert desc.handles == () + assert desc.fds == () alias = Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) assert alias.size == 0 assert alias.is_mapped @@ -542,17 +933,17 @@ def test_descriptor_owns_its_file_descriptors(self, vmm_ipc_device): mr = _resource(device) buf = mr.allocate(1) desc = buf.ipc_descriptor - fd = int(desc._handles[0]) + fd = desc.fds[0] os.fstat(fd) # open with pytest.raises(TypeError): pickle.dumps(desc) # file descriptors need multiprocessing with pytest.raises(TypeError): pickle.dumps(buf) - buf.close() # drops the cached descriptor - del desc + del desc # the buffer keeps no descriptor, so this closes the file descriptor gc.collect() with pytest.raises(OSError): os.fstat(fd) + buf.close() @pytest.mark.agent_authored(model="claude-fable-5-1") def test_errors(self, vmm_ipc_device): From de85c4cd4c081ac7b04e8ae7f1251124e8d09fdc Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Tue, 6 Oct 2026 15:58:38 -0700 Subject: [PATCH 06/11] docs(cuda.core): record the rejected retention design for VMM IPC forwarding Write down in VMM_DESIGN.md the opt-in design that would let an imported VirtualMemoryResource buffer be forwarded (keep the received handles with the imported buffer and duplicate them on export), why the option would belong on the resource, and its cost, so that it does not have to be re-derived. Not implemented; nothing needs it yet. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index b9c0bdbd9b9..5f2c9e452be 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -270,8 +270,25 @@ exception surfaces from `pickle`/`multiprocessing` (a `Queue` feeder thread reports it through `Queue._on_queue_feeder_error` and drops the item; the queue stays usable). PyTorch applies the same contract to received CUDA tensors ("Attempted to send CUDA tensor received from another process"). There is no -retention option. A `modify_allocation` config must keep the resource's -`handle_type`: every chunk of a buffer must be exportable the same way. +retention option. + +Forwarding an imported buffer would need the received file descriptors, because +they are the only exportable form of imported memory. The design, not +implemented because nothing needs it yet: a `VirtualMemoryResourceOptions` +field that makes imports through that resource keep their received +`IPCAllocationHandle`s with the imported buffer, and an export loop that uses +the provenance mark to duplicate the kept handle for imported allocations while +exporting created ones through the driver. The option belongs on the resource +rather than on `from_ipc_descriptor`, because a pickled buffer imports through +the resource it carries and the shared `Buffer.from_ipc_descriptor` signature +stays unchanged; the consequence is that the sender's options decide whether a +receiver can forward a pickled buffer. The cost is one open file descriptor per +imported allocation for the imported buffer's lifetime, the same pressure on +the process limit that the no-cache export rule avoids, which is why it would +stay opt-in. + +A `modify_allocation` config must keep the resource's `handle_type`: every +chunk of a buffer must be exportable the same way. ## Why `modify_allocation` returns a new buffer From 2fb30c3ac374c032fe03f927c62bedd37b2fbe62 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Tue, 6 Oct 2026 16:29:52 -0700 Subject: [PATCH 07/11] test(cuda.core): import resource inside the fd-limit test The module-level import of the POSIX-only resource module made pytest fail to collect tests/memory_ipc/test_vmm_ipc.py on Windows, so the Windows rows ran no cuda.core tests at all. Every test in the module is skipped on Windows by its fixture; only the import has to move. Co-Authored-By: Claude Fable 5.1 --- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index 25b22a870a8..1b3753a4bee 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -18,7 +18,6 @@ import multiprocessing.queues import os import pickle -import resource import socket import threading @@ -635,6 +634,8 @@ class TestFileDescriptorLimit: @pytest.mark.thread_unsafe(reason="changes the process-wide file descriptor limit") def test_export_near_the_limit(self, vmm_ipc_device): """Under a lowered RLIMIT_NOFILE an export either succeeds or fails with a readable error and no leak.""" + import resource # POSIX-only; a module-level import breaks collection on Windows + device = vmm_ipc_device stream = device.default_stream mr = _resource(device) From 2cf07d916b1c3cf09d9d8b54b929d651fbcd381e Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 7 Oct 2026 13:01:55 -0700 Subject: [PATCH 08/11] cuda.core: VMM IPC review fixes; no sharing of host-located virtual memory Take ownership of each exported or duplicated file descriptor before appending its handle, so a failed append cannot close the descriptor twice. Validate imported chunk sizes against the minimum granularity that cuMemMap requires rather than the resource's granularity option. Name both causes in the note on a failed import: a handle that is not a CUDA allocation of the expected type, or an owning device that this process cannot use. Report is_ipc_enabled False for host-located resources: host NUMA virtual memory has no verified allocation path, so it is not shared yet; the docs and the design note say so. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 9 ++-- cuda_core/cuda/core/_memory/_ipc.pyx | 4 +- .../core/_memory/_virtual_memory_resource.pyi | 16 ++++--- .../core/_memory/_virtual_memory_resource.pyx | 44 ++++++++++++------- cuda_core/docs/source/release/1.3.0-notes.rst | 3 +- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 7 ++- 6 files changed, 53 insertions(+), 30 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index 5f2c9e452be..6efe07a2bb0 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -144,8 +144,8 @@ what makes that safe: the allocation is released exactly once, when its last map `location_type="host"` with a handle type other than `None`, which the driver rejects, and a request for GPUDirect RDMA on a device without support. `__init__` runs it after the VMM-support check; `modify_allocation` runs it on a per-call configuration, which must also name the - resource's location. The resource reports `is_ipc_enabled = False`, which - `Buffer.ipc_descriptor` reads. + resource's location. The resource reports `is_ipc_enabled` (True for `posix_fd` on Linux with + device-located memory), which `Buffer.ipc_descriptor` and `Buffer.from_ipc_descriptor` check. - `cdef class VirtualMemoryBuffer(Buffer)` carries no extra state. It is created with `Buffer_from_deviceptr_handle(h_ptr, size, self, cls=VirtualMemoryBuffer)` and documented in `api.rst` like `ManagedBuffer`. It overrides `close(stream=None)` to reject a capturing stream @@ -203,8 +203,9 @@ 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; -`is_ipc_enabled` reports it). The exporter calls `cuMemExportToShareableHandle` +`handle_type` can travel between processes (`posix_fd` 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 descriptor carries the handles with each allocation's size. Each file descriptor is owned by an `IPCAllocationHandle` and closed when the descriptor diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index 81251bc595a..55b054761f0 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -293,10 +293,12 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): continue dup = os.dup(operator.index(item)) try: - handles.append(IPCAllocationHandle._init(dup, None)) + handle = IPCAllocationHandle._init(dup, None) except: # noqa: E722 rollback-then-raise: the duplicate is not owned yet os.close(dup) raise + # The handle owns the duplicate now; if append fails, the handle closes it once. + handles.append(handle) return VirtualMemoryIPCBufferDescriptor._from_exports(driver_type, sizes_t, tuple(handles), size) @property diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index ade6e744067..b127218fe06 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -329,18 +329,20 @@ 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; - ``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. + 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 + 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 + ``"host_numa_current"``) is not shared yet. """ def _import_ipc_buffer(self, desc: VirtualMemoryIPCBufferDescriptor, 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 - type and checked to live where this resource allocates (its device, - or its host location). The allocations are then mapped in order into + type and checked to live on this resource's device. The allocations + are then mapped in order into one new address reservation with the access descriptors this resource's options produce for its device and peers, and the deallocation stream is recorded as :meth:`allocate` does. The diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index a647b3a7c79..4b3ca193503 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -885,20 +885,25 @@ 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; - ``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. + 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 + 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 + ``"host_numa_current"``) is not shared yet. """ - return self.config.handle_type in _IPC_HANDLE_TYPES + return ( + self.config.handle_type in _IPC_HANDLE_TYPES + and self.config.location_type not in _HOST_LOCATION_TYPES + ) def _import_ipc_buffer(self, VirtualMemoryIPCBufferDescriptor 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 - type and checked to live where this resource allocates (its device, - or its host location). The allocations are then mapped in order into + type and checked to live on this resource's device. The allocations + are then mapped in order into one new address reservation with the access descriptors this resource's options produce for its device and peers, and the deallocation stream is recorded as :meth:`allocate` does. The @@ -919,7 +924,7 @@ cdef class VirtualMemoryResource(MemoryResource): cdef VmmRangeHandle rng cdef DevicePtrHandle h_ptr cdef Buffer buf - cdef size_t gran, addr_align, chunk, total = 0, offset = 0 + 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 @@ -955,11 +960,17 @@ cdef class VirtualMemoryResource(MemoryResource): self._fill_prop(cfg, &prop) self._fill_access(cfg, &prop, descs) gran = self._granularity(cfg, &prop) + # cuMemMap requires the minimum granularity, whatever this resource's + # granularity option says; the exporter may have used either. + with nogil: + HANDLE_RETURN(cydriver.cuMemGetAllocationGranularity( + &gran_min, &prop, cydriver.CU_MEM_ALLOC_GRANULARITY_MINIMUM)) for chunk_obj in desc._chunk_sizes: chunk = chunk_obj - if chunk == 0 or chunk % gran != 0: + if chunk == 0 or chunk % gran_min != 0: raise ValueError( - f"descriptor chunk size {chunk} is not a positive multiple of the {gran}-byte granularity" + f"descriptor chunk size {chunk} is not a positive multiple of the {gran_min}-byte minimum " + "granularity" ) if chunk > -1 - total: raise OverflowError("the descriptor's chunk sizes do not fit in size_t") @@ -988,9 +999,10 @@ cdef class VirtualMemoryResource(MemoryResource): 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; the " - "shared handle may not refer to a CUDA allocation exported with " - f"handle_type={str(cfg.handle_type)!r}" + 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)" ) # The exporter chose where the memory lives. This resource must map # memory of that kind on that device, or its access descriptors @@ -1106,9 +1118,11 @@ cdef VirtualMemoryIPCBufferDescriptor _export_ipc_descriptor(Buffer buf): "allocation takes one file descriptor, so a process near its file descriptor limit fails here", ) try: - handles.append(IPCAllocationHandle._init(fd, None)) + handle = IPCAllocationHandle._init(fd, None) except: # noqa: E722 rollback-then-raise: the fd is not owned yet os.close(fd) raise + # The handle owns the fd now; if append fails, the handle closes it once. + handles.append(handle) sizes.append(chunk) return VirtualMemoryIPCBufferDescriptor._from_exports(handle_type, tuple(sizes), tuple(handles), buf._size) 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 af61df694f4..a8612d3e391 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -86,7 +86,8 @@ New features 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. A grown buffer + 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 diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index 1b3753a4bee..09cf19fecbe 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -181,11 +181,14 @@ def _on_queue_feeder_error(e, _obj): 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({"location_type": "host_numa", "handle_type": "posix_fd"}, True, id="host_numa-posix_fd"), + pytest.param({"location_type": "host_numa", "handle_type": "posix_fd"}, False, id="host_numa-posix_fd"), + pytest.param( + {"location_type": "host_numa_current", "handle_type": "posix_fd"}, False, id="host_numa_current-posix_fd" + ), ], ) def test_is_ipc_enabled(vmm_ipc_device, options, expected): - """Only POSIX file descriptors have a transport; the location does not matter.""" + """Only 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 From 85f1d6d896e9d21b5ef95678ab45aa824321d10e Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Tue, 13 Oct 2026 12:44:42 -0700 Subject: [PATCH 09/11] cuda.core: VMM IPC review fixes: posix_fd-only from_fds, scoped docs, no test skips - from_fds() takes no handle_type; a file descriptor can only carry posix_fd - The descriptor, resource, and release-note text scope the pin, pickle, and same-device rules to POSIX file descriptors, and the fds docstrings state the descriptor's lifetime rule - The 1.3.0 notes record that a per-call modify_allocation config must keep the resource's handle_type - VMM_DESIGN.md records the fabric direction: a descriptor class of its own - TestManyChunkBuffer and TestReexport no longer skip when the driver refuses the decoy's address hint; the grow relocates in that case too Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 23 ++++++-- cuda_core/cuda/core/_memory/_ipc.pyi | 45 +++++++++------- cuda_core/cuda/core/_memory/_ipc.pyx | 52 +++++++++---------- .../core/_memory/_virtual_memory_resource.pyi | 44 +++++++++------- .../core/_memory/_virtual_memory_resource.pyx | 44 +++++++++------- cuda_core/docs/source/release/1.3.0-notes.rst | 14 ++--- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 48 +++++++---------- 7 files changed, 144 insertions(+), 126 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index 6efe07a2bb0..9cd7994b6b4 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -144,8 +144,9 @@ what makes that safe: the allocation is released exactly once, when its last map `location_type="host"` with a handle type other than `None`, which the driver rejects, and a request for GPUDirect RDMA on a device without support. `__init__` runs it after the VMM-support check; `modify_allocation` runs it on a per-call configuration, which must also name the - resource's location. The resource reports `is_ipc_enabled` (True for `posix_fd` on Linux with - device-located memory), which `Buffer.ipc_descriptor` and `Buffer.from_ipc_descriptor` check. + resource's location and `handle_type`. The resource reports `is_ipc_enabled` (True for + `posix_fd` on Linux with device-located memory), which `Buffer.ipc_descriptor` and + `Buffer.from_ipc_descriptor` check. - `cdef class VirtualMemoryBuffer(Buffer)` carries no extra state. It is created with `Buffer_from_deviceptr_handle(h_ptr, size, self, cls=VirtualMemoryBuffer)` and documented in `api.rst` like `ManagedBuffer`. It overrides `close(stream=None)` to reject a capturing stream @@ -253,9 +254,9 @@ descriptor either; it records only that it was imported (`is_mapped`). A descriptor parked in a `Queue` or sent to a `Pool` pins the memory in the sender until the receiver unpickles it. Besides `multiprocessing`, a process with its own file descriptor passing sends the descriptor's `fds`, `sizes`, -`handle_type`, and `size` and rebuilds it with -`VirtualMemoryIPCBufferDescriptor.from_fds()`, which duplicates the integers -it is given. +and `size` and rebuilds it with `VirtualMemoryIPCBufferDescriptor.from_fds()`, +which duplicates the integers it is given and always builds a `posix_fd` +descriptor. Re-export. The driver exports only allocations this process created with the requested handle type; an imported allocation has none, and the driver reports @@ -291,6 +292,18 @@ stay opt-in. 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 `VirtualMemoryIPCBufferDescriptor`, 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). + ## Why `modify_allocation` returns a new buffer The input buffer stays open and aliases the result. The chunks the input already mapped are diff --git a/cuda_core/cuda/core/_memory/_ipc.pyi b/cuda_core/cuda/core/_memory/_ipc.pyi index a85c15363de..8e1cce43296 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyi +++ b/cuda_core/cuda/core/_memory/_ipc.pyi @@ -61,12 +61,13 @@ class IPCAllocationHandle: class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. - The descriptor holds one exported handle per physical allocation that - backs the buffer, in address order (``handles``), with each allocation's - size (``sizes``) and the buffer's size (``size``, which may be smaller - than their sum). With POSIX file descriptors the handles are - :class:`IPCAllocationHandle` objects: the descriptor owns them and closes - them when it is released. + The descriptor holds one exported POSIX file descriptor per physical + allocation that backs the buffer, in address order (``handles``, each an + :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. A descriptor pins the physical memory while it exists, in every process that holds one, and nothing else does: the exporting buffer keeps no @@ -81,8 +82,8 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): ``Process`` arguments, or a ``Pool``) pickles the descriptor, or a buffer, and duplicates the file descriptors into the receiving process; plain ``pickle`` cannot carry them and raises. A process with its own file - descriptor passing sends ``fds``, ``sizes``, ``handle_type``, and ``size`` - itself, and the receiver rebuilds the descriptor with :meth:`from_fds`. + descriptor passing sends ``fds``, ``sizes``, and ``size`` itself, and the + receiver rebuilds the descriptor with :meth:`from_fds`. Note ---- @@ -98,13 +99,14 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): @staticmethod def _from_exports(handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple[Any, ...], size: int) -> VirtualMemoryIPCBufferDescriptor: ... @classmethod - def from_fds(cls, fds: Iterable[int | IPCAllocationHandle], sizes: Iterable[int], *, handle_type: VirtualMemoryHandleType | str='posix_fd', size: int | None=None) -> VirtualMemoryIPCBufferDescriptor: - """Build a descriptor from file descriptors received outside ``multiprocessing``. + def from_fds(cls, fds: Iterable[int | IPCAllocationHandle], sizes: Iterable[int], *, size: int | None=None) -> VirtualMemoryIPCBufferDescriptor: + """Build a descriptor from POSIX file descriptors received outside ``multiprocessing``. - The exporter sends ``fds``, ``sizes``, ``handle_type``, and ``size`` - of its descriptor through a transport of its own (a Unix socket with - ``SCM_RIGHTS``, for example); the receiver rebuilds the descriptor - here and imports it with :meth:`Buffer.from_ipc_descriptor`. + The exporter sends ``fds``, ``sizes``, and ``size`` of its descriptor + through a transport of its own (a Unix socket with ``SCM_RIGHTS``, for + example); the receiver rebuilds the descriptor here and imports it + with :meth:`Buffer.from_ipc_descriptor`. The result's ``handle_type`` + is ``"posix_fd"``, the only handle type a file descriptor can carry. Parameters ---------- @@ -116,8 +118,6 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): holder. sizes : Iterable[int] The exporter's ``sizes``, one per handle. - handle_type : VirtualMemoryHandleType | str, optional - The exporter's ``handle_type``; ``"posix_fd"`` by default. size : int, optional The exporter's ``size``, at most the sum of ``sizes``; the sum by default. @@ -125,8 +125,8 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): Raises ------ ValueError - If the counts differ, a size is not positive, ``size`` exceeds the - sum of ``sizes``, or ``handle_type`` cannot be shared. + If the counts differ, a size is not positive, or ``size`` exceeds + the sum of ``sizes``. RuntimeError If an :class:`IPCAllocationHandle` is closed. OSError @@ -140,13 +140,18 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): """Size of each physical allocation, in address order; the sum is at least ``size``.""" @property def handles(self) -> tuple[IPCAllocationHandle, ...]: - """The exported handles, one per allocation, in address order; the descriptor owns them.""" + """The exported handles, one :class:`IPCAllocationHandle` per allocation, in address order. + + The descriptor owns them and closes them when it is released. + """ @property def fds(self) -> tuple[int, ...]: """The exported file descriptors as integers, for a transport of your own. They belong to this descriptor: send or duplicate them while it is - alive. Each one is released when the descriptor is. + alive, because each one is closed when the descriptor is released. + Keep the descriptor in a variable while you use them; the numbers + from ``buf.ipc_descriptor.fds`` alone may already be closed. """ def __reduce__(self) -> tuple[object, ...]: ... diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index 55b054761f0..5bc35fa84b0 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -180,12 +180,13 @@ multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handl cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): """Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes. - The descriptor holds one exported handle per physical allocation that - backs the buffer, in address order (``handles``), with each allocation's - size (``sizes``) and the buffer's size (``size``, which may be smaller - than their sum). With POSIX file descriptors the handles are - :class:`IPCAllocationHandle` objects: the descriptor owns them and closes - them when it is released. + The descriptor holds one exported POSIX file descriptor per physical + allocation that backs the buffer, in address order (``handles``, each an + :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. A descriptor pins the physical memory while it exists, in every process that holds one, and nothing else does: the exporting buffer keeps no @@ -200,8 +201,8 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): ``Process`` arguments, or a ``Pool``) pickles the descriptor, or a buffer, and duplicates the file descriptors into the receiving process; plain ``pickle`` cannot carry them and raises. A process with its own file - descriptor passing sends ``fds``, ``sizes``, ``handle_type``, and ``size`` - itself, and the receiver rebuilds the descriptor with :meth:`from_fds`. + descriptor passing sends ``fds``, ``sizes``, and ``size`` itself, and the + receiver rebuilds the descriptor with :meth:`from_fds`. Note ---- @@ -230,15 +231,15 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): fds: Iterable[int | IPCAllocationHandle], sizes: Iterable[int], *, - handle_type: VirtualMemoryHandleType | str = "posix_fd", size: int | None = None, ) -> VirtualMemoryIPCBufferDescriptor: - """Build a descriptor from file descriptors received outside ``multiprocessing``. + """Build a descriptor from POSIX file descriptors received outside ``multiprocessing``. - The exporter sends ``fds``, ``sizes``, ``handle_type``, and ``size`` - of its descriptor through a transport of its own (a Unix socket with - ``SCM_RIGHTS``, for example); the receiver rebuilds the descriptor - here and imports it with :meth:`Buffer.from_ipc_descriptor`. + The exporter sends ``fds``, ``sizes``, and ``size`` of its descriptor + through a transport of its own (a Unix socket with ``SCM_RIGHTS``, for + example); the receiver rebuilds the descriptor here and imports it + with :meth:`Buffer.from_ipc_descriptor`. The result's ``handle_type`` + is ``"posix_fd"``, the only handle type a file descriptor can carry. Parameters ---------- @@ -250,8 +251,6 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): holder. sizes : Iterable[int] The exporter's ``sizes``, one per handle. - handle_type : VirtualMemoryHandleType | str, optional - The exporter's ``handle_type``; ``"posix_fd"`` by default. size : int, optional The exporter's ``size``, at most the sum of ``sizes``; the sum by default. @@ -259,20 +258,14 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): Raises ------ ValueError - If the counts differ, a size is not positive, ``size`` exceeds the - sum of ``sizes``, or ``handle_type`` cannot be shared. + If the counts differ, a size is not positive, or ``size`` exceeds + the sum of ``sizes``. RuntimeError If an :class:`IPCAllocationHandle` is closed. OSError If an ``int`` is not an open file descriptor. """ - # Lazy import: this module is loaded while _virtual_memory_resource imports it. - from cuda.core._memory._virtual_memory_resource import VirtualMemoryResourceOptions - from cuda.core.typing import VirtualMemoryHandleType - - if handle_type is None: - raise ValueError("handle_type=None cannot be shared; the exporter's handle_type is required") - driver_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(VirtualMemoryHandleType(handle_type))) + cdef int driver_type = cydriver.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR cdef tuple sizes_t = tuple(operator.index(n) for n in sizes) cdef list given = list(fds) if len(given) != len(sizes_t): @@ -318,7 +311,10 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): @property def handles(self) -> tuple[IPCAllocationHandle, ...]: - """The exported handles, one per allocation, in address order; the descriptor owns them.""" + """The exported handles, one :class:`IPCAllocationHandle` per allocation, in address order. + + The descriptor owns them and closes them when it is released. + """ return self._handles @property @@ -326,7 +322,9 @@ cdef class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor): """The exported file descriptors as integers, for a transport of your own. They belong to this descriptor: send or duplicate them while it is - alive. Each one is released when the descriptor is. + alive, because each one is closed when the descriptor is released. + Keep the descriptor in a variable while you use them; the numbers + from ``buf.ipc_descriptor.fds`` alone may already be closed. """ return tuple(int(handle) for handle in self._handles) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi index b127218fe06..3f630f09834 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyi @@ -102,9 +102,12 @@ class VirtualMemoryBuffer(Buffer): Every access exports the physical allocations that back the buffer again and returns a new descriptor that owns one file descriptor per allocation; the file descriptors close when the descriptor is - released. A descriptor pins the physical memory while it exists, in - every process that holds one, and nothing else does: the buffer - itself keeps no descriptor, so it pins nothing beyond its own mapping. + released. Keep the descriptor in a variable while you use its + ``fds``: the numbers from ``buf.ipc_descriptor.fds`` alone may + already be closed. A descriptor pins the physical memory while it + exists, in every process that holds one, and nothing else does: the + buffer itself keeps no descriptor, so it pins nothing beyond its own + mapping. Pickling the buffer through ``multiprocessing`` builds a transient descriptor, so the sender holds no file descriptors between sends; a buffer passed as a ``Process`` argument is the exception, because @@ -186,23 +189,24 @@ class VirtualMemoryResource(MemoryResource): :meth:`Buffer.set_deallocation_stream`. Buffers can be shared with other processes when ``config.handle_type`` - is ``"posix_fd"`` (Linux); ``handle_type`` is the switch and - :attr:`is_ipc_enabled` the query. :attr:`Buffer.ipc_descriptor` exports - the physical allocations that back a buffer, one file descriptor each, - into a new descriptor on every 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``, ``handle_type``, 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 it exists, in every process that holds one, and - nothing else does. 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. + is ``"posix_fd"`` (Linux), the only handle type cuda.core shares today; + ``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 + a buffer, one file descriptor each, into a new descriptor on every + 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 + ``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 + it exists, in every process that holds one, and nothing else does. 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. """ device: object config: object diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx index 4b3ca193503..9a4bb083922 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx @@ -296,9 +296,12 @@ cdef class VirtualMemoryBuffer(Buffer): Every access exports the physical allocations that back the buffer again and returns a new descriptor that owns one file descriptor per allocation; the file descriptors close when the descriptor is - released. A descriptor pins the physical memory while it exists, in - every process that holds one, and nothing else does: the buffer - itself keeps no descriptor, so it pins nothing beyond its own mapping. + released. Keep the descriptor in a variable while you use its + ``fds``: the numbers from ``buf.ipc_descriptor.fds`` alone may + already be closed. A descriptor pins the physical memory while it + exists, in every process that holds one, and nothing else does: the + buffer itself keeps no descriptor, so it pins nothing beyond its own + mapping. Pickling the buffer through ``multiprocessing`` builds a transient descriptor, so the sender holds no file descriptors between sends; a buffer passed as a ``Process`` argument is the exception, because @@ -416,23 +419,24 @@ cdef class VirtualMemoryResource(MemoryResource): :meth:`Buffer.set_deallocation_stream`. Buffers can be shared with other processes when ``config.handle_type`` - is ``"posix_fd"`` (Linux); ``handle_type`` is the switch and - :attr:`is_ipc_enabled` the query. :attr:`Buffer.ipc_descriptor` exports - the physical allocations that back a buffer, one file descriptor each, - into a new descriptor on every 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``, ``handle_type``, 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 it exists, in every process that holds one, and - nothing else does. 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. + is ``"posix_fd"`` (Linux), the only handle type cuda.core shares today; + ``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 + a buffer, one file descriptor each, into a new descriptor on every + 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 + ``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 + it exists, in every process that holds one, and nothing else does. 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. """ cdef: 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 a8612d3e391..ff7b9c931df 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -91,10 +91,9 @@ New features 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``, - ``handle_type``, and ``size``, and - ``VirtualMemoryIPCBufferDescriptor.from_fds()`` rebuilds it in the receiver. - Fabric handles are not shared yet. + transport of your own, the descriptor exposes ``fds``, ``sizes``, and + ``size``, and ``VirtualMemoryIPCBufferDescriptor.from_fds()`` rebuilds it + in the receiver. Fabric handles are not shared yet. (`#2980 `__) Fixes and enhancements @@ -259,8 +258,11 @@ Fixes and enhancements :attr:`Buffer.size` after a grow is the aligned total. The ``config`` argument applies to the chunk that call adds and is no longer stored on the resource (`#2909 `__); - it must keep the resource's location and passes the constructor's option - checks. It never changes the access of memory that is already mapped, so a + it must keep the resource's location and ``handle_type`` and passes the + constructor's option checks. Every chunk of a buffer must be exportable + the same way, so a resource created with ``handle_type=None`` needs + ``handle_type=None`` in a per-call ``config`` as well. It never changes + the access of memory that is already mapped, so a request the buffer already covers has no effect even with ``config``. Buffers from ``allocate()`` free themselves when they close and no longer call :meth:`VirtualMemoryResource.deallocate`, which now serves pointers diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index 09cf19fecbe..55d7779e09e 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -223,15 +223,17 @@ def test_main(self, vmm_ipc_device): gran = buf.size buf = mr.modify_allocation(buf, 2 * gran) # extends in place or moves buffers.append(buf) - # Take the range after the buffer so that this grow must remap the - # existing chunks at a new address. + # Take the range right after the buffer so that this grow must + # remap the existing chunks at a new address. The grow probes that + # range with the same driver call; when the driver places the decoy + # elsewhere, it places the grow's probe elsewhere too, and the grow + # moves without the decoy. decoy = _reserve_at(int(buf.handle) + buf.size, gran) - if decoy is None: - pytest.skip("the driver did not grant a reservation right after the buffer") try: moved = mr.modify_allocation(buf, 3 * gran) finally: - handle_return(driver.cuMemAddressFree(decoy, gran)) + if decoy is not None: + handle_return(driver.cuMemAddressFree(decoy, gran)) buffers.append(moved) assert int(moved.handle) != int(buf.handle) assert moved.size == 3 * gran @@ -438,19 +440,21 @@ class TestReexport: @staticmethod def _derived(mr, imported, gran): - """Buffers modify_allocation derives from an import: an alias, a grow, and (if the driver allows) a move.""" + """Buffers modify_allocation derives from an import: an alias, a grow, and a move.""" derived = { "alias": mr.modify_allocation(imported, imported.size), "grow": mr.modify_allocation(imported, 2 * gran), } base = derived["grow"] + # The decoy takes the range the grow would extend into; when the driver + # refuses the decoy's hint, it refuses the grow's too, and the grow moves. decoy = _reserve_at(int(base.handle) + base.size, gran) - if decoy is not None: - try: - derived["move"] = mr.modify_allocation(base, base.size + gran) - finally: + try: + derived["move"] = mr.modify_allocation(base, base.size + gran) + finally: + if decoy is not None: handle_return(driver.cuMemAddressFree(decoy, gran)) - assert int(derived["move"].handle) != int(base.handle) + assert int(derived["move"].handle) != int(base.handle) return derived @pytest.mark.agent_authored(model="claude-fable-5-1") @@ -695,7 +699,7 @@ def test_export_near_the_limit(self, vmm_ipc_device): class TestRawFileDescriptorTransport: - """The descriptor's fds, sizes, handle_type, and size travel over any file-descriptor-passing transport.""" + """The descriptor's fds, sizes, and size travel over any file-descriptor-passing transport.""" @pytest.mark.agent_authored(model="claude-fable-5-1") @pytest.mark.thread_unsafe(reason="counts process-global file descriptors") @@ -718,15 +722,12 @@ def test_round_trip_with_os_dup(self, vmm_ipc_device): wire = { "fds": [os.dup(fd) for fd in desc.fds], "sizes": list(desc.sizes), - "handle_type": str(desc.handle_type), "size": desc.size, } del desc gc.collect() assert _open_fds() == fds0 + 2 - rebuilt = VirtualMemoryIPCBufferDescriptor.from_fds( - wire["fds"], wire["sizes"], handle_type=wire["handle_type"], size=wire["size"] - ) + rebuilt = VirtualMemoryIPCBufferDescriptor.from_fds(wire["fds"], wire["sizes"], size=wire["size"]) assert _open_fds() == fds0 + 4, "from_fds duplicates the integers it is given" for fd in wire["fds"]: os.close(fd) # the caller keeps ownership of its own @@ -757,6 +758,7 @@ def test_from_fds_shares_handle_objects_and_rejects_bad_input(self, vmm_ipc_devi shared = from_fds(desc.handles, desc.sizes) assert shared.handles[0] is desc.handles[0] # an IPCAllocationHandle is shared, not duplicated assert shared.size == gran + assert shared.handle_type == "posix_fd" Buffer.from_ipc_descriptor(mr, shared, stream=stream).close() with pytest.raises(ValueError, match="one size per handle"): @@ -765,20 +767,12 @@ def test_from_fds_shares_handle_objects_and_rejects_bad_input(self, vmm_ipc_devi from_fds(desc.fds, (0,)) with pytest.raises(ValueError, match="sum of the allocation sizes"): from_fds(desc.fds, desc.sizes, size=gran + 1) - with pytest.raises(ValueError): - from_fds(desc.fds, desc.sizes, handle_type="win32") - with pytest.raises(ValueError, match="handle_type=None"): - from_fds(desc.fds, desc.sizes, handle_type=None) with pytest.raises(OSError): from_fds([1 << 20], desc.sizes) closed = IPCAllocationHandle._init(os.dup(desc.fds[0]), None) closed.close() with pytest.raises(RuntimeError, match="closed"): from_fds([closed], desc.sizes) - # A handle type other than the resource's is refused at import, before any driver call. - other = from_fds(desc.fds, desc.sizes, handle_type="fabric") - with pytest.raises(ValueError, match="handle type does not match"): - Buffer.from_ipc_descriptor(mr, other, stream=stream) buf.close() @pytest.mark.agent_authored(model="claude-fable-5-1") @@ -797,7 +791,7 @@ def test_round_trip_through_a_unix_socket(self, vmm_ipc_device): process.start() child_sock.close() desc = buf.ipc_descriptor - header = json.dumps({"sizes": list(desc.sizes), "handle_type": str(desc.handle_type), "size": desc.size}) + header = json.dumps({"sizes": list(desc.sizes), "size": desc.size}) socket.send_fds(parent_sock, [header.encode()], list(desc.fds)) del desc # the kernel duplicated the descriptors into the message gc.collect() @@ -819,9 +813,7 @@ def child_main(device_id, sock, queue): with sock: msg, fds, _flags, _addr = socket.recv_fds(sock, 4096, 16) header = json.loads(msg) - desc = VirtualMemoryIPCBufferDescriptor.from_fds( - fds, header["sizes"], handle_type=header["handle_type"], size=header["size"] - ) + desc = VirtualMemoryIPCBufferDescriptor.from_fds(fds, header["sizes"], size=header["size"]) for fd in fds: os.close(fd) # from_fds duplicated them; the received ones are ours to close imported = Buffer.from_ipc_descriptor(_resource(device), desc, stream=device.default_stream) From f87e5c515358d39b1f252f0c6bf100174763af4a Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Tue, 13 Oct 2026 14:27:10 -0700 Subject: [PATCH 10/11] test(cuda.core): collect garbage before each file-descriptor baseline in the VMM IPC tests init_cuda no longer runs gc.collect() at every teardown (#3053), so objects left by earlier tests (child pipes, queues, descriptors) can be collected in the middle of a test that counts process-wide file descriptors, and the count then drops below its baseline. Collect before taking the baseline. Co-Authored-By: Claude Fable 5.1 --- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 24 ++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index 55d7779e09e..c2cfecbe0db 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -88,6 +88,18 @@ def _open_fds(): return len(os.listdir("/proc/self/fd")) +def _settled_open_fds(): + """The open file descriptor count after a garbage collection. + + Objects left by earlier tests may still hold file descriptors (child + pipes, queues, descriptors), and ``init_cuda`` no longer collects at + every teardown. Collecting here keeps a baseline stable against a + collection that the test itself triggers later. + """ + gc.collect() + return _open_fds() + + def _has_note(exc, text): """Whether ``text`` is in one of the exception's notes, or in its message on Python 3.10.""" return any(text in note for note in getattr(exc, "__notes__", ())) or text in str(exc) @@ -241,7 +253,7 @@ def test_main(self, vmm_ipc_device): for i, value in enumerate(values): _fill(moved, value, offset=i * gran, size=gran) - fds_before = _open_fds() + fds_before = _settled_open_fds() desc = moved.ipc_descriptor assert desc.size == 3 * gran assert desc._chunk_sizes == (gran, gran, gran) @@ -551,7 +563,7 @@ def test_exporter_holds_none_without_a_descriptor(self, vmm_ipc_device): first = mr.allocate(1) gran = first.size buf = mr.modify_allocation(first, 2 * gran) - fds0 = _open_fds() + fds0 = _settled_open_fds() desc = buf.ipc_descriptor assert desc.sizes == (gran, gran) assert _open_fds() == fds0 + 2 @@ -579,7 +591,7 @@ def test_process_argument_releases_with_the_process_object(self, vmm_ipc_device) device = vmm_ipc_device mr = _resource(device) buf = mr.allocate(1) - fds0 = _open_fds() + fds0 = _settled_open_fds() process = mp.Process(target=_close_argument, args=(buf,)) process.start() process.join(timeout=CHILD_TIMEOUT_SEC) @@ -624,7 +636,7 @@ def child_main(device_id, to_child, from_child): desc = to_child.get(timeout=CHILD_TIMEOUT_SEC) # A first import lets the driver finish any lazy initialization before counting. Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream).close() - counts = {"with_descriptor": _open_fds()} + counts = {"with_descriptor": _settled_open_fds()} imported = Buffer.from_ipc_descriptor(mr, desc, stream=device.default_stream) counts["after_import"] = _open_fds() del desc @@ -665,7 +677,7 @@ def test_export_near_the_limit(self, vmm_ipc_device): os.close(fd) break fillers.append(fd) - baseline = _open_fds() + baseline = _settled_open_fds() # Room for two more descriptors: a three-allocation export cannot complete. resource.setrlimit(resource.RLIMIT_NOFILE, (highest + 1 + 2, hard)) failure = None @@ -712,7 +724,7 @@ def test_round_trip_with_os_dup(self, vmm_ipc_device): buf = mr.modify_allocation(first, 2 * gran) _fill(buf, 0x71, size=gran) _fill(buf, 0x72, offset=gran, size=gran) - fds0 = _open_fds() + fds0 = _settled_open_fds() desc = buf.ipc_descriptor assert desc.handle_type == "posix_fd" assert desc.sizes == (gran, gran) From b5827efd6f1cf6d7ab9b7030a3222a79457e7bb4 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 14 Oct 2026 13:18:23 -0700 Subject: [PATCH 11/11] test(cuda.core): mark the VMM re-export test thread-unsafe TestReexport._derived asserts that a grow moved the buffer after a decoy took the adjacent range. Under pytest-run-parallel another thread can free that range between the decoy attempt and the grow, so the grow extends in place and the assertion fails. Mark the test thread_unsafe like the other address-probing tests in the file. Also record in VMM_DESIGN.md the numbers behind the no-cache export rule and the age of the multiprocessing machinery the Process path relies on. Co-Authored-By: Claude Fable 5.1 --- cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md | 9 +++++++-- cuda_core/tests/memory_ipc/test_vmm_ipc.py | 1 + 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md index 9cd7994b6b4..2797af2a5c8 100644 --- a/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md +++ b/cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md @@ -239,14 +239,19 @@ exists, in every process that holds one, and nothing else does. The exporting buffer keeps no descriptor: `ipc_descriptor` exports again on every access and returns a new descriptor that owns one file descriptor per allocation, and `VirtualMemoryBuffer.__reduce__` builds a transient one while `multiprocessing` -pickles the buffer, so a sender holds no file descriptors between sends. One +pickles the buffer, so a sender holds no file descriptors between sends. The +numbers set that rule: PyTorch's expandable segments grow in 2 MB to 20 MB +chunks, so a buffer of several GB holds several hundred allocations, against a +default process limit of 1024 file descriptors. One transport needs more: a `Queue`, `Pipe`, or `Pool` duplicates the file descriptors while pickling (the resource sharer), but a spawned `Process` records their numbers during pickling and passes them to the child only when it is created (`spawnv_passfds`), so `__reduce__` keeps the transient descriptor on the spawning `Popen` object, which dies with the `Process` object. Without that, the numbers are reused by the spawn's own pipes -("bad value(s) in fds_to_keep") or the child imports a stranger's fd. The +("bad value(s) in fds_to_keep") or the child imports a stranger's fd. That +machinery (`spawnv_passfds`, `get_spawning_popen()`) has not changed since +Python 3.4 (2014). The driver copies each shared allocation into a handle of its own during the import and keeps no reference to the file descriptor, so the imported buffer does not keep the diff --git a/cuda_core/tests/memory_ipc/test_vmm_ipc.py b/cuda_core/tests/memory_ipc/test_vmm_ipc.py index c2cfecbe0db..b42f1eb4623 100644 --- a/cuda_core/tests/memory_ipc/test_vmm_ipc.py +++ b/cuda_core/tests/memory_ipc/test_vmm_ipc.py @@ -470,6 +470,7 @@ def _derived(mr, imported, gran): return derived @pytest.mark.agent_authored(model="claude-fable-5-1") + @pytest.mark.thread_unsafe(reason="checks mapping state by address; another thread can reuse a freed address") def test_export_and_pickling_raise_before_any_driver_call(self, vmm_ipc_device): device = vmm_ipc_device stream = device.default_stream