fix(sc): keep the watchdog armed while the train pump drains

nemo_rl/algorithms/single_controller.py

@@ -563,42 +563,48 @@ async def run(self) -> dict[str, Any]:
         if probe_task is not None:
             tasks.append(probe_task)
         try:
-            done, _ = await asyncio.wait(
-                set(tasks), return_when=asyncio.FIRST_COMPLETED
-            )
-            stop_after_rollout_checkpoint = False
-            if rollout_checkpoint_task is not None and rollout_checkpoint_task in done:
-                await rollout_checkpoint_task
-                if not self._rollout_checkpoint_stop_requested.is_set():
-                    raise RuntimeError(
-                        "rollout checkpoint pump exited without requesting stop"
-                    )
-                stop_after_rollout_checkpoint = True
-            if stop_after_rollout_checkpoint:
-                # FIRST_COMPLETED may return several tasks. Do not let the
-                # orderly pre-step checkpoint stop hide a rollout/train failure
-                # that completed in the same event-loop turn.
-                for task in done:
-                    if task is not rollout_checkpoint_task:
-                        await task
-            if (
-                not stop_after_rollout_checkpoint
-                and probe_task is not None
-                and probe_task in done
-            ):
-                # Loops forever like the watchdog, so finishing at all means it raised.
-                await probe_task
-            if not stop_after_rollout_checkpoint and watchdog_task in done:
-                # The watchdog loops forever, so finishing at all means it raised --
-                # a stall or an unhealthy environment. Surface that ahead of the
-                # pumps, whose own symptom would just be "waiting".
-                await watchdog_task
-            if not stop_after_rollout_checkpoint and rollout_task in done:
-                # Propagate rollout failures immediately. A normally exhausted
-                # rollout pump leaves the train pump to drain committed groups.
-                await rollout_task
-            if not stop_after_rollout_checkpoint:
-                await train_task
+            # Keep supervising every task until training drains, or until the
+            # rollout-checkpoint pump requests its orderly pre-step stop. A
+            # single wait would stop watching the monitors as soon as rollout
+            # generation exhausts normally, so a watchdog failure during the
+            # remaining train drain could otherwise be discarded by teardown.
+            pending = set(tasks)
+            while train_task in pending:
+                done, pending = await asyncio.wait(
+                    pending, return_when=asyncio.FIRST_COMPLETED
+                )
+                stop_after_rollout_checkpoint = False
+                if (
+                    rollout_checkpoint_task is not None
+                    and rollout_checkpoint_task in done
+                ):
+                    await rollout_checkpoint_task
+                    if not self._rollout_checkpoint_stop_requested.is_set():
+                        raise RuntimeError(
+                            "rollout checkpoint pump exited without requesting stop"
+                        )
+                    stop_after_rollout_checkpoint = True
+                if stop_after_rollout_checkpoint:
+                    # FIRST_COMPLETED may return several tasks. Do not let an
+                    # orderly checkpoint stop hide a rollout/train failure that
+                    # completed in the same event-loop turn.
+                    for task in done:
+                        if task is not rollout_checkpoint_task:
+                            await task
+                    break
+                if probe_task is not None and probe_task in done:
+                    # Loops forever like the watchdog, so finishing means it raised.
+                    await probe_task
+                if watchdog_task in done:
+                    # Surface a stall or unhealthy environment ahead of pumps
+                    # whose own symptom would only be "waiting".
+                    await watchdog_task
+                if rollout_task in done:
+                    # Propagate rollout failures immediately. Normal exhaustion
+                    # leaves the train pump and monitors running through the drain.
+                    await rollout_task
+                if train_task in done:
+                    await train_task
         finally:
             for task in tasks:
                 task.cancel()

tests/unit/single_controller/test_run_supervision.py

@@ -0,0 +1,200 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""run() must keep supervising the watchdog for as long as the train pump lives.
+
+The rollout pump finishing first is the *normal* end-of-data path, not a
+failure -- the train pump then drains the groups already committed. The
+watchdog has to stay armed across that drain, because a wedged collective
+during it is exactly the kind of stall nothing else detects.
+"""
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from nemo_rl.algorithms.single_controller import SingleControllerActor
+from nemo_rl.experience.failures import RolloutStall
+
+
+def _bare_actor():
+    """An actor with __init__ skipped, stubbed down to what run() touches."""
+    controller_cls = SingleControllerActor.__ray_metadata__.modified_class
+    ctrl = object.__new__(controller_cls)
+    ctrl._gen_fleet = None  # no fleet -> run() creates no probe task
+    ctrl._train_steps = 0
+    ctrl._trainer_version = 0
+    ctrl._master_config = SimpleNamespace(
+        rollout_checkpointing=SimpleNamespace(snapshot_attempt_interval_s=None)
+    )
+    ctrl._rollout_checkpoint_stop_requested = asyncio.Event()
+    ctrl._weight_synchronizer = SimpleNamespace(is_stale=False, shutdown=lambda: None)
+    ctrl._logger = SimpleNamespace(finish=lambda: None)
+    ctrl._checkpointer = SimpleNamespace(shutdown=lambda: None)
+    ctrl._finalizer_actors = []
+    # run() stamps the rollout manager with the starting weight version before
+    # it launches any pump.
+    ctrl._rollout_manager = SimpleNamespace(set_weight_version=lambda _v: None)
+
+    async def _noop():
+        return None
+
+    ctrl._sync_weights = _noop
+    ctrl._maybe_restore_replay_buffer = _noop
+
+    async def _noop_restore_recovery(*, restored_replay_groups):
+        return None
+
+    ctrl._maybe_restore_rollout_recovery = _noop_restore_recovery
+    ctrl._maybe_restore_replacement_reserve = _noop
+    return ctrl
+
+
+async def _exhausts():
+    """A rollout pump that reaches the end of the data and returns."""
+    await asyncio.sleep(0)
+
+
+async def _wedged():
+    """A train pump blocked in a collective that never returns."""
+    await asyncio.Event().wait()
+
+
+def _stalls_after(delay):
+    async def _watchdog():
+        await asyncio.sleep(delay)
+        raise RolloutStall("simulated stall")
+
+    return _watchdog
+
+
+def test_watchdog_still_aborts_after_the_rollout_pump_exhausts():
+    """The drain phase is the window this regressed in: rollout done, train wedged."""
+    ctrl = _bare_actor()
+    ctrl._rollout_pump = _exhausts
+    ctrl._train_pump = _wedged
+    ctrl._stall_watchdog_pump = _stalls_after(0.02)
+
+    with pytest.raises(RolloutStall):
+        asyncio.run(asyncio.wait_for(ctrl.run(), timeout=5.0))
+
+
+def test_watchdog_aborts_while_the_rollout_pump_is_still_running():
+    """The pre-existing path, kept as a guard against fixing one and breaking the other."""
+    ctrl = _bare_actor()
+    ctrl._rollout_pump = _wedged
+    ctrl._train_pump = _wedged
+    ctrl._stall_watchdog_pump = _stalls_after(0.02)
+
+    with pytest.raises(RolloutStall):
+        asyncio.run(asyncio.wait_for(ctrl.run(), timeout=5.0))
+
+
+def test_a_clean_run_still_returns_its_summary():
+    """Both pumps finish, the watchdog never fires: run() returns normally."""
+    ctrl = _bare_actor()
+    ctrl._train_steps = 7
+    ctrl._trainer_version = 7
+    ctrl._rollout_pump = _exhausts
+    ctrl._train_pump = _exhausts
+
+    async def _quiet_watchdog():
+        await asyncio.Event().wait()
+
+    ctrl._stall_watchdog_pump = _quiet_watchdog
+
+    result = asyncio.run(asyncio.wait_for(ctrl.run(), timeout=5.0))
+    assert result == {"train_steps": 7, "trainer_version": 7}
+
+
+class _Boom(RuntimeError):
+    """Distinct from RolloutStall so a test cannot pass on the wrong task."""
+
+
+def _fails_after(delay):
+    async def _pump():
+        await asyncio.sleep(delay)
+        raise _Boom("simulated failure")
+
+    return _pump
+
+
+def test_a_rollout_failure_propagates_while_the_train_pump_is_still_working():
+    """The loop's comment says rollout failures propagate immediately. The
+    three tests above only ever have the rollout pump exhaust cleanly or never
+    finish, so nothing exercised the raise. Under a loop that stops watching,
+    a failed rollout task sits in `pending` unawaited while run() parks on the
+    train pump -- the job then holds its GPUs until the scheduler kills it.
+    """
+    ctrl = _bare_actor()
+    ctrl._rollout_pump = _fails_after(0.02)
+    ctrl._train_pump = _wedged
+    ctrl._stall_watchdog_pump = _wedged
+
+    with pytest.raises(_Boom):
+        asyncio.run(asyncio.wait_for(ctrl.run(), timeout=5.0))
+
+
+def test_the_fleet_probe_still_aborts_after_the_rollout_pump_exhausts():
+    """`_bare_actor` sets `_gen_fleet = None`, so no test above creates the
+    probe task at all and the `probe_task in done` branch is unreached. Fleet
+    health is the seconds-scale liveness signal; the probe pump loops forever,
+    so it finishing means it raised."""
+    ctrl = _bare_actor()
+    ctrl._gen_fleet = object()  # truthy -> run() creates the probe task
+    ctrl._rollout_pump = _exhausts
+    ctrl._train_pump = _wedged
+    ctrl._stall_watchdog_pump = _wedged
+    ctrl._gen_fleet_probe_pump = _fails_after(0.02)
+
+    with pytest.raises(_Boom):
+        asyncio.run(asyncio.wait_for(ctrl.run(), timeout=5.0))
+
+
+def test_rollout_checkpoint_stop_finishes_without_waiting_for_train_drain():
+    """The new supervision loop must preserve main's orderly pre-step stop."""
+    ctrl = _bare_actor()
+    ctrl._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 1.0
+    ctrl._rollout_pump = _wedged
+    ctrl._train_pump = _wedged
+    ctrl._stall_watchdog_pump = _wedged
+
+    async def _checkpoint_stop():
+        await asyncio.sleep(0)
+        ctrl._rollout_checkpoint_stop_requested.set()
+
+    ctrl._rollout_checkpoint_pump = _checkpoint_stop
+
+    result = asyncio.run(asyncio.wait_for(ctrl.run(), timeout=5.0))
+
+    assert result == {"train_steps": 0, "trainer_version": 0}
+
+
+def test_rollout_checkpoint_exit_without_stop_request_is_an_error():
+    ctrl = _bare_actor()
+    ctrl._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 1.0
+    ctrl._rollout_pump = _wedged
+    ctrl._train_pump = _wedged
+    ctrl._stall_watchdog_pump = _wedged
+
+    async def _unexpected_checkpoint_exit():
+        await asyncio.sleep(0)
+
+    ctrl._rollout_checkpoint_pump = _unexpected_checkpoint_exit
+
+    with pytest.raises(
+        RuntimeError, match="rollout checkpoint pump exited without requesting stop"
+    ):
+        asyncio.run(asyncio.wait_for(ctrl.run(), timeout=5.0))