diff --git a/miles/rollout/filter_hub/dynamic_sampling_filters.py b/miles/rollout/filter_hub/dynamic_sampling_filters.py index 0c0c6233ae5..798c5d79546 100644 --- a/miles/rollout/filter_hub/dynamic_sampling_filters.py +++ b/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: diff --git a/tests/fast/rollout/test_dynamic_sampling_filters.py b/tests/fast/rollout/test_dynamic_sampling_filters.py new file mode 100644 index 00000000000..9186f093b84 --- /dev/null +++ b/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