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

[yueming-yuan]

Motivation

#1716 moved fully_async_rollout.py from examples/ into miles/rollout/ unchanged. What landed 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).

This PR rewrites it on the class-based rollout API and gives it a real flag.

What this PR does

FullyAsyncRolloutFn, selected by train_async.py --fully-async (requires MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1):

• 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 (a full queue pauses submission).
• New-stack primitives: inference_rollout_common.generate_and_rm_group with an instance-owned GenerateState, so custom generate functions 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; 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.

Rollout features the drain used to silently ignore. The drain reimplements the collection loop, so everything the standard path does around it was lost:

• --dynamic-sampling-filter-path and --rollout-sample-filter-path are now applied. A dynamic-filtered group is dropped rather than recycled, matching the standard path — continuous submission already replaces it, so no over-sampling machinery is needed. This mattered in practice: examples/swe-agent/run-glm47-flash-agentic-async.py, one of the scripts this PR migrates, passes --dynamic-sampling-filter-path.
• --partial-rollout, --recompute-logprobs-via-prefill and --rollout-all-samples-process-path now fail fast instead of being ignored. Partial rollout is actively defeated by the recycle path (reset_for_retry() drops the partial response), and a group can cross a weight update before it is drained, so a drain-time logprob recompute would score it under newer weights than generated it — the correct fix there is a per-group recompute at completion time, which is out of scope here.

Argument resolution. --rollout-function-path loses its string default in favor of None, so resolve_rollout_function_paths() is the single place that answers "which rollout and eval function did these arguments select". --fully-async asserts against the flags it conflicts with (colocate, multi-LoRA, an explicit --rollout-function-path, the missing refactor env var), and eval is resolved before the fully-async override, so "fully async does not serve eval" holds by construction — eval keeps the standard InferenceRolloutFn unless --eval-function-path is set. train.py rejects --fully-async outright (it needs the async driver).

Weight-version cache. _CachedWeightVersion only stamped _last_query on success, so an unreachable router made every group in the drain pay the full 2s timeout instead of one per TTL — silently, at logger.debug. Stamped in a finally on completion now, which also fixes a router slower than the TTL defeating the cache entirely.

Scripts and docs move to the --fully-async path; the qwen scripts gain the env flag. _submit_one_group is shaped to receive #1673's sample-completion backfill with minimal conflict.

Testing

New e2e case — tests/e2e/megatron/test_qwen3_30B_A3B/test_fully_async.py, stage-c-8-gpu-h100. Fully-async on the disaggregated topology (it cannot colocate); three rollouts exercise the states that only exist with a persistent worker: a cold start, a drain from an already-warm queue, and a drain across a weight update (which pauses generation and makes the worker recycle aborted groups).

New CPU tests — tests/fast/rollout/test_fully_async_rollout.py, stage-a-cpu, 12 cases on a FakeGenerateState / FakeDataSource harness:

drain_collects_batch_sorted_with_metrics	sorted batch + metrics, worker persists across calls
eval_raises	eval is rejected without starting a worker
aborted_group_recycled / stale_group_recycled	both recycle paths return the prompt group to the data source
worker_error_propagates / worker_failure_beats_queued_groups	a dead worker fails the step, even with a backlog queued
worker_bounds_in_flight_groups / async_max_concurrent_samples_caps_in_flight_groups	both in-flight bounds
nested_group_recycles_the_flat_prompt_group	multi-sample generate functions resubmit the flat prompt group
dynamic_filter_drops_group_without_recycling	dropped, not recycled; drop reason reaches metrics
sample_filter_marks_samples_without_shrinking_the_batch	remove_sample set, batch size unchanged
weight_version_throttles_failed_queries	one HTTP call per TTL when the router is down, not one per group

resolve_rollout_function_paths was checked to be behavior-identical to the two functions it replaces across all 10 valid argument combinations. pre-commit run --all-files passes.

[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]

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]

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).

[guapisolo]

Pushed two small fixes from review:

• a793f52 — docs: the async-rollout snippet in docs/user-guide/training-script-walkthrough.md still said "two changes" and was missing MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1. Since --fully-async hard-asserts that env var, the launch command failed at startup as written. Aligned the snippet with docs/user-guide/fully-async.md.
• 4470cd9 — refactor: extracted a _first_sample(group) helper next to _iter_samples for the group[0][0] if isinstance(group[0], list) else group[0] unwrap that appeared three times in _drain (first-sample log, finish log, sort key).

tests/fast/rollout/test_fully_async_rollout.py (12 cases) passes locally; pre-commit hooks pass on both commits.

[guapisolo]

LGTM. clean design.

[guapisolo]

CC said eval has some conflict with current fully async code path. But I think we can leave this PR as-is.