diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 06fad292132..d5d5ec58233 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1817,7 +1817,13 @@ async def _train_pump(self) -> None: ) if groups_dispatched == 0 and self._gen is not None: - await asyncio.to_thread(self._gen.snapshot_step_metrics) + # Raise here for observability. + try: + await asyncio.to_thread(self._gen.snapshot_step_metrics) + except RayActorError as error: + log.warning( + "Skipping generation snapshot metrics: %s", error + ) # ---- 2. Prepare the batch ---- # Compute prev_logprobs / ref_logprobs @@ -2032,9 +2038,12 @@ async def _train_pump(self) -> None: aggregate_rollout_metrics(per_group_rollout_metrics) ) if self._gen is not None: - step_metrics.update( - await asyncio.to_thread(self._gen.get_step_metrics) - ) + try: + step_metrics.update( + await asyncio.to_thread(self._gen.get_step_metrics) + ) + except RayActorError as error: + log.warning("Skipping generation step metrics: %s", error) self._step_log_dict = {k: [] for k in self._step_log_dict} step_metrics.update( _pooled_opd_metrics( diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index 21f15127f5a..b3760607dee 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -38,12 +38,12 @@ run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller run_test fast uv run --no-sync bash ./tests/functional/ppo_async_single_controller.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_gym_single_controller.sh -# Full mode only (~10 min): SIGKILLs a generation worker and asserts the job fails fast +# Fast mode too (~10 min): SIGKILLs a generation worker and asserts the job fails fast # and attributably instead of wedging. This is the ONLY end-to-end check of the # containment behaviour -- without it, a regression that restores the silent wedge is # caught by nothing, because a wedged job produces no exception and no failing assertion # anywhere else. -run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_chaos.sh # Full mode only: the same Gym run, but with NeMo-Gym pointed at the NeMo-RL-owned router. # Without this the router has no functional coverage at all -- the default Gym run above # leaves it disabled, so a regression in the proxy would ship silently. diff --git a/tests/functional/grpo_dp_single_controller_chaos.sh b/tests/functional/grpo_dp_single_controller_chaos.sh index 719b16e607c..01bed187223 100755 --- a/tests/functional/grpo_dp_single_controller_chaos.sh +++ b/tests/functional/grpo_dp_single_controller_chaos.sh @@ -8,7 +8,7 @@ # spun, with no exception raised anywhere. The pass condition here is therefore not # "training succeeds"; it is "the job stops, quickly, with an attributable error". # -# Registered in the SingleController L1 lane (full mode). It was originally kept out as +# Registered in the SingleController L1 lane (fast mode too). It was originally kept out as # "timing-sensitive", but that no longer justifies exclusion: the death deadline is now # 600s against an observed 222s, and the victim selection is asserted rather than assumed. # (The shard-recovery harness added in part 3/4 of this series kills processes the same diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 1793a26beca..5f3f41d471e 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -21,6 +21,7 @@ import pytest import torch +from ray.exceptions import ActorDiedError from tensordict import TensorDict import nemo_rl.algorithms.single_controller as single_controller @@ -1222,14 +1223,20 @@ def finish_train_step(self) -> dict: class _StepMetricRecordingGeneration: requires_kv_scale_sync = False - def __init__(self, events: list[str]) -> None: + def __init__(self, events: list[str], dies_in: str | None = None) -> None: self._events = events + self._dies_in = dies_in + + def _record(self, method: str) -> None: + self._events.append(method) + if method == self._dies_in: + raise ActorDiedError() def snapshot_step_metrics(self) -> None: - self._events.append("snapshot_step_metrics") + self._record("snapshot_step_metrics") def get_step_metrics(self) -> dict[str, float]: - self._events.append("get_step_metrics") + self._record("get_step_metrics") return {"vllm/spec_acceptance_rate": 0.8} @@ -1670,9 +1677,12 @@ def test_train_pump_aggregates_selected_rollout_metrics_across_chunks( assert "histogram/gen_tokens_length" not in capsys.readouterr().out +@pytest.mark.parametrize("dies_in", [None, "snapshot_step_metrics", "get_step_metrics"]) def test_train_pump_collects_generation_metrics_at_step_boundaries( - monkeypatch, + monkeypatch, dies_in ) -> None: + """A shard killed mid-step (grpo_dp_single_controller_chaos) must not end the pump + from the metrics fan-out; the typed failure belongs to the probe/refit paths.""" meta = KVBatchMeta( partition_id="rollout_data", task_name="train", @@ -1684,7 +1694,7 @@ def test_train_pump_collects_generation_metrics_at_step_boundaries( events: list[str] = [] ctrl = _train_pump_controller(sampler=_ChunkedSampler(meta, chunks=2)) ctrl._trainer = _StepMetricRecordingTrainer(events) - ctrl._gen = _StepMetricRecordingGeneration(events) + ctrl._gen = _StepMetricRecordingGeneration(events, dies_in=dies_in) ctrl._sync_weights = AsyncMock(return_value=0) ctrl._logger = MagicMock() monkeypatch.setattr(single_controller.ray, "cluster_resources", lambda: {}) @@ -1700,7 +1710,10 @@ def test_train_pump_collects_generation_metrics_at_step_boundaries( "get_step_metrics", ] train_metrics = ctrl._logger.log_metrics.call_args_list[0].args[0] - assert train_metrics["vllm/spec_acceptance_rate"] == pytest.approx(0.8) + if dies_in == "get_step_metrics": + assert "vllm/spec_acceptance_rate" not in train_metrics + else: + assert train_metrics["vllm/spec_acceptance_rate"] == pytest.approx(0.8) def test_train_pump_skips_generation_metrics_without_generation_handle(