fix: stop the fully-async rollout dropping completed groups

.github/workflows/pr-test.yml

@@ -637,6 +637,10 @@ jobs:
               "num_gpus": 0,
               "test_file": "test_rollout_validation.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_fully_async_rollout.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "test_rollout_sample_hooks.py"

.github/workflows/pr-test.yml.j2

@@ -85,6 +85,7 @@
         {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0},
         {'test_file': 'test_sample.py', 'num_gpus': 0},
         {'test_file': 'test_rollout_validation.py', 'num_gpus': 0},
+        {'test_file': 'test_fully_async_rollout.py', 'num_gpus': 0},
         {'test_file': 'test_rollout_sample_hooks.py', 'num_gpus': 0},
         {'test_file': 'test_hf_to_megatron.py', 'num_gpus': 0},
         {'test_file': 'test_qwen3_5_vl_native.py', 'num_gpus': 0},

slime/rollout/fully_async_rollout.py

@@ -82,7 +82,13 @@ def __init__(self, args, data_buffer, concurrency: int = 10):
         self.data_buffer = data_buffer
         self.concurrency = concurrency
         self.running = True
-        self.output_queue: queue.Queue[tuple[int, list[Sample]]] = queue.Queue(maxsize=1000)
+        # Unbounded on purpose: put() runs inside the event-loop thread (task
+        # done-callback), so a bounded queue that fills up would block the loop
+        # and freeze every in-flight generation. Backpressure lives in _loop()
+        # instead, which stops topping up while a full pool of completed groups
+        # is already waiting to be consumed.
+        self.output_queue: queue.Queue[tuple[int, list[Sample]]] = queue.Queue()
+        self.poll_interval = 1.0
         self.worker_thread: threading.Thread | None = None
         self.state = GenerateState(args)
 
@@ -98,9 +104,16 @@ def stop(self) -> None:
         if self.worker_thread and self.worker_thread.is_alive():
             self.worker_thread.join(timeout=5)
 
-    def get_completed_groups(self) -> list[tuple[int, list[Sample]]]:
+    def get_completed_groups(self, limit: int | None = None) -> list[tuple[int, list[Sample]]]:
+        """Pop up to ``limit`` completed groups (all of them when ``None``).
+
+        Callers that only need a fixed number of groups must pass ``limit`` —
+        anything popped beyond it would otherwise have to be thrown away, and
+        these groups are fully generated and reward-scored, with their prompts
+        already consumed from ``data_buffer``.
+        """
         completed: list[tuple[int, list[Sample]]] = []
-        while True:
+        while limit is None or len(completed) < limit:
             try:
                 completed.append(self.output_queue.get_nowait())
             except queue.Empty:
@@ -132,8 +145,12 @@ async def _loop(self) -> None:
                             logger.warning("fully-async task crashed: %r", e)
                     active_tasks -= done
 
-                # Top up.
-                while len(active_tasks) < max_concurrent and self.running:
+                # Top up. The qsize gate is the queue's backpressure: once a
+                # full pool of completed groups is waiting, stop pulling new
+                # prompts until the training side drains some.
+                while (
+                    len(active_tasks) < max_concurrent and self.output_queue.qsize() < max_concurrent and self.running
+                ):
                     groups = self.data_buffer.get_samples(1)
                     if not groups:
                         break
@@ -151,10 +168,10 @@ async def _loop(self) -> None:
                         task.add_done_callback(self._make_done_cb(gid))
                         active_tasks.add(task)
 
-                await asyncio.sleep(1)
+                await asyncio.sleep(self.poll_interval)
             except Exception as e:  # noqa: BLE001
                 logger.exception("fully-async loop iteration error: %s", e)
