Skip to content

Commit b804e18

Browse files
Andy-Jostclaude
Andy-Jost
and
claude
committed
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 <[email protected]>
1 parent ef9b712 commit b804e18

13 files changed

Lines changed: 846 additions & 161 deletions

File tree

‎cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.md‎

Lines changed: 40 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -232,21 +232,46 @@ process yields an alias of the exporter's memory at a new address. A
232232
`VirtualMemoryResource` pickles as (device, options), which is what lets a
233233
`Buffer` pickle as (resource, descriptor).
234234

235-
File descriptors and memory. The driver copies each shared allocation into a
236-
handle of its own during the import and keeps no reference to the file
237-
descriptor, so the imported buffer does not keep the descriptor; it records
238-
only that it was imported (`is_mapped`). The descriptor's file descriptors live
239-
as long as the descriptor object, on the exporter (where the buffer caches it)
240-
and in every process that unpickled a copy, and each open descriptor keeps the
241-
physical memory allocated. A descriptor on a `Queue` pins the memory in the
242-
sender until the receiver unpickles it.
243-
244-
Re-export. The driver exports only allocations created with the requested
245-
handle type, and an imported allocation carries none, so an imported buffer,
246-
and any buffer grown from one, cannot be exported again; `ipc_descriptor` says
247-
so instead of surfacing the driver's INVALID_VALUE. A `modify_allocation`
248-
config must keep the resource's `handle_type` for the same reason: every chunk
249-
of a buffer must be exportable the same way.
235+
File descriptors and memory. A descriptor pins the physical memory while it
236+
exists, in every process that holds one, and nothing else does. The exporting
237+
buffer keeps no descriptor: `ipc_descriptor` exports again on every access and
238+
returns a new descriptor that owns one file descriptor per allocation, and
239+
`VirtualMemoryBuffer.__reduce__` builds a transient one while `multiprocessing`
240+
pickles the buffer, so a sender holds no file descriptors between sends. One
241+
transport needs more: a `Queue`, `Pipe`, or `Pool` duplicates the file
242+
descriptors while pickling (the resource sharer), but a spawned `Process`
243+
records their numbers during pickling and passes them to the child only when
244+
it is created (`spawnv_passfds`), so `__reduce__` keeps the transient
245+
descriptor on the spawning `Popen` object, which dies with the `Process`
246+
object. Without that, the numbers are reused by the spawn's own pipes
247+
("bad value(s) in fds_to_keep") or the child imports a stranger's fd. The
248+
driver copies
249+
each shared allocation into a handle of its own during the import and keeps no
250+
reference to the file descriptor, so the imported buffer does not keep the
251+
descriptor either; it records only that it was imported (`is_mapped`). A
252+
descriptor parked in a `Queue` or sent to a `Pool` pins the memory in the
253+
sender until the receiver unpickles it. Besides `multiprocessing`, a process
254+
with its own file descriptor passing sends the descriptor's `fds`, `sizes`,
255+
`handle_type`, and `size` and rebuilds it with
256+
`VirtualMemoryIPCBufferDescriptor.from_fds()`, which duplicates the integers
257+
it is given.
258+
259+
Re-export. The driver exports only allocations this process created with the
260+
requested handle type; an imported allocation has none, and the driver reports
261+
INVALID_VALUE. cuda.core does not rely on that: `MemAllocationBox.imported`
262+
records provenance per allocation (`mem_allocation_is_imported`), and since
263+
`modify_allocation` reuses the allocation handles of the input, the mark
264+
travels into every alias, in-place grow, and move. `_export_ipc_descriptor`
265+
checks the allocations that cover the buffer before any driver call and raises
266+
a RuntimeError that says the buffer contains memory imported from another
267+
process, that such memory cannot be exported again, and that the caller should
268+
forward the descriptor it imported from or copy into a buffer it owns. The same
269+
exception surfaces from `pickle`/`multiprocessing` (a `Queue` feeder thread
270+
reports it through `Queue._on_queue_feeder_error` and drops the item; the queue
271+
stays usable). PyTorch applies the same contract to received CUDA tensors
272+
("Attempted to send CUDA tensor received from another process"). There is no
273+
retention option. A `modify_allocation` config must keep the resource's
274+
`handle_type`: every chunk of a buffer must be exportable the same way.
250275

251276
## Why `modify_allocation` returns a new buffer
252277

‎cuda_core/cuda/core/_cpp/rt/api.hpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,13 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat
262262
// Size of the allocation; the only size cuMemMap accepts for it.
263263
size_t mem_allocation_size(const MemAllocationHandle& h) noexcept;
264264

265+
// Whether cuMemImportFromShareableHandle produced the allocation. The driver
266+
// exports only allocations this process created with a handle type, so an
267+
// imported allocation cannot be exported again; callers check this before
268+
// cuMemExportToShareableHandle instead of relying on the driver's error. The
269+
// mark travels with the allocation into every range that maps it.
270+
bool mem_allocation_is_imported(const MemAllocationHandle& h) noexcept;
271+
265272
// Import a physical allocation that another process exported with
266273
// cuMemExportToShareableHandle. `os_handle` is the shareable handle in the
267274
// form cuMemImportFromShareableHandle expects for `handle_type` (a POSIX file

