From 2e61d9643e70b015cbd3c9816bbaa8c1bf7a180d Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 2 Jun 2026 02:38:21 +0000 Subject: [PATCH] feat: promote fully_async rollout into core (slime #1920) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port of THUDM/slime#1920 — moves the fully-async rollout from an example into the core package so it's usable via --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async. - vime/rollout/fully_async_rollout.py (new): promoted from examples/fully_async/fully_async_rollout.py and translated to vime's vLLM seam (vime.rollout.vllm_rollout, args.vllm_server_concurrency). Carries slime's promotion refinements: docstring, logging, __all__, atexit cleanup, num_engines concurrency scaling, ABORTED->data_buffer redirect, fan-out handling. - vime/rollout/vllm_rollout.py: generate_and_rm_group return type -> list[Sample] | list[list[Sample]] + comment (analog of slime's sglang_rollout change; supports --custom-generate-function-path fan-out into multiple samples). - examples/fully_async/: remove the now-duplicated fully_async_rollout.py (moved to core); repoint run-qwen3-4b-fully_async.sh + README at the core --rollout-function-path. (vime keeps the 4B example rather than swapping to 0.5B as upstream did; the CI test below covers 0.5B.) - tests/test_qwen2.5_0.5B_fully_async_short.py (new) + registered in the run-ci-short matrix (pr-test.yml + .j2 regenerated). Mirrors test_qwen2.5_0.5B_async_short, only flipping the rollout-function-path. Refs: THUDM/slime#1920, #107 Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: aoshen02 --- .github/workflows/pr-test.yml | 2 +- .github/workflows/pr-test.yml.j2 | 1 + examples/fully_async/README.md | 5 +- examples/fully_async/fully_async_rollout.py | 262 ------------------ .../fully_async/run-qwen3-4b-fully_async.sh | 2 +- tests/test_qwen2.5_0.5B_fully_async_short.py | 134 +++++++++ vime/rollout/fully_async_rollout.py | 256 +++++++++++++++++ vime/rollout/vllm_rollout.py | 9 +- 8 files changed, 404 insertions(+), 267 deletions(-) delete mode 100644 examples/fully_async/fully_async_rollout.py create mode 100644 tests/test_qwen2.5_0.5B_fully_async_short.py create mode 100644 vime/rollout/fully_async_rollout.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 5d9db6a23..efbfe8636 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -56,7 +56,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}] + info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}] defaults: run: working-directory: ${{ github.workspace }} diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index def8dc6bf..2f0f34d95 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -5,6 +5,7 @@ {'test_file': 'test_qwen3.5_0.8B_gsm8k_async_short.py', 'num_gpus': 4}, {'test_file': 'test_qwen3.5_0.8B_gsm8k_short.py', 'num_gpus': 4}, {'test_file': 'test_qwen2.5_0.5B_ppo_critic_only_short.py', 'num_gpus': 4}, + {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4}, ], }, 'e2e-test-vllm-config': { diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index 36a36cb7d..2545ada98 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -3,7 +3,8 @@ This example shows a simple way to make rollout generation **fully asynchronous**: a single global worker is created once and then keeps running in the background, continuously pulling prompts and launching generation tasks. Training only needs to fetch already finished results. This removes the per‑step wait that happens in the normal synchronous style. ## Files -* `fully_async_rollout.py`: global async worker + `generate_rollout_fully_async` entry. +The fully-async worker has been **promoted from this example into the core package** — it now lives in +`vime/rollout/fully_async_rollout.py`. This directory keeps only the launch script: * `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B. ## Prerequisite @@ -37,7 +38,7 @@ To enable the fully async pattern there are only two changes compared to a norma 1. Use the async training driver: `train_async.py` (not `train.py`). 2. Set the rollout function path: ```bash - --rollout-function-path fully_async_rollout.generate_rollout_fully_async + --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async ``` Why is it still "fully" async although `train_async.py` itself schedules rollouts step‑by‑step? diff --git a/examples/fully_async/fully_async_rollout.py b/examples/fully_async/fully_async_rollout.py deleted file mode 100644 index 484370057..000000000 --- a/examples/fully_async/fully_async_rollout.py +++ /dev/null @@ -1,262 +0,0 @@ -import asyncio -import atexit -import queue -import threading -import time - -from vime.rollout.vllm_rollout import GenerateState, generate_and_rm_group -from vime.utils.async_utils import run -from vime.utils.types import Sample - -# Global worker manager -_global_worker = None -_worker_lock = threading.Lock() - - -def get_global_worker(args, data_buffer): - """Get or create global worker""" - global _global_worker - with _worker_lock: - if _global_worker is None or not _global_worker.worker_thread.is_alive(): - print("Creating new global async worker...") - _global_worker = AsyncRolloutWorker(args, data_buffer, concurrency=args.vllm_server_concurrency) - _global_worker.start() - return _global_worker - - -def stop_global_worker(): - """Stop global worker""" - global _global_worker - with _worker_lock: - if _global_worker is not None: - _global_worker.stop() - _global_worker = None - - -class AsyncRolloutWorker: - """ - Simplified asynchronous rollout worker, using threads instead of processes - Supports continuous running, independent of rollout function lifecycle - """ - - def __init__(self, args, data_buffer, concurrency=10): - self.args = args - self.data_buffer = data_buffer # Directly save data_buffer reference - self.concurrency = concurrency - self.running = True - self.output_queue = queue.Queue(maxsize=1000) # Continuous output queue - self.worker_thread = None - self.state = GenerateState(args) - - async def continuous_worker_loop(self): - """Continuous work loop - constantly get data from data_buffer and process""" - print("Continuous async rollout worker started") - - active_tasks = set() - max_concurrent_tasks = self.args.rollout_batch_size - group_id_counter = 0 - - while self.running: - try: - # Clean up completed tasks - if active_tasks: - done_tasks = {task for task in active_tasks if task.done()} - for task in done_tasks: - try: - task.result() # Results are already handled in callbacks - except Exception as e: - 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 - - # Brief sleep to avoid busy waiting - await asyncio.sleep(1) - - except Exception as e: - print(f"Error in continuous worker loop: {e}") - await asyncio.sleep(1) - - if active_tasks: - print(f"Waiting for {len(active_tasks)} continuous tasks to complete...") - await asyncio.wait(active_tasks) - - print("Continuous async rollout worker stopped") - - def worker_thread_func(self): - """Worker function running in independent thread""" - asyncio.run(self.continuous_worker_loop()) - - def start(self): - """Start continuous work mode""" - if self.worker_thread is None or not self.worker_thread.is_alive(): - self.worker_thread = threading.Thread(target=self.worker_thread_func, daemon=True) - self.worker_thread.start() - print("Started continuous async worker thread") - - def stop(self): - """Stop worker thread""" - self.running = False - if self.worker_thread and self.worker_thread.is_alive(): - self.worker_thread.join(timeout=5) - print("Stopped async worker thread") - - def get_completed_groups(self) -> list[tuple]: - """Get completed sample groups""" - completed = [] - while True: - try: - result = self.output_queue.get_nowait() - completed.append(result) - except queue.Empty: - break - return completed - - def get_queue_size(self) -> int: - """Get current output queue size""" - return self.output_queue.qsize() - - -async def generate_rollout_async(args, rollout_id: int, data_buffer) -> list[list[Sample]]: - """ - Simplified asynchronous rollout generation - using global continuous worker - """ - assert args.rollout_global_dataset - - # Get global worker, which will run continuously - worker = get_global_worker(args, data_buffer) - - # Simplified: directly use rollout_batch_size as target - target_data_size = args.rollout_batch_size - - data = [] - completed_groups = {} - do_print = True - - print(f"Starting async rollout generation for {target_data_size} groups") - print(f"Global worker queue size: {worker.get_queue_size()}") - - # Main loop: collect results from global worker's output queue - start_time = time.time() - last_progress_time = start_time - no_progress_timeout = 30.0 # Warn if no progress for 30 seconds - - while len(data) < target_data_size: - # Collect completed results - completed = worker.get_completed_groups() - - made_progress = False - for group_id, group in completed: - completed_groups[group_id] = group - made_progress = True - - if made_progress: - last_progress_time = time.time() - - # Process completed groups in order (try to maintain order, but not strict requirement) - processed_any = False - - # Process all available completed groups - available_ids = list(completed_groups.keys()) - for group_id in available_ids: - if len(data) >= target_data_size: - break - - group = completed_groups.pop(group_id) - - # If any sample in the group was aborted, return the whole group to the data buffer - # and do not forward it to the training engine. - try: - any_aborted = any([sample.status == Sample.Status.ABORTED for sample in group]) - except Exception: - any_aborted = False - - if any_aborted: - try: - # add back to buffer so it can be retried or handled by buffer policy - data_buffer.add_samples([group]) - print(f"Returned aborted group {group_id} to data buffer", flush=True) - except Exception as e: - print(f"Failed to return aborted group {group_id} to buffer: {e}", flush=True) - # don't count as processed for training - continue - - if do_print: - print( - f"First rollout sample: {[group[0].prompt + group[0].response]}, " - f"label: {group[0].label}, reward: {group[0].reward}", - flush=True, - ) - do_print = False - - # Simplified: directly add samples, no filters used - data.append(group) - processed_any = True - - # Check progress - current_time = time.time() - if current_time - last_progress_time > no_progress_timeout: - print( - f"Warning: No progress for {no_progress_timeout}s. " - f"Queue size: {worker.get_queue_size()}, " - f"Collected: {len(data)}/{target_data_size}" - ) - last_progress_time = current_time - - # If no results were processed, brief sleep to avoid busy waiting - if not processed_any: - await asyncio.sleep(0.01) - - duration = time.time() - start_time - print(f"Rollout completed in {duration:.2f}s! Global worker queue size: {worker.get_queue_size()}") - - if data: - print( - f"Finish rollout: {[data[-1][0].prompt + data[-1][0].response]}, " - f"label: {data[-1][0].label}, reward: {data[-1][0].reward}", - flush=True, - ) - - data = sorted(data, key=lambda group: group[0].index) - return data - - -def generate_rollout_fully_async(args, rollout_id, data_buffer, evaluation=False): - if evaluation: - raise ValueError("Evaluation mode not supported in simple async rollout") - - completed_samples = run(generate_rollout_async(args, rollout_id, data_buffer)) - return completed_samples - - -# Register exit cleanup function - -atexit.register(stop_global_worker) diff --git a/examples/fully_async/run-qwen3-4b-fully_async.sh b/examples/fully_async/run-qwen3-4b-fully_async.sh index c293e115c..6535e08fa 100644 --- a/examples/fully_async/run-qwen3-4b-fully_async.sh +++ b/examples/fully_async/run-qwen3-4b-fully_async.sh @@ -38,7 +38,7 @@ CKPT_ARGS=( PROMPT_SET=/path/to/dapo-math-17k.jsonl ROLLOUT_ARGS=( - --rollout-function-path fully_async_rollout.generate_rollout_fully_async + --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async --prompt-data ${PROMPT_SET} --input-key prompt --label-key label diff --git a/tests/test_qwen2.5_0.5B_fully_async_short.py b/tests/test_qwen2.5_0.5B_fully_async_short.py new file mode 100644 index 000000000..b055949f6 --- /dev/null +++ b/tests/test_qwen2.5_0.5B_fully_async_short.py @@ -0,0 +1,134 @@ +"""CI smoke test for the fully-async rollout path. + +Mirrors ``test_qwen2.5_0.5B_async_short`` (Qwen2.5-0.5B + dapo-math-17k + +3 rollouts of GRPO) but flips the rollout function over to +``vime.rollout.fully_async_rollout.generate_rollout_fully_async`` so the +fully-async worker path gets exercised end-to-end. + +Kept intentionally minimal so it runs in the same time budget as the +existing 0.5B short tests. +""" + +import os +import vime.utils.external_utils.command_utils as U + +TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 4 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + + rollout_args = ( + # The only line that differs from test_qwen2.5_0.5B_async_short.py: + # use the public fully-async rollout function. + "--rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async " + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 3 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 8192 " + "--rollout-temperature 0.8 " + "--global-batch-size 32 " + "--balance-data " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 9216 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + vllm_args = ( + "--rollout-num-gpus-per-engine 1 " + f"--vllm-gpu-memory-utilization {0.55 if TIGHT_DEVICE_MEMORY else 0.65} " + "--vllm-max-cudagraph-capture-size 32 " + ) + + ci_args = "--ci-test " + + fault_tolerance_args = ( + "--use-fault-tolerance " + "--rollout-health-check-interval 5 " + "--rollout-health-check-timeout 10 " + "--rollout-health-check-first-wait 0 " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 1 " + "--rollout-num-gpus 3 " + "--megatron-to-hf-mode bridge " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{vllm_args} " + f"{ci_args} " + f"{fault_tolerance_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + train_script="train_async.py", + ) + + +if __name__ == "__main__": + prepare() + os.environ.pop("http_proxy", None) + os.environ.pop("https_proxy", None) + os.environ.pop("HTTP_PROXY", None) + os.environ.pop("HTTPS_PROXY", None) + execute() diff --git a/vime/rollout/fully_async_rollout.py b/vime/rollout/fully_async_rollout.py new file mode 100644 index 000000000..3ba24b74a --- /dev/null +++ b/vime/rollout/fully_async_rollout.py @@ -0,0 +1,256 @@ +"""Fully-async rollout for vime. + +Decouples ``max_concurrent_tasks`` from ``rollout_batch_size``: a background +asyncio worker keeps a fixed pool of in-flight trajectories across rollout +boundaries, so the next training step doesn't have to wait for the slowest +in-flight sample to finish. + +Use with ``--rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async``. +Plug in per-sample logic via ``--custom-generate-function-path`` and +per-sample reward via ``--custom-rm-path`` — the worker calls vime's stock +:func:`generate_and_rm_group` which dispatches to those. + +Concurrency is sourced from ``args.vllm_server_concurrency`` and scaled by +the number of vLLM engines (``rollout_num_gpus // rollout_num_gpus_per_engine``) +to match the per-sample semaphore cap in :mod:`vime.rollout.vllm_rollout`. + +The worker is intentionally oblivious to vime's higher-level pause / +weight-update signalling (e.g. ``GenerateState.aborted``). Each in-flight +generation short-circuits on those signals on its own and surfaces +:data:`Sample.Status.ABORTED`; the only piece the worker owns is +**redirecting ABORTED groups back to ``data_buffer``** instead of shipping +them to training, so the next rollout (with refreshed weights) can pick +them up. +""" + +from __future__ import annotations + +import asyncio +import atexit +import logging +import queue +import threading +import time + +from vime.rollout.vllm_rollout import GenerateState, generate_and_rm_group +from vime.utils.async_utils import run +from vime.utils.types import Sample + +__all__ = [ + "AsyncRolloutWorker", + "generate_rollout_fully_async", +] + +logger = logging.getLogger("vime.rollout.fully_async") + + +# Global worker, shared across rollout calls so the queue stays warm. +_global_worker: AsyncRolloutWorker | None = None +_worker_lock = threading.Lock() + + +def _get_global_worker(args, data_buffer) -> AsyncRolloutWorker: + global _global_worker + with _worker_lock: + if _global_worker is None or not _global_worker.worker_thread.is_alive(): + logger.info("starting fully-async rollout worker") + num_engines = max(1, args.rollout_num_gpus // args.rollout_num_gpus_per_engine) + _global_worker = AsyncRolloutWorker( + args, data_buffer, concurrency=args.vllm_server_concurrency * num_engines + ) + _global_worker.start() + return _global_worker + + +def _stop_global_worker() -> None: + global _global_worker + with _worker_lock: + if _global_worker is not None: + _global_worker.stop() + _global_worker = None + + +atexit.register(_stop_global_worker) + + +class AsyncRolloutWorker: + """Background thread + asyncio loop that continuously consumes groups + from ``data_buffer`` and runs :func:`generate_and_rm_group` on each.""" + + def __init__(self, args, data_buffer, concurrency: int = 10): + self.args = args + self.data_buffer = data_buffer + self.concurrency = concurrency + self.running = True + self.output_queue: queue.Queue[tuple[int, list[Sample]]] = queue.Queue(maxsize=1000) + self.worker_thread: threading.Thread | None = None + self.state = GenerateState(args) + + # -- public -------------------------------------------------------------- + + def start(self) -> None: + if self.worker_thread is None or not self.worker_thread.is_alive(): + self.worker_thread = threading.Thread(target=self._thread_main, name="fully-async-rollout", daemon=True) + self.worker_thread.start() + + def stop(self) -> None: + self.running = False + 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]]]: + completed: list[tuple[int, list[Sample]]] = [] + while True: + try: + completed.append(self.output_queue.get_nowait()) + except queue.Empty: + break + return completed + + def queue_size(self) -> int: + return self.output_queue.qsize() + + # -- internals ----------------------------------------------------------- + + def _thread_main(self) -> None: + asyncio.run(self._loop()) + + async def _loop(self) -> None: + active_tasks: set[asyncio.Task] = set() + max_concurrent = self.concurrency + gid_counter = 0 + + while self.running: + try: + # Reap done tasks + if active_tasks: + done = {t for t in active_tasks if t.done()} + for t in done: + try: + t.result() # results already handled in callback + except Exception as e: # noqa: BLE001 + logger.warning("fully-async task crashed: %r", e) + active_tasks -= done + + # Top up. + while len(active_tasks) < max_concurrent and self.running: + groups = self.data_buffer.get_samples(1) + if not groups: + break + for group in groups: + gid = gid_counter + gid_counter += 1 + task = asyncio.create_task( + generate_and_rm_group( + self.args, + group, + sampling_params=self.state.sampling_params.copy(), + evaluation=False, + ) + ) + task.add_done_callback(self._make_done_cb(gid)) + active_tasks.add(task) + + await asyncio.sleep(1) + except Exception as e: # noqa: BLE001 + logger.exception("fully-async loop iteration error: %s", e) + await asyncio.sleep(1) + + if active_tasks: + logger.info( + "fully-async: waiting for %d in-flight tasks to drain", + len(active_tasks), + ) + try: + await asyncio.wait(active_tasks, timeout=30) + except Exception: # noqa: BLE001 + pass + + def _make_done_cb(self, gid: int): + def _cb(done_task: asyncio.Task) -> None: + try: + result = done_task.result() + except Exception: # noqa: BLE001 + logger.exception("fully-async: process task raised") + return + if not isinstance(result, list): + logger.warning( + "fully-async: generate_and_rm_group returned %r, expected list[Sample]; dropping", + type(result).__name__, + ) + return + # Aborted group → requeue, don't ship to training. + if any(getattr(s, "status", None) == Sample.Status.ABORTED for s in result): + try: + self.data_buffer.add_samples([result]) + except Exception: # noqa: BLE001 + logger.exception("fully-async: failed to requeue aborted group") + return + self.output_queue.put((gid, result)) + + return _cb + + +async def _generate_rollout_async(args, rollout_id: int, data_buffer) -> list[list[Sample]]: + assert args.rollout_global_dataset + worker = _get_global_worker(args, data_buffer) + + target = args.rollout_batch_size + logger.info( + "fully-async rollout %d: target=%d queue_warm=%d", + rollout_id, + target, + worker.queue_size(), + ) + + collected: dict[int, list[Sample]] = {} + started = time.time() + last_log = started + LOG_EVERY = 30.0 + + while len(collected) < target: + # Pull whatever's done. + drained = 0 + for gid, group in worker.get_completed_groups(): + collected[gid] = group + drained += 1 + + if not drained: + await asyncio.sleep(0.05) + + now = time.time() + if now - last_log > LOG_EVERY: + logger.info( + "fully-async rollout %d: collected %d/%d, queue=%d, elapsed=%.1fs", + rollout_id, + len(collected), + target, + worker.queue_size(), + now - started, + ) + last_log = now + + # Order by sample.index for determinism (vime convention). + def _key(group: list[Sample]) -> int: + for s in group: + idx = getattr(s, "index", None) + if idx is not None: + return int(idx) + return 0 + + out = sorted(collected.values(), key=_key)[:target] + logger.info( + "fully-async rollout %d: done in %.1fs, queue_left=%d", + rollout_id, + time.time() - started, + worker.queue_size(), + ) + return out + + +def generate_rollout_fully_async(args, rollout_id, data_buffer, evaluation: bool = False): + """vime ``--rollout-function-path`` entrypoint.""" + + if evaluation: + raise ValueError("fully-async rollout doesn't support evaluation mode") + return run(_generate_rollout_async(args, rollout_id, data_buffer)) diff --git a/vime/rollout/vllm_rollout.py b/vime/rollout/vllm_rollout.py index 0b1b3ec56..ff99776d6 100644 --- a/vime/rollout/vllm_rollout.py +++ b/vime/rollout/vllm_rollout.py @@ -663,7 +663,14 @@ async def generate_and_rm( ) async def generate_and_rm_group( args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False -) -> list[Sample]: +) -> list[Sample] | list[list[Sample]]: + # ``generate_and_rm`` may return either a ``Sample`` or a ``list[Sample]`` + # depending on whether the ``--custom-generate-function-path`` callable + # emits one trainable sample or several (e.g. multi-turn agent rollouts + # that fan out into multiple prefix-chained samples). The asyncio.gather + # below preserves whichever shape each task produced, so the group is + # ``list[Sample]`` for plain rollouts and ``list[list[Sample]]`` for + # the fan-out case. state = GenerateState(args) if state.aborted: