Drop stale async rollout results

[hamishivi]

Summary

• Add max_result_age_steps to accumulate_inference_batches: in async GRPO, a rollout whose generating policy is more than that many steps behind the current trainer step (training_step - result.model_step > max_result_age_steps) is dropped instead of trained on.
• When a stale result is dropped, a fresh prompt is replenished into the generator (when replenish_prompts is on) so throughput is unaffected, and the result is not kept for requeue-on-timeout.
• Logs a stale_results_dropped reward metric per accumulation.
• Wired from DataPreparationActor using async_steps as the staleness threshold (results more than one async window behind are dropped).

Notes

• model_step tracking and the model_step_min/max/mean / num_steps_off_policy metrics already existed on main; this only adds the drop policy on top.
• The staleness predicate is factored into a small pure result_is_stale(...) helper for testability.

Test plan

• New unit tests TestResultIsStale cover: disabled when max_result_age_steps/training_step/model_step is None, stale when lag exceeds the threshold, not stale at the threshold boundary, and not stale when fresh.
• ruff, formatting, and whole-project ty check pass.

GPU_TESTS=bypass

Made with Cursor

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to drop stale async rollout results whose generating policy lags too far behind the trainer, replenishing them with fresh prompts and logging a stale_results_dropped metric. Unit tests were also added to verify the staleness logic. The reviewer identified a critical issue where dropping stale results when replenish_prompts is disabled could lead to a silent deadlock or hang due to a drained pipeline, and suggested adding a validation check to prevent this configuration.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If max_result_age_steps is configured but replenish_prompts is False, dropping stale results without replenishing them will permanently reduce the number of prompts in flight. Over time, this will drain the pipeline and cause a silent deadlock/hang as the accumulator waits indefinitely for num_prompts results that will never arrive.

To prevent this, we should explicitly validate that replenish_prompts is enabled when staleness checking is active.

        if result_is_stale(result.model_step, training_step, max_result_age_steps):
            if not replenish_prompts:
                raise ValueError("max_result_age_steps requires replenish_prompts to be True to avoid deadlocks.")
            stale_results_dropped += 1
            logger.warning(
                "[accumulate_inference_batches] Dropping stale result for index=%s at training_step=%s: "
                "model_step=%s lag=%s max_result_age_steps=%s",
                result.index,
                training_step,
                result.model_step,
                training_step - result.model_step,
                max_result_age_steps,
            )
            assert iter_dataloader is not None and param_prompt_Q is not None
            add_prompt_to_generator(
                next(iter_dataloader),
                iter_dataloader._epoch,
                param_prompt_Q,
                generation_config,
                is_eval=False,
                base_env_config=base_env_config,
                ground_truth_overrides=ground_truth_overrides,
            )
            continue

[hamishivi]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to drop stale async rollout results whose generating policy lags too far behind the trainer (exceeding max_result_age_steps). When a stale result is detected, it is discarded, a fresh prompt is replenished to prevent the pipeline from draining, and a stale_results_dropped metric is logged. Unit tests have been added to verify the staleness detection logic and the associated configuration validation. There are no review comments, so no additional feedback is provided.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[farhatkevin]

looks good!