Skip to content

Persistent async checkpoint worker failures are not propagated to training ranks #6970

Description

@zupengwang

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:

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions