Allow zero loss_mask for remove_sample'd rollouts

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(

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

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: