[TransferEngine] Propagate batch memory operation errors

[feichai0017]

Description

Batch memory operations could lose transport errors and leave partial registration state. This PR:

• Propagates nonzero per-buffer errors from batch register/unregister implementations.
• Makes batch unregister best-effort across all addresses and installed transports, returning the first error while publishing successful removals.
• Rolls back every attempted transport when batch registration fails.

Regression tests cover unknown addresses, continuation after errors, and partial-registration rollback.

Module

• Transfer Engine (`mooncake-transfer-engine`)

Type of Change

• Bug fix

How Has This Been Tested?

Test commands:
```bash
cmake -S . -B build -G Ninja -DBUILD_UNIT_TESTS=ON -DWITH_TE=ON -DWITH_STORE=OFF -DWITH_STORE_RUST=OFF -DWITH_EP=OFF -DWITH_P2P_STORE=OFF
cmake --build build --target transport_uint_test -j4
./build/mooncake-transfer-engine/tests/transport_uint_test
```

Test results:

• Unit tests pass (18/18 on Ubuntu)
• Integration tests pass (not applicable)
• Manual testing done (not applicable)

Checklist

• I have performed a self-review of my own code
• I have formatted my code using `./scripts/code_format.sh`
• I have run `pre-commit run --all-files` and all hooks pass (touched-file hooks pass)
• I have updated the documentation (not applicable; no user-facing API change)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue (not applicable)

AI Assistance Disclosure

• AI tools were used: Codex assisted with implementation and test preparation.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[gemini-code-assist]

Code Review

This pull request improves error handling and propagation in memory registration and unregistration batch operations across various transport implementations (such as HCCL, UBShmem, Barex, CXI, CXL, EFA, HIP, NVLink, RDMA, Sunrise Link, and TCP). It ensures that errors from individual registration or unregistration calls are properly captured and returned (either by returning early or tracking the first error encountered), rather than being ignored or mapped to generic errors. A unit test has also been added to verify that unregisterLocalMemoryBatch correctly propagates transport errors. I have no feedback to provide.

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.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 73.37662% with 41 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...e/src/transport/kunpeng_transport/ub_transport.cpp	0.00%	11 Missing ⚠️
...ne/src/transport/rdma_transport/rdma_transport.cpp	0.00%	11 Missing ⚠️
...gine/src/transport/cxl_transport/cxl_transport.cpp	0.00%	9 Missing ⚠️
...cake-transfer-engine/tests/transport_uint_test.cpp	91.91%	8 Missing ⚠️
...ncake-transfer-engine/src/transfer_engine_impl.cpp	86.66%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[he-yufeng]

Reviewed as a maintainer: traced the per-buffer error flow through the backends and the aggregation loop, not just the diff. This is a real fix and the pattern is applied consistently.

What I checked:

The bug is real. Every classic backend's registerLocalMemoryBatch / unregisterLocalMemoryBatch called registerLocalMemory(...) / unregisterLocalMemory(...) in a loop and discarded the return value, then fell through to updateLocalSegmentDesc() and returned that. So a per-buffer transport failure was logged-or-dropped and the batch still reported success. The fix captures each per-buffer return and short-circuits with if (ret) return ret;, so the first failing buffer propagates and — because the check sits before updateLocalSegmentDesc() — metadata is published only when every buffer registered. That gating is the important half and it's correct.

The if (ret < 0) → if (ret) change in transfer_engine_impl.cpp is the other half, not cosmetic. Some backends surface positive status codes (e.g. the CANN/HCCL family, where success is 0 and errors are positive), so the old ret < 0 silently let a positive error through as success. Switching to "any nonzero is an error" is what actually catches those. It's sound as long as no backend's batch call returns a positive non-error value — and it doesn't: each batch function returns either a propagated registerLocalMemory code (0 / error) or updateLocalSegmentDesc() (0 / error), never a positive count. Worth a one-line confirmation from you that this holds for every backend you touched, since it's the one assumption the change rests on.

The parallel-backend "wait for all in-flight ops, return the first error" behavior stays inside each backend, and the TCP regression (unregister an address that was never registered) pins the previously-swallowed path. Consistent across the backend list.

One note in passing: the HcclTransport hunk also changes the trailing registerLocalMemory(..., true, -1) to ..., true, false) — an int -1 where a bool is expected. Looks like a correct cleanup, just calling it out so it's intentional and not swept in.

LGTM.

[feichai0017]

Confirmed: every touched batch API returns 0 on success and nonzero only on error; none returns a positive non-error value. The HCCL -1 to false change is intentional because the final argument is update_metadata, so batch registration defers publication to the final updateLocalSegmentDesc(). I also rebased onto current main and preserved #2870’s reservation cleanup on error.

[feichai0017]

The current build (3.10) failure is runner infrastructure: the job annotation reports No space left on device in the Actions worker log after the project build completed. I do not have permission to rerun upstream Actions; could a maintainer rerun the failed job once this workflow finishes?

[he-yufeng]

Read the batch memory paths across the backends. This is a real fix: the classic / hccl / ubshmem batch register and unregister discarded per-buffer transport failures and returned only the metadata-update result, so a failed buffer registration surfaced as success. Two things are done well: the parallel backends (barex, cxi) waiting for every in-flight op and returning the first error is the right shape, and switching the per-buffer calls to update_metadata=false so the single batch-end updateLocalSegmentDesc() publishes once (instead of a redundant per-buffer update) is a nice cleanup.

A few points, none blocking:

1. Partial-batch register can leave orphans. In the sequential backends (e.g. HcclTransport::registerLocalMemoryBatch), if (ret) return ret early-returns on the first buffer's failure. Since the per-buffer calls now pass update_metadata=false and the batch-end updateLocalSegmentDesc() is skipped by the early return, buffers [0..i-1] end up registered in the transport but not reflected in metadata. What's the contract on a partial-batch failure - does the caller tear down / roll back, or are those registrations reconciled later? If neither, unregistering the already-registered buffers before returning (best-effort), or documenting the partial-failure state, would avoid an untracked registration.

2. The if (ret < 0) -> if (ret) widening now treats any nonzero as an error. Is the register/unregister return contract strictly 0=success / nonzero=error across every transport, i.e. can no transport ever return a positive non-error (a count, handle, or size)? If one can, this would turn a positive non-error into a spurious failure. If the convention is already 0-or-negative everywhere, the widening is harmless - just worth confirming.

3. Sequential unregister should be best-effort like the parallel ones. The parallel backends' unregister loops correctly collect first_error and keep going, but the sequential unregisterLocalMemoryBatch (hccl) early-returns on the first failure, leaving the remaining addresses registered. On a teardown path that is a leak, and it is inconsistent with the parallel backends. Making the sequential unregister best-effort (keep unregistering, return the first error) would match them.

[feichai0017]

Addressed in two follow-up commits: 86f72ec3 makes unregister attempt every address and installed transport, preserve the first error, and publish successful removals; 516b9495 rolls back the failing and previously successful transports before releasing reservations when registration fails. I also confirmed the touched batch APIs use 0 for success and nonzero for failure; none returns a positive success value. Added regression coverage, and all 18 transport_uint_test cases passed.