‎cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,10 @@ using namespace detail;
2828
namespace {
2929

3030
struct MemAllocationBox {
31-
MemAllocationValue resource; // from cuMemCreate
31+
MemAllocationValue resource; // from cuMemCreate or cuMemImportFromShareableHandle
3232
size_t size; // the only size cuMemMap accepts for it
3333
std::vector<CUmemAccessDesc> access; // applied to every mapping of this allocation
34+
bool imported; // another process created it; it cannot be exported again
3435
};
3536

3637
struct VaReservationBox {
@@ -56,9 +57,9 @@ const Box* box_of(const Handle& h) noexcept {
5657
// cuMemCreate or cuMemImportFromShareableHandle produced it. The last
5758
// reference releases it; the memory is freed once no mapping remains.
5859
MemAllocationHandle wrap_mem_allocation(CUmemGenericAllocationHandle handle, size_t size,
59-
std::vector<CUmemAccessDesc>&& access) {
60+
std::vector<CUmemAccessDesc>&& access, bool imported) {
6061
auto box = std::shared_ptr<const MemAllocationBox>(
61-
new MemAllocationBox{{handle}, size, std::move(access)},
62+
new MemAllocationBox{{handle}, size, std::move(access), imported},
6263
[](const MemAllocationBox* b) {
6364
GILReleaseGuard gil;
6465
pw_cuMemRelease(b->resource.raw);
@@ -84,7 +85,7 @@ MemAllocationHandle create_mem_allocation_handle(size_t size, const CUmemAllocat
8485
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemCreate, &handle, size, &prop, 0))) {
8586
return {};
8687
}
87-
return wrap_mem_allocation(handle, size, std::move(access));
88+
return wrap_mem_allocation(handle, size, std::move(access), false);
8889
}
8990

9091
MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocationHandleType handle_type,
@@ -97,13 +98,17 @@ MemAllocationHandle import_mem_allocation_handle(void* os_handle, CUmemAllocatio
9798
if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemImportFromShareableHandle, &handle, os_handle, handle_type))) {
9899
return {};
99100
}
100-
return wrap_mem_allocation(handle, size, std::move(access));
101+
return wrap_mem_allocation(handle, size, std::move(access), true);
101102
}
102103

103104
size_t mem_allocation_size(const MemAllocationHandle& h) noexcept {
104105
return h ? box_of<MemAllocationBox>(h)->size : 0;
105106
}
106107

108+
bool mem_allocation_is_imported(const MemAllocationHandle& h) noexcept {
109+
return h ? box_of<MemAllocationBox>(h)->imported : false;
110+
}
111+
107112
// ============================================================================
108113
// Address reservations
109114
// ============================================================================

‎cuda_core/cuda/core/_memory/_buffer.pyi‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,11 @@ class Buffer:
109109
"""
110110
@property
111111
def ipc_descriptor(self) -> IPCBufferDescriptor:
112-
"""Descriptor for sharing this buffer with other processes."""
112+
"""Descriptor for sharing this buffer with other processes.
113+
114+
A pool-backed buffer caches its descriptor; a
115+
:class:`VirtualMemoryBuffer` exports a new one on every access.
116+
"""
113117
def close(self, stream: Stream | GraphBuilder | None=None) -> None:
114118
"""Deallocate this buffer asynchronously on the given stream.
115119

‎cuda_core/cuda/core/_memory/_buffer.pyx‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -421,7 +421,11 @@ cdef class Buffer:
421421
@property
422422
@cython.critical_section
423423
def ipc_descriptor(self) -> IPCBufferDescriptor:
424-
"""Descriptor for sharing this buffer with other processes."""
424+
"""Descriptor for sharing this buffer with other processes.
425+
426+
A pool-backed buffer caches its descriptor; a
427+
:class:`VirtualMemoryBuffer` exports a new one on every access.
428+
"""
425429
Buffer_check_open(self)
426430
cdef object ipc_data
427431
if self._ipc_data is None:

‎cuda_core/cuda/core/_memory/_ipc.pyi‎

Lines changed: 75 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
# This file was generated by stubgen-pyx v0.2.22 from cuda_core/cuda/core/_memory/_ipc.pyx
22

33
import uuid
4-
from typing import Any
4+
from typing import Any, Iterable
5+
6+
from cuda.core.typing import VirtualMemoryHandleType
57

68
__all__ = []
79