-                await asyncio.sleep(1)
+                await asyncio.sleep(self.poll_interval)
 
         if active_tasks:
             logger.info(
@@ -209,9 +226,10 @@ async def _generate_rollout_async(args, rollout_id: int, data_buffer) -> list[li
     LOG_EVERY = 30.0
 
     while len(collected) < target:
-        # Pull whatever's done.
+        # Pull only what this rollout still needs; the surplus stays queued for
+        # the next rollout (that is the "queue stays warm" contract).
         drained = 0
-        for gid, group in worker.get_completed_groups():
+        for gid, group in worker.get_completed_groups(limit=target - len(collected)):
             collected[gid] = group
             drained += 1
 
@@ -238,7 +256,7 @@ def _key(group: list[Sample]) -> int:
                 return int(idx)
         return 0
 
-    out = sorted(collected.values(), key=_key)[:target]
+    out = sorted(collected.values(), key=_key)
     logger.info(
         "fully-async rollout %d: done in %.1fs, queue_left=%d",
         rollout_id,

tests/test_fully_async_rollout.py

@@ -0,0 +1,171 @@
+"""CPU unit tests for the fully-async rollout worker's queue contract.
+
+The module docstring of ``slime.rollout.fully_async_rollout`` promises that the
+worker's output queue "stays warm" across ``generate_rollout`` calls: each call
+takes ``rollout_batch_size`` completed groups and leaves the rest queued.
+
+Three behaviours are pinned here:
+
+  1. ``_generate_rollout_async`` consumes exactly ``rollout_batch_size`` groups
+     and leaves the surplus in the queue. (It used to drain the whole queue and
+     slice — throwing away fully generated, reward-scored groups whose prompts
+     had already been consumed from the data buffer.)
+  2. The task done-callback never blocks. It runs on the event-loop thread, so
+     a bounded queue that filled up would freeze every in-flight generation.
+  3. Backpressure exists anyway: ``_loop`` stops pulling new prompts while a
+     full pool of completed groups is already waiting to be consumed.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+import threading
+import time
+import types
+from collections import deque
+from types import SimpleNamespace
+
+# ``fully_async_rollout`` imports ``sglang_rollout``, which needs sglang_router
+# and (transitively) transformers — both deliberately absent from the CPU CI
+# env. The tests below never dial a server or touch a tokenizer, so stub the
+# imports, same as tests/test_agent/test_agent_rollout_cpu.py.
+if "sglang_router" not in sys.modules:
+    _router_stub = types.ModuleType("sglang_router")
+    _router_stub.__version__ = "0.2.3"
+    sys.modules["sglang_router"] = _router_stub
+if "transformers" not in sys.modules:
+    _tf_stub = types.ModuleType("transformers")
+    for _name in ("AutoProcessor", "AutoTokenizer", "PreTrainedTokenizerBase", "ProcessorMixin"):
+        setattr(_tf_stub, _name, type(_name, (), {}))
+    sys.modules["transformers"] = _tf_stub
+
+import pytest
+
+import slime.rollout.fully_async_rollout as fa
+from slime.utils.types import Sample
+
+
+NUM_GPUS = 0
+
+
+class _FakeGenerateState:
+    def __init__(self, args):
+        self.sampling_params = {}
+
+
+class _FakeDataBuffer:
+    """Finite fuel: one group per ``get_samples`` call until exhausted."""
+
+    def __init__(self, groups):
+        self._groups = deque(groups)
+        self.requeued = []
+
+    def get_samples(self, n):
+        assert n == 1
+        if not self._groups:
+            return []
+        return [self._groups.popleft()]
+
+    def add_samples(self, groups):
+        self.requeued.extend(groups)
+
+
+def _make_group(index: int) -> list[Sample]:
+    sample = Sample(index=index, prompt=f"p{index}")
+    sample.status = Sample.Status.COMPLETED
+    return [sample]
+
+
+def _make_worker(monkeypatch, data_buffer=None, concurrency=4) -> fa.AsyncRolloutWorker:
+    monkeypatch.setattr(fa, "GenerateState", _FakeGenerateState)
+    args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4)
+    return fa.AsyncRolloutWorker(args, data_buffer or _FakeDataBuffer([]), concurrency=concurrency)
+
+
+@pytest.mark.unit
+def test_rollout_takes_target_groups_and_leaves_surplus_queued(monkeypatch):
+    worker = _make_worker(monkeypatch)
+    for gid in range(10):
+        worker.output_queue.put((gid, _make_group(gid)))
+    monkeypatch.setattr(fa, "_get_global_worker", lambda args, data_buffer: worker)
+
+    args = SimpleNamespace(rollout_global_dataset=True, rollout_batch_size=4)
+    out = asyncio.run(fa._generate_rollout_async(args, rollout_id=0, data_buffer=None))
+
+    assert len(out) == 4
+    # FIFO: the oldest four groups ship first.
+    assert [group[0].index for group in out] == [0, 1, 2, 3]
+    # The other six are still queued for the next rollout, not thrown away.
+    assert worker.queue_size() == 6
+    assert [gid for gid, _ in worker.get_completed_groups()] == [4, 5, 6, 7, 8, 9]
+
+
+@pytest.mark.unit
+def test_get_completed_groups_limit(monkeypatch):
+    worker = _make_worker(monkeypatch)
+    for gid in range(5):
+        worker.output_queue.put((gid, _make_group(gid)))
+
+    assert [gid for gid, _ in worker.get_completed_groups(limit=2)] == [0, 1]
+    assert [gid for gid, _ in worker.get_completed_groups()] == [2, 3, 4]
+    assert worker.get_completed_groups(limit=3) == []
+
+
+@pytest.mark.unit
+def test_done_callback_never_blocks_event_loop_thread(monkeypatch):
+    """The callback runs on the loop thread; blocking there freezes every
+    in-flight generation. Push more results than the old bounded-queue cap
+    (1000) through it and require completion."""
+    worker = _make_worker(monkeypatch)
+
+    class _DoneTask:
+        def __init__(self, gid):
+            self._result = _make_group(gid)
+
+        def result(self):
+            return self._result
+
+    def _push_all():
+        for gid in range(1001):
+            worker._make_done_cb(gid)(_DoneTask(gid))
+
+    pusher = threading.Thread(target=_push_all, daemon=True)
+    pusher.start()
+    pusher.join(timeout=30)
+
+    assert not pusher.is_alive(), "done-callback blocked on a full output queue"
+    assert worker.queue_size() == 1001
+
+
+@pytest.mark.unit
+def test_loop_backpressure_stops_topping_up_when_queue_is_full(monkeypatch):
+    """With instantly-completing generations and plenty of fuel, the queue must
+    plateau around ``concurrency`` instead of absorbing the whole dataset."""
+    concurrency = 3
+    fuel = 60
+    data_buffer = _FakeDataBuffer([_make_group(i) for i in range(fuel)])
+
+    async def _instant_generate(args, group, sampling_params, evaluation):
+        return group
+
+    monkeypatch.setattr(fa, "generate_and_rm_group", _instant_generate)
+    worker = _make_worker(monkeypatch, data_buffer=data_buffer, concurrency=concurrency)
+    worker.poll_interval = 0.01
+
+    worker.start()
+    try:
+        # Give the loop ample iterations to overshoot if it is going to.
+        deadline = time.time() + 3.0
+        max_seen = 0
+        while time.time() < deadline:
+            max_seen = max(max_seen, worker.queue_size())
+            if max_seen > 2 * concurrency:
+                break
+            time.sleep(0.02)
+    finally:
+        worker.stop()
+
+    # In-flight tasks may still land after the gate check, so allow one pool
+    # beyond the gate — but nothing near the unthrottled fuel size.
+    assert 0 < max_seen <= 2 * concurrency, f"queue grew to {max_seen} with concurrency={concurrency}"