feat(connector): share sleep-mode (cumem/VMM) KV via exported POSIX-FD handles#394
feat(connector): share sleep-mode (cumem/VMM) KV via exported POSIX-FD handles#394moreWax wants to merge 2 commits into
Conversation
…D handles vLLM's --enable-sleep-mode allocates KV with cuMemCreate (VMM); legacy CUDA IPC (storage._share_cuda_()) has no handles for such memory, so registration died with 'CUDA error: invalid argument' and PegaFlow could not be used together with sleep mode at all. This adds a VMM sharing path with no vLLM changes required: - worker: vLLM's cumem extension already requests exportable handle types wherever the driver reports fabric/posix support, so the covering allocation's generic handle exports as a POSIX FD (cuMemExportToShareableHandle). One export per cumem allocation, shared by every KV tensor inside it. - handoff: abstract unix socket + SCM_RIGHTS, served by a per-worker thread. 128-bit random tokens gate each FD (abstract socket names are world-visible in /proc/net/unix). pidfd_getfd was rejected: yama ptrace_scope=1 forbids it between sibling processes. - server: imports + maps the FD in the embedded interpreter (cuMemImportFromShareableHandle / AddressReserve / Map / SetAccess via ctypes) and returns a duck-typed view satisfying the registry contract (data_ptr / device.index / untyped_storage().nbytes). Mapping lifetime is tied to view GC: registry drop_instance already decrefs held views and gc.collect()s, which releases the import — critical, because an imported handle refcounts the exporter's physical memory. - sleep/wake lifecycle: unregister_for_sleep (drop server mappings, close exported FDs — required or sleep frees no VRAM) and reregister_after_wake (wake maps fresh physical chunks into the SAME virtual addresses, so the cached tensor dict re-registers; handles are re-read at export time). Exposed on the connector for invocation via vLLM's /collective_rpc through a thin Worker passthrough. Validated live (vLLM 0.22.1, RTX 3090): engine boots with sleep mode + connector (previously an instant init crash); sleep L1 drops the GPU from ~21GB to ~1GB; wake + re-register in seconds; a 5.3k-token prefix replay after wake hits 335/335 blocks from the server at 0.2s — the KV that sleep discarded is refilled from the external cache. Unit tests cover the FD handoff (real SCM_RIGHTS over an abstract socket), import-mapping cache/reuse, GC-driven release including FD close, release idempotence, and fail-safe cumem detection — all CPU-only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Adds a Python-side CUDA VMM (cuMemCreate) IPC path so the PegaFlow vLLM connector can register/share sleep-mode KV caches by exporting POSIX-FD handles per allocation, plus worker lifecycle hooks to safely unregister before sleep and reregister after wake.
Changes:
- Introduces
pegaflow.vmm_ipcimplementing worker-side FD export + SCM_RIGHTS handoff and server-side VMM import/mapping + view duck-typing. - Extends the worker connector to choose VMM-IPC vs legacy CUDA IPC per tensor, and adds
unregister_for_sleep()/reregister_after_wake()lifecycle methods with save-drain + failure surfacing. - Installs vLLM
Workerpassthrough RPC methods so an orchestrator can drive sleep/wake ordering via/collective_rpc, and adds unit tests for the VMM IPC layer and failure cleanup paths.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| python/pegaflow/vmm_ipc.py | New VMM-IPC implementation: POSIX-FD export + SCM_RIGHTS handoff and ctypes-based cuMem import/mapping with mapping cache. |
| python/pegaflow/connector/worker.py | Uses VMM-IPC wrapper for cumem tensors, closes exported FDs on all teardown paths, and adds sleep/wake unregister/reregister helpers. |
| python/pegaflow/connector/init.py | Installs vLLM Worker RPC passthroughs and exposes connector methods callable via /collective_rpc. |
| python/tests/test_vmm_ipc.py | New CPU-only unit tests for FD handoff behavior and import/mapping lifecycle using a fake CUDA driver. |
| python/tests/test_connector_fault_tolerance.py | Adds tests ensuring sleep unregister failure is raised and that registration failure closes exported FDs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| conn.settimeout(5.0) | ||
| token = conn.recv(64).decode(errors="replace") | ||
| with self._exports_lock: |
There was a problem hiding this comment.
Fixed in 38a8302 — the token read is now framed (loop to exactly 32 bytes; uuid4().hex length) before the lookup.
| with self._exports_lock: | ||
| token = self._token_by_dmem.get(d_mem) | ||
| if token is not None and token in self._exports: | ||
| return token | ||
| fd = export_fd() | ||
| with self._exports_lock: | ||
| token = uuid.uuid4().hex | ||
| self._exports[token] = _Export(fd=fd, d_mem=d_mem, size=size) | ||
| self._token_by_dmem[d_mem] = token | ||
| return token |
There was a problem hiding this comment.
Fixed in 38a8302 — publish_allocation re-checks under the lock after minting and closes the losing FD.
| s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) | ||
| try: | ||
| s.connect(uds_path) | ||
| s.sendall(token.encode()) | ||
| msg, ancillary, _, _ = s.recvmsg(1, socket.CMSG_SPACE(4)) | ||
| if msg != b"F" or not ancillary: | ||
| raise RuntimeError(f"FD handoff failed for token {token!r}") | ||
| fds = array.array("i") | ||
| fds.frombytes(ancillary[0][2][:4]) | ||
| fd = fds[0] |
There was a problem hiding this comment.
Fixed in 38a8302 — 10s socket timeout on the handoff and array('i').itemsize-sized ancillary buffer.
| cuda.cuMemImportFromShareableHandle( | ||
| ctypes.byref(handle), ctypes.c_int(fd), _CU_MEM_HANDLE_TYPE_POSIX_FD | ||
| ), |
There was a problem hiding this comment.
Fixed in 38a8302 — fd now passed as c_void_p (osHandle is pointer-sized).
| mapping = _Mapping(va.value, size, handle.value, fd) | ||
| with _mappings_lock: | ||
| _mappings[token] = weakref.ref(mapping) |
There was a problem hiding this comment.
Fixed in 38a8302 — a weakref callback prunes the token entry when the mapping is collected.
| def fd_server(): | ||
| server = _FdServer() | ||
| yield server | ||
| server.close_all() |
There was a problem hiding this comment.
Fixed in 38a8302 — added _FdServer.close() (closes the listening socket, ending the accept loop) and the fixture now uses it.
- frame the 32-byte token read on the SOCK_STREAM handoff (recv is not message-framed; a partial read produced spurious unknown-token errors) - close the loser's freshly minted FD when two threads race publish_allocation for the same allocation (a duplicate export pins memory) - timeout + itemsize-sized ancillary buffer on the server-side handoff recvmsg (a wedged peer must fail registration, not hang the embedded interpreter) - pass the fd as c_void_p to cuMemImportFromShareableHandle (osHandle is pointer-sized; c_int is ABI-fragile off x86-64) - weakref callback prunes dead mapping-cache entries (unbounded growth across sleep/wake cycles) - test fixture closes the per-test _FdServer listening socket (thread/fd leak across the suite) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Addressed all review feedback in 38a8302: framed 32-byte token read; mint race closes the losing FD; handoff recvmsg has a timeout and itemsize-sized ancillary buffer; |
|
Thanks for the detailed PR. Before we dive into a full review — what's your concrete use case for sleep mode? We don't use |
What
Lets
PegaKVConnectorwork with vLLM's sleep mode. Sleep mode allocates KV under the cumem (CUDA VMM /cuMemCreate) allocator, which legacy CUDA IPC cannot share —storage._share_cuda_()fails withCUDA error: invalid argument, so today PegaFlow and--enable-sleep-modeare mutually exclusive. With this PR the combination gives the best of both worlds: seconds-scale sleep/wake and KV prefixes that survive in the external cache (which sleep otherwise discards).No vLLM changes are needed: vLLM's cumem extension already requests exportable handle types wherever the driver reports fabric/POSIX handle support (which includes consumer GPUs — verified on RTX 3090), so the allocations are exportable as POSIX FDs out of the box.
Design
New module
pegaflow/vmm_ipc.py, plus a wrapper-selection branch at registration and two lifecycle methods:cuMemExportToShareableHandle(POSIX_FD). One export per allocation, shared by every KV tensor inside it. Detection is fail-safe: anything not provably cumem-managed takes the existing legacy-IPC path unchanged.SCM_RIGHTS, served by one per-worker thread. Each FD is gated by a 128-bit random token (abstract socket names are world-visible in/proc/net/unix, so the token is the secret).pidfd_getfdwas considered and rejected: yamaptrace_scope=1forbids it between sibling processes, which is exactly the worker↔server relationship.to_tensor()runs in the server's embedded interpreter and imports/maps via ctypes (cuMemImportFromShareableHandle→cuMemAddressReserve→cuMemMap→cuMemSetAccess), returning a duck-typed view that satisfies the registry contract (data_ptr/device.index/untyped_storage().nbytes). One mapping per allocation, shared across layer views.registry.drop_instancealready decrefs held views and runsgc.collect(), which unmaps/releases the import. Symmetrically, the worker exposes:unregister_for_sleep()— unregister + close all exported FDs. Must run before vLLM sleeps, or sleep frees no VRAM.reregister_after_wake()— wake maps fresh physical chunks into the same virtual addresses, so the cached kv-cache dict re-registers the same tensor objects; generic handles are re-read at export time.Both are reachable through vLLM's
/collective_rpc: the connector installs thinWorker.pegaflow_sleep_unregister/Worker.pegaflow_wake_reregisterpassthroughs at worker-role construction (the Worker class is already imported there; non-PegaFlow processes pay nothing), so an orchestrator drives:collective_rpc(pegaflow_sleep_unregister)→/sleep→/wake_up→collective_rpc(pegaflow_wake_reregister).unregister_for_sleepraises when saves haven't drained or the server unregister RPC fails — a non-OK response the caller must treat as "do not sleep" (sleeping with live mappings frees no VRAM while reporting success).unregister_context, registration-failure rollback), not just the sleep path; one export per cumem allocation is shared by all layer tensors inside it; the FD-handoff thread is wedge-proof (per-connection timeout, broad exception guard — abstract sockets are connectable by any local process) and TOCTOU-safe (dup-under-lock so a concurrently reused fd number can never hand out the wrong descriptor); a failed import releases the handle/VA/fd instead of pinning the exporter's memory.The Rust tree already contains a complete-but-unwired VMM wrapper (
pegaflow-transfer/src/cuda_lib/cumem.rs); this PR deliberately stays in Python to avoid touching the data plane — the imported mapping is an ordinary device VA, soMemcpyBackend/KernelBackendwork unchanged. Migrating the import intoCudaTensorRegistryon top of that module would be a natural follow-up.Validation
close_allinvalidation, per-allocation export reuse, garbage/stalled-client resilience of the handoff thread, import-flow + mapping-cache reuse against a call-recording fake driver, import-failure cleanup, GC-driven release including FD close, release idempotence, fail-safe cumem detection, registration-failure FD rollback, and unregister-failure raising.test_vllm_e2e_correctness.pygate (no/data/models/Qwen3-4Bhere); happy to iterate if CI or a reviewer run surfaces anything.Limitations
cuMemCreatetime (vLLM requests it whenCU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTEDprobes true; on drivers where it probes false, allocations are unexportable and this path cannot engage — the connector then fails registration exactly as before this PR)./sleep//wake_up; the PR documents the contract but cannot enforce ordering from inside the connector (vLLM has no connector sleep hooks).🤖 Generated with Claude Code