From b98a2bf6ce660eb773064243cd556ddd5f267fa5 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 1 Jul 2026 02:55:53 +0000 Subject: [PATCH] Add infra-failure/truncation filters; allow zero loss_mask for masked samples Co-Authored-By: Claude Opus 4.8 --- miles/backends/megatron_utils/actor.py | 4 +- .../filter_hub/dynamic_sampling_filters.py | 43 +++++++++++++++++++ miles/rollout/filter_hub/rollout_filters.py | 27 +++++++++++- 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index c138ea8746..5bb4600c57 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -411,7 +411,9 @@ def _basic_batch_summary(): mask_sum = _sum_float(mask) if mask_sum <= 0: - _add_error(f"loss_masks[{i}] has no active tokens, sum={mask_sum}, response_len={resp}") + # Warning-only: rollout sample filters zero a sample's loss_mask + # (remove_sample) to drop it from the gradient; the reducer clamps. + _add_warning(f"loss_masks[{i}] has no active tokens, sum={mask_sum}, response_len={resp}") if mask_sum > resp: # Warning-only: float/weighted masks can legitimately have sum > resp. _add_warning( diff --git a/miles/rollout/filter_hub/dynamic_sampling_filters.py b/miles/rollout/filter_hub/dynamic_sampling_filters.py index f39ff8c9c1..0c0c6233ae 100644 --- a/miles/rollout/filter_hub/dynamic_sampling_filters.py +++ b/miles/rollout/filter_hub/dynamic_sampling_filters.py @@ -6,10 +6,34 @@ __all__ = [ "check_reward_nonzero_std", "check_no_aborted", + "check_no_infra_failures", "drop_zero_std_groups_and_extreme_pass_rate", "drop_truncated_or_extreme_pass_rate", ] +# Harbor exit_status values (from _extract_exit_status, after _TIMEOUT_EXCEPTION_MAP) +# for infra / non-policy failures: the rollout was severed by the environment, +# engine, or verifier rather than the policy, so it isn't valid training signal. +INFRA_FAILURE_EXIT_STATUSES = frozenset( + { + "AgentTimeout", + "AgentTimeoutError", + "HealthcheckError", + "_K8sInternalInfraError", + "Cancelled", + "RewardFileNotFoundError", + "AgentSetupTimeout", + "AgentSetupTimeoutError", + "SqsConsumerError", + "VerifierTimeout", + "VerifierTimeoutError", + "EnvStartTimeout", + "EnvironmentStartTimeoutError", + "TimeoutError", + "AddTestsDirError", + } +) + def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): rewards = [sample.get_reward_value(args) for sample in samples] @@ -46,6 +70,25 @@ def check_no_aborted(args, samples: list[Sample], **kwargs): return DynamicFilterOutput(keep=True) +def check_no_infra_failures(args, samples: list[Sample], **kwargs) -> DynamicFilterOutput: + """Reject the group if any rollout aborted or failed for an infra reason. + + 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. + + --dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_infra_failures + """ + for sample in _flatten_samples(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) + + def drop_zero_std_groups_and_extreme_pass_rate(args, samples: list[Sample], **kwargs) -> DynamicFilterOutput: """Filter groups with near-zero reward std or extreme mean rewards. For 0/1 rewards, mean reward is equivalent to pass rate --- so this function can be used to filter diff --git a/miles/rollout/filter_hub/rollout_filters.py b/miles/rollout/filter_hub/rollout_filters.py index 72282d4c64..30be747d71 100644 --- a/miles/rollout/filter_hub/rollout_filters.py +++ b/miles/rollout/filter_hub/rollout_filters.py @@ -1,7 +1,32 @@ from miles.rollout.filter_hub.dynamic_sampling_filters import _flatten_samples from miles.utils.types import Sample -__all__ = ["mask_truncated", "mask_truncated_and_llm_judge_failed"] +__all__ = ["mask_truncated", "mask_truncated_and_llm_judge_failed", "mask_token_truncated"] + +# Harbor exit_status values for rollouts truncated by running out of output +# tokens or context. Real attempts (verifier reward is meaningful) but the +# trajectory shouldn't get gradient — masked, reward kept in the group baseline. +TRUNCATION_EXIT_STATUSES = frozenset( + { + "BadRequestError", + "ContextWindowExceededError", + "OutputLengthExceededError", + } +) + + +def mask_token_truncated(args, samples: list[Sample]) -> None: + """Mask token/context-truncated samples (zero gradient, reward kept in baseline). + + Keyed on Harbor ``exit_status`` (BadRequest/ContextWindow record COMPLETED + turns, so a status check misses them); also covers TRUNCATED status. + + --rollout-sample-filter-path miles.rollout.filter_hub.rollout_filters.mask_token_truncated + """ + for sample in _flatten_samples(samples): + exit_status = (sample.metadata or {}).get("exit_status", "") + if exit_status in TRUNCATION_EXIT_STATUSES or sample.status == Sample.Status.TRUNCATED: + sample.remove_sample = True def mask_truncated(args, samples: list[Sample]) -> None: