[Reshard] Add model-weight Transfer Engine adapters

[Bo-Vincent]

Description

This PR adds model-weight adapters that execute validated reshard plans through
Mooncake Transfer Engine.

The adapters consume an attested TransferPlan, fresh runtime binding
manifests, and framework-provided allocation guards. They validate the selected
runtime fragments, lower compact N-D transfer regions into bounded TE batches,
and execute either target-initiated reads or source-initiated writes. The
resource-neutral TE executor remains independent of model names and framework
tensor layouts.

The execution lifecycle is fail-closed. GPU allocations and memory
registrations remain pinned until completion reaches a known terminal state.
Unknown completion or cleanup outcomes are retained by the pending-transfer
manager, fence further submissions on the same engine, and expose a pending ID
for status and recovery handling. Reader and sink paths share the same
completion, registration, allocation-lifetime, and operation-budget contracts.

This PR also documents the execution boundary and public APIs. Store
persistence, framework runtime discovery, model-specific conversion, and
quantization or packed-layout transformation remain separate layers.

Related RFC: #3111

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake Conductor (mooncake-conductor)
• Reshard (mooncake-reshard)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Test commands:

PYTHONPATH=mooncake-reshard/python \
  python -m pytest -q mooncake-reshard/tests

python -m pyright --project mooncake-reshard/pyrightconfig.json

./scripts/code_format.sh --changed-lines --check -b origin/main

pre-commit run --files $(git diff --name-only origin/main...HEAD)

make -C docs clean html

git diff --check origin/main...HEAD

The 22 commits in this PR were also checked individually by running the full
mooncake-reshard/tests suite at each commit.

Test results:

• Unit tests pass
• Integration tests pass (if applicable)
• Manual testing done (described below)

Exact-head results:

• 786 passed, 4 skipped, 2 subtests passed
• Pyright: 0 errors, 0 warnings, 0 informations
• All 22 semantic commits passed their individual test gate
• Changed-lines formatting, PR-scoped pre-commit hooks, documentation build,
  and git diff --check passed
• Public-path fault injection covered completion unknown, interrupted waits,
  registration cleanup, allocation-token release failure, duplicate token IDs,
  pending ownership, restart-required fencing, and reader/sink symmetry

