feat: promote fully_async rollout into core (slime #1920)

[aoshen02]

What

Sync of THUDM/slime#1920 ("Move fully_async example to main codebase") into vime (RFC #107). 🔧 PORT.

Promotes the fully-async rollout from an example into the core package, usable via
--rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async.

Changes

• vime/rollout/fully_async_rollout.py (new) — promoted from examples/fully_async/fully_async_rollout.py, 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).
• examples/fully_async/ — removed the now-duplicated fully_async_rollout.py (moved to core); repointed run-qwen3-4b-fully_async.sh + README at the core --rollout-function-path.
• tests/test_qwen2.5_0.5B_fully_async_short.py (new) + registered in run-ci-short (pr-test.yml + .j2 regenerated). Mirrors test_qwen2.5_0.5B_async_short, only flipping the rollout-function-path.

vime notes / divergence from upstream

• vime's example was already vLLM-ported, so this is "promote the (already-vLLM) example + apply slime's core-integration deltas", not a fresh sglang→vLLM port.
• vime keeps the 4B example (just repointed) rather than swapping it for 0.5B as upstream did — the new CI test already covers 0.5B, and keeping a working example is better than deleting it.

Validation

• pre-commit (pinned: ruff/black/isort) PASS; commit signed off (DCO).
• cpu gate in a vime cu129 image: import vime.rollout.fully_async_rollout OK; test_megatron_argument_validation + 4 plugin_contracts PASS; new test file valid (structurally identical to the passing async_short base, only the rollout-function-path differs).
• GPU e2e pending: the 4-GPU test_qwen2.5_0.5B_fully_async_short.py (exercises the async-worker path end-to-end) will run once an H200 node frees from the in-flight [Docker] upgrade torch_memory_saver to a193d9dd + reduce host memory (slime #1916, #1924, #1932) #114 CI — or via the run-ci-short label.

Refs: THUDM/slime#1920, #107

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request promotes the fully-async rollout worker from an example into the core package (vime/rollout/fully_async_rollout.py), updates the example launch scripts, and adds a CI smoke test. It also updates generate_and_rm_group to support fan-out multi-turn agent rollouts. Feedback on the changes highlights several critical issues: an event loop mismatch with the GenerateState semaphore in the background thread, potential loss of excess completed groups when draining the queue, and broken aborted checks and sorting logic in the fan-out case due to nested lists.

[gemini-code-assist]

The GenerateState class is a singleton that initializes self.semaphore during its init. Since AsyncRolloutWorker is instantiated in the main thread, self.semaphore is bound to the main thread's event loop. However, the background thread runs its own event loop via asyncio.run(self._loop()). When generate_and_rm is called inside the background thread, attempting to acquire state.semaphore will raise a RuntimeError due to the event loop mismatch. Re-creating the semaphore inside _loop binds it to the background thread's event loop and resolves this issue.

Suggested change

	async def _loop(self) -> None:
	active_tasks: set[asyncio.Task] = set()
	async def _loop(self) -> None:
	self.state.semaphore = asyncio.Semaphore(self.concurrency)
	active_tasks: set[asyncio.Task] = set()

[aoshen02]

Valid. GenerateState is a SingletonMeta, so its asyncio.Semaphore is shared with the main-thread rollout path and can end up bound to a different event loop than the bg thread's asyncio.run(self._loop()). Rebinding the semaphore inside _loop() is the correct fix. Will apply.

[gemini-code-assist]

When _generate_rollout_async collects completed groups, worker.get_completed_groups() drains all completed groups from the queue. If the number of drained groups exceeds the target rollout batch size, the excess groups are discarded after sorting, resulting in lost training samples and wasted GPU resources. Adding an optional limit parameter to get_completed_groups ensures we only drain up to the remaining target count, leaving any excess completed groups in the persistent queue for the next rollout step.

Suggested change

	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 get_completed_groups(self, limit: int | None = None) -> list[tuple[int, list[Sample]]]:
	completed: list[tuple[int, list[Sample]]] = []
	while limit is None or len(completed) < limit:
	try:
	completed.append(self.output_queue.get_nowait())
	except queue.Empty:
	break
	return completed

[aoshen02]

Valid. get_completed_groups() draining all completed groups can discard the excess beyond target. Adding the optional limit param to leave the surplus in the queue for the next step. Will apply.

[gemini-code-assist]

To prevent discarding excess completed groups when the drained count exceeds the target, limit the number of groups pulled from the worker to the remaining target count.

Suggested change

	while len(collected) < target:
	# Pull whatever's done.
	drained = 0
	for gid, group in worker.get_completed_groups():
	while len(collected) < target:
	# Pull whatever's done.
	drained = 0
	for gid, group in worker.get_completed_groups(limit=target - len(collected)):

[aoshen02]

Valid — pairs with the get_completed_groups(limit=...) change: cap the pull at target - len(collected). Will apply.

[gemini-code-assist]

In the fan-out case (e.g., multi-turn agent rollouts), generate_and_rm_group returns a nested list[list[Sample]] instead of a flat list[Sample]. The current aborted check iterates over result and checks getattr(s, 'status', None). Since s is a list in the fan-out case, this check always evaluates to False, allowing aborted samples to be shipped to training. Additionally, wrapping result in a list ([result]) when calling add_samples results in incorrect nesting for the fan-out case. Flattening the samples for the aborted check and conditionally wrapping result based on its structure resolves both issues.

Suggested change

	# 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
	# Aborted group → requeue, don't ship to training.
	flat_samples = []
	for item in result:
	if isinstance(item, list):
	flat_samples.extend(item)
	else:
	flat_samples.append(item)
	if any(getattr(s, "status", None) == Sample.Status.ABORTED for s in flat_samples):
	try:
	if result and isinstance(result[0], list):
	self.data_buffer.add_samples(result)
	else:
	self.data_buffer.add_samples([result])
	except Exception: # noqa: BLE001
	logger.exception("fully-async: failed to requeue aborted group")
	return

[aoshen02]

Valid. In the fan-out case generate_and_rm_group returns list[list[Sample]], so the aborted check must flatten before inspecting .status, and the requeue must not double-nest. Will apply the flatten + conditional-wrap fix.

[gemini-code-assist]

The sorting key _key assumes group is a flat list[Sample]. In the fan-out case where group is list[list[Sample]], s is a list, so getattr(s, 'index', None) returns None, causing all groups to default to an index of 0 and breaking deterministic sorting. Updating _key to recursively search nested lists ensures correct sorting for both plain and fan-out rollouts.

    def _key(group: list[Sample] | list[list[Sample]]) -> int:
        for item in group:
            if isinstance(item, list):
                for s in item:
                    idx = getattr(s, "index", None)
                    if idx is not None:
                        return int(idx)
            else:
                idx = getattr(item, "index", None)
                if idx is not None:
                    return int(idx)
        return 0

[aoshen02]

Valid. _key assumes a flat list[Sample]; for the fan-out shape it must recurse into nested lists or every group sorts to index 0. Will apply.

[CalvinXKY]

LGTM