AsyncGRPO: cancel stale in-flight rollouts and survive generation fai…

[AmineDiro]

Fresh take on #5320, which went stale. Both problems are still there on main: a single failed generation kills the rollout worker, and in-flight rollouts of a group the trainer will drop as stale keep running on vLLM until they finish.

Cancel stale in-flight rollouts

When the policy version advances, the worker cancels the in-flight generations of every group that started more than max_staleness versions ago, and frees their slots. RolloutQueueDataset would drop those samples anyway, so finishing them only burns vLLM compute. Cancelling the request makes vLLM abort it.

The earlier PR could strand a partially dispatched group: cancelling it deleted its state, then _repeat_iterator recreated it with too few rows, so it never reached num_generations. The iterator now yields the row index alongside the group id. A group created at index k had its first k rows cancelled, so it starts with k rows already counted and is scored as a smaller group at the new version.

Survive generation failures

A failed rollout is logged and dropped, and the group is scored with the rollouts that succeeded. Reward kwargs are trimmed to match. A group with a single survivor is dropped, since a group-relative advantage needs at least two. A group where every rollout failed still re-raises: that points at a broken server or setup, and swallowing it would leave the trainer waiting forever on an empty queue.

Metrics

Three counters, documented in the metrics table: rollout/failed_total, rollout/dropped_groups_total, rollout/stale_groups_total.

Tests

The new tests run the real generate and score loops against a scripted _generate_one and assert on what reaches the rollout buffer: a failed rollout, a single-survivor group, an all-failed group, stale cancellation on a version bump, and the partially dispatched case above. All five fail on main.

---

Note

Medium Risk
Changes async rollout scheduling, cancellation, and group completion semantics; incorrect edge-case handling could drop valid groups or waste vLLM work, but training logic is isolated to the rollout worker with targeted tests.

Overview
The async GRPO rollout worker now honors max_staleness on the generation side: when the policy version bumps, it cancels in-flight vLLM work for groups that lag too far behind, frees slots, and tracks rollout/stale_groups_total. _repeat_iterator yields (group_id, index, row) so a group that lost early rollouts to cancellation still finishes as a smaller group at the new version instead of hanging forever.

Generation failures no longer tear down the worker: a failed rollout is logged and skipped (rollout/failed_total), the rest of the group is scored with reward kwargs sliced to surviving completions, groups with only one success are dropped (rollout/dropped_groups_total), and a group where every rollout fails still propagates the error. The trainer passes max_staleness into AsyncRolloutWorker, docs describe the behavior and new metrics, and TestGenerateLoop covers failure handling, stale cancellation, and partial groups.

^{Reviewed by Cursor Bugbot for commit 692f774. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 692f774. Configure here.}

[cursor]

Cancel skips OpenEnv harness sessions

Medium Severity

Stale cancel only cancels the asyncio Task wrapping _generate_one. _HarnessRolloutLoop runs each session with run_in_executor, so task.cancel() does not stop _run_session. Freed slots then queue more sessions onto a pool already sized to max_inflight_tasks, so stale OpenEnv work keeps running and can stall fresh rollouts after a policy bump.

 

^{Reviewed by Cursor Bugbot for commit 692f774. Configure here.}

[cursor]

Distillation loop not kept in sync

Medium Severity

The generate-loop change that drops a failed rollout instead of re-raising, and that cancels stale in-flight work, was applied only in async_grpo. The copied loop in async_distillation still does raise task.exception() on a single failure and still lets stale in-flight generations run to completion.

 

^{Triggered by project rule: ../.ai/AGENTS.md}

^{Reviewed by Cursor Bugbot for commit 692f774. Configure here.}