Rewrite fully-async rollout as FullyAsyncRolloutFn on the class-based rollout API

[yueming-yuan]

Motivation

Stacked on #1716 (pure move; this PR's diff shows only the rewrite). The moved module was still example-quality code on the legacy stack: a module-global worker (_global_worker + threading.Lock + atexit), a private thread event loop, the GenerateState singleton from sglang_rollout, print everywhere, and broad try/except that silently drops data. The singleton's semaphore binds to the worker's private loop, which is why fully-async + eval was structurally impossible (cross-loop RuntimeError).

What this PR does

Rewrites the module as FullyAsyncRolloutFn on the class-based rollout API (requires MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1, which the swe-agent async scripts already run):

• Worker = long-lived task on the shared rollout event loop, lazy-started on the first train call. No thread, no private loop, no globals, no atexit; asyncio.Queue(maxsize=1000) keeps the same backpressure (full queue pauses submission).
• New-stack primitives: inference_rollout_common.generate_and_rm_group with an instance-owned GenerateState. Custom generate functions now receive the state type GenerateFnInput declares.
• Errors are loud: a failed generation task kills the worker and the next drain raises the original exception instead of hanging; the recycle paths no longer swallow. (_iter_samples also fixes a latent bug where multi-sample groups could never be abort-recycled — the old broad except hid an AttributeError.)
• Behavior preserved: in-flight bound = rollout_batch_size groups; ABORTED-group recycle on weight-update pauses; --max-weight-staleness filter (/model_info query now via http_utils with a 2s timeout); final sort by index. New: assert len(group) == n_samples_per_prompt and per-step metrics (rollout/fully_async/queue_size, recycle counts, staleness stats) via RolloutFnTrainOutput.metrics.
• Eval raises with guidance (set --eval-function-path to the standard InferenceRolloutFn); RolloutManager now reuses the rollout-fn instance when eval_function_path == rollout_function_path instead of constructing twice.
• Scripts/docs updated to the FullyAsyncRolloutFn path; the qwen scripts gain the env flag; _submit_one_group is shaped to receive Backfill async rollouts on sample completions #1673's sample-completion backfill with minimal conflict.

Testing

New CPU CI tests (tests/fast/rollout/test_fully_async_rollout.py, FakeGenerateState pattern): drain collects a sorted batch with metrics and the worker persists across calls; eval raises; ABORTED and stale groups recycle to the data source; worker exceptions propagate to the caller; in-flight submissions stay bounded. All 6 pass; pre-commit run --all-files passes.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request refactors the fully asynchronous rollout implementation by replacing the thread-based global worker with a class-based FullyAsyncRolloutFn that runs a persistent worker task on the shared rollout event loop. The feedback highlights several important issues to address: reusing the rollout function instance for both training and evaluation in RolloutManager introduces state-sharing and concurrency bugs; parsing the engine weight version outside the try-except block could crash the training process on malformed responses; a pending task leak may occur if _next_group is cancelled; and assert should be replaced with ValueError for configuration validation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Reusing the same rollout function instance for both training and evaluation introduces concurrency and state-sharing bugs. Specifically, InferenceRolloutFn holds a single GenerateState instance (self.state), which contains mutable state such as self.state.aborted and self.state.generate_fn_semaphore.

Since training rollouts and evaluation rollouts can run concurrently on the RolloutManager actor, sharing this state means:

1. If a training rollout is aborted (setting state.aborted = True), any concurrent evaluation rollout will also be silently aborted.
2. The concurrency semaphore is shared, which can cause unexpected blocking or rate-limiting between training and evaluation.

Since FullyAsyncRolloutFn explicitly raises an error on evaluation anyway, there is no need to share the instance to prevent it from being constructed twice for eval. It is much safer to always construct separate instances for training and evaluation.

            self.eval_generate_rollout = load_rollout_function(input, self.args.eval_function_path)

[gemini-code-assist]

The current implementation catches httpx.HTTPError and asyncio.TimeoutError but performs the integer parsing of data["weight_version"] outside the try-except block. If the router returns a malformed response (e.g., missing the "weight_version" key, or containing a non-integer value), this will raise a KeyError or ValueError and crash the entire training process.

Since the staleness filter is best-effort, we should defensively handle any parsing or key errors and fall back to the cached/previous value instead of raising a hard error.

Suggested change

	try:
	async with aiohttp.ClientSession() as session:
	async with session.get(url, timeout=aiohttp.ClientTimeout(total=2)) as resp:
	if resp.status == 200:
	data = await resp.json()
	self._value = int(data["weight_version"])
	self._last_query = now
	except Exception as e:
	data = await asyncio.wait_for(get(url), timeout=WEIGHT_VERSION_QUERY_TIMEOUT_SECS)
	except (httpx.HTTPError, asyncio.TimeoutError) as e:
	# Transient router unavailability; the staleness filter is best-effort.
	logger.debug(f"Failed to query engine weight version: {e}")
	return self._value
	self._value = int(data["weight_version"])
	self._last_query = now
	return self._value
	try:
	data = await asyncio.wait_for(get(url), timeout=WEIGHT_VERSION_QUERY_TIMEOUT_SECS)
	self._value = int(data["weight_version"])
	self._last_query = now
	except (httpx.HTTPError, asyncio.TimeoutError, KeyError, ValueError, TypeError) as e:
	# Transient router unavailability or malformed response; the staleness filter is best-effort.
	logger.debug(f"Failed to query engine weight version: {e}")
	return self._value

References

1. Avoid raising hard errors (such as ValueError) for non-critical metrics or data produced solely by internal components. Raising errors in these cases can turn a minor, non-critical issue (like a missing or malformed metric) into a hard application failure. Instead, use defensive fallback behavior (e.g., silent no-ops).

[gemini-code-assist]

If _next_group is cancelled or exits early, the background task queue_get (which is fetching from the queue) will remain pending, potentially causing resource leaks or 'Task was destroyed but it is pending!' warnings. Wrapping the wait loop in a try...finally block ensures that queue_get is always cancelled when the coroutine exits.

Suggested change

	queue_get = asyncio.create_task(self._output.get())
	while True:
	done, _ = await asyncio.wait(
	{queue_get, self._worker},
	return_when=asyncio.FIRST_COMPLETED,
	timeout=NO_PROGRESS_WARN_SECS,
	)
	if queue_get in done:
	return queue_get.result()
	if self._worker in done:
	queue_get.cancel()
	self._worker.result()
	raise RuntimeError("fully-async rollout worker exited without an exception")
	logger.warning(
	f"No completed rollout groups for {NO_PROGRESS_WARN_SECS}s " f"(queued: {self._output.qsize()})"
	)
	queue_get = asyncio.create_task(self._output.get())
	try:
	while True:
	done, _ = await asyncio.wait(
	{queue_get, self._worker},
	return_when=asyncio.FIRST_COMPLETED,
	timeout=NO_PROGRESS_WARN_SECS,
	)
	if queue_get in done:
	return queue_get.result()
	if self._worker in done:
	self._worker.result()
	raise RuntimeError("fully-async rollout worker exited without an exception")
	logger.warning(
	f"No completed rollout groups for {NO_PROGRESS_WARN_SECS}s " f"(queued: {self._output.qsize()})"
	)
	finally:
	if not queue_get.done():
	queue_get.cancel()

References

1. To prevent resource leaks (e.g., counters that are not decremented), use constructs like try...finally or a with statement to ensure cleanup logic is always executed, even in the case of exceptions or early returns.

[gemini-code-assist]

Use ValueError instead of assert for validating configuration or function arguments, as assertions can be globally disabled in Python when run with optimization flags (-O), leading to silent bypasses of critical configuration guards.

Suggested change

	assert args.rollout_global_dataset
	if not args.rollout_global_dataset:
	raise ValueError("FullyAsyncRolloutFn requires rollout_global_dataset to be True")

References

1. Use ValueError instead of assert for validating function or constructor arguments (such as checking for positive or non-negative values).