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 assembly on the critical training data path; behavior shifts when rollouts fail or policy versions move quickly, though covered by new tests.

Overview
Async rollout workers now enforce max_staleness on in-flight work, not only when samples are dequeued for training. When the shared policy version advances, AsyncGRPO cancels whole in-flight groups that are too stale, frees vLLM slots, and records rollout/stale_groups_total; async distillation does the same per sample with rollout/stale_samples_total. Trainers pass max_staleness into both workers, and config/docs describe cancellation plus queue discard.

AsyncGRPO generation is more fault-tolerant. A failed rollout is logged and omitted while the rest of the group is still scored; reward kwargs are trimmed to survivors. Groups with fewer than two successful rollouts are dropped (rollout/dropped_groups_total); if every rollout in a group fails, the worker still raises. After stale cancellation, _repeat_iterator yields (group_id, index, row) so partially dispatched groups can finish as smaller groups at the new version instead of hanging forever.

OpenEnv harness rollouts cooperate with asyncio cancellation via a threading token that stops pool-thread sessions and in-flight vLLM sampling futures.

Tests cover stale cancellation, partial groups, failure handling, and distillation stale samples; metrics tables in the async GRPO/distillation docs are updated.

^{Reviewed by Cursor Bugbot for commit a15bbe0. 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]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

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

^{Reviewed by Cursor Bugbot for commit 44c45d9. Configure here.}