diff --git a/examples/fully_async/fully_async_rollout.py b/examples/fully_async/fully_async_rollout.py index e4c23cc120..d515fc6183 100644 --- a/examples/fully_async/fully_async_rollout.py +++ b/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) diff --git a/miles/rollout/inference_rollout/inference_rollout_common.py b/miles/rollout/inference_rollout/inference_rollout_common.py index 9f1cc603b0..9e957168d6 100644 --- a/miles/rollout/inference_rollout/inference_rollout_common.py +++ b/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") diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py index eed15557bd..95e2139326 100644 --- a/miles/rollout/inference_rollout/inference_rollout_train.py +++ b/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 diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index 64662ed522..e55161a65e 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -306,7 +306,11 @@ async def generate_and_rm( async def generate_and_rm_group( - args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False + args: Namespace, + group: list[Sample], + sampling_params: dict[str, Any], + evaluation: bool = False, + sample_done_callback: Callable[[], None] | None = None, ) -> list[Sample]: state = GenerateState(args) @@ -325,9 +329,13 @@ async def generate_and_rm_group( if getattr(args, "sglang_enable_deterministic_inference", False): seed = state.group_sampling_seeds[idx] current_sampling_params["sampling_seed"] = seed - tasks.append( - asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation)) - ) + task = asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation)) + if sample_done_callback is not None: + # Fires for every sample task regardless of outcome; see the note in + # inference_rollout_common.generate_and_rm_group -- credits must conserve + # in-flight sample concurrency exactly. + task.add_done_callback(lambda _task: sample_done_callback()) + tasks.append(task) group = await asyncio.gather(*tasks) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 53dcc34f8e..ea66a1ef5b 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -458,6 +458,17 @@ def add_rollout_arguments(parser): "You could use `miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std` as an example." ), ) + parser.add_argument( + "--rollout-sample-completion-backfill", + action="store_true", + default=False, + help=( + "If set, training rollout replenishes one new prompt group after " + "n_samples_per_prompt individual samples complete, instead of waiting " + "for a full group task to return. Disabled by default to preserve the " + "legacy group-level scheduling behavior." + ), + ) # partial rollout parser.add_argument( diff --git a/tests/fast/rollout/inference_rollout/test_sample_completion_backfill.py b/tests/fast/rollout/inference_rollout/test_sample_completion_backfill.py new file mode 100644 index 0000000000..15e8ee9304 --- /dev/null +++ b/tests/fast/rollout/inference_rollout/test_sample_completion_backfill.py @@ -0,0 +1,155 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu", labels=[]) + +import asyncio +import sys +from argparse import Namespace +from types import ModuleType + +import pytest + +if "sglang_router" not in sys.modules: + sglang_router_stub = ModuleType("sglang_router") + sglang_router_stub.__version__ = "0.0.0" + sys.modules["sglang_router"] = sglang_router_stub + +import miles.rollout.inference_rollout.inference_rollout_train as train +from miles.rollout.filter_hub.base_types import MetricGatherer +from miles.utils.types import Sample + + +class FakeGenerateState: + def __init__(self, args: Namespace): + self.args = args + self.sampling_params = {} + self.aborted = False + self.reset_count = 0 + + def reset(self) -> None: + self.aborted = False + self.reset_count += 1 + + +def make_group(group_index: int, group_size: int) -> list[Sample]: + return [ + Sample( + group_index=group_index, + index=group_index * 100 + sample_index, + prompt=f"prompt {group_index}", + response="ok", + response_length=1, + label="ok", + reward=1, + status=Sample.Status.COMPLETED, + ) + for sample_index in range(group_size) + ] + + +@pytest.mark.asyncio +async def test_generate_rollout_without_backfill_flag_uses_legacy_group_scheduler(monkeypatch): + group = make_group(group_index=1, group_size=2) + args = Namespace( + rollout_global_dataset=True, + rollout_batch_size=1, + rollout_sample_filter_path=None, + rollout_all_samples_process_path=None, + dynamic_sampling_filter_path=None, + sglang_router_ip="127.0.0.1", + sglang_router_port=30000, + ) + state = FakeGenerateState(args) + called = [] + + async def noop_configure_sglang(_args): + return None + + async def noop_recompute(*_args, **_kwargs): + return None + + async def fake_group_level(*_args, **_kwargs): + called.append("group_level") + return [group], [group], [] + + async def unexpected_sample_completion_backfill(*_args, **_kwargs): + raise AssertionError("sample-completion backfill should be disabled by default") + + monkeypatch.setattr(train.dumper_utils, "configure_sglang", noop_configure_sglang) + monkeypatch.setattr(train, "recompute_samples_rollout_logprobs_via_prefill", noop_recompute) + monkeypatch.setattr(train, "load_function", lambda _path: None) + monkeypatch.setattr(train, "_generate_rollout_group_level_async", fake_group_level) + monkeypatch.setattr( + train, + "_generate_rollout_sample_completion_backfill_async", + unexpected_sample_completion_backfill, + ) + + output, aborted_samples = await train.generate_rollout_async(state, rollout_id=0, data_source=lambda _n: []) + + assert called == ["group_level"] + assert output.samples == [group] + assert aborted_samples == [] + assert state.reset_count == 1 + + +@pytest.mark.asyncio +async def test_sample_completion_backfill_submits_group_after_enough_samples_finish(monkeypatch): + args = Namespace(rollout_batch_size=1, n_samples_per_prompt=2) + state = FakeGenerateState(args) + data_source_calls = [] + submitted_group_indices = [] + next_group_index = 0 + + def data_source(num_groups: int) -> list[list[Sample]]: + nonlocal next_group_index + data_source_calls.append(num_groups) + groups = [] + for _ in range(num_groups): + next_group_index += 1 + groups.append(make_group(group_index=next_group_index, group_size=args.n_samples_per_prompt)) + return groups + + async def never_complete(): + await asyncio.Future() + + async def complete_group(group: list[Sample]) -> list[Sample]: + await asyncio.sleep(0) + return group + + def fake_submit_generate_tasks(_state, samples, sample_done_callback=None): + tasks = [] + for group in samples: + submitted_group_indices.append(group[0].group_index) + if len(submitted_group_indices) == 1: + assert sample_done_callback is not None + for _ in group: + sample_done_callback() + tasks.append(asyncio.create_task(never_complete())) + else: + tasks.append(asyncio.create_task(complete_group(group))) + return tasks + + async def fake_abort(_state, pendings, _rollout_id): + _state.aborted = True + for task in pendings: + task.cancel() + await asyncio.gather(*pendings, return_exceptions=True) + return [] + + monkeypatch.setattr(train, "submit_generate_tasks", fake_submit_generate_tasks) + monkeypatch.setattr(train, "abort", fake_abort) + + data, all_data, aborted_samples = await train._generate_rollout_sample_completion_backfill_async( + state, + rollout_id=0, + data_source=data_source, + dynamic_filter=None, + metric_gatherer=MetricGatherer(), + ) + + assert data_source_calls == [1, 1] + assert submitted_group_indices == [1, 2] + assert data == [make_group(group_index=2, group_size=args.n_samples_per_prompt)] + assert all_data == data + assert aborted_samples == []