Backfill async rollouts on sample completions

examples/fully_async/fully_async_rollout.py

@@ -96,6 +96,54 @@ async def continuous_worker_loop(self):
         active_tasks = set()
         max_concurrent_tasks = self.args.rollout_batch_size
         group_id_counter = 0
+        sample_completion_backfill = getattr(self.args, "rollout_sample_completion_backfill", False)
+        samples_per_group = max(1, int(getattr(self.args, "n_samples_per_prompt", 1) or 1))
+        completed_sample_credits = 0
+        sample_backfill_initialized = False
+
+        if sample_completion_backfill:
+            print(
+                "Sample-completion backfill enabled: "
+                f"initial_groups={max_concurrent_tasks}, samples_per_group={samples_per_group}"
+            )
+
+        def on_sample_done():
+            nonlocal completed_sample_credits
+            completed_sample_credits += 1
+
+        def submit_one_group() -> bool:
+            nonlocal group_id_counter
+            samples = self.data_buffer.get_samples(1)
+
+            for group in samples:
+                group_id = group_id_counter
+                group_id_counter += 1
+
+                generate_kwargs = dict(
+                    args=self.args,
+                    group=group,
+                    sampling_params=self.state.sampling_params.copy(),
+                    evaluation=False,
+                )
+                if sample_completion_backfill:
+                    generate_kwargs["sample_done_callback"] = on_sample_done
+
+                # Create new async task
+                task = asyncio.create_task(generate_and_rm_group(**generate_kwargs))
+
+                # Add completion callback
+                def make_callback(gid):
+                    def task_done_callback(done_task):
+                        result = done_task.result()
+                        self.output_queue.put((gid, result))
+
+                    return task_done_callback
+
+                task.add_done_callback(make_callback(group_id))
+                active_tasks.add(task)
+                return True
+
+            return False
 
         while self.running:
             try:
@@ -109,35 +157,27 @@ async def continuous_worker_loop(self):
                             print(f"Task failed with exception: {e}")
                     active_tasks -= done_tasks
 
-                # If active task count hasn't reached limit, try to get new data and start tasks
-                while len(active_tasks) < max_concurrent_tasks and self.running:
-                    samples = self.data_buffer.get_samples(1)
-
-                    for group in samples:
-                        group_id = group_id_counter
-                        group_id_counter += 1
-
-                        # Create new async task
-                        task = asyncio.create_task(
-                            generate_and_rm_group(
-                                self.args,
-                                group,
-                                sampling_params=self.state.sampling_params.copy(),
-                                evaluation=False,
-                            )
-                        )
-
-                        # Add completion callback
-                        def make_callback(gid):
-                            def task_done_callback(done_task):
-                                result = done_task.result()
-                                self.output_queue.put((gid, result))
-
-                            return task_done_callback
-
-                        task.add_done_callback(make_callback(group_id))
-                        active_tasks.add(task)
-                        break
+                if sample_completion_backfill:
+                    if not sample_backfill_initialized:
+                        while len(active_tasks) < max_concurrent_tasks and self.running:
+                            if not submit_one_group():
+                                break
+                        sample_backfill_initialized = len(active_tasks) >= max_concurrent_tasks
+
+                    while (
+                        sample_backfill_initialized and completed_sample_credits >= samples_per_group and self.running
+                    ):
+                        if not submit_one_group():
+                            break
+                        completed_sample_credits -= samples_per_group
+
+                    if not active_tasks and self.running:
+                        submit_one_group()
+                else:
+                    # If active task count hasn't reached limit, try to get new data and start tasks.
+                    while len(active_tasks) < max_concurrent_tasks and self.running:
+                        if not submit_one_group():
+                            break
 
                 # Brief sleep to avoid busy waiting
                 await asyncio.sleep(1)

miles/rollout/inference_rollout/inference_rollout_common.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 from argparse import Namespace
+from collections.abc import Callable
 from copy import deepcopy
 from typing import Any
 
@@ -118,7 +119,11 @@ async def generate_and_rm(
 
 
 async def generate_and_rm_group(
-    state: GenerateState, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False
+    state: GenerateState,
+    group: list[Sample],
+    sampling_params: dict[str, Any],
+    evaluation: bool = False,
+    sample_done_callback: Callable[[], None] | None = None,
 ) -> list[Sample]:
     args = state.args
 
@@ -132,9 +137,15 @@ async def generate_and_rm_group(
         current_sampling_params = sampling_params.copy()
         if getattr(args, "sglang_enable_deterministic_inference", False):
             current_sampling_params["sampling_seed"] = args.rollout_seed + idx
-        tasks.append(
-            asyncio.create_task(generate_and_rm(state, sample, current_sampling_params, evaluation=evaluation))
-        )
+        task = asyncio.create_task(generate_and_rm(state, sample, current_sampling_params, evaluation=evaluation))
+        if sample_done_callback is not None:
+            # Fires for every sample task regardless of outcome (success, exception,
+            # or cancellation): each submitted sample yields exactly one completion
+            # credit, so in-flight sample concurrency is conserved exactly. Do not
+            # condition this on task success -- skipping failed samples would leak
+            # scheduling credits and decay concurrency over time.
+            task.add_done_callback(lambda _task: sample_done_callback())
+        tasks.append(task)
 
     group = await asyncio.gather(*tasks)
     logger.debug(f"{log_prefix} [group] All {len(group)} samples completed")

miles/rollout/inference_rollout/inference_rollout_train.py

@@ -2,6 +2,7 @@
 import logging
 from argparse import Namespace
 from collections.abc import Callable
+from contextlib import suppress
 
 import sglang_router
 from packaging.version import parse
@@ -56,7 +57,11 @@ async def get_worker_urls(args: Namespace):
         return [worker["url"] for worker in response["workers"]]
 
 
-def submit_generate_tasks(state: GenerateState, samples: list[list[Sample]]):
+def submit_generate_tasks(
+    state: GenerateState,
+    samples: list[list[Sample]],
+    sample_done_callback: Callable[[], None] | None = None,
+):
     return [
         asyncio.create_task(
             # submit a group of samples as a single task.
@@ -65,6 +70,7 @@ def submit_generate_tasks(state: GenerateState, samples: list[list[Sample]]):
                 group,
                 sampling_params=state.sampling_params.copy(),
                 evaluation=False,
+                sample_done_callback=sample_done_callback,
             )
         )
         for group in samples
@@ -84,6 +90,57 @@ async def generate_rollout_async(
 
     metric_gatherer = MetricGatherer()
 
+    if getattr(args, "rollout_sample_completion_backfill", False):
+        data, all_data, aborted_samples = await _generate_rollout_sample_completion_backfill_async(
+            state,
+            rollout_id,
+            data_source,
+            dynamic_filter,
+            metric_gatherer,
+        )
+    else:
+        data, all_data, aborted_samples = await _generate_rollout_group_level_async(
+            state,
+            rollout_id,
+            data_source,
+            dynamic_filter,
+            metric_gatherer,
+        )
+
+    assert len(data) == args.rollout_batch_size, f"Got {len(data)} samples, expected {args.rollout_batch_size}"
+    data = sorted(data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index)
+    all_samples = sorted(
+        all_data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index
+    )
+
+    # reset the global state to prevent effects on the next rollout or eval.
+    state.reset()
+
+    if f := load_function(args.rollout_sample_filter_path):
+        f(args, data)
+    # There can be circumstances where users want to process all samples including filtered ones.
+    if f := load_function(args.rollout_all_samples_process_path):
+        f(args, all_samples, data_source)
+
+    await recompute_samples_rollout_logprobs_via_prefill(
+        args,
+        [sample for group in data for sample in group],
+        url=f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate",
+        sampling_params=state.sampling_params,
+    )
+
+    return RolloutFnTrainOutput(samples=data, metrics=metric_gatherer.collect()), aborted_samples
+
+
+async def _generate_rollout_group_level_async(
+    state: GenerateState,
+    rollout_id: int,
+    data_source: Callable[[int], list[list[Sample]]],
+    dynamic_filter,
+    metric_gatherer: MetricGatherer,
+) -> tuple[list[list[Sample]], list[list[Sample]], list[list[Sample]]]:
+    args = state.args
+
     # target_data_size is the total number of valid samples to get
     target_data_size = args.rollout_batch_size
 
@@ -138,26 +195,117 @@ async def generate_rollout_async(
     # there are still some unfinished requests, abort them
     aborted_samples = await abort(state, pendings, rollout_id)
 
-    assert len(data) == args.rollout_batch_size, f"Got {len(data)} samples, expected {args.rollout_batch_size}"
-    data = sorted(data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index)
-    all_samples = sorted(
-        all_data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index
+    return data, all_data, aborted_samples
+
+
+async def _generate_rollout_sample_completion_backfill_async(
+    state: GenerateState,
+    rollout_id: int,
+    data_source: Callable[[int], list[list[Sample]]],
+    dynamic_filter,
+    metric_gatherer: MetricGatherer,
+) -> tuple[list[list[Sample]], list[list[Sample]], list[list[Sample]]]:
+    args = state.args
+    target_data_size = args.rollout_batch_size
+    group_size = args.n_samples_per_prompt
+    sample_done_queue: asyncio.Queue[int] = asyncio.Queue()
+    accept_sample_done = True
+
+    def on_sample_done() -> None:
+        if accept_sample_done and not state.aborted:
+            sample_done_queue.put_nowait(1)
+
+    pendings = set()
+    data = []
+    all_data = []
+    do_print = True
+    sample_done_credit = 0
+    pbar = tqdm(total=target_data_size * group_size, desc="Rollout generation")
+
+    def submit_groups(num_groups: int) -> int:
+        if num_groups <= 0:
+            return 0
+        samples = data_source(num_groups)
+        new_tasks = submit_generate_tasks(state, samples, sample_done_callback=on_sample_done)
+        pendings.update(new_tasks)
+        return len(new_tasks)
+
+    logger.info(
+        "[rollout] sample-completion backfill enabled: target_groups=%s group_size=%s",
+        target_data_size,
+        group_size,
     )
+    submit_groups(target_data_size)
 
-    # reset the global state to prevent effects on the next rollout or eval.
-    state.reset()
+    while len(data) < target_data_size:
+        if not pendings:
+            # Defensive fallback for group-level task exceptions. Normal flow keeps
+            # pending sample slots replenished from sample completion credits. If the
+            # data source is exhausted and nothing is in flight, stop instead of
+            # blocking forever on the sample-completion queue.
+            if submit_groups(max(1, target_data_size - len(data))) == 0:
+                break
+
+        sample_done_task = asyncio.create_task(sample_done_queue.get())
+        done, _ = await asyncio.wait(pendings | {sample_done_task}, return_when=asyncio.FIRST_COMPLETED)
+
+        sample_done_count = 0
+        if sample_done_task in done:
+            sample_done_count += sample_done_task.result()
+        else:
+            sample_done_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await sample_done_task
+
+        while True:
+            try:
+                sample_done_count += sample_done_queue.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+        sample_done_credit += sample_done_count
 
-    if f := load_function(args.rollout_sample_filter_path):
-        f(args, data)
-    # There can be circumstances where users want to process all samples including filtered ones.
-    if f := load_function(args.rollout_all_samples_process_path):
-        f(args, all_samples, data_source)
+        group_done = done & pendings
+        if group_done:
+            pendings.difference_update(group_done)
 
-    await recompute_samples_rollout_logprobs_via_prefill(
-        args,
-        [sample for group in data for sample in group],
-        url=f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate",
-        sampling_params=state.sampling_params,
+        for task in group_done:
+            try:
+                group: list[Sample] = task.result()
+            except Exception as e:
+                logger.error(f"[rollout] Task raised exception: {e!r}", exc_info=True)
+                continue
+
+            if do_print:
+                sample = group[0][0] if isinstance(group[0], list) else group[0]
+                logger.info(
+                    f"First rollout sample: {[str(sample.prompt) + sample.response]}, label: {sample.label}, reward: {sample.reward}",
+                )
+                do_print = False
+
+            assert len(group) == group_size
+            all_data.append(group)
+            dynamic_filter_output = call_dynamic_filter(dynamic_filter, args, group)
+            if not dynamic_filter_output.keep:
+                metric_gatherer.on_dynamic_filter_drop(reason=dynamic_filter_output.reason)
+                continue
+
+            if len(data) < target_data_size:
+                data.append(group)
+                pbar.update(group_size)
+
+        while sample_done_credit >= group_size and len(data) < target_data_size:
+            submitted = submit_groups(1)
+            if submitted <= 0:
+                break
+            sample_done_credit -= group_size
+
+    pbar.close()
+    sample = data[-1][0][0] if isinstance(data[-1][0], list) else data[-1][0]
+    logger.info(
+        f"Finish rollout: {[str(sample.prompt) + sample.response]}, label: {sample.label}, reward: {sample.reward}",
     )
 
-    return RolloutFnTrainOutput(samples=data, metrics=metric_gatherer.collect()), aborted_samples
+    accept_sample_done = False
+    aborted_samples = await abort(state, pendings, rollout_id)
+
+    return data, all_data, aborted_samples