Persistent async checkpoint worker failures are not propagated to training ranks
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**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**
1. Initialize a distributed process group with at least two ranks.
2. Create `AsyncCallsQueue(persistent=True)`.
3. Schedule an `AsyncRequest` whose `async_fn` raises on rank 0.
4. 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:
```text
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.
Contributor guide
Assessment
This issue has not been assessed yet.