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

cosmos_rl/tools/slurm/cosmos_rl_job_multi_node.sh

@@ -241,6 +241,41 @@ submit_autoresume() {
     return ${status}
 }
 
+# Is a non-zero controller exit actually a deliberate, successful shutdown?
+#
+# 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 it surfaces
+# here as 128+SIGTERM and was indistinguishable from a crash -- so every
+# successful multi-node run burned its whole retry budget re-running completed
+# training, then reported FAILED to SLURM.
+#
+# Mirrors ``_is_coordinated_controller_exit`` in launcher/launch_all.py, which
+# fixed the same misreading for the single-node CLI path.  Deliberately narrow,
+# so a SIGTERM from anywhere else still fails loudly:
+#
+#   * only the controller -- the caller passes ${exit_code_controller};
+#   * only 128+SIGTERM;
+#   * only when the feature that produces it is on.  Truthiness matches
+#     utils/constant.py, which accepts 1/true/yes;
+#   * only when this script was NOT itself signalled.  ``scancel`` and the
+#     ``--signal=B:SIGUSR1@...`` pre-timeout both set ``received_signal``, and
+#     a SLURM time-limit SIGTERM reaches the controller as 143 too -- without
+#     this clause a job killed at its wall-clock limit would report success.
+is_coordinated_controller_exit() {
+    local code=$1
+
+    [[ ${code} -eq $((128 + 15)) ]] || return 1
+    [[ -z "${received_signal}" ]] || return 1
+
+    local flag
+    flag=$(printf '%s' "${COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS:-0}" | tr '[:upper:]' '[:lower:]')
+    case "${flag}" in
+        1 | true | yes) return 0 ;;
+        *) return 1 ;;
+    esac
+}
+
 # Retry the job on transient failures (decrements remaining-retries counter).
 handle_auto_retry() {
     local status=$1
@@ -579,7 +614,16 @@ while true; do
     fi
 
     if [[ -n "${exit_code_controller}" ]] && [[ ${exit_code_controller} -ne 0 ]]; then
-        log "Controller failed with exit code ${exit_code_controller}. Terminating other processes."
+        if is_coordinated_controller_exit "${exit_code_controller}"; then
+            log "Controller exited ${exit_code_controller} via coordinated shutdown (COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS). Treating as success."
+            status=0
+        else
+            log "Controller failed with exit code ${exit_code_controller}. Terminating other processes."
+            status=${exit_code_controller}
+        fi
+        # Reap peers either way.  On a coordinated shutdown the policy replicas
+        # have already exited -- that is what triggered it -- so these are
+        # normally no-ops, but a rollout can still be draining.
         if [[ "${policy_waited}" == "false" ]]; then
             kill "$pid_policy" 2>/dev/null || true
             wait "$pid_policy" 2>/dev/null || true
@@ -588,7 +632,6 @@ while true; do
             kill "$pid_rollout" 2>/dev/null || true
             wait "$pid_rollout" 2>/dev/null || true
         fi
-        status=${exit_code_controller}
         break
     fi
 

tests/run_test.sh

@@ -129,6 +129,8 @@ run python tests/test_ucxx_fetch_engine.py
 run python tests/test_ucxx_rollout_mixin.py
 run python tests/test_ucxx_transport.py
 run python tests/test_launcher_shutdown.py
+# Same predicate as above, for the multi-node sbatch template.
+run python tests/test_slurm_multinode_exit.py
 # Guards the wait/teardown helper the GPU suites below rely on to stay bounded.
 run python tests/test_subprocess_helpers.py
 run python tests/test_process_flow.py

tests/test_slurm_multinode_exit.py

@@ -0,0 +1,223 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The multi-node sbatch template must not retry a completed run.
+
+``COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS`` makes the controller SIGTERM *itself*
+once the last policy replica unregisters.  That is how a successful run ends,
+but it reaches the batch script as 128+SIGTERM.  Read as a crash, it sent the
+job through its whole retry budget re-running finished training and then
+reported FAILED -- observed on a real 2-node run: four full 500-step runs, then
+FAILED.
+
+``tests/test_launcher_shutdown.py`` covers the same predicate for the
+single-node CLI launcher.  This one covers the sbatch template, which is shell
+and has no import surface, so the functions are extracted and driven directly.
+
+Excusing too much would be worse than the original bug, so the negative cases
+matter as much as the positive one: a SIGTERM from ``scancel``, from
+pre-emption, or from the wall-clock limit must still fail loudly.
+"""
+
+import re
+import signal
+import subprocess
+import unittest
+from pathlib import Path
+
+from cosmos_rl.tools import slurm as slurm_tools
+
+# Resolve the template next to the *installed* module, not by repo layout.
+# CI copies only tests/ into the image and installs the package into
+# site-packages, so a repo-relative path resolves to nothing there.  The .sh
+# ships alongside the module in both cases.
+TEMPLATE = Path(slurm_tools.__file__).resolve().parent / "cosmos_rl_job_multi_node.sh"
+
+SIGTERM_RC = 128 + signal.SIGTERM  # 143
+
+
+def _extract_monitor_loop() -> str:
+    """Pull the process-monitoring loop out of the template.
+
+    Testing the predicate alone is not enough: it would still pass with the
+    call site deleted, i.e. with the fix reverted.  This drives the loop that
+    has to consult it.
+    """
+    src = TEMPLATE.read_text()
+    start = src.index("\nwhile true; do\n")
+    end = src.index("\ndone\n", start)
+    return src[start : end + len("\ndone\n")]
+
+
+def _run_monitor(controller_rc, policy_rc=0, rollout_rc=0, *, flag="1"):
+    """Drive the real loop with real children and return its final status."""
+    script = "\n".join(
+        [
+            "set -u",
+            'received_signal=""',
+            f'COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS="{flag}"',
+            "log() { :; }",
+            "status=",
+            "policy_waited=false; rollout_waited=false; controller_waited=false",
+            "exit_code_policy=; exit_code_rollout=; exit_code_controller=",
+            f"bash -c 'exit {policy_rc}' & pid_policy=$!",
+            f"bash -c 'exit {rollout_rc}' & pid_rollout=$!",
+            f"bash -c 'exit {controller_rc}' & pid_controller=$!",
+            "sleep 0.3",
+            _extract("is_coordinated_controller_exit"),
+            _extract_monitor_loop(),
+            'echo "status=${status}"',
+        ]
+    )
+    out = subprocess.run(
+        ["bash", "-c", script], capture_output=True, text=True, timeout=120
+    )
+    assert out.returncode == 0, f"harness failed: {out.stderr}"
+    line = [x for x in out.stdout.splitlines() if x.startswith("status=")][-1]
+    return int(line.split("=", 1)[1])
+
+
+def _extract(func_name: str) -> str:
+    """Pull one shell function out of the template.
+
+    The template is an sbatch script: sourcing it would run a job.  These
+    functions are self-contained, so lifting them by text is enough to drive
+    them, and it keeps the test honest about which definition it is checking.
+    """
+    src = TEMPLATE.read_text()
+    match = re.search(
+        rf"^{re.escape(func_name)}\(\) \{{.*?^\}}$",
+        src,
+        re.MULTILINE | re.DOTALL,
+    )
+    assert match, f"{func_name} not found in {TEMPLATE}"
+    return match.group(0)
+
+
+def _classify(code, *, flag=None, received_signal=""):
+    """Run ``is_coordinated_controller_exit`` and return whether it excused."""
+    script = "\n".join(
+        [
+            "set -u",
+            f'received_signal="{received_signal}"',
+            (
+                f'COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS="{flag}"'
+                if flag is not None
+                else "unset COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS || true"
+            ),
+            _extract("is_coordinated_controller_exit"),
+            f'if is_coordinated_controller_exit "{code}"; then echo YES; else echo NO; fi',
+        ]
+    )
+    out = subprocess.run(
+        ["bash", "-c", script], capture_output=True, text=True, timeout=60
+    )
+    assert out.returncode == 0, f"harness failed: {out.stderr}"
+    return out.stdout.strip() == "YES"
+
+
+class TestCoordinatedControllerExit(unittest.TestCase):
+    def test_self_sigterm_with_feature_on_is_success(self):
+        """The case that made every completed multi-node run report FAILED."""
+        self.assertTrue(_classify(SIGTERM_RC, flag="1"))
+
+    def test_truthiness_matches_the_python_constant(self):
+        """utils/constant.py accepts 1/true/yes, case-insensitively."""
+        for value in ("1", "true", "TRUE", "yes", "Yes"):
+            self.assertTrue(_classify(SIGTERM_RC, flag=value), value)
+        for value in ("0", "", "no", "false", "off", "2"):
+            self.assertFalse(_classify(SIGTERM_RC, flag=value), value)
+
+    def test_feature_off_or_unset_still_fails(self):
+        self.assertFalse(_classify(SIGTERM_RC, flag="0"))
+        self.assertFalse(_classify(SIGTERM_RC, flag=None))
+
+    def test_other_exit_codes_are_never_excused(self):
+        # 137 = SIGKILL (OOM reaper), 1 = ordinary crash, 0 handled upstream.
+        for code in (1, 2, 137, 139, 255):
+            self.assertFalse(_classify(code, flag="1"), code)
+
+    def test_scancel_still_fails(self):
+        """``scancel`` traps SIGTERM on the batch script itself."""
+        self.assertFalse(_classify(SIGTERM_RC, flag="1", received_signal="SIGTERM"))
+
+    def test_wall_clock_pre_timeout_still_fails(self):
+        """``--signal=B:SIGUSR1@...`` fires before the time limit.
+
+        Without the received_signal clause a job killed at its wall-clock
+        limit would report success, trading one false verdict for another.
+        """
+        self.assertFalse(_classify(SIGTERM_RC, flag="1", received_signal="SIGUSR1"))
+
+
+class TestAutoRetryGate(unittest.TestCase):
+    """A success classification must actually stop the requeue."""
+
+    @staticmethod
+    def _retry_outcome(status, received_signal=""):
+        script = "\n".join(
+            [
+                "set -u",
+                f'received_signal="{received_signal}"',
+                'latest_part_dir="/nonexistent"',
+                'log() { echo "$@"; }',
+                "scontrol() { echo SCONTROL_CALLED; }",
+                _extract("handle_auto_retry"),
+                f"handle_auto_retry {status}",
+                'echo "rc=$?"',
+            ]
+        )
+        out = subprocess.run(
+            ["bash", "-c", script], capture_output=True, text=True, timeout=60
+        )
+        return out.stdout
+
+    def test_status_zero_never_requeues(self):
+        out = self._retry_outcome(0)
+        self.assertNotIn("SCONTROL_CALLED", out)
+        self.assertIn("rc=0", out)
+
+    def test_nonzero_without_retries_left_does_not_requeue(self):
+        # remaining-retries file is absent -> treated as 0 remaining.
+        out = self._retry_outcome(1)
+        self.assertNotIn("SCONTROL_CALLED", out)
+        self.assertIn("No retries remaining", out)
+
+
+class TestMonitorLoopWiring(unittest.TestCase):
+    """The predicate must actually be consulted by the loop.
+
+    Without these, deleting the call site -- reverting the fix entirely --
+    leaves every predicate test still green.
+    """
+
+    def test_controller_self_sigterm_yields_success_status(self):
+        self.assertEqual(_run_monitor(SIGTERM_RC, flag="1"), 0)
+
+    def test_controller_self_sigterm_with_feature_off_still_fails(self):
+        self.assertEqual(_run_monitor(SIGTERM_RC, flag="0"), SIGTERM_RC)
+
+    def test_controller_crash_still_fails(self):
+        self.assertEqual(_run_monitor(1, flag="1"), 1)
+
+    def test_policy_failure_is_untouched_by_the_controller_path(self):
+        self.assertEqual(_run_monitor(0, policy_rc=7, flag="1"), 7)
+
+    def test_all_clean_exits_are_success(self):
+        self.assertEqual(_run_monitor(0, flag="1"), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()