Native RDMA/CUDA execution and process-restart recovery were not rerun on the
final exact head.

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 on the files changed in this PR and all hooks pass
• I have updated the documentation (if applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue ([RFC]: Manifest-driven heterogeneous model weight resharding and storage #3111)

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specified below)

AI tools assisted with implementation, test generation, fault-injection review,
and documentation. The changes were independently reviewed and validated on the
exact PR head; the human submitter remains responsible for the design and code.

[zxpdemonio]

Thanks for the thorough lifecycle work and test coverage. I found several issues that should be addressed before merging:

1. Acquisition rollback can lose allocation pins. In weight/_lifetime.py, if validation of a later acquired binding fails and one token's release_after_terminal() raises, _release_tokens() stops immediately. Earlier tokens are neither released nor quarantined, and the caller never receives them. Constructing the temporary AllocationTokenSet can also fail on duplicate token IDs before cleanup starts. Please make rollback best-effort and recoverable, preserving every unresolved token in quarantine, and add a multi-token release-failure test.

2. Unregistration progress is not interruption-safe. In both normal registration cleanup and PendingTransferManager.drain_pending_transfer(), progress is only recorded after the unregister loop. An interruption after one successful unregister_memory() return can leave remaining registrations live while the outer adapter releases allocation tokens; the pending path can also retry an address that was already successfully unregistered. Please track an unresolved-address set, checkpoint each successful unregister immediately, and catch BaseException around the complete cleanup loop.

3. Completed physical I/O can be reported as ABORTED. Reader and sink only switch to FAILED_DRAINED for TransferCompletionFailedError. If a native batch completes and an interruption occurs before the adapter returns, all tokens are finalized with ABORTED even though target memory was modified. Please track whether physical I/O has started/completed and classify subsequent non-pending failures as FAILED_DRAINED, as the Store paths already do.

4. Per-participant execution acquires the global peer set. A source that only writes a subset of targets still requires guards and executor bindings for every target in the plan; the reader has the symmetric source-side issue. For example, source 0 in TP2-to-TP4 only touches target 0/1, but execution fails if the unrelated target 2 guard is absent. Please derive required participants/fragments from the selected executor's operation indices and validate completeness against that subset.

5. The current execution loop serializes every batch and endpoint. execute_batch() waits for a terminal ticket before the next batch or peer is submitted, leaving only one transfer in flight. This prevents multi-peer/NIC aggregation from approaching line rate. Please introduce a bounded in-flight window while retaining each ticket's registrations and lifetime resources until terminal completion.

6. The weight adapter tests do not exercise the production scatter-ticket path. The model-weight fake engine primarily exposes the flat compatibility API, whereas the real binding selects scatter_transfer_sync_*_with_ticket. Please add reader/sink tests covering lowering through a scatter-capable engine (bases, capacities, nested offsets, direction mapping and completion), plus a native CPU/GPU or RDMA smoke test.

I reproduced the first three lifecycle failures locally. The PR-scoped TE tests pass (173 passed), Pyright reports no errors, and git diff --check passes, but these failure paths are not covered by the current suite.

[Bo-Vincent]

Thanks for the thorough lifecycle work and test coverage. I found several issues that should be addressed before merging:

1. Acquisition rollback can lose allocation pins. In weight/_lifetime.py, if validation of a later acquired binding fails and one token's release_after_terminal() raises, _release_tokens() stops immediately. Earlier tokens are neither released nor quarantined, and the caller never receives them. Constructing the temporary AllocationTokenSet can also fail on duplicate token IDs before cleanup starts. Please make rollback best-effort and recoverable, preserving every unresolved token in quarantine, and add a multi-token release-failure test.
2. Unregistration progress is not interruption-safe. In both normal registration cleanup and PendingTransferManager.drain_pending_transfer(), progress is only recorded after the unregister loop. An interruption after one successful unregister_memory() return can leave remaining registrations live while the outer adapter releases allocation tokens; the pending path can also retry an address that was already successfully unregistered. Please track an unresolved-address set, checkpoint each successful unregister immediately, and catch BaseException around the complete cleanup loop.
3. Completed physical I/O can be reported as ABORTED. Reader and sink only switch to FAILED_DRAINED for TransferCompletionFailedError. If a native batch completes and an interruption occurs before the adapter returns, all tokens are finalized with ABORTED even though target memory was modified. Please track whether physical I/O has started/completed and classify subsequent non-pending failures as FAILED_DRAINED, as the Store paths already do.
4. Per-participant execution acquires the global peer set. A source that only writes a subset of targets still requires guards and executor bindings for every target in the plan; the reader has the symmetric source-side issue. For example, source 0 in TP2-to-TP4 only touches target 0/1, but execution fails if the unrelated target 2 guard is absent. Please derive required participants/fragments from the selected executor's operation indices and validate completeness against that subset.
5. The current execution loop serializes every batch and endpoint. execute_batch() waits for a terminal ticket before the next batch or peer is submitted, leaving only one transfer in flight. This prevents multi-peer/NIC aggregation from approaching line rate. Please introduce a bounded in-flight window while retaining each ticket's registrations and lifetime resources until terminal completion.
6. The weight adapter tests do not exercise the production scatter-ticket path. The model-weight fake engine primarily exposes the flat compatibility API, whereas the real binding selects scatter_transfer_sync_*_with_ticket. Please add reader/sink tests covering lowering through a scatter-capable engine (bases, capacities, nested offsets, direction mapping and completion), plus a native CPU/GPU or RDMA smoke test.

I reproduced the first three lifecycle failures locally. The PR-scoped TE tests pass (173 passed), Pyright reports no errors, and git diff --check passes, but these failure paths are not covered by the current suite.

Thanks for the detailed review. I addressed issues 1–4 and 6.

1. Acquisition rollback

Rollback now uses the executor’s recoverable resource finalizer for every acquired token. A release failure retains the unresolved token in pending quarantine, while the remaining tokens are still finalized. Duplicate token IDs can no longer prevent cleanup from starting. Multi-token release-failure and duplicate-ID regression tests were added.

2. Interruption-safe unregistration

Registration cleanup now tracks an unresolved-address set and checkpoints every successful unregister immediately. Both normal cleanup and pending-transfer draining catch BaseException around the complete loop. An interruption with uncertain progress marks the cleanup as restart-required and does not replay an address whose outcome is unknown.

3. Physical I/O terminal state

TransferSubmission now records when physical I/O enters the native completion fence. Any subsequent non-pending failure is finalized as FAILED_DRAINED rather than ABORTED, preserving the fact that target memory may already have been modified. Reader and sink regression tests cover this interruption boundary.

4. Per-participant execution scope

Reader and sink now derive required peers and fragments from the selected local executor’s operation indices. Binding, guard acquisition, and executor-completeness validation are limited to that operation scope, so unrelated plan participants are no longer required. Both target-initiated and source-initiated subset cases are covered.

5. Bounded in-flight execution

I confirmed that the current loop keeps one transfer in flight. I have not implemented concurrency in this PR because a safe bounded window requires a native non-draining submission API, composite pending-ticket recovery, and per-batch registration and lease ownership. Wrapping the synchronous API in a thread pool would weaken the current fail-closed lifecycle. I propose handling this as a focused performance follow-up while keeping this PR’s execution contract synchronous and recoverable.

6. Scatter-ticket coverage

Reader and sink tests now exercise the production-shaped scatter_transfer_sync_read_with_ticket and scatter_transfer_sync_write_with_ticket paths, including allocation bases, capacities, nested offsets, direction mapping, and completion. A native CPU/TCP Scatter write/read smoke test also passed.

Validation on the exact head:

• 799 passed, 4 skipped, 2 subtests passed
• Pyright: 0 errors
• all 22 semantic commits passed their individual full test gate
• PR-scoped pre-commit, changed-lines formatting, Sphinx documentation build, and git diff --check passed

[zxpdemonio]

Thanks for the update. I rechecked exact head 7f6e502c68d0ec445535fbd20413228f072aedd9. The rollback, interruption-safe cleanup, post-I/O terminal-state handling, and scatter-ticket coverage from items 1–3 and 6 now look addressed. I still found two correctness blockers in the operation-scoping change:

1. The peer scope expands back to every worker executor in the same participant. The new preflight correctly derives source_requirements / target_requirements from the local operation indices, but _execute_reserved() then calls resolve_runtime_executors() for each peer binding. That resolver returns all executor snapshots for the (instance_id, participant_id), not just the operation-scoped workers. Reader then compares that expanded set with expected_source_executors at reader.py:289-320; sink has the symmetric check at sink.py:280-311.

   I reproduced this with one valid WeightRuntimeBindingManifest containing two fragments assigned to two worker IDs under the same participant. When local target 0 touches only source worker 0, reader fails with TransferEngineError: source executor set is incomplete; swapping the topology makes sink fail with target executor set is incomplete. No unrelated participant is present—the extra executor is only another worker in the same participant. Please resolve/revalidate the fresh binding against the selected executor/fragment scope rather than resolving the complete participant, and add reader/sink regressions where one participant contains multiple worker groups but the local operation subset uses only one.

2. An explicit unknown worker selector now silently succeeds with no work. After filtering by target_worker_id / source_worker_id, reader.py:125-132 and sink.py:122-129 return () when nothing matches. A typo, an empty string, or a non-string selector therefore looks like a successful zero-byte transfer. The Store adapters already reject invalid selectors and unknown workers, and the reserved TE paths still require exactly one selected local executor. Please preserve that fail-closed contract: only return empty when the binding itself is outside the plan (if that is intentional), but reject an explicitly supplied selector that does not identify a planned worker. Add reader/sink tests for invalid selector type/value and unknown worker ID.

The bounded in-flight item remains unresolved. I agree that a thread-pool wrapper over the synchronous API would weaken the lifecycle contract; please at least link a concrete native-async follow-up and document that this version serializes all endpoint/batch submissions.

Validation on this head: the focused TE/weight-adapter suite passes (186 passed), Pyright reports 0 errors, and git diff --check passes. The full reshard suite reached 799 passed, 3 skipped, 2 subtests passed; its only failure on this host is an unrelated installed native mooncake.store that does not expose begin_weight_snapshot.

[Bo-Vincent]

Thanks for the update. I rechecked exact head 7f6e502c68d0ec445535fbd20413228f072aedd9. The rollback, interruption-safe cleanup, post-I/O terminal-state handling, and scatter-ticket coverage from items 1–3 and 6 now look addressed. I still found two correctness blockers in the operation-scoping change:

1. The peer scope expands back to every worker executor in the same participant. The new preflight correctly derives source_requirements / target_requirements from the local operation indices, but _execute_reserved() then calls resolve_runtime_executors() for each peer binding. That resolver returns all executor snapshots for the (instance_id, participant_id), not just the operation-scoped workers. Reader then compares that expanded set with expected_source_executors at reader.py:289-320; sink has the symmetric check at sink.py:280-311.
   I reproduced this with one valid WeightRuntimeBindingManifest containing two fragments assigned to two worker IDs under the same participant. When local target 0 touches only source worker 0, reader fails with TransferEngineError: source executor set is incomplete; swapping the topology makes sink fail with target executor set is incomplete. No unrelated participant is present—the extra executor is only another worker in the same participant. Please resolve/revalidate the fresh binding against the selected executor/fragment scope rather than resolving the complete participant, and add reader/sink regressions where one participant contains multiple worker groups but the local operation subset uses only one.
2. An explicit unknown worker selector now silently succeeds with no work. After filtering by target_worker_id / source_worker_id, reader.py:125-132 and sink.py:122-129 return () when nothing matches. A typo, an empty string, or a non-string selector therefore looks like a successful zero-byte transfer. The Store adapters already reject invalid selectors and unknown workers, and the reserved TE paths still require exactly one selected local executor. Please preserve that fail-closed contract: only return empty when the binding itself is outside the plan (if that is intentional), but reject an explicitly supplied selector that does not identify a planned worker. Add reader/sink tests for invalid selector type/value and unknown worker ID.

The bounded in-flight item remains unresolved. I agree that a thread-pool wrapper over the synchronous API would weaken the lifecycle contract; please at least link a concrete native-async follow-up and document that this version serializes all endpoint/batch submissions.

Validation on this head: the focused TE/weight-adapter suite passes (186 passed), Pyright reports 0 errors, and git diff --check passes. The full reshard suite reached 799 passed, 3 skipped, 2 subtests passed; its only failure on this host is an unrelated installed native mooncake.store that does not expose begin_weight_snapshot.

Thanks for the follow-up review. Both operation-scoping issues have been fixed on the current head (8e9109d0992181094281a2229d93f3482adc06bb).

1. Runtime executor validation is now limited to the workers and fragments referenced by the selected operation indices. Other workers under the same participant are no longer pulled back into the required executor set. Reader and sink regressions cover this multi-worker case.

2. Explicit worker selectors are now fail-closed. Empty, non-string, and unknown worker IDs raise TransferEngineError instead of returning an empty successful result. Reader and sink negative tests cover these cases.

The focused regression suite passes (9 passed), and git diff --check passes.

Bounded in-flight execution is being handled as a separate performance follow-up because it requires native non-draining submission and recoverable multi-ticket lifetime management.