From cb3dc1d08aa3b8c8b5fa733b94238a99a298e42e Mon Sep 17 00:00:00 2001 From: xiaoming <1294892474@qq.com> Date: Sun, 26 Jul 2026 18:39:48 +0800 Subject: [PATCH] fix: stop the fully-async rollout dropping completed groups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two coupled defects in the fully-async worker's output queue: 1. `_generate_rollout_async` drained the entire queue on every poll but returned only `[:rollout_batch_size]`. Everything past the slice was discarded — fully generated, reward-scored groups whose prompts were already consumed from the data buffer. With the shipped example config the worker keeps `sglang_server_concurrency * num_engines` (1536) groups in flight against a target of 8, so nearly every completed group was thrown away. The module's own contract says the opposite: "each generate_rollout call drains until it has rollout_batch_size groups", "the queue stays warm", and `queue_left=` is logged as if leftovers persisted. 2. The task done-callback called blocking `Queue.put` on a bounded queue (maxsize=1000) from the event-loop thread. When the queue filled, the callback blocked the loop itself, freezing every in-flight generation. They have to be fixed together: keeping the surplus queued (1) makes a standing backlog normal, which turns the blocking put (2) from an occasional stall into a guaranteed freeze. - `get_completed_groups` takes a `limit`; the rollout pulls only what it still needs, and the surplus stays queued for the next call (FIFO). - The queue is unbounded so the callback can never block the loop; real backpressure moves to `_loop`, which stops pulling new prompts while a full pool of completed groups is already waiting. - The loop's sleep becomes `self.poll_interval` so the CPU test can drive iterations quickly. Adds tests/test_fully_async_rollout.py (CPU CI): surplus retention + FIFO, `limit` semantics, callback completing past the old 1000-item cap, and the queue plateauing near `concurrency` instead of absorbing the whole dataset. --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + slime/rollout/fully_async_rollout.py | 38 ++++-- tests/test_fully_async_rollout.py | 171 +++++++++++++++++++++++++++ 4 files changed, 204 insertions(+), 10 deletions(-) create mode 100644 tests/test_fully_async_rollout.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 45606dd532..35be769cf3 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -629,6 +629,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_reloadable_process_group_world.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index f58b064830..d3c1cf1412 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -82,6 +82,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_reloadable_process_group_world.py', 'num_gpus': 0}, {'test_file': 'test_placement_group.py', 'num_gpus': 0}, {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0}, diff --git a/slime/rollout/fully_async_rollout.py b/slime/rollout/fully_async_rollout.py index c301075c5c..6b94fe7acd 100644 --- a/slime/rollout/fully_async_rollout.py +++ b/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, diff --git a/tests/test_fully_async_rollout.py b/tests/test_fully_async_rollout.py new file mode 100644 index 0000000000..7c5a595b28 --- /dev/null +++ b/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}"