@@ -60,18 +62,27 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor):
6062
"""Serializable object describing a :class:`VirtualMemoryBuffer` that can be shared between processes.
6163
6264
The descriptor holds one exported handle per physical allocation that
63-
backs the buffer, in address order, with each allocation's size. With
64-
POSIX file descriptors the handles are :class:`IPCAllocationHandle`
65-
objects: the descriptor owns them and closes them when it is released, and
66-
``multiprocessing`` duplicates them into the receiving process. Plain
67-
``pickle`` cannot carry file descriptors; send the descriptor, or the
68-
buffer, through ``multiprocessing`` (a ``Queue``, a ``Pipe``, ``Process``
69-
arguments, or a ``Pool``).
70-
71-
A live descriptor costs one file descriptor per allocation and keeps the
72-
physical memory allocated, in every process that holds a copy, until it is
73-
released. A descriptor placed on a ``Queue`` or sent to a ``Pool`` pins the
74-
memory in the sender until the receiver has unpickled it.
65+
backs the buffer, in address order (``handles``), with each allocation's
66+
size (``sizes``) and the buffer's size (``size``, which may be smaller
67+
than their sum). With POSIX file descriptors the handles are
68+
:class:`IPCAllocationHandle` objects: the descriptor owns them and closes
69+
them when it is released.
70+
71+
A descriptor pins the physical memory while it exists, in every process
72+
that holds one, and nothing else does: the exporting buffer keeps no
73+
descriptor, and an imported buffer does not keep the one it came from. A
74+
descriptor parked in a ``multiprocessing`` ``Queue`` or sent to a ``Pool``
75+
pins the memory in the sender until the receiver has unpickled it; a
76+
buffer passed as a ``Process`` argument pins it until the ``Process``
77+
object is released, because the child receives the file descriptors when
78+
it is spawned.
79+
80+
Two transports exist. ``multiprocessing`` (a ``Queue``, a ``Pipe``,
81+
``Process`` arguments, or a ``Pool``) pickles the descriptor, or a buffer,
82+
and duplicates the file descriptors into the receiving process; plain
83+
``pickle`` cannot carry them and raises. A process with its own file
84+
descriptor passing sends ``fds``, ``sizes``, ``handle_type``, and ``size``
85+
itself, and the receiver rebuilds the descriptor with :meth:`from_fds`.
7586
7687
Note
7788
----
@@ -86,6 +97,57 @@ class VirtualMemoryIPCBufferDescriptor(IPCBufferDescriptor):
8697

8798
@staticmethod
8899
def _from_exports(handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple[Any, ...], size: int) -> VirtualMemoryIPCBufferDescriptor: ...
100+
@classmethod
101+
def from_fds(cls, fds: Iterable[int | IPCAllocationHandle], sizes: Iterable[int], *, handle_type: VirtualMemoryHandleType | str='posix_fd', size: int | None=None) -> VirtualMemoryIPCBufferDescriptor:
102+
"""Build a descriptor from file descriptors received outside ``multiprocessing``.
103+
104+
The exporter sends ``fds``, ``sizes``, ``handle_type``, and ``size``
105+
of its descriptor through a transport of its own (a Unix socket with
106+
``SCM_RIGHTS``, for example); the receiver rebuilds the descriptor
107+
here and imports it with :meth:`Buffer.from_ipc_descriptor`.
108+
109+
Parameters
110+
----------
111+
fds : Iterable[int | IPCAllocationHandle]
112+
One handle per physical allocation, in address order, as the
113+
exporter's ``fds`` listed them. An ``int`` is duplicated: the
114+
caller keeps ownership of the original and closes it. An
115+
:class:`IPCAllocationHandle` is shared and closes with its last
116+
holder.
117+
sizes : Iterable[int]
118+
The exporter's ``sizes``, one per handle.
119+
handle_type : VirtualMemoryHandleType | str, optional
120+
The exporter's ``handle_type``; ``"posix_fd"`` by default.
121+
size : int, optional
122+
The exporter's ``size``, at most the sum of ``sizes``; the sum by
123+
default.
124+
125+
Raises
126+
------
127+
ValueError
128+
If the counts differ, a size is not positive, ``size`` exceeds the
129+
sum of ``sizes``, or ``handle_type`` cannot be shared.
130+
RuntimeError
131+
If an :class:`IPCAllocationHandle` is closed.
132+
OSError
133+
If an ``int`` is not an open file descriptor.
134+
"""
135+
@property
136+
def handle_type(self) -> VirtualMemoryHandleType:
137+
"""The handle type the exporter used; the importing resource must be configured with the same."""
138+
@property
139+
def sizes(self) -> tuple[int, ...]:
140+
"""Size of each physical allocation, in address order; the sum is at least ``size``."""
141+
@property
142+
def handles(self) -> tuple[IPCAllocationHandle, ...]:
143+
"""The exported handles, one per allocation, in address order; the descriptor owns them."""
144+
@property
145+
def fds(self) -> tuple[int, ...]:
146+
"""The exported file descriptors as integers, for a transport of your own.
147+
148+
They belong to this descriptor: send or duplicate them while it is
149+
alive. Each one is released when the descriptor is.
150+
"""
89151
def __reduce__(self) -> tuple[object, ...]: ...
90152

91153
def _reduce_allocation_handle(alloc_handle: IPCAllocationHandle) -> tuple[object, ...]: ...

0 commit comments

Comments
 (0)