From efbcc387d7175d395d08075b12cfa0e6a0e3d4a7 Mon Sep 17 00:00:00 2001
From: Ashton Chew <44445391+ashtonchew@users.noreply.github.com>
Date: Thu, 30 Jul 2026 21:21:34 +0000
Subject: [PATCH] fix(rollout): preserve fully async concurrency control

Carry the existing sample-concurrency bound into the class-based fully async rollout path and align its evaluation fixtures.
---
 miles/rollout/fully_async_rollout.py          | 19 ++++++--
 tests/fast/rollout/test_checkpoint_eval.py    |  1 +
 .../fast/rollout/test_fully_async_rollout.py  | 45 +++++++++++++++++++
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/miles/rollout/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py
index 0162a8aac81..ca8c73acba7 100644
--- a/miles/rollout/fully_async_rollout.py
+++ b/miles/rollout/fully_async_rollout.py
@@ -1,7 +1,7 @@
 """Fully asynchronous rollout generation.
 
-A persistent background worker keeps up to ``rollout_batch_size`` prompt groups in
-flight at all times; each training step only drains already-completed groups from the
+A persistent background worker keeps prompt groups in flight up to the configured
+concurrency limit. Each training step only drains already-completed groups from the
 worker's output queue. Rollout production and training consumption run in parallel,
 so per-iteration wall time moves from ``rollout_time + train_time`` toward
 ``max(rollout_time, train_time)``.
@@ -97,6 +97,17 @@ class FullyAsyncRolloutFn:
 
     def __init__(self, input: RolloutFnConstructorInput):
         self.args = input.args
+        if self.args.async_max_concurrent_samples is not None:
+            client_capacity = (
+                self.args.sglang_server_concurrency
+                * self.args.rollout_num_gpus
+                // self.args.rollout_num_gpus_per_engine
+            )
+            if self.args.async_max_concurrent_samples > client_capacity:
+                logger.warning(
+                    f"--async-max-concurrent-samples ({self.args.async_max_concurrent_samples}) exceeds the "
+                    f"client concurrency cap ({client_capacity}); the excess queues on the semaphore"
+                )
         self.data_source = input.data_source
         self.state = GenerateState(input.args)
         self._weight_version = _CachedWeightVersion()
@@ -132,7 +143,9 @@ async def _call_eval(self, input: RolloutFnEvalInput) -> RolloutFnOutput:
     # -------------------------- producer --------------------------
 
     def _max_in_flight_groups(self) -> int:
-        return self.args.rollout_batch_size
+        if self.args.async_max_concurrent_samples is None:
+            return self.args.rollout_batch_size
+        return max(1, self.args.async_max_concurrent_samples // self.args.n_samples_per_prompt)
 
     def _submit_one_group(self) -> asyncio.Task:
         [group] = self.data_source.get_samples(1)
diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py
index b698b0a6a28..913f52d3ed3 100644
--- a/tests/fast/rollout/test_checkpoint_eval.py
+++ b/tests/fast/rollout/test_checkpoint_eval.py
@@ -61,6 +61,7 @@ def _eval_dataset_env(monkeypatch, generate):
         chat_template_path=None,
         reward_key=None,
         eval_reward_key=None,
+        sglang_router_policy="round_robin",
     )
     dataset_cfg = SimpleNamespace(
         name="ds",
diff --git a/tests/fast/rollout/test_fully_async_rollout.py b/tests/fast/rollout/test_fully_async_rollout.py
index ac3ea6dc8fb..9d6bc95eaf8 100644
--- a/tests/fast/rollout/test_fully_async_rollout.py
+++ b/tests/fast/rollout/test_fully_async_rollout.py
@@ -70,6 +70,10 @@ def make_args(**overrides) -> Namespace:
         rollout_batch_size=2,
         n_samples_per_prompt=N_SAMPLES_PER_PROMPT,
         max_weight_staleness=None,
+        async_max_concurrent_samples=None,
+        sglang_server_concurrency=8,
+        rollout_num_gpus=4,
+        rollout_num_gpus_per_engine=1,
         sglang_router_ip="127.0.0.1",
         sglang_router_port=30000,
         eval_num_gpus=0,
@@ -253,3 +257,44 @@ async def blocking_generate(state, group, sampling_params, evaluation=False):
     release.set()
     output = await drain
     assert len(output.samples) == 2
+
+
+async def test_legacy_concurrency_flag_bounds_in_flight_groups(monkeypatch):
+    release = asyncio.Event()
+
+    async def blocking_generate(state, group, sampling_params, evaluation=False):
+        await release.wait()
+        return group
+
+    data_source = FakeDataSource()
+    fn = make_fn(
+        monkeypatch,
+        make_args(rollout_batch_size=4, async_max_concurrent_samples=5),
+        data_source,
+        generate=blocking_generate,
+    )
+
+    drain = asyncio.create_task(fn(RolloutFnTrainInput(rollout_id=0)))
+    await asyncio.sleep(0.05)
+    assert data_source.num_get_calls == 2
+
+    release.set()
+    output = await drain
+    assert len(output.samples) == 4
+
+
+def test_legacy_concurrency_flag_warns_above_client_capacity(monkeypatch, caplog):
+    args = make_args(
+        async_max_concurrent_samples=17,
+        sglang_server_concurrency=4,
+        rollout_num_gpus=4,
+        rollout_num_gpus_per_engine=1,
+    )
+
+    with caplog.at_level("WARNING"):
+        make_fn(monkeypatch, args, FakeDataSource())
+
+    assert caplog.messages == [
+        "--async-max-concurrent-samples (17) exceeds the client concurrency cap (16); "
+        "the excess queues on the semaphore"
+    ]