Flaky + undiagnosable: test_full_dcp_handoff_is_exact_on_same_dtensor_topology[muon_normuon] reds main

[thad0ctor]

Summary

tests/test_gefen_fsdp2_checkpoint.py::test_full_dcp_handoff_is_exact_on_same_dtensor_topology[muon_normuon] is flaky on CI and, worse, is written so that a failure is undiagnosable from the logs — it can only ever report assert False, never why a worker died.

This matters more now: #87 removed the duplicate workflow runs, so a single flaky run reds main outright (it already did — the #83 merge run 29538690201 failed on main), and a release PR (#88, v0.5.0) is in flight.

Evidence it's flaky (not a code defect)

• Run 29538418197 (pull_request event, bd47e80) — passed.
• Run 29538416045 (push event, same SHA bd47e80) — failed.
• Run 29538690201 (push, df3a615 = the Compat: DCP resharding for FSDP2 Gefen state #83 merge on main) — failed.
• It hit py3.12 in one occurrence and py3.11 in another, while the other Python versions in the same run passed — i.e. nondeterministic across versions, not version-specific.
• Locally it passes 15/15 on a large box (CPU/Gloo), so it appears to be a resource/timing sensitivity on small CI runners.
• The full suite on the same code passes locally (822 passed / 0 failed, 4 GPUs).

The observed failure signature

        rank_checks = result_queue.get(timeout=180)
        ...
        assert rank_checks is not None, "distributed checkpoint workers timed out"
>       assert all(process.exitcode == 0 for process in processes)
E       assert False

Note rank_checks is not None passed — the result was retrieved — and only the exitcode assertion failed. So the test logic completed, but a worker process exited nonzero.

Why it can't be diagnosed (the actual problem)

In test_full_dcp_handoff_is_exact_on_same_dtensor_topology:

1. Only rank 0 reports. The worker puts its result on the queue only from rank 0, so a failure on rank 1 is completely silent — the parent still gets rank 0's rank_checks and then trips the exitcode assertion with no explanation.
2. No error propagation. _worker has no try/except; any exception becomes a subprocess traceback on a stderr nobody captures, surfacing only as exitcode 1.
3. terminate() conflates slow with broken. After join(timeout=...), a still-alive (merely slow) worker is terminate()d, giving exitcode -15, which then fails assert all(exitcode == 0) identically to a real crash.

So every CI occurrence yields assert False with zero signal about the root cause.

Suggested fix

Adopt the hardened worker pattern already used in tests/test_fsdp2_cpu_offload.py and tests/test_dcp_resharding.py (added during #80/#83):

• Wrap each worker body in try/except BaseException and result_queue.put({"rank": rank, "fatal_error": traceback.format_exc()}) so a failure on any rank surfaces the real traceback in the CI log.
• Have every rank report (and drain one result per process), rather than rank 0 only, so rank 1 failures aren't silent.
• Fail fast when a worker exits early with a nonzero code instead of blocking on a long queue timeout.
• Reap workers in finally (join, then terminate only as a last resort), and distinguish "terminated because slow" from "crashed" in the assertion message.

That won't by itself remove whatever resource/timing sensitivity trips it on small runners, but it converts an opaque assert False into an actionable traceback — after which the underlying cause can actually be fixed.

Other spawn-based distributed tests in test_gefen_fsdp2_checkpoint.py share the same rank-0-only/no-try-except shape and would benefit from the same treatment.

Scope note

This test is pre-existing and unrelated to the CPU-offload (#80) and DCP-resharding (#83) work — those merges just made the flake visible on main.

[thad0ctor]

Correction to the analysis above, and the actual mechanism.

My original claim that _worker "has no try/except" is wrong — it does:

except BaseException:
    result_queue.put({"rank": rank, "traceback": traceback.format_exc()})
    raise

So error propagation exists. Reading the whole worker, the real reasons a failure is undiagnosable are more specific:

1. dist.init_process_group ran outside the try. A rendezvous/init failure — a real possibility on a loaded CI runner, and a plausible root cause of this very flake — exits nonzero with no traceback at all, because the handler that would report it hasn't been entered yet.
2. The traceback was put but never read. Only rank 0 put a result, and the parent drained exactly one item. When a non-zero rank fails, it does put its traceback — but the parent takes rank 0's success off the queue instead and never looks again, so the real error sits unread and the failure surfaces as a bare assert all(process.exitcode == 0) → assert False. That exactly matches the observed CI signature (rank_checks retrieved fine, only the exitcode assertion failing).
3. terminate() conflates slow with broken — a merely-slow worker gets exitcode -15, failing the same assertion as a genuine crash.

Fixed in the follow-up PR: init moved inside the try (and the finally only tears down a group that was actually initialized), every rank reports, the parent drains one result per rank and raises any worker traceback in the assertion, and the timeout/exitcode messages distinguish terminated-because-slow from crashed.

Verified by injecting a rank-1 failure: the test now fails in ~4s with the worker's real traceback (RuntimeError: INJECTED_RANK1_FAILURE) instead of a bare assert False.

This makes the next occurrence diagnosable; it does not by itself remove whatever timing/resource sensitivity trips it on small runners (it does not reproduce locally — 15/15 green on a large box).