Skip to content

[TransferEngine] Make single unregisterLocalMemory best-effort#2962

Open
he-yufeng wants to merge 1 commit into
kvcache-ai:mainfrom
he-yufeng:fix/unregister-local-memory-best-effort
Open

[TransferEngine] Make single unregisterLocalMemory best-effort#2962
he-yufeng wants to merge 1 commit into
kvcache-ai:mainfrom
he-yufeng:fix/unregister-local-memory-best-effort

Conversation

@he-yufeng

@he-yufeng he-yufeng commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Description

TransferEngineImpl::unregisterLocalMemory iterates every transport and returned on the first transport whose unregisterLocalMemory failed:

for (auto& transport : multi_transports_->listTransports()) {
    int ret = transport->unregisterLocalMemory(addr, update_metadata);
    if (ret) return ret;   // aborts before the remaining transports
}

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

  • Transfer Engine (mooncake-transfer-engine)

Type of Change

  • Bug fix

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:

OLD: attempts=1 ret=5 -> transports left registered=2 (LEAK)
NEW: attempts=3 ret=5 -> all transports attempted, error still propagated

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:

  • Unit tests pass
  • Integration tests pass (if applicable)
  • Manual testing done (standalone control-flow harness, above)

Checklist

  • I have performed a self-review of my own code
  • I have formatted my code using ./scripts/code_format.sh — clang-format was not available on this host; I hand-matched the existing style (Google, 80 col) and CI will verify
  • I have run pre-commit run --all-files and all hooks pass — toolchain unavailable locally; relying on CI
  • I have updated the documentation (if applicable) — not applicable
  • I have added tests to prove my changes are effective — no unit-test seam exists for this transport-loop path without a full multi-transport engine; verified with the standalone harness above
  • For changes >500 LOC: I have filed an RFC issue — not applicable (5 LOC)

AI Assistance Disclosure

  • AI tools were used (specify below)

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_error best-effort pattern #2869 introduced for the batch/Ascend/HCCL paths, applied to the single-memory path in the same file.

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]>

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +626 to 629
if (first_error) return first_error;

std::unique_lock<std::shared_mutex> lock(mutex_);
eraseMemoryRegionLocked(addr);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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:

  1. 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 with ERR_ADDRESS_OVERLAPPED in tryReserveMemoryRegions, even if the physical memory was freed and reallocated.
  2. Impossible Recovery: If the user retries unregisterLocalMemory, the transports that did succeed in the first attempt might now return an error (such as ERR_ADDRESS_NOT_REGISTERED), which will again set first_error and 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.

Suggested change
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-commenter

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants