Repository navigation
[Bugfix][Store] Register ShmHelper memory with SPDK for NoF zero-copy transfers - #3717
Conversation
NUABO
changed the title
SPDK's RDMA transport can only DMA from SPDK-registered memory. The host pool allocated via ShmHelper (memfd+mmap) was never registered, so NoF (NVMe-oF) put/get passing user buffers directly fails with "spdk_rdma_get_translation: No translation for ptr". Register the mapping with spdk_mem_register in allocate() (USE_NOF only) and unregister before munmap in free()/cleanup(). Registration failure is non-fatal so non-NoF paths are unaffected. A dedicated ShmSegment::spdk_registered flag tracks SPDK registration, keeping the existing `registered` flag (used by DummyClient for IPC registration) untouched.
ShmHelper::allocate pads shm->size up to the hugepage/2MB boundary when MC_STORE_REGISTER_SPDK=1 is set, but DummyClient::register_buffer's strict `size != shm->size` check (with only an is_hugepage() compensation) then rejected every non-2MB-multiple buffer (e.g. the SGLang host KV pool), so registration returned -1 even though the buffer was valid. The caller-supplied size is only used for this validation (register_shm_via_ipc sends shm->size, which the receiver maps), so the fix records the caller's original request in a new ShmSegment::requested_size field and compares against it instead of re-deriving the alignment policy. This is robust to any future change in the alignment (2MB / 1GB / none on SPDK >= 26.09).
ShmHelper and SpdkWrapper are both function-local statics. The first host-pool allocation constructs SpdkWrapper AFTER ShmHelper (via RegisterMemory inside ShmHelper::allocate), so at process exit SpdkWrapper was destroyed FIRST (~SpdkWrapper -> Cleanup() -> spdk_env_fini()) and then ~ShmHelper -> cleanup() -> SpdkWrapper::GetInstance().UnregisterMemory() ran on the destroyed singleton, after env fini (UB; NULL-deref of g_mem_reg_map on SPDK >= 26.09). Fix: - ShmHelper's constructor forces SpdkWrapper to complete construction before it does (SpdkWrapper::GetInstance(); under USE_NOF), so ShmHelper is destroyed first and its unregister runs while the SPDK env is still alive. Constructing SpdkWrapper here is side-effect free: its ctor is `= default` and the env is only initialized lazily via InitializeEnv(). - UnregisterMemory now early-returns 0 when the env was never initialized or has already been finalized by Cleanup() -> spdk_env_fini(), instead of calling spdk_mem_unregister on a torn-down or never-created global mem_map.
…e/dummy) In the dummy/real split, the host pool is allocated in the SGLang/dummy process (ShmHelper registers ITS mapping with SPDK under MC_STORE_REGISTER_SPDK=1), but the NoF submit runs in the separate real-client process on ITS OWN mapping of the shared memfd (resolved via shm_addr_offset). That receiver mapping was never spdk_mem_registered, so NoF transfers from the shared pool still failed with "No translation for ptr". Register the receiver-side mapping in RealClient::map_shm_internal_with_device and unregister it before munmap on every teardown path: - map_shm_internal_with_device: when MC_STORE_REGISTER_SPDK=1 and shm_size is a 2MB multiple, map the fd at a 2MB-aligned base (mmap_shm_2mb_aligned, now shared in utils.h) and spdk_mem_register it; registration failure is non-fatal. Falls back to a plain mmap when the aligned MAP_FIXED fails for HugeTLB fds (e.g. 1GB hugepages, whose base must be hugepage-aligned; the kernel then picks a 2MB/1GB-aligned base). - Unregister before munmap in unmap_shm_internal, tearDownAll_internal, ascend_unmap_shm_internal (POSIX branch) and unregister_shm_buffer_internal, guarded by a new MappedShm::spdk_registered flag. - ShmHelper::is_register_spdk_enabled() is the single source of truth for the env semantics, shared by the sender (constructor) and the receiver. Also moves mmap_shm_2mb_aligned from shm_helper.cpp (static) to utils.h (inline) so both sides share it, and notes in comments that on SPDK >= 26.09 the 2MB alignment can be removed.
Three follow-up fixes to the receiver-side SPDK registration: - unmap_shm_internal: on a partial unregisterLocalMemory failure the old code munmapped only the failing segment and dropped the sibling segments via mapped_shms.clear() without unregistering or unmapping them, leaking their mappings and the newly-added SPDK registrations. Track the failure with a flag and munmap every segment (each still gets its unregisterLocalMemory + SPDK unregister), then report the error at the end. - map_shm_internal_with_device: reserve vector capacity before push_back so a bad_alloc from reallocation cannot happen after RegisterMemory succeeded, which would leak the SPDK registration and the mapping. If reserve itself throws it happens before registration. - unregister_shm_buffer_internal: move the SPDK unregister block before the `if (!rc)` early return so it runs on the failure path too, matching the ordering in unmap_shm_internal (SPDK unregistered before any munmap).
Fix: when MC_STORE_REGISTER_SPDK=1, ShmHelper now forces HugeTLB backing for its allocations so registration succeeds and NoF RDMA can DMA to/from host-pool buffers (no SGLang-side staging needed). Defaults to 2MB hugepages; MC_STORE_HUGEPAGE_SIZE=1GB is honored. If the host has no hugepages reserved the first host-pool allocation fails loudly with a message naming the hugepage size and nr_hugepages to configure (or to unset MC_STORE_REGISTER_SPDK), instead of silently losing NoF zero-copy. When the env is unset the allocation path is byte-for-byte the old flow. Follow-ups to the same receiver-side registration work, also folded in: - unregister_shm_buffer_internal unmaps/erases even when unregisterLocalMemory fails and only then reports INTERNAL_ERROR (symmetric with unmap_shm_internal); DummyClient clears its registered flag on INTERNAL_ERROR/INVALID_PARAMS, keeps it on RPC_FAIL/RPC_TIMEOUT. - DummyClient::register_buffer accepts both the original requested_size and the padded mapping size, so padded (hugepage-aligned) segments register cleanly. Signed-off-by: tan changzhi <544463199@qq.com>
github-actions
Bot
added
the
documentation
|
The teardown path is unsafe when |
Thank you for the review. I will review the code and fix the issue. |
free()/cleanup() and the RealClient teardown paths previously munmapped the region and cleared spdk_registered even when spdk_mem_unregister() failed, leaving SPDK's RDMA mem_map holding a translation to an unmapped or later reused VA and risking silent NoF data corruption. Only munmap and clear spdk_registered after a successful unregister. On failure the mapping is retained (never munmapped, so the VA cannot be reused) and quarantined: ShmHelper keeps the segment for a later free() retry, and RealClient moves the MappedShm into quarantine_shms_ and reports INTERNAL_ERROR. quarantine_shms_ is re-attempted at each teardown entry point (RetryQuarantinedShmsLocked) so transient unregister failures do not leak the memory until process exit. Signed-off-by: tan changzhi <544463199@qq.com>
@catyans Thanks for the review — the underlying concern is valid and I've fixed it in commit ("Retain SPDK-registered shm mappings when spdk_mem_unregister fails"). What changed All teardown paths (ShmHelper::free()/cleanup(), RealClient::unmap_shm_internal, ascend_unmap_shm_internal, unregister_shm_buffer_internal, tearDownAll_internal) now only munmap/clear spdk_registered after spdk_mem_unregister() succeeds. On failure the mapping is retained (never unmapped, so the VA cannot be reused while SPDK still holds a translation to it) and quarantined: ShmHelper keeps the segment for a later free() retry; RealClient moves it into quarantine_shms_, reports INTERNAL_ERROR, and retries it at each teardown entry point (RetryQuarantinedShmsLocked) so a transient unregister failure does not leak memory until process exit. Why I did not add a failure-injection regression test I considered it carefully but it does not fit the current test infrastructure: No injection seam. SpdkWrapper is a non-virtual singleton that links the real SPDK symbols directly. Forcing spdk_mem_unregister to fail from a test would require adding a test-only fault-injection hook to production code, which I'd rather avoid. Environment prerequisites. Reaching the guarded state (spdk_registered == true) requires a real allocation path: ShmHelper::allocate under MC_STORE_REGISTER_SPDK=1 forces HugeTLB-backed mappings, and the receiver path needs an fd over IPC. A fault hook can bypass the SPDK calls but not mmap(MAP_HUGETLB); CI runners without a configured hugepage pool cannot run it. No NoF CI coverage today. There is currently no USE_NOF build/test job in CI (the existing nof_heartbeat_test is also built only locally). A new USE_NOF test would not be exercised by CI, so it would neither prevent regressions nor justify its maintenance cost. The invariant is instead guaranteed structurally: munmap is only reachable after a successful unregister, and a failed one leaves the mapping alive and quarantined. I verified the change on the Linux + USE_NOF setup by exercising teardown/unmap/unregister paths and confirming no mapping is released while its SPDK registration is still live. If a USE_NOF CI job with hugepages becomes available, I'm happy to add the fault-injection test on top of it. |
|
hi @zxpdemonio @catyans I have already updated the code as required. The CI test errors are unrelated to my modifications. Please review it again, Thank you. |
|
hi @fcczzz could you help review this pull request? thanks |
Sure, I’m still reviewing it. I’ll leave some feedback once I’m done. Thanks! |
Great, thank you for your review! I’ll keep an eye out for your feedback and address it promptly. |
RealClient teardown paths munmapped a received shm even when Client::unregisterLocalMemory() failed. The transfer engine keeps the region in that case: RdmaTransport::unregisterLocalMemoryInternal returns before deregistering the MR when the metadata update fails, and TransferEngineImpl::unregisterLocalMemory() only erases its region record after every transport succeeded. A later mmap reusing that address then fails to register with ERR_ADDRESS_OVERLAPPED, or NoF DMAs into unmapped memory. Track this with MappedShm::te_unregister_pending and, in unmap_shm_internal, ascend_unmap_shm_internal and unregister_shm_buffer_internal (both the USE_NOF and the non-NoF arm), only munmap once the transfer engine AND SPDK unregistrations succeeded; otherwise retain and quarantine the mapping and report INTERNAL_ERROR. RetryQuarantinedShmsLocked() retries TE before SPDK for quarantined segments, but only for segments whose TE unregister actually failed: Client::unregisterLocalMemory() collapses every TE error to INVALID_PARAMS, so retrying unconditionally could never converge. Also check the SPDK rollback result in ShmHelper::allocate() and RealClient::map_shm_internal_with_device(): spdk_mem_register() can fail after partially registering the range, and when the compensating UnregisterMemory() also failed, spdk_registered stayed false so the teardown guard was skipped and the mapping was munmapped anyway. A rollback result other than 0 / -EINVAL now marks the range registered, so teardown attempts the unregister again and quarantines the mapping on failure. Signed-off-by: tan changzhi <544463199@qq.com>
A mapping quarantined because unregisterLocalMemory() failed kept spdk_registered=true even when the SPDK unregister had already succeeded. RetryQuarantinedShmsLocked() then called spdk_mem_unregister() on the already released range, which returns the benign -EINVAL no-op (no segment in g_mem_reg_map is REGISTERED), and the retry treats any non-zero rc as failure -- so the entry never converged: it stayed quarantined, the mapping was never munmapped, and every later teardown entry point logged an ERROR and retried again. Record the actual SPDK side result (`!spdk_ok`) in the three quarantine branches (unmap_shm_internal, ascend_unmap_shm_internal, unregister_shm_buffer_internal), so the flag keeps meaning "SPDK still holds a translation for this range" -- the same invariant the success paths already maintain when they clear it before munmap. Signed-off-by: tan changzhi <544463199@qq.com>
unregister_shm_buffer_internal() returns INVALID_PARAMS for two unrelated situations: the receiver holds no such mapping (the two early returns before any teardown), and set_context_if_needed() failed before the teardown ran, leaving the mapping, its transfer engine region and its SPDK registration all live and the entry still in mapped_shms. DummyClient::unregister_buffer() drops shm->registered on INVALID_PARAMS, so the second case made the dummy forget a mapping the receiver still holds. Every later unregister attempt is then rejected locally with "Buffer is not registered with RealClient" and never reaches the receiver, so the mapping is leaked for the lifetime of the process; the reconnect loop, which re-registers only the segments still flagged as registered, also skips it. unregister_device_buffer_for_reconnect() tracks device buffers the same way through the same RPC. Return RPC_FAIL there instead: the dummy already treats anything outside OK / INTERNAL_ERROR / INVALID_PARAMS as "the operation did not complete, keep the bookkeeping and retry". The two not-found paths keep INVALID_PARAMS, so their clearing behavior and the meaning the callers already assume for it are unchanged. This branch is compiled only under USE_ASCEND_DIRECT, and only on the Ascend context setup failure path, so non-Ascend builds are unaffected. Signed-off-by: tan changzhi <544463199@qq.com>
Rolling back a failed spdk_mem_register() treated -EINVAL as proof that nothing had been registered, so free()/cleanup() went on to munmap. That is unsound in SPDK v23.01.1: spdk_mem_register() marks the range in g_mem_reg_map before running its notify callbacks and does not roll back on failure (memory.c:370-384), and in iova=va the vtophys notify installs the IOMMU DMA mapping before populating the translation table (memory.c:1089-1108). An unregister then returns -EINVAL while the DMA mapping is still live -- for a mid-loop failure that mapping covers the whole range, because buffer_len != len is detected before vtophys_iommu_unmap_dma (memory.c:1217-1224) -- so the region could be munmapped and reused while SPDK still translates to it. Only rc == 0 proves the range is clean, so every unregister site now retains/quarantines the mapping on any non-zero return, on both the sender (ShmHelper) and the receiver (RealClient::map_shm_internal_with_device). -EINVAL covers both a failure before anything was mapped and one after the DMA mapping was installed, and the codes cannot be told apart, so retaining a mapping that turns out to have been clean is the accepted cost. To preserve the munmap for ranges SPDK rejects before marking anything, add SpdkWrapper::IsRegistrableRange(), which mirrors the up-front -EINVAL checks (memory.c:339-348), and gate the registration on it. -EBUSY is handled separately: the range is already registered to another owner, so rolling back would clear their registration instead (memory.c:358-366, 445-466). OpenNofSegment() now initializes the env itself instead of relying on an earlier RegisterMemory() to have done so.
|
hi, no further blocking findings from my review. Please resolve the merge conflicts with main. |
Thank you very much! I will resolve the conflict as soon as possible. |
Conflicts resolved:
- mooncake-store/src/dummy_client.cpp
* unregister_device_buffer_for_reconnect(): keep main's reference-counted
release_buffer_registration() and function-scope lock, and keep our wider
erase condition (drop local bookkeeping on INTERNAL_ERROR/INVALID_PARAMS
too). Removed the inner lock_guard our side had: it would re-lock the
non-recursive registered_device_buffers_mutex_ and deadlock.
* register_buffer(): keep main's hugepage size normalization and our
requested_size/size bounds check (accepts both the caller's original
request and the padded mapping size).
- mooncake-store/include/utils.h (deleted by main): accept the deletion and
move our mmap_shm_2mb_aligned() into include/common/mmap_aligned.h, matching
main's split-utility-code-by-function layout. SZ_2MB/align_up/
get_hugepage_size_from_env come from common/client_buffer_allocation.h.
- mooncake-store/src/shm_helper.cpp: keep main's includes/HugepageConfig and
our SPDK registration + teardown-quarantine logic; point the
mmap_shm_2mb_aligned comment at the new header.
No behavior change intended beyond the deadlock avoidance above.
Signed-off-by: tan changzhi <544463199@qq.com>
Hi @fcczzz, I've resolved the merge conflicts with main. Could you please take another look when you get a chance? Thanks! |
Resolve conflicts between the NoF/SPDK decoupling and mainline work: - Env-var structuration PRs (kvcache-ai#4166/kvcache-ai#4170/kvcache-ai#4203): keep the neutral NoF naming/refactor while adopting the client-owned config structs (NoFDebugConfig, NoFWorkerPoolConfig, FilereadWorkerPoolConfig); QoS byte-limit vars stay in transfer_task.cpp, matching main (open kvcache-ai#4171). - kvcache-ai#3665 embedded master + P2PHANDSHAKE metadata resolution compose with the added nof_initiator injection in Client::Create. - kvcache-ai#4070 TENT native API layered onto the refactored TransferSubmitter. - kvcache-ai#3717 (ShmHelper/receiver shm SPDK registration) is ported from the deleted SpdkWrapper singleton onto the NVMeoFInitiator interface and the process-global NofPageRegistry: registrations normalize to 2MB pages with per-page refcounts, -EBUSY pages are absorbed as external, and a failed Register owns no registration. ShmHelper's singleton is now deliberately leaked so exit-time cleanup can never touch a torn-down registry/env; the SpdkEnvGuard process pin makes the old construction-ordering hack unnecessary.
Description
fix: #3736
Implement NOF integration with SGLang, SPDK's RDMA transport can only DMA from SPDK-registered memory. The host pool allocated via ShmHelper (memfd+mmap) was never registered, so NoF (NVMe-oF) put/get passing user buffers directly fails with "spdk_rdma_get_translation: No translation for ptr".
Register the mapping with spdk_mem_register in allocate() (USE_NOF only) and unregister before munmap in free()/cleanup(). Registration failure is non-fatal so non-NoF paths are unaffected. A dedicated ShmSegment::spdk_registered flag tracks SPDK registration, keeping the existing
registeredflag (used by DummyClient for IPC registration) untouched.This PR fixes an issue where SGLang's MooncakeHostTensorAllocator triggers a No translation for ptr error under Mooncake NoF mode. The allocator uses MooncakeHostMemAllocator → ShmHelper (memfd+mmap), but the allocated memory was not registered with SPDK, so RDMA transport cannot access the buffer. The patch adds spdk_mem_register()/spdk_mem_unregister() to ShmHelper's allocate/free paths to enable NoF zero-copy transfers for SGLang.
Module
mooncake-transfer-engine)mooncake-store)mooncake-reshard)mooncake-ep)mooncake-pg)mooncake-integration)mooncake-p2p-store)mooncake-wheel)mooncake-common)mooncake-rl)Type of Change
How Has This Been Tested?
Test commands:
# Example: bash scripts/run_ci_test.shTest results:
Checklist
./scripts/code_format.shAI Assistance Disclosure