Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 38 additions & 11 deletions cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ what makes that safe: the allocation is released exactly once, when its last map

A `VirtualMemoryBuffer` is shared the way pool-backed buffers are, through
`Buffer.ipc_descriptor` and `Buffer.from_ipc_descriptor`, when the resource's
`handle_type` can travel between processes (`posix_fd` on Linux) and the memory
`handle_type` can travel between processes (`posix_fd` or `fabric` on Linux) and the memory
lives on a device (`is_ipc_enabled` reports both; host-located memory is not
shared yet). The exporter calls `cuMemExportToShareableHandle`
once per mapping that covers the buffer's size, in address order, and the
Expand Down Expand Up @@ -298,16 +298,43 @@ A `modify_allocation` config must keep the resource's `handle_type`: every
chunk of a buffer must be exportable the same way.

Fabric handles. `CU_MEM_HANDLE_TYPE_FABRIC` exports a `CUmemFabricHandle`, 64
bytes of plain data that travel through `pickle`, MPI, or a socket, need no
close, and cannot hold a reference to the allocation while in flight. The
importer needs access to the exporter's IMEX channel rather than a matching
device index, and the memory may live on another node. The file-descriptor
rules above (ownership, pinning, `pickle` raising, the same-device check) do
not hold for it, so fabric sharing gets a descriptor class of its own instead
of a widened `VirtualMemoryFdBufferDescriptor`, whose `handles` stay
`IPCAllocationHandle`s. `from_fds()` is POSIX-only by construction for the
same reason and takes no handle type; a fabric descriptor is plain data and
needs no out-of-band constructor. Tracked on the VMM epic (#2980).
opaque bytes cached by the driver at `cuMemCreate`. Repeated exports of one
allocation return the same bytes once the handle struct is zeroed before the
export (the driver fills its fields and leaves the padding, which otherwise
made two exports of one buffer compare unequal on B300); export creates no OS
object and takes no reference. `VirtualMemoryFabricBufferDescriptor` carries one `bytes` per
allocation, its size, the buffer size, and the owner's device UUID. It
compares and hashes by (sizes, handles, size); the UUID is diagnostic metadata.
`from_handles()` validates externally transported handles and sizes. Plain
`pickle`, `multiprocessing`, MPI, or a socket can carry the descriptor, with
no reduction registration or spawning-Popen retention. The fd descriptor
keeps its ownership and identity semantics and its POSIX-only `from_fds()`.

Nothing pins the allocation between export and import. The exporter must keep
the buffer alive until every importer has imported it; a descriptor parked
in a queue holds no driver reference. Import takes references of its own,
so the imported buffer is independent of the exporter after import. A stale
or forged handle fails with a driver error; an import failure carries a note
about stale handles, forgery, and differing IMEX channels. Imported memory
cannot be re-exported; the same provenance check covers both transports.

Both concrete descriptors derive from `_VirtualMemoryDescriptor`, whose
`handle_ptr` supplies an fd cast to `void*` or fills a caller-provided fabric
scratch struct and returns its address. The import loop checks the handle
type before driver calls and uses the same location check, reservation,
mappings, and range for both transports. An imported fabric allocation
reports DEVICE and the owner's device ordinal as numbered by the importer.
A location mismatch includes the owner's UUID and its local ordinal, or says
that the UUID is not visible in this process.

Construction checks `Device.properties.handle_type_fabric_supported` once
for a device-located fabric resource. It requires an NVLink fabric and an
accessible IMEX channel (`/dev/nvidia-caps-imex-channels/channel*`); otherwise
construction raises a useful `RuntimeError` before an opaque allocation
failure. `is_ipc_enabled` remains a configuration query without driver calls.
The types are available in both supported bindings majors; no version gate
is needed. Sharing is supported on one node; multi-node IMEX sharing is out
of scope and untested. Tracked on the VMM epic (#2980).

## Why `modify_allocation` returns a new buffer

Expand Down
12 changes: 11 additions & 1 deletion cuda_core/cuda/core/_memory/_ipc.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,20 @@ cdef class IPCAllocationHandle:

# Descriptor of a VirtualMemoryResource buffer: one exported handle per
# physical allocation that backs the buffer, in address order.
cdef class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor):
cdef class _VirtualMemoryDescriptor(IPCBufferDescriptor):
cdef readonly int _handle_type
cdef readonly tuple _chunk_sizes

cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL


cdef class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor):
cdef readonly tuple _handles


cdef class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor):
cdef readonly tuple _handles
cdef object _owner_uuid


# Buffer IPC Implementation
Expand Down
80 changes: 69 additions & 11 deletions cuda_core/cuda/core/_memory/_ipc.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -58,16 +58,27 @@ class IPCAllocationHandle:
@property
def uuid(self) -> uuid.UUID: ...

class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor):
class _VirtualMemoryDescriptor(IPCBufferDescriptor):
"""Common allocation layout and driver-handle interface for VMM descriptors."""
_handle_type: int
_chunk_sizes: tuple[Any, ...]

@property
def handle_type(self) -> VirtualMemoryHandleType:
"""The export handle type; the importing resource must use the same type."""
@property
def sizes(self) -> tuple[int, ...]:
"""Size of each physical allocation, in address order; the sum is at least ``size``."""

class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor):
"""Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes.

The descriptor holds one exported POSIX file descriptor per physical
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.
``"posix_fd"``; the rules below describe file descriptors.

A descriptor pins the physical memory while it exists, in every process
that holds one, and nothing else does: the exporting buffer keeps no
Expand All @@ -92,8 +103,6 @@ class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor):
which checks what it can and fails instead of mapping a size the driver
rejects.
"""
_handle_type: int
_chunk_sizes: tuple[Any, ...]
_handles: tuple[Any, ...]

@staticmethod
Expand Down Expand Up @@ -133,12 +142,6 @@ class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor):
If an ``int`` is not an open file descriptor.
"""
@property
def handle_type(self) -> VirtualMemoryHandleType:
"""The handle type the exporter used; the importing resource must be configured with the same."""
@property
def sizes(self) -> tuple[int, ...]:
"""Size of each physical allocation, in address order; the sum is at least ``size``."""
@property
def handles(self) -> tuple[IPCAllocationHandle, ...]:
"""The exported handles, one :class:`IPCAllocationHandle` per allocation, in address order.

Expand All @@ -155,5 +158,60 @@ class VirtualMemoryFdBufferDescriptor(IPCBufferDescriptor):
"""
def __reduce__(self) -> tuple[object, ...]: ...

class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor):
"""Plain-data descriptor of a :class:`VirtualMemoryBuffer` exported with fabric handles.

Each physical allocation contributes one opaque 64-byte handle, in address
order. Plain ``pickle``, ``multiprocessing``, or a transport of your own
can carry these bytes. Use :meth:`from_handles` to rebuild a descriptor
received outside cuda.core, then :meth:`Buffer.from_ipc_descriptor` to
import it. ``handle_type`` is always ``"fabric"``.

Nothing is pinned: unlike a file descriptor, a fabric descriptor parked
in a queue does not retain the memory. The exporter must keep the buffer
alive until every importer has imported it. After import, the imported
allocation holds its own driver references and is independent of the
exporter. Importing a stale or forged handle fails with a driver error.

Sharing is supported on one node; multi-node sharing is untested. Fabric
handles require an NVLink fabric and an accessible IMEX channel. Treat
handles and sizes received from another process as untrusted and import
only through :meth:`Buffer.from_ipc_descriptor`.
"""
_handles: tuple[Any, ...]

@staticmethod
def _from_exports(sizes: tuple[int, ...], handles: tuple[bytes, ...], size: int, owner_uuid: str | None) -> VirtualMemoryFabricBufferDescriptor: ...
@classmethod
def from_handles(cls, handles: Iterable[bytes], sizes: Iterable[int], *, size: int | None=None, owner_uuid: str | None=None) -> VirtualMemoryFabricBufferDescriptor:
"""Build a descriptor from fabric handles received through a transport of your own.

Parameters
----------
handles : Iterable[bytes]
One opaque 64-byte handle per physical allocation, in address order.
sizes : Iterable[int]
One positive allocation size per handle, in the same order.
size : int, optional
Buffer size, between zero and the sum of ``sizes``; the sum by default.
owner_uuid : str, optional
Exporting device's UUID, used to identify the owner in error messages.

Raises
------
ValueError
If counts differ, a handle is not 64 bytes, an allocation size is
not a positive integer, or the buffer size is outside the allocation sizes.
"""
@property
def handles(self) -> tuple[bytes, ...]:
"""One opaque 64-byte fabric handle per allocation, in address order."""
@property
def owner_uuid(self) -> str | None:
"""The exporting device's UUID, if supplied by the exporter."""
def __reduce__(self) -> tuple[object, ...]: ...
def __eq__(self, other) -> bool: ...
def __hash__(self) -> int: ...

def _reduce_allocation_handle(alloc_handle: IPCAllocationHandle) -> tuple[object, ...]: ...
def _reconstruct_allocation_handle(cls: type, df: object, uuid: uuid.UUID | None) -> IPCAllocationHandle: ...
Loading
Loading