fix: stop the fully-async rollout dropping completed groups

[keepkeen]

What

The fully-async rollout throws away almost every completed group, and its output queue can freeze the worker's event loop. The two defects are coupled, so this fixes both.

Bug 1: completed groups are dropped

_generate_rollout_async drains the entire output queue on every poll, then returns only a slice:

for gid, group in worker.get_completed_groups():   # drains everything
    collected[gid] = group
...
out = sorted(collected.values(), key=_key)[:target]  # surplus silently dropped

The dropped groups are fully generated and reward-scored, and their prompts were already consumed from data_buffer (never requeued) — burned out of the epoch, GPU work wasted.

This is not a rare edge: the worker is sized at sglang_server_concurrency * num_engines in-flight groups. With the shipped examples/fully_async/run-qwen2.5-0.5B-fully_async.sh that is 512 × 3 = 1536 in flight against --rollout-batch-size 8. The module's own documentation states the opposite contract three times: "The worker is shared across all subsequent generate_rollout calls so its queue stays warm", "each generate_rollout call drains until it has rollout_batch_size groups and returns them", and queue_left= is logged at the end as if leftovers persisted.

Reproduction (real _generate_rollout_async): 10 groups ready, target=4 → 4 returned, queue left 0, groups 4–9 gone.

Bug 2: blocking put on the event-loop thread

The task done-callback runs on the worker's event-loop thread and calls blocking Queue.put on a bounded queue (maxsize=1000). Once the queue fills — sized 1536-in-flight vs cap 1000, this happens on the first rollout of the example config — the callback blocks the loop itself, freezing every in-flight generate_and_rm_group coroutine until the training side happens to drain.

Why one PR

Fixing bug 1 alone makes bug 2 strictly worse: keeping the surplus queued makes a standing backlog the normal state, turning the blocking put from an occasional stall into a guaranteed freeze.

Fix

• get_completed_groups(limit=None): the rollout pulls only target - len(collected) per poll; the surplus stays queued FIFO for the next rollout (the documented warm-queue behaviour). The now-dead [:target] slice is removed.
• The queue becomes unbounded, so the done-callback can never block the loop. Real backpressure moves to _loop's top-up gate: no new prompts are pulled while a full pool (concurrency) of completed groups is already waiting to be consumed. Memory stays bounded at ~2 × concurrency groups, same order as today's in-flight pool.
• _loop's sleep becomes self.poll_interval (default 1.0, unchanged) so the CPU test can drive iterations quickly.

Test

New CPU test tests/test_fully_async_rollout.py, registered in the cpu-unittest job (sglang_router / transformers stubbed the same way as tests/test_agent/test_agent_rollout_cpu.py):

1. _generate_rollout_async returns exactly target groups FIFO and leaves the surplus queued — fails on main with assert 0 == 6.
2. get_completed_groups limit semantics.
3. The done-callback completes past the old 1000-item cap instead of blocking.
4. With instantly-completing generations and 60 groups of fuel, the queue plateaus at ≤ 2×concurrency instead of absorbing the dataset — fails on main with queue grew to 9 with concurrency=3.

[shinytang6]

@keepkeen thx for the fix, we met the similar problem, please help take a look @zhuzilin