fix(slurm): stop the multi-node template retrying a completed run

[shuangwu]

The bug

COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS makes the controller SIGTERM itself once the last policy replica unregisters, so the allocation is released instead of idling to wall-clock. That is how a successful run ends — but the batch script sees 128+SIGTERM and cannot tell it from a crash, so it requeues and re-runs finished training until the retry budget is gone, then reports FAILED.

Observed on a 2-node GCP run (job 1781412) — four full 500/500-step runs across ~1.5 h of allocation, every one of them logging a clean shutdown:

run	last step	coordinated shutdown logged	outcome
run_0	500/500	yes	retried
run_1	499/500	yes	retried
run_2	500/500	yes	retried
run_3	500/500	yes	FAILED

[cosmos-rl-sbatch]: Main loop exited with status: 143
[cosmos-rl-sbatch]: Job failed with status 143. No retries remaining.

The cost isn't only wasted allocation: sacct reports FAILED for a run that completed, so there is no way to tell a finished job from a broken one without reading logs.

Why the existing guard didn't catch it

handle_auto_retry already skips retry on SIGTERM:

if [[ "${received_signal}" == "SIGTERM" ]]; then
    log "Job was manually cancelled (SIGTERM). Skipping auto-retry."

But received_signal is set only by this script's own trap — it answers "was I cancelled?", not "how did the child exit?". A controller self-SIGTERM never signals the batch script, so the guard was never consulted.

The fix

Classify at the point where the monitor already knows which process exited and with what code:

if is_coordinated_controller_exit "${exit_code_controller}"; then
    log "Controller exited ... via coordinated shutdown. Treating as success."
    status=0
else
    status=${exit_code_controller}
fi

handle_auto_retry needs no change — it already returns early on status -eq 0. So one classification fixes the requeue, the FAILED state, and the misleading sacct output together.

Peer reaping now runs on both paths; it previously sat under the failure log only. On a coordinated shutdown the policy replicas are already gone (that is what triggered it), but a rollout can still be draining.

Kept deliberately narrow

Mirrors _is_coordinated_controller_exit (launcher/launch_all.py:397), which fixed the same misreading for the single-node CLI path in #714. Excusing too much would be worse than the original bug, so:

• only the controller — the caller passes ${exit_code_controller}
• only 128+SIGTERM
• only when the feature that produces it is on, with truthiness matching utils/constant.py (1/true/yes, case-insensitive) rather than a non-empty test, so =0 stays off
• only when the script was not itself signalled

That last clause is new relative to the single-node version, and it matters: SLURM's wall-clock SIGTERM also reaches the controller as 143. Without it, a job killed at its time limit would now report success — trading one false verdict for another. scancel and the --signal=B:SIGUSR1@... pre-timeout both set received_signal, so requiring it empty excludes both.

Tests

tests/test_slurm_multinode_exit.py, 13 cases, registered in run_test.sh. The template is an sbatch script with no import surface, so the functions are extracted and driven directly.

The first draft tested the predicate in isolation and was not sufficient — deleting the call site, i.e. reverting the fix entirely, left all 8 predicate tests green. Confirmed by mutation. So the suite now also drives the real monitoring loop with real child processes and asserts on the resulting status; the same mutation fails it with 143 != 0.

Covered in both directions:

• controller self-SIGTERM with the feature on → status=0; with it off → still 143
• an ordinary crash, a SIGKILL (137), a policy-side failure → unaffected
• scancel and the pre-timeout SIGUSR1 → still fail loudly
• the 1/true/yes truthiness utils/constant.py accepts

CPU-only, ~1.7 s.

Verification

Unit-tested and reasoned against a real failure. Not re-run on SLURM — the fix is in the template that job 1781412 exercised, but confirming end-to-end would need another multi-node allocation.