fix(ci): Fix #3589's test broken by #3883

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(

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.

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

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(