Describe the bug
When --use-persistent-ckpt-worker is enabled, an exception in a persistent checkpoint worker's preload or async function can terminate the child process without publishing a completion record. The training process then waits for a completion that will never arrive, and a single failed rank can strand its distributed peers during checkpoint finalization.
This affects ordinary Python exceptions as well as hard worker exits. The parent currently polls the completion queue without turning the worker traceback or exit code into a distributed error.
Steps/Code to reproduce bug
- Initialize a distributed process group with at least two ranks.
- Create
AsyncCallsQueue(persistent=True).
- Schedule an
AsyncRequest whose async_fn raises on rank 0.
- Call
maybe_finalize_async_calls(blocking=True) on every rank.
A bounded regression harness on current main reaches its assertion that the worker failure was not propagated on both ranks. The same behavior can be reproduced when preload_fn raises or calls os._exit(17).
Tested base:
eb26b52878b217ae19963a96ac35f9f1569629b3
Expected behavior
The training ranks should detect that a persistent worker failed, receive a rank-tagged traceback or exit code, clean up failed worker state without an unbounded queue wait, and raise a consistent error on every rank. A subsequent persistent worker should be able to start cleanly.
Additional context
A focused local fix has been validated with controlled failure injection:
- 2 ranks: async function exception, preload exception, and
os._exit(17) all propagated successfully.
- 6 ranks: the same three scenarios passed on all ranks on 6x RTX 3090.
- Existing async caller shutdown tests passed.
The test node uses PCIe GPUs without NVLink; the affected logic is process lifecycle and distributed error propagation rather than topology-specific checkpoint throughput.
Describe the bug
When
--use-persistent-ckpt-workeris enabled, an exception in a persistent checkpoint worker's preload or async function can terminate the child process without publishing a completion record. The training process then waits for a completion that will never arrive, and a single failed rank can strand its distributed peers during checkpoint finalization.This affects ordinary Python exceptions as well as hard worker exits. The parent currently polls the completion queue without turning the worker traceback or exit code into a distributed error.
Steps/Code to reproduce bug
AsyncCallsQueue(persistent=True).AsyncRequestwhoseasync_fnraises on rank 0.maybe_finalize_async_calls(blocking=True)on every rank.A bounded regression harness on current
mainreaches its assertion that the worker failure was not propagated on both ranks. The same behavior can be reproduced whenpreload_fnraises or callsos._exit(17).Tested base:
Expected behavior
The training ranks should detect that a persistent worker failed, receive a rank-tagged traceback or exit code, clean up failed worker state without an unbounded queue wait, and raise a consistent error on every rank. A subsequent persistent worker should be able to start cleanly.
Additional context
A focused local fix has been validated with controlled failure injection:
os._exit(17)all propagated successfully.The test node uses PCIe GPUs without NVLink; the affected logic is process lifecycle and distributed error propagation rather than topology-specific checkpoint throughput.