Filter zero-variance groups after infra checks

miles/rollout/filter_hub/dynamic_sampling_filters.py

@@ -71,22 +71,26 @@ def check_no_aborted(args, samples: list[Sample], **kwargs):
 
 
 def check_no_infra_failures(args, samples: list[Sample], **kwargs) -> DynamicFilterOutput:
-    """Reject the group if any rollout aborted or failed for an infra reason.
+    """Reject the group if any rollout aborted or failed for an infra reason,
+    or if the group has zero reward std (all-correct / all-incorrect).
 
     Superset of ``check_no_aborted``: also drops samples whose Harbor
     ``exit_status`` is an infra / non-policy failure. Keyed on ``exit_status``
     because such rollouts may still record COMPLETED turns; aborted samples
-    (zero records, no exit_status) are caught by the status check.
+    (zero records, no exit_status) are caught by the status check. Groups that
+    survive the infra checks are then passed through ``check_reward_nonzero_std``
+    so all-same-reward groups (advantage 0, no gradient) are dropped too.
 
         --dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_infra_failures
     """
-    for sample in _flatten_samples(samples):
+    flat_samples = list(_flatten_samples(samples))
+    for sample in flat_samples:
         if sample.status == Sample.Status.ABORTED:
             return DynamicFilterOutput(keep=False, reason="group_has_aborted")
         exit_status = (sample.metadata or {}).get("exit_status", "")
         if exit_status in INFRA_FAILURE_EXIT_STATUSES:
             return DynamicFilterOutput(keep=False, reason=f"group_has_{exit_status}")
-    return DynamicFilterOutput(keep=True)
+    return check_reward_nonzero_std(args, flat_samples, **kwargs)
 
 
 def drop_zero_std_groups_and_extreme_pass_rate(args, samples: list[Sample], **kwargs) -> DynamicFilterOutput:

tests/fast/rollout/test_dynamic_sampling_filters.py

@@ -0,0 +1,73 @@
+from argparse import Namespace
+
+import pytest
+
+from miles.rollout.filter_hub.base_types import DynamicFilterOutput
+from miles.rollout.filter_hub.dynamic_sampling_filters import INFRA_FAILURE_EXIT_STATUSES, check_no_infra_failures
+from miles.utils.types import Sample
+
+ARGS = Namespace(reward_key=None)
+
+
+def _sample(
+    reward: float | None,
+    *,
+    status: Sample.Status = Sample.Status.COMPLETED,
+    exit_status: str | None = None,
+) -> Sample:
+    metadata = {} if exit_status is None else {"exit_status": exit_status}
+    return Sample(reward=reward, status=status, metadata=metadata)
+
+
+@pytest.mark.parametrize("exit_status", sorted(INFRA_FAILURE_EXIT_STATUSES))
+def test_check_no_infra_failures_rejects_every_infra_exit_status(exit_status: str) -> None:
+    # The rewards are mixed, so this group would otherwise pass the std check.
+    samples = [_sample(0, exit_status=exit_status), _sample(1)]
+
+    output = check_no_infra_failures(ARGS, samples)
+
+    assert output == DynamicFilterOutput(keep=False, reason=f"group_has_{exit_status}")
+
+
+def test_check_no_infra_failures_rejects_aborted_before_reward_check() -> None:
+    samples = [_sample(None, status=Sample.Status.ABORTED), _sample(1)]
+
+    output = check_no_infra_failures(ARGS, samples)
+
+    assert output == DynamicFilterOutput(keep=False, reason="group_has_aborted")
+
+
+def test_check_no_infra_failures_keeps_nested_mixed_reward_group() -> None:
+    samples = [[_sample(0)], [_sample(1)]]
+
+    output = check_no_infra_failures(ARGS, samples)
+
+    assert bool(output.keep)
+    assert output.reason is None
+
+
+@pytest.mark.parametrize(
+    ("reward", "expected_reason"),
+    [
+        (1, "zero_std_1"),
+        (0, "zero_std_0"),
+    ],
+    ids=["all-correct", "all-incorrect"],
+)
+def test_check_no_infra_failures_rejects_zero_std_reward_groups(
+    reward: int,
+    expected_reason: str,
+) -> None:
+    output = check_no_infra_failures(ARGS, [_sample(reward), _sample(reward)])
+
+    assert not bool(output.keep)
+    assert output.reason == expected_reason
+
+
+def test_check_no_infra_failures_allows_policy_failure_with_mixed_rewards() -> None:
+    samples = [_sample(0, exit_status="TestsFailed"), _sample(1)]
+
+    output = check_no_infra_failures(ARGS, samples)
+
+    assert bool(output.keep)
+    assert output.reason is None