[Bug] P2P transfer failures do not stop weight-update finalization

[ashtonchew]

Summary

P2PTransferManager.wait_transfers() catches transfer exceptions and timeouts, logs them, clears every tracked future, and returns normally. The enclosing weight update therefore continues to publish the new weight version and resume generation even when one or more P2P writes failed or may still be running.

Current behavior

On main at f7d5190729658b5c141a3349c3ce2de320567929:

1. P2PTransferManager.wait_transfers() calls Future.result(timeout=...), catches every exception, and clears the future list.
2. UpdateWeightP2P treats that return as successful completion.
3. The updater then publishes the new version and finalizes the update, after which generation resumes.
4. The underlying write path already raises when Mooncake reports a negative result, but that exception is stored in the worker future and subsequently suppressed: _do_p2p_write_one_session().

A timeout is especially unsafe. Python documents that Future.result(timeout=...) raises TimeoutError without stopping a running call, and a running thread cannot be cancelled. Clearing that future loses ownership while the write may continue against the shared transfer buffer: Python concurrent.futures documentation.

Minimal reproduction

This reproduces the swallowed-failure boundary without RDMA hardware:

from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.p2p_transfer_utils import (
    P2PTransferManager,
)


def fail_transfer() -> None:
    raise RuntimeError("synthetic transfer failure")


manager = P2PTransferManager(num_workers=1, transfer_timeout=1.0)
manager.submit(fail_transfer)
manager.wait_transfers()
print("returned successfully")

Observed behavior:

[P2P] Transfer future failed: synthetic transfer failure
returned successfully

The same behavior can be tested deterministically for timeouts by placing an unresolved Future in the manager with transfer_timeout=0: wait_transfers() returns and clears the unresolved future.

Impact

• A rollout engine may be resumed with stale or partially updated parameters.
• A timed-out background write may continue modifying parameters after finalization.
• The published weight-version label can claim an update completed when its writes did not.
• Raising only on the source rank is not sufficient because peer trainer ranks can block at the following collective.

Expected behavior

• Any P2P write failure or timeout makes the complete weight update fail.
• Every trainer rank reaches a common Gloo success/failure agreement before any rank raises.
• Failed updates skip version publication, post-processing, and generation resume.
• Unresolved timed-out futures remain owned; their shared buffers are never reused.
• A failed updater performs no further writes before the enclosing trainer actor is stopped or recreated.

Suggested direction

1. Record transfer exceptions instead of logging and forgetting them.
2. Apply the configured timeout to both immediate and background transfer batches.
3. Stop loading new values into shared transfer buffers after the first failure or timeout.
4. At the common post-transfer rendezvous, use the existing collective_bool_and() so every trainer rank observes the failure before raising.
5. Add deterministic CPU regression coverage for task exceptions, unresolved timeouts, remote-rank failure consensus, and suppression of finalization.

This should remain a fail-closed correction to the P2P publication boundary; automatic retry/recovery policy can be handled separately once failed target sessions have an explicit reset contract.

Related pull request

• fix(weight-update): propagate P2P transfer failures #1732

[saisharan0103]

The repro details point to wait_transfers() as the part I would inspect first: it catches exceptions from Future.result(timeout=...), clears the tracked futures, and returns normally. That lets UpdateWeightP2P continue into version publication even though the write path can already raise on a negative Mooncake result. For a PR, I would start by preserving the failure state and making finalization abort or propagate when any transfer fails or times out.