diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 36dd5ffd103..0eb0cdeea6d 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/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() diff --git a/tests/unit/single_controller/test_run_supervision.py b/tests/unit/single_controller/test_run_supervision.py new file mode 100644 index 00000000000..4def2538d45 --- /dev/null +++ b/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))