Repository navigation
cuda.core: transfer VMM IPC allocations one at a time #3110
Description
Summary
#3028 shares VirtualMemoryResource buffers across processes with one POSIX file descriptor per physical allocation. A descriptor holds all N descriptors at once, and each transport multiplies that count. With chunks of a few MB, a buffer of several GB holds several hundred allocations, so 2N reaches the default limit of 1024 at about 500 chunks. When the receiver is the side at its limit, the kernel drops descriptors from the message and multiprocessing reports only received 0 items of ancdata (pytorch/pytorch#34768).
Analysis by @juenglin in #3028 (review) and #3028 (comment).
Where the count scales
- Export:
ipc_descriptoropens one descriptor per chunk before it returns. Queue,Pipe,Pool: pickling runsDupFdonce per handle, andresource_sharerholds each copy until the receiver fetches it. The sender holds up to 2N during the transfer.- Receive: unpickling fetches all N before the import starts.
from_fds: it duplicates every integer it is given, so the receiver holds 2N until the caller closes the originals.- Raw socket: one
SCM_RIGHTSmessage carries at most 253 descriptors (SCM_MAX_FD), so the documented one-callsend_fdspattern fails above 253 chunks.
Proposed design
- Pickling registers one callback per chunk with a sender-side server and sends only chunk ids and sizes. No descriptor is opened yet.
- Each callback exports its chunk when the receiver asks, sends that single descriptor, and closes it.
- Unpickling reserves the whole address range first, then fetches, imports, maps, and closes one chunk at a time. If a chunk fails, it releases everything mapped so far.
Prior art: NCCL moves cuMem handles one at a time through a proxy thread over a Unix socket (p2p.cc, proxy.cc). PyTorch expandable segments pull each descriptor with pidfd_getfd, which needs ptrace permission and fails in containers (pytorch/pytorch#165685), so that part should not be copied. CPython's resource_sharer already sends one descriptor per connection; all N pile up only because one unpickle fetches every one.
Consequences
queue.put(buf)keeps working unchanged, and each side holds about one descriptor at a time.- A spawned child fetches its chunks from the parent's sharer like any other receiver, so the
get_spawning_popen()stash inVirtualMemoryBuffer.__reduce__becomes unnecessary. The rule that the memory stays pinned until theProcessobject is released goes with it. VirtualMemoryIPCBufferDescriptor.from_fds()already takes anIterable, so it can consume a generator and close each descriptor after its import without a signature change. The export side needs a new generator that yields one handle per chunk next tofds.- Fabric handles are 64 bytes of plain data and do not have this problem. The design stays specific to
posix_fd.
Activity