[TransferEngine] Make single unregisterLocalMemory best-effort#2962
[TransferEngine] Make single unregisterLocalMemory best-effort#2962he-yufeng wants to merge 1 commit into
Conversation
TransferEngineImpl::unregisterLocalMemory returned on the first transport whose unregister failed, leaving the region registered on the remaining transports and skipping the local bookkeeping erase. Collect the first error but attempt every transport, mirroring the batch path made best-effort in kvcache-ai#2869. Refs kvcache-ai#2869 Signed-off-by: Yufeng He <[email protected]>
There was a problem hiding this comment.
Code Review
This pull request updates unregisterLocalMemory in TransferEngineImpl to adopt a best-effort approach, attempting to unregister memory from all transports even if one fails. However, returning early when an error is encountered prevents local memory region cleanup (eraseMemoryRegionLocked), which can cause permanent address blocking and prevent future registrations. The reviewer's suggestion to perform local cleanup before returning the error is highly recommended to ensure state consistency.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if (first_error) return first_error; | ||
|
|
||
| std::unique_lock<std::shared_mutex> lock(mutex_); | ||
| eraseMemoryRegionLocked(addr); |
There was a problem hiding this comment.
If any transport fails to unregister, first_error is set to a non-zero value. In the current implementation, if first_error is non-zero, the function returns early on line 626, skipping the local bookkeeping cleanup (eraseMemoryRegionLocked(addr)).
This has two severe consequences:
- Permanent Address Block: The memory region remains in the
local_memory_regions_map. Any future attempt to register a memory region at the same address (or an overlapping address) will fail withERR_ADDRESS_OVERLAPPEDintryReserveMemoryRegions, even if the physical memory was freed and reallocated. - Impossible Recovery: If the user retries
unregisterLocalMemory, the transports that did succeed in the first attempt might now return an error (such asERR_ADDRESS_NOT_REGISTERED), which will again setfirst_errorand prevent the local bookkeeping from ever being cleaned up.
By performing the local eraseMemoryRegionLocked(addr) regardless of first_error (while still returning the error to the caller), we ensure the local engine's state remains clean and future registrations are not permanently blocked.
| if (first_error) return first_error; | |
| std::unique_lock<std::shared_mutex> lock(mutex_); | |
| eraseMemoryRegionLocked(addr); | |
| std::unique_lock<std::shared_mutex> lock(mutex_); | |
| eraseMemoryRegionLocked(addr); | |
| if (first_error) return first_error; |
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Description
TransferEngineImpl::unregisterLocalMemoryiterates every transport and returned on the first transport whoseunregisterLocalMemoryfailed:In a multi-transport setup, if the first transport's unregister fails the remaining transports keep the region registered (a resource leak) and the local
eraseMemoryRegionLocked(addr)bookkeeping is skipped.This makes the single-memory path best-effort — attempt every transport, keep the first error, and return it — exactly the pattern #2869 just established for
registerLocalMemoryBatch/unregisterLocalMemoryBatch(and the Ascend/HCCL unregister loops). Behavior on the all-success path is unchanged.Refs #2869.
Module
mooncake-transfer-engine)Type of Change
How Has This Been Tested?
I could not build the full stack on this host (Linux-only deps), so I verified the control-flow change with a standalone harness that mirrors both the old and new loops with a fake 3-transport list where the first transport fails:
So the old code stops after the first failure (leaving 2 transports registered), while the new code unregisters from all three and still returns the first error. Full compile/tests rely on CI.
Test results:
Checklist
./scripts/code_format.sh— clang-format was not available on this host; I hand-matched the existing style (Google, 80 col) and CI will verifypre-commit run --all-filesand all hooks pass — toolchain unavailable locally; relying on CIAI Assistance Disclosure
Claude Code helped locate this sibling of #2869 and build the standalone harness above. The change is small and was reviewed line by line: it mirrors the
first_errorbest-effort pattern #2869 introduced for the batch/Ascend/HCCL paths, applied to the single-memory path in the same file.