From be30dc8b52f37879c5ec85a7a1791ccdd778448f Mon Sep 17 00:00:00 2001 From: Linyu Wu Date: Tue, 4 Aug 2026 15:24:52 +0800 Subject: [PATCH] refactor(rollout): move the driver-side async engines to rollout/manager/ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit They hold no model and cannot generate, so calling them engines made 'engine' mean three things at once: the ABC, a concrete backend, and these. They are the layer that owns admission, acceptance and disposal over time — a manager. engine/asynchronous.py -> manager/{buffers,batch,agentic}.py AsyncBatchRolloutEngine -> BatchManager AsyncAgenticRolloutEngine -> AgenticManager 'manager' rather than 'scheduler' because this tree already calls three unrelated things a scheduler: the LR scheduler (utils/scheduler_utils.py), the diffusion noise scheduler, and SGLang's own scheduler subprocesses. Adds manager/protocol.py recording the consumer surface both implementations share, and splits the mechanisms (VersionedBuffer, PendingGroups, root_of) out of the file that used to hold both engines and their machinery. launch_ceiling stays with BatchManager, whose two trainers are its only callers, so its ownership note remains accurate. Pure move + rename; no behaviour change. Test plan: an AST comparison against the pre-move file proves VersionedBuffer, PendingGroups, InflightPool, _InflightJob, root_of and BatchManager byte-identical modulo the rename. compileall, check-recipe-targets (2434 paths) and check-experimental-boundaries all pass. --- unirl/rollout/README.md | 10 +- unirl/rollout/engine/__init__.py | 15 +- unirl/rollout/engine/agentic/engine.py | 2 +- unirl/rollout/engine/asynchronous.py | 444 ------------------------- unirl/rollout/manager/__init__.py | 23 ++ unirl/rollout/manager/agentic.py | 133 ++++++++ unirl/rollout/manager/batch.py | 205 ++++++++++++ unirl/rollout/manager/buffers.py | 116 +++++++ unirl/rollout/manager/protocol.py | 39 +++ unirl/trainer/README.md | 4 +- unirl/trainer/agentic_async.py | 6 +- unirl/trainer/agentic_partial.py | 6 +- unirl/trainer/async_ar.py | 6 +- unirl/trainer/async_diffusion.py | 6 +- 14 files changed, 543 insertions(+), 472 deletions(-) delete mode 100644 unirl/rollout/engine/asynchronous.py create mode 100644 unirl/rollout/manager/__init__.py create mode 100644 unirl/rollout/manager/agentic.py create mode 100644 unirl/rollout/manager/batch.py create mode 100644 unirl/rollout/manager/buffers.py create mode 100644 unirl/rollout/manager/protocol.py diff --git a/unirl/rollout/README.md b/unirl/rollout/README.md index d06a0478..a6f4c09d 100644 --- a/unirl/rollout/README.md +++ b/unirl/rollout/README.md @@ -63,13 +63,13 @@ wrong objective. ratio is 1 on the first update; *separate* — a dedicated engine on its own GPUs plus a `sync:` block; *colocate* — a dedicated engine sharing GPUs with train, plus offload/onload and `sync:`. -- **Driver-side async engines** (`engine/asynchronous.py`, the driver-side half next +- **Driver-side async engines** (`manager/`, the driver-side half next to `engine/synchronous.py`'s worker-side sync contracts). Both engines expose the same consumer verbs the async trainers program against: `poll` / `drain_freshest` / - `pop_evicted` / `quiesce` + engine-owned `weight_version`. `AsyncBatchRolloutEngine` + `pop_evicted` / `quiesce` + engine-owned `weight_version`. `BatchManager` (batch granularity; non-blocking `Handle.launch_nowait` generations, stamps versions at launch, used by `AsyncARTrainer`/`AsyncDiffusionTrainer`) and - `AsyncAgenticRolloutEngine` (trajectory granularity over the agentic rank-0 + `AgenticManager` (trajectory granularity over the agentic rank-0 coordinator; normalizes the `[0]` unwraps, assembles n-sibling GRPO groups, stamps versions at completion, used by the partial/async agentic trainers). @@ -93,10 +93,10 @@ implements its weight-receive method and a matching `sync:` handler in - **Direct sampling forbids a `sync:` block; dedicated requires one.** The trainside engine also can't live on a `layout: separate` slab — `_build_rollout` raises. - **Quiesce before weight sync / eval / checkpoint on the batch async path** — - `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a + `BatchManager.quiesce()` drains every in-flight generation; a weight + KV update corrupts one mid-flight. The agentic quiesce is a turn-boundary `abort` + final poll, folded into - `AsyncAgenticRolloutEngine.quiesce()`; its `sync_weights()` rejects a live + `AgenticManager.quiesce()`; its `sync_weights()` rejects a live drive, then pairs the weight push with the version bump and logs the sync. Reap-vs-launch ordering is trainer statement order (diffusion polls before topping up; see its `_next_step`). diff --git a/unirl/rollout/engine/__init__.py b/unirl/rollout/engine/__init__.py index 28bb1ab2..1a5b3002 100644 --- a/unirl/rollout/engine/__init__.py +++ b/unirl/rollout/engine/__init__.py @@ -1,12 +1,11 @@ """Rollout engines over the canonical ``Sample`` request type. -Two halves of one design: ``synchronous.py`` records the worker-side sync contracts -(``BaseRolloutEngine`` — the broad ABC including coordinator engines — and -``SyncRolloutEngine``, the ``Sample`` → ``Sample`` refinement the per-backend -subpackages implement); ``asynchronous.py`` records the driver side (the -batch/agentic async engines and their mechanisms). +``synchronous.py`` records the worker-side contracts (``BaseRolloutEngine`` — the +broad ABC including coordinator engines — and ``SyncRolloutEngine``, the ``Sample`` +→ ``Sample`` refinement the per-backend subpackages implement). The driver side +lives in ``../manager/``, which owns admission, acceptance and disposal over time +and holds no model. -Deliberately empty otherwise: importing the driver-side ``asynchronous`` module -must stay ray/torch-free, so this init imports nothing and consumers import the -halves directly. +Deliberately empty otherwise: this init imports nothing, so consumers import the +module they need directly and the manager package stays ray/torch-free. """ diff --git a/unirl/rollout/engine/agentic/engine.py b/unirl/rollout/engine/agentic/engine.py index 60fc5490..b09583bc 100644 --- a/unirl/rollout/engine/agentic/engine.py +++ b/unirl/rollout/engine/agentic/engine.py @@ -12,7 +12,7 @@ (how many to over-sample, staleness, when to sync). The coordinator exposes a **submit / poll / finalize / abort** interface over a **background** drain the trainer reaps and interrupts — consumed driver-side through -:class:`~unirl.rollout.engine.asynchronous.AsyncAgenticRolloutEngine` (rank-0 unwrap + group assembly + +:class:`~unirl.rollout.manager.AgenticManager` (rank-0 unwrap + group assembly + versioned buffering): - ``submit(request)`` — enqueue a pool (fresh prompts and/or carried partials) and fire the drain diff --git a/unirl/rollout/engine/asynchronous.py b/unirl/rollout/engine/asynchronous.py deleted file mode 100644 index 20018a8c..00000000 --- a/unirl/rollout/engine/asynchronous.py +++ /dev/null @@ -1,444 +0,0 @@ -"""Driver-side async rollout engines and their mechanisms (LIN-631). - -The async half of the engine design: ``synchronous.py`` records the worker-side sync -contracts (``BaseRolloutEngine`` / ``SyncRolloutEngine``); this module records -the driver side — everything is single-threaded, lock-free, and ray-free -(non-blocking dispatch is ``Handle.launch_nowait``). - -Mechanisms (policy-free — launch ceilings, reap/launch ordering, and step -loops live in the trainers): - -- :class:`VersionedBuffer` — payload-agnostic freshness/staleness buffer. -- :class:`InflightPool` — non-blocking pool of distributed ``generate`` calls. - -Engines share one consumer surface (``poll`` / ``drain_freshest`` / -``pop_evicted`` / ``quiesce`` + engine-owned ``weight_version``): - -- :class:`AsyncBatchRolloutEngine` — batch granularity over a single-turn - engine slab; one ``submit`` is one non-blocking distributed ``generate``. - ``(weight_version, gen_id)`` are stamped at LAUNCH. -- :class:`AsyncAgenticRolloutEngine` — trajectory granularity over the - ``AgenticRolloutEngine`` rank-0 coordinator; ``submit`` fires a task-pool - drive and completions stream in via ``poll``. ``(weight_version, gen_id)`` - are stamped at COMPLETION; the per-turn version spread inside a carried - trajectory is corrected per-token by each gen Part's own ``weight_version``. - -Submission is deliberately engine-specific (incompatible signatures and -stamping semantics); the consumer verbs above are what the async trainers -program against. The colocate barrier path (``AgenticTrainer``) keeps -calling ``rollout.generate(sample)[0]`` directly. -""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - Generic, - Iterable, - List, - Optional, - Tuple, - TypeVar, -) - -if TYPE_CHECKING: - from unirl.types.sample import Sample - -logger = logging.getLogger(__name__) - -T = TypeVar("T") - - -class VersionedBuffer(Generic[T]): - """Payload-agnostic freshness buffer of ``(payload, weight_version, gen_id)`` items. - - Unifies the batch path's per-``Sample`` buffering and the agentic path's - per-group (``List[Sample]``) buffering; stamping semantics belong to the - caller (batch stamps at launch, agentic at completion). - """ - - def __init__(self) -> None: - self._items: List[Tuple[T, int, int]] = [] - self._evicted: List[T] = [] - - def put(self, payload: T, *, weight_version: int, gen_id: int) -> None: - self._items.append((payload, int(weight_version), int(gen_id))) - - def size(self) -> int: - return len(self._items) - - def drain_freshest( - self, - n: int, - *, - current_version: Optional[int] = None, - max_staleness: Optional[int] = None, - ) -> Optional[List[T]]: - """Pop the ``n`` freshest eligible payloads, carrying leftovers forward. - - Over-stale items are evicted first (retrievable via :meth:`pop_evicted`), - then remaining items are sorted by descending ``gen_id`` (stable — ties - keep insertion order). Returns ``None`` without consuming anything when - fewer than ``n`` remain after eviction. - """ - if max_staleness is not None and current_version is not None: - kept: List[Tuple[T, int, int]] = [] - for item in self._items: - if current_version - item[1] <= max_staleness: - kept.append(item) - else: - self._evicted.append(item[0]) - self._items = kept - if len(self._items) < n: - return None - self._items.sort(key=lambda item: item[2], reverse=True) - picked, self._items = self._items[:n], self._items[n:] - return [payload for payload, _, _ in picked] - - def pop_evicted(self) -> List[T]: - """Return and clear payloads rejected by the latest staleness checks.""" - evicted, self._evicted = self._evicted, [] - return evicted - - -Complete = Callable[[int, int, Any], None] - - -@dataclass(frozen=True) -class _InflightJob: - gen_id: int - weight_version: int - pending: Any - - -class InflightPool: - """Non-blocking pool of distributed ``generate`` launches on a rollout Handle. - - Mechanism only: launch ceilings, reap/launch ordering, and step loops are - caller policy. Jobs are launched via ``Handle.launch_nowait`` and completed - through ``complete(gen_id, weight_version, payload)`` — all of ``complete``'s - fallible work must happen before it mutates caller state, because a job - whose completion raises stays in flight for retry. - """ - - def __init__(self, rollout: Any, *, start_gen_id: int = 0) -> None: - self._rollout = rollout - self._next_gen_id = int(start_gen_id) - self._jobs: List[_InflightJob] = [] - - @property - def next_gen_id(self) -> int: - return self._next_gen_id - - def __len__(self) -> int: - return len(self._jobs) - - def launch(self, sample: Any, *, weight_version: int) -> int: - gen_id = self._next_gen_id - pending = self._rollout.launch_nowait("generate", sample) - self._jobs.append(_InflightJob(gen_id, int(weight_version), pending)) - self._next_gen_id += 1 - return gen_id - - def reap_ready(self, complete: Complete) -> int: - """Complete every ready job; leave unresolved and failed jobs in flight. - - A job whose ``result()``/``complete`` raises stays in flight for retry; - the first error re-raises after the sweep, later errors are logged. - KeyboardInterrupt/SystemExit propagate immediately (not deferred behind - remaining ready jobs). Returns the number completed. - """ - still: List[_InflightJob] = [] - first_error: Optional[Exception] = None - completed = 0 - for job in self._jobs: - if not job.pending.ready(): - still.append(job) - continue - try: - complete(job.gen_id, job.weight_version, job.pending.result()) - completed += 1 - except Exception as exc: - still.append(job) - if first_error is None: - first_error = exc - else: - logger.error("reap_ready: additional failure for gen_id=%s", job.gen_id, exc_info=exc) - self._jobs = still - if first_error is not None: - raise first_error - return completed - - def drain_all(self, complete: Complete) -> int: - """Quiesce: complete every job, blocking as needed. Same error contract as - :meth:`reap_ready`.""" - jobs, self._jobs = self._jobs, [] - first_error: Optional[Exception] = None - completed = 0 - for job in jobs: - try: - complete(job.gen_id, job.weight_version, job.pending.result()) - completed += 1 - except Exception as exc: - self._jobs.append(job) - if first_error is None: - first_error = exc - else: - logger.error("drain_all: additional failure for gen_id=%s", job.gen_id, exc_info=exc) - if first_error is not None: - raise first_error - return completed - - def wait_oldest(self) -> None: - """Block until the oldest in-flight generation resolves, without collecting.""" - if self._jobs: - self._jobs[0].pending.wait() - - -class AsyncBatchRolloutEngine: - """Batch-granular async engine over a ``SyncRolloutEngine`` slab Handle; buffers ``Sample`` groups. - - ``complete(gen_id, completed) -> groups`` runs at reap time — scoring must - precede training, and on transfer-sensitive backends the next launch. All - of ``complete``'s fallible work happens before any buffer mutation, so a - failed job stays in flight for retry without double-inserting groups. - - ``quiesce()`` (drain everything) is MANDATORY before a weight sync, eval, or - checkpoint: a weight + KV update corrupts an in-flight generation. - """ - - def __init__( - self, - rollout: Any, - *, - complete: Callable[[int, "Sample"], List["Sample"]], - start_gen_id: int = 0, - ) -> None: - self._complete = complete - self._pool = InflightPool(rollout, start_gen_id=start_gen_id) - self._buffer: VersionedBuffer["Sample"] = VersionedBuffer() - self._weight_version = 0 - - @property - def weight_version(self) -> int: - return self._weight_version - - def bump_weight_version(self) -> int: - self._weight_version += 1 - return self._weight_version - - @property - def next_gen_id(self) -> int: - """The gen_id the next ``submit`` will get (1:1 with rollout ids).""" - return self._pool.next_gen_id - - @property - def inflight(self) -> int: - return len(self._pool) - - def submit(self, sample: "Sample") -> int: - """Launch one non-blocking distributed ``generate``; stamps the CURRENT version.""" - return self._pool.launch(sample, weight_version=self._weight_version) - - def poll(self) -> int: - return self._pool.reap_ready(self._on_complete) - - def drain_freshest(self, n: int, *, max_staleness: int) -> Optional[List["Sample"]]: - return self._buffer.drain_freshest(n, current_version=self._weight_version, max_staleness=max_staleness) - - def pop_evicted(self) -> List["Sample"]: - return self._buffer.pop_evicted() - - def quiesce(self) -> List["Sample"]: - self._pool.drain_all(self._on_complete) - return [] - - def wait_oldest(self) -> None: - """Block until the oldest in-flight generation resolves (reap via ``poll``).""" - self._pool.wait_oldest() - - def _on_complete(self, gen_id: int, weight_version: int, completed: "Sample") -> None: - groups = self._complete(gen_id, completed) # fallible (scoring) before any buffer put - for group in groups: - self._buffer.put(group, weight_version=weight_version, gen_id=gen_id) - - -def root_of(traj: "Sample") -> str: - """Root id shared by a prompt's ``n`` sibling trajectories.""" - return traj.parts[0].sample_ids[0] - - -class PendingGroups: - """Bucket a flat stream of terminal trajectories into complete GRPO groups. - - Poll returns variable-depth trajectory ``Sample``s; a prompt's ``n`` - siblings share its slash-free root id. A group is complete once all ``n`` - of a root's siblings are terminal. Variable-depth trajectories are NOT - concatenated — a group stays a ``List[Sample]`` (the trainer flattens their - gen Parts at train time). - """ - - def __init__(self, n: int) -> None: - self._n = int(n) - self._by_root: Dict[str, List["Sample"]] = {} - - def add_completed(self, trajs: List["Sample"]) -> None: - for t in trajs: - self._by_root.setdefault(root_of(t), []).append(t) - - def pop_complete_groups(self) -> List[List["Sample"]]: - ready = [root for root, sibs in self._by_root.items() if len(sibs) >= self._n] - out: List[List["Sample"]] = [] - for root in ready: - sibs = self._by_root.pop(root) - out.append(sibs[: self._n]) - return out - - def discard_roots(self, roots: Iterable[str]) -> int: - """Drop incomplete buckets for abandoned roots; returns the number of - terminal sibling trajectories that were being held for them.""" - discarded = 0 - for root in set(roots): - discarded += len(self._by_root.pop(root, [])) - return discarded - - def size(self) -> int: - return len(self._by_root) - - -class AsyncAgenticRolloutEngine: - """Trajectory-granular async engine over the ``AgenticRolloutEngine`` - rank-0 coordinator Handle; buffers ``List[Sample]`` sibling groups. - - Normalizes the coordinator's BROADCAST+RANK_ZERO returns (every value - unwraps ``[0]``). Groups are stamped at COMPLETION: ``weight_version`` is - the engine's counter when a root's last sibling lands, ``gen_id`` a - monotonic completed-group counter. - - ``submit`` requires the prior drive to be finalized or quiesced — two live - drains would double-pull the coordinator queue. - """ - - def __init__(self, rollout: Any, *, group_size: int, start_gen_id: int = 0) -> None: - self._rollout = rollout - self._pending = PendingGroups(group_size) - self._buffer: VersionedBuffer[List["Sample"]] = VersionedBuffer() - self._gen_id = int(start_gen_id) - self._weight_version = 0 - self._drive_live = False - - @property - def weight_version(self) -> int: - return self._weight_version - - def sync_weights(self, weight_sync: Any) -> int: - """Push train weights via *weight_sync* and advance the version ledger. - - The only sanctioned weight-push path — pairing the push with the bump - is what keeps the ledger truthful. Raises while a drive is active (a - weight push must be decode-idle); a joined ``finalize_if_drained`` or - ``quiesce`` ends the drive. - """ - if self._drive_live: - raise RuntimeError("sync_weights with a drive active; finalize or quiesce() first") - weight_sync.sync() - self._weight_version += 1 - logger.info("sync_weights: pushed train weights; weight_version -> %d", self._weight_version) - return self._weight_version - - def submit(self, tasks: List["Sample"]) -> None: - """Fire a background drive over a flat task list (fresh siblings + carried partials). - - Enforced double-pull guard: two live drains would double-pull the - coordinator queue, so a second ``submit`` before ``finalize_if_drained`` - reported the drive done (or before ``quiesce``) raises instead of - silently corrupting the drive. - """ - if self._drive_live: - raise RuntimeError( - "AsyncAgenticRolloutEngine.submit: prior drive still live — wait for " - "finalize_if_drained() to report it done or quiesce() first (a second " - "drain would double-pull the coordinator queue)." - ) - # Set before RPC so ambiguous submit failures remain guarded. - self._drive_live = True - self._rollout.submit(tasks) - - def poll(self) -> int: - return self._ingest(self._rollout.poll()[0]) - - def finalize_if_drained(self) -> Optional[int]: - """``None`` while the in-flight drive is still running; otherwise join it - and ingest its final completions — atomic with the readiness check (a - separate poll would race the next ``submit``'s worker-buffer reset).""" - completed = self._rollout.finalize_if_drained()[0] - if completed is None: - return None - self._drive_live = False - return self._ingest(completed) - - def drain_freshest(self, n: int, *, max_staleness: int) -> Optional[List[List["Sample"]]]: - return self._buffer.drain_freshest(n, current_version=self._weight_version, max_staleness=max_staleness) - - def pop_evicted(self) -> List[List["Sample"]]: - return self._buffer.pop_evicted() - - def quiesce(self) -> List["Sample"]: - """Turn-boundary stop: abort, then one final poll for trajectories that - completed DURING the quiesce (before the next ``submit`` resets worker - buffers). Call before ``sync_weights`` so those groups carry the - version they completed under.""" - carried = self._rollout.abort()[0] - self.poll() - self._drive_live = False - return carried - - def discard_roots(self, roots: Iterable[str]) -> int: - """Drop abandoned roots' incomplete pending buckets (tail-drop policy).""" - return self._pending.discard_roots(roots) - - def pending_groups(self) -> int: - """Roots with some-but-not-all siblings terminal (the pending backlog).""" - return self._pending.size() - - def buffered_groups(self) -> int: - return self._buffer.size() - - def _ingest(self, completed: List["Sample"]) -> int: - if completed: - self._pending.add_completed(completed) - for group in self._pending.pop_complete_groups(): - self._buffer.put(group, weight_version=self._weight_version, gen_id=self._gen_id) - self._gen_id += 1 - return len(completed) - - -def launch_ceiling(rollout_id: int, *, sync_interval: int, max_staleness: int, num_rollouts: int) -> int: - """The batch trainers' on-policy launch clamp — trainer POLICY, defined once. - - A generation launched now is consumed later, so how far ahead the gen_id - allocator may run is bounded to ``max_staleness`` weight-sync windows: - ``max_staleness=0`` ⇒ never launch into a future sync-window ⇒ no - generation crosses a sync ⇒ ``ratio≈1`` (on-policy). - - OWNERSHIP: this is trainer-side POLICY, not engine surface — its vocabulary - (``rollout_id`` / ``sync_interval`` / ``num_rollouts``) is the trainers', - the engine classes never call it, and it must never become an engine - method. It is hosted in this module only because it is the two batch - trainers' one shared torch-free home; the step loops that use it stay in - the trainers as visible statement order. - """ - return min(num_rollouts, ((rollout_id // sync_interval) + 1 + max_staleness) * sync_interval) - - -__all__ = [ - "AsyncAgenticRolloutEngine", - "AsyncBatchRolloutEngine", - "launch_ceiling", - "root_of", -] diff --git a/unirl/rollout/manager/__init__.py b/unirl/rollout/manager/__init__.py new file mode 100644 index 00000000..35e46bd1 --- /dev/null +++ b/unirl/rollout/manager/__init__.py @@ -0,0 +1,23 @@ +"""Driver-side rollout managers (LIN-693). See ``unirl/rollout/README.md``. + +The layer between the trainer and the rollout engines: it owns admission, +acceptance and disposal over time, and holds no model. Named *manager* rather than +*scheduler* because this tree already calls three unrelated things a scheduler — +the LR scheduler, the diffusion noise scheduler, and SGLang's own subprocesses. +""" + +from unirl.rollout.manager.agentic import AgenticManager, Carried +from unirl.rollout.manager.batch import BatchManager, InflightPool +from unirl.rollout.manager.buffers import PendingGroups, VersionedBuffer, root_of +from unirl.rollout.manager.protocol import RolloutManager + +__all__ = [ + "AgenticManager", + "BatchManager", + "Carried", + "InflightPool", + "PendingGroups", + "RolloutManager", + "VersionedBuffer", + "root_of", +] diff --git a/unirl/rollout/manager/agentic.py b/unirl/rollout/manager/agentic.py new file mode 100644 index 00000000..11c81ad1 --- /dev/null +++ b/unirl/rollout/manager/agentic.py @@ -0,0 +1,133 @@ +"""AgenticManager — trajectory-granular driver-side rollout manager. + +Sits over the ``AgenticRolloutEngine`` rank-0 coordinator and normalizes its +BROADCAST+RANK_ZERO returns (every value unwraps ``[0]``). Groups are stamped at +COMPLETION: ``weight_version`` is the counter when a root's last sibling lands, +``gen_id`` a monotonic completed-group counter. The per-turn version spread inside a +carried trajectory is corrected per-token by each gen Part's own ``weight_version``. + +Moved verbatim from ``engine/asynchronous.py``; renamed because it holds no model and +cannot generate, so calling it an engine was misleading. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional + +from unirl.rollout.manager.buffers import PendingGroups, VersionedBuffer, root_of + +if TYPE_CHECKING: + from unirl.types.sample import Sample + +logger = logging.getLogger(__name__) + + +class AgenticManager: + """Trajectory-granular async engine over the ``AgenticRolloutEngine`` + rank-0 coordinator Handle; buffers ``List[Sample]`` sibling groups. + + Normalizes the coordinator's BROADCAST+RANK_ZERO returns (every value + unwraps ``[0]``). Groups are stamped at COMPLETION: ``weight_version`` is + the engine's counter when a root's last sibling lands, ``gen_id`` a + monotonic completed-group counter. + + ``submit`` requires the prior drive to be finalized or quiesced — two live + drains would double-pull the coordinator queue. + """ + + def __init__(self, rollout: Any, *, group_size: int, start_gen_id: int = 0) -> None: + self._rollout = rollout + self._pending = PendingGroups(group_size) + self._buffer: VersionedBuffer[List["Sample"]] = VersionedBuffer() + self._gen_id = int(start_gen_id) + self._weight_version = 0 + self._drive_live = False + + @property + def weight_version(self) -> int: + return self._weight_version + + def sync_weights(self, weight_sync: Any) -> int: + """Push train weights via *weight_sync* and advance the version ledger. + + The only sanctioned weight-push path — pairing the push with the bump + is what keeps the ledger truthful. Raises while a drive is active (a + weight push must be decode-idle); a joined ``finalize_if_drained`` or + ``quiesce`` ends the drive. + """ + if self._drive_live: + raise RuntimeError("sync_weights with a drive active; finalize or quiesce() first") + weight_sync.sync() + self._weight_version += 1 + logger.info("sync_weights: pushed train weights; weight_version -> %d", self._weight_version) + return self._weight_version + + def submit(self, tasks: List["Sample"]) -> None: + """Fire a background drive over a flat task list (fresh siblings + carried partials). + + Enforced double-pull guard: two live drains would double-pull the + coordinator queue, so a second ``submit`` before ``finalize_if_drained`` + reported the drive done (or before ``quiesce``) raises instead of + silently corrupting the drive. + """ + if self._drive_live: + raise RuntimeError( + "AgenticManager.submit: prior drive still live — wait for " + "finalize_if_drained() to report it done or quiesce() first (a second " + "drain would double-pull the coordinator queue)." + ) + # Set before RPC so ambiguous submit failures remain guarded. + self._drive_live = True + self._rollout.submit(tasks) + + def poll(self) -> int: + return self._ingest(self._rollout.poll()[0]) + + def finalize_if_drained(self) -> Optional[int]: + """``None`` while the in-flight drive is still running; otherwise join it + and ingest its final completions — atomic with the readiness check (a + separate poll would race the next ``submit``'s worker-buffer reset).""" + completed = self._rollout.finalize_if_drained()[0] + if completed is None: + return None + self._drive_live = False + return self._ingest(completed) + + def drain_freshest(self, n: int, *, max_staleness: int) -> Optional[List[List["Sample"]]]: + return self._buffer.drain_freshest(n, current_version=self._weight_version, max_staleness=max_staleness) + + def pop_evicted(self) -> List[List["Sample"]]: + return self._buffer.pop_evicted() + + def quiesce(self) -> List["Sample"]: + """Turn-boundary stop: abort, then one final poll for trajectories that + completed DURING the quiesce (before the next ``submit`` resets worker + buffers). Call before ``sync_weights`` so those groups carry the + version they completed under.""" + carried = self._rollout.abort()[0] + self.poll() + self._drive_live = False + return carried + + def discard_roots(self, roots: Iterable[str]) -> int: + """Drop abandoned roots' incomplete pending buckets (tail-drop policy).""" + return self._pending.discard_roots(roots) + + def pending_groups(self) -> int: + """Roots with some-but-not-all siblings terminal (the pending backlog).""" + return self._pending.size() + + def buffered_groups(self) -> int: + return self._buffer.size() + + def _ingest(self, completed: List["Sample"]) -> int: + if completed: + self._pending.add_completed(completed) + for group in self._pending.pop_complete_groups(): + self._buffer.put(group, weight_version=self._weight_version, gen_id=self._gen_id) + self._gen_id += 1 + return len(completed) + + +__all__ = ["AgenticManager"] diff --git a/unirl/rollout/manager/batch.py b/unirl/rollout/manager/batch.py new file mode 100644 index 00000000..c6379eb9 --- /dev/null +++ b/unirl/rollout/manager/batch.py @@ -0,0 +1,205 @@ +"""BatchManager — batch-granular async rollout over a single-turn engine slab. + +One ``submit`` is one non-blocking distributed ``generate``: a batch IS one logical +unit, so its all-or-nothing completion is the right semantics and it stays on the +slab-wide ``Handle.launch_nowait`` path. ``(weight_version, gen_id)`` are stamped at +LAUNCH. Moved verbatim from the former ``engine/asynchronous.py``. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Iterable, List, Optional, Tuple, TypeVar + +if TYPE_CHECKING: + from unirl.types.sample import Sample + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + +from unirl.rollout.manager.buffers import VersionedBuffer + + +Complete = Callable[[int, int, Any], None] + + +@dataclass(frozen=True) +class _InflightJob: + gen_id: int + weight_version: int + pending: Any + + +class InflightPool: + """Non-blocking pool of distributed ``generate`` launches on a rollout Handle. + + Mechanism only: launch ceilings, reap/launch ordering, and step loops are + caller policy. Jobs are launched via ``Handle.launch_nowait`` and completed + through ``complete(gen_id, weight_version, payload)`` — all of ``complete``'s + fallible work must happen before it mutates caller state, because a job + whose completion raises stays in flight for retry. + """ + + def __init__(self, rollout: Any, *, start_gen_id: int = 0) -> None: + self._rollout = rollout + self._next_gen_id = int(start_gen_id) + self._jobs: List[_InflightJob] = [] + + @property + def next_gen_id(self) -> int: + return self._next_gen_id + + def __len__(self) -> int: + return len(self._jobs) + + def launch(self, sample: Any, *, weight_version: int) -> int: + gen_id = self._next_gen_id + pending = self._rollout.launch_nowait("generate", sample) + self._jobs.append(_InflightJob(gen_id, int(weight_version), pending)) + self._next_gen_id += 1 + return gen_id + + def reap_ready(self, complete: Complete) -> int: + """Complete every ready job; leave unresolved and failed jobs in flight. + + A job whose ``result()``/``complete`` raises stays in flight for retry; + the first error re-raises after the sweep, later errors are logged. + KeyboardInterrupt/SystemExit propagate immediately (not deferred behind + remaining ready jobs). Returns the number completed. + """ + still: List[_InflightJob] = [] + first_error: Optional[Exception] = None + completed = 0 + for job in self._jobs: + if not job.pending.ready(): + still.append(job) + continue + try: + complete(job.gen_id, job.weight_version, job.pending.result()) + completed += 1 + except Exception as exc: + still.append(job) + if first_error is None: + first_error = exc + else: + logger.error("reap_ready: additional failure for gen_id=%s", job.gen_id, exc_info=exc) + self._jobs = still + if first_error is not None: + raise first_error + return completed + + def drain_all(self, complete: Complete) -> int: + """Quiesce: complete every job, blocking as needed. Same error contract as + :meth:`reap_ready`.""" + jobs, self._jobs = self._jobs, [] + first_error: Optional[Exception] = None + completed = 0 + for job in jobs: + try: + complete(job.gen_id, job.weight_version, job.pending.result()) + completed += 1 + except Exception as exc: + self._jobs.append(job) + if first_error is None: + first_error = exc + else: + logger.error("drain_all: additional failure for gen_id=%s", job.gen_id, exc_info=exc) + if first_error is not None: + raise first_error + return completed + + def wait_oldest(self) -> None: + """Block until the oldest in-flight generation resolves, without collecting.""" + if self._jobs: + self._jobs[0].pending.wait() + + +class BatchManager: + """Batch-granular async engine over a ``SyncRolloutEngine`` slab Handle; buffers ``Sample`` groups. + + ``complete(gen_id, completed) -> groups`` runs at reap time — scoring must + precede training, and on transfer-sensitive backends the next launch. All + of ``complete``'s fallible work happens before any buffer mutation, so a + failed job stays in flight for retry without double-inserting groups. + + ``quiesce()`` (drain everything) is MANDATORY before a weight sync, eval, or + checkpoint: a weight + KV update corrupts an in-flight generation. + """ + + def __init__( + self, + rollout: Any, + *, + complete: Callable[[int, "Sample"], List["Sample"]], + start_gen_id: int = 0, + ) -> None: + self._complete = complete + self._pool = InflightPool(rollout, start_gen_id=start_gen_id) + self._buffer: VersionedBuffer["Sample"] = VersionedBuffer() + self._weight_version = 0 + + @property + def weight_version(self) -> int: + return self._weight_version + + def bump_weight_version(self) -> int: + self._weight_version += 1 + return self._weight_version + + @property + def next_gen_id(self) -> int: + """The gen_id the next ``submit`` will get (1:1 with rollout ids).""" + return self._pool.next_gen_id + + @property + def inflight(self) -> int: + return len(self._pool) + + def submit(self, sample: "Sample") -> int: + """Launch one non-blocking distributed ``generate``; stamps the CURRENT version.""" + return self._pool.launch(sample, weight_version=self._weight_version) + + def poll(self) -> int: + return self._pool.reap_ready(self._on_complete) + + def drain_freshest(self, n: int, *, max_staleness: int) -> Optional[List["Sample"]]: + return self._buffer.drain_freshest(n, current_version=self._weight_version, max_staleness=max_staleness) + + def pop_evicted(self) -> List["Sample"]: + return self._buffer.pop_evicted() + + def quiesce(self) -> List["Sample"]: + self._pool.drain_all(self._on_complete) + return [] + + def wait_oldest(self) -> None: + """Block until the oldest in-flight generation resolves (reap via ``poll``).""" + self._pool.wait_oldest() + + def _on_complete(self, gen_id: int, weight_version: int, completed: "Sample") -> None: + groups = self._complete(gen_id, completed) # fallible (scoring) before any buffer put + for group in groups: + self._buffer.put(group, weight_version=weight_version, gen_id=gen_id) + + +def launch_ceiling(rollout_id: int, *, sync_interval: int, max_staleness: int, num_rollouts: int) -> int: + """The batch trainers' on-policy launch clamp — trainer POLICY, defined once. + + A generation launched now is consumed later, so how far ahead the gen_id + allocator may run is bounded to ``max_staleness`` weight-sync windows: + ``max_staleness=0`` ⇒ never launch into a future sync-window ⇒ no + generation crosses a sync ⇒ ``ratio≈1`` (on-policy). + + OWNERSHIP: this is trainer-side POLICY, not engine surface — its vocabulary + (``rollout_id`` / ``sync_interval`` / ``num_rollouts``) is the trainers', + the engine classes never call it, and it must never become an engine + method. It is hosted in this module only because it is the two batch + trainers' one shared torch-free home; the step loops that use it stay in + the trainers as visible statement order. + """ + return min(num_rollouts, ((rollout_id // sync_interval) + 1 + max_staleness) * sync_interval) + + +__all__ = ["BatchManager", "InflightPool", "launch_ceiling"] diff --git a/unirl/rollout/manager/buffers.py b/unirl/rollout/manager/buffers.py new file mode 100644 index 00000000..afaf17a9 --- /dev/null +++ b/unirl/rollout/manager/buffers.py @@ -0,0 +1,116 @@ +"""Payload-agnostic buffering mechanisms shared by the rollout managers. + +Policy-free: freshness/staleness bookkeeping and GRPO group assembly, with no +opinion about admission, placement, or when a batch is due. Moved verbatim from +the former ``engine/asynchronous.py``. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Iterable, List, Optional, Tuple, TypeVar + +if TYPE_CHECKING: + from unirl.types.sample import Sample + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + + +class VersionedBuffer(Generic[T]): + """Payload-agnostic freshness buffer of ``(payload, weight_version, gen_id)`` items. + + Unifies the batch path's per-``Sample`` buffering and the agentic path's + per-group (``List[Sample]``) buffering; stamping semantics belong to the + caller (batch stamps at launch, agentic at completion). + """ + + def __init__(self) -> None: + self._items: List[Tuple[T, int, int]] = [] + self._evicted: List[T] = [] + + def put(self, payload: T, *, weight_version: int, gen_id: int) -> None: + self._items.append((payload, int(weight_version), int(gen_id))) + + def size(self) -> int: + return len(self._items) + + def drain_freshest( + self, + n: int, + *, + current_version: Optional[int] = None, + max_staleness: Optional[int] = None, + ) -> Optional[List[T]]: + """Pop the ``n`` freshest eligible payloads, carrying leftovers forward. + + Over-stale items are evicted first (retrievable via :meth:`pop_evicted`), + then remaining items are sorted by descending ``gen_id`` (stable — ties + keep insertion order). Returns ``None`` without consuming anything when + fewer than ``n`` remain after eviction. + """ + if max_staleness is not None and current_version is not None: + kept: List[Tuple[T, int, int]] = [] + for item in self._items: + if current_version - item[1] <= max_staleness: + kept.append(item) + else: + self._evicted.append(item[0]) + self._items = kept + if len(self._items) < n: + return None + self._items.sort(key=lambda item: item[2], reverse=True) + picked, self._items = self._items[:n], self._items[n:] + return [payload for payload, _, _ in picked] + + def pop_evicted(self) -> List[T]: + """Return and clear payloads rejected by the latest staleness checks.""" + evicted, self._evicted = self._evicted, [] + return evicted + +def root_of(traj: "Sample") -> str: + """Root id shared by a prompt's ``n`` sibling trajectories.""" + return traj.parts[0].sample_ids[0] + + +class PendingGroups: + """Bucket a flat stream of terminal trajectories into complete GRPO groups. + + Poll returns variable-depth trajectory ``Sample``s; a prompt's ``n`` + siblings share its slash-free root id. A group is complete once all ``n`` + of a root's siblings are terminal. Variable-depth trajectories are NOT + concatenated — a group stays a ``List[Sample]`` (the trainer flattens their + gen Parts at train time). + """ + + def __init__(self, n: int) -> None: + self._n = int(n) + self._by_root: Dict[str, List["Sample"]] = {} + + def add_completed(self, trajs: List["Sample"]) -> None: + for t in trajs: + self._by_root.setdefault(root_of(t), []).append(t) + + def pop_complete_groups(self) -> List[List["Sample"]]: + ready = [root for root, sibs in self._by_root.items() if len(sibs) >= self._n] + out: List[List["Sample"]] = [] + for root in ready: + sibs = self._by_root.pop(root) + out.append(sibs[: self._n]) + return out + + def discard_roots(self, roots: Iterable[str]) -> int: + """Drop incomplete buckets for abandoned roots; returns the number of + terminal sibling trajectories that were being held for them.""" + discarded = 0 + for root in set(roots): + discarded += len(self._by_root.pop(root, [])) + return discarded + + def size(self) -> int: + return len(self._by_root) + + +__all__ = ["PendingGroups", "VersionedBuffer", "root_of"] diff --git a/unirl/rollout/manager/protocol.py b/unirl/rollout/manager/protocol.py new file mode 100644 index 00000000..fdc38c8b --- /dev/null +++ b/unirl/rollout/manager/protocol.py @@ -0,0 +1,39 @@ +"""The manager surface the async trainers program against. + +Mechanism only. Every knob a trainer sets — batch size, over-sample width, tail +policy, sync cadence — stays in the trainer; a manager answers *what is ready* and +*what is in flight*, never *how much to train on*. + +Two implementations, differing only where they genuinely must: +:class:`~unirl.rollout.manager.batch.BatchManager` (one batch per call, slab-wide +dispatch, versions stamped at launch) and +:class:`~unirl.rollout.manager.agentic.AgenticManager` (one trajectory per call, +point-to-point dispatch, versions stamped at completion). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, List, Optional, Protocol + +if TYPE_CHECKING: + from unirl.types.sample import Sample + + +class RolloutManager(Protocol): + """Consumer verbs shared by the batch and agentic managers.""" + + @property + def weight_version(self) -> int: ... + + def sync_weights(self, weight_sync: Any) -> int: ... + + def poll(self) -> int: ... + + def drain_freshest(self, n: int, *, max_staleness: Optional[int]) -> Optional[List[Any]]: ... + + def pop_evicted(self) -> List[Any]: ... + + def quiesce(self) -> List["Sample"]: ... + + +__all__ = ["RolloutManager"] diff --git a/unirl/trainer/README.md b/unirl/trainer/README.md index 842ff191..48690854 100644 --- a/unirl/trainer/README.md +++ b/unirl/trainer/README.md @@ -68,8 +68,8 @@ The current trainer surface is: | `AsyncAgenticTrainer` / `AsyncAgenticEnvTrainer` | buffered complete trajectory groups → concatenated turn `Part` | Disaggregated train/rollout slabs, resident agentic drive, weight-version staleness control, and the same explicit `carry`/`drop` tail policy. | The async variants program against the driver-side async engines in -`unirl/rollout/engine/asynchronous.py`: `AsyncBatchRolloutEngine` (AR/diffusion — non-blocking -batched generations, launch-time version stamps) and `AsyncAgenticRolloutEngine` +`unirl/rollout/manager/`: `BatchManager` (AR/diffusion — non-blocking +batched generations, launch-time version stamps) and `AgenticManager` (partial/async agentic — trajectory drives, group assembly, completion-time stamps). The trainers keep the policy: launch ceilings, reap-vs-launch order, quiesce points, and tail carry/drop. diff --git a/unirl/trainer/agentic_async.py b/unirl/trainer/agentic_async.py index eb940da8..e4269ee0 100644 --- a/unirl/trainer/agentic_async.py +++ b/unirl/trainer/agentic_async.py @@ -10,7 +10,7 @@ Mechanism vs policy (LIN-531): the **engine** exposes a ``submit`` / ``poll`` / ``finalize_if_drained`` / ``abort`` interface over a background drain (see :class:`~unirl.rollout.engine.agentic.engine.AgenticRolloutEngine`), consumed -through the driver-side :class:`~unirl.rollout.engine.asynchronous.AsyncAgenticRolloutEngine` +through the driver-side :class:`~unirl.rollout.manager.AgenticManager` (group assembly + versioned buffering); this **trainer** owns the *policy* — * **Producer** — keep the rollout slab saturated: ``submit`` a pool of fresh prompt @@ -51,7 +51,7 @@ from omegaconf import DictConfig from unirl.distributed.group.placement import placement, remote -from unirl.rollout.engine.asynchronous import AsyncAgenticRolloutEngine, root_of +from unirl.rollout.manager import AgenticManager, root_of from unirl.train.stack import TrainStepResult from unirl.trainer.agentic import AgenticTrainer from unirl.trainer.base import BaseTrainer, build_sampling_dict @@ -353,7 +353,7 @@ def train( }, ) - self._engine = AsyncAgenticRolloutEngine(self.rollout, group_size=self._n, start_gen_id=start_rollout) + self._engine = AgenticManager(self.rollout, group_size=self._n, start_gen_id=start_rollout) if start_rollout < num_rollouts and start_rollout and self.weight_sync is not None: self._engine.sync_weights(self.weight_sync) # push restored weights into the fresh engine diff --git a/unirl/trainer/agentic_partial.py b/unirl/trainer/agentic_partial.py index e2b5d003..b4a01915 100644 --- a/unirl/trainer/agentic_partial.py +++ b/unirl/trainer/agentic_partial.py @@ -16,7 +16,7 @@ Motivation (LIN-531 ALFWorld comparison): the fully-async trainer *lost* on ALFWorld because disaggregation halves the generation GPUs; colocate+partial keeps all GPUs for generation while still cutting the straggler tail — the slime "over-sample + abort + recycle" / verl `bypass_mode` -pattern. The driver-side :class:`~unirl.rollout.engine.asynchronous.AsyncAgenticRolloutEngine` +pattern. The driver-side :class:`~unirl.rollout.manager.AgenticManager` (submit/poll/finalize_if_drained/quiesce + group assembly and versioned buffering) is shared with ``AsyncAgenticTrainer``; only the colocate wake/sync/sleep choreography is new. @@ -35,7 +35,7 @@ from collections import Counter from typing import Dict, List, Literal, Optional -from unirl.rollout.engine.asynchronous import AsyncAgenticRolloutEngine, root_of +from unirl.rollout.manager import AgenticManager, root_of from unirl.trainer.agentic import AgenticTrainer from unirl.trainer.agentic_env import _EnvRewardSource from unirl.types.sample import Part, Sample @@ -204,7 +204,7 @@ def train( }, ) - self._engine = AsyncAgenticRolloutEngine(self.rollout, group_size=self._n, start_gen_id=start_rollout) + self._engine = AgenticManager(self.rollout, group_size=self._n, start_gen_id=start_rollout) self._carried = [] try: diff --git a/unirl/trainer/async_ar.py b/unirl/trainer/async_ar.py index 85b72d33..20cdf00c 100644 --- a/unirl/trainer/async_ar.py +++ b/unirl/trainer/async_ar.py @@ -19,7 +19,7 @@ regime). ``>0`` = **off-policy continuous buffer**: generations may run ahead across syncs, bounded by eviction; the rollout-anchored DRPO ratio absorbs it. -Generation runs through :class:`~unirl.rollout.engine.asynchronous.AsyncBatchRolloutEngine` +Generation runs through :class:`~unirl.rollout.manager.BatchManager` (non-blocking Ray futures over the rollout Handle) on the single driver thread — no producer thread, no locks; the trainer's ``_next_step`` loop owns the policy (launch ceiling, launch-then-reap order). Draining all in-flight generations @@ -45,7 +45,7 @@ from unirl.distributed.group.placement import placement, remote from unirl.distributed.tensor import hydrate from unirl.models.qwen3_5.validation import validate_qwen3_5_training_contract -from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, launch_ceiling +from unirl.rollout.manager import BatchManager, launch_ceiling from unirl.train.stack import TrainStepResult from unirl.trainer.ar import ARTrainer from unirl.trainer.base import BaseTrainer, build_sampling_dict @@ -308,7 +308,7 @@ def train( }, ) - self._async_engine = AsyncBatchRolloutEngine( + self._async_engine = BatchManager( self.rollout, complete=self._score_completed, start_gen_id=start_rollout, diff --git a/unirl/trainer/async_diffusion.py b/unirl/trainer/async_diffusion.py index 6b41db9f..100a1136 100644 --- a/unirl/trainer/async_diffusion.py +++ b/unirl/trainer/async_diffusion.py @@ -9,7 +9,7 @@ checkpoint / FlowGRPO ``stack.train_track``). The async loop runs over the shared -:class:`~unirl.rollout.engine.asynchronous.AsyncBatchRolloutEngine` (the same engine +:class:`~unirl.rollout.manager.BatchManager` (the same engine ``AsyncARTrainer`` drives) — one single-threaded driver loop over non-blocking Ray dispatch, no producer thread and no locks. This trainer supplies only the diffusion hooks: @@ -50,7 +50,7 @@ import torch from unirl.distributed.tensor import hydrate -from unirl.rollout.engine.asynchronous import AsyncBatchRolloutEngine, launch_ceiling +from unirl.rollout.manager import BatchManager, launch_ceiling from unirl.train.stack import TrainStepResult from unirl.trainer.diffusion import DiffusionTrainer from unirl.types.sample import Sample @@ -167,7 +167,7 @@ def train( }, ) - self._async_engine = AsyncBatchRolloutEngine( + self._async_engine = BatchManager( self.rollout, complete=self._score_completed, start_gen_id=start_rollout,