refactor(rollout): engine-level sync/async contracts for async rollout by celve · Pull Request #287 · Tencent-Hunyuan/UniRL

celve · 2026-07-31T07:08:23Z

Summary

Unifies the two async rollout paths behind engine-level contracts under unirl/rollout/engine/, replacing the over-layered async_runtime.py (scheduler + one-implementation dispatcher Protocol + callback types) and the agentic trainers' duplicated producer plumbing.

engine/synchronous.py (was base.py): worker-side sync contracts. BaseSingleTurnRolloutEngine is renamed SyncRolloutEngine; chunked_engine_generate moves in from the package __init__.
engine/asynchronous.py (new): the driver side — VersionedBuffer + InflightPool mechanisms, the AsyncRolloutEngine protocol (poll / drain_freshest / pop_evicted / quiesce, engine-owned weight_version), and two concretes: AsyncBatchRolloutEngine (AR/diffusion; launch-time version stamps) and AsyncAgenticRolloutEngine (facade over the agentic rank-0 coordinator; [0] unwraps, PendingGroups sibling assembly, completion-time stamps, quiesce folds abort + the mandatory follow-up poll).
Handle gains a public non-blocking seam (launch_nowait → PendingHandleCall.ready()/wait()/result()), removing RayGenerationDispatcher's documented "update in lockstep" coupling to the private _execute_all/_rebind_tree; handle_fn is untouched.
The four async trainers keep policy as explicit loops (launch ceiling, reap-vs-launch as statement order — diffusion polls before topping up, preserving the measured ~~150s→~~8s BAGEL constraint; quiesce before sync; tail carry/drop). The reap_before_launch flag, dead drained()/pending_roots()/_pending_carried, and trainer-side _weight_version counters are removed.
engine/__init__.py becomes a lazy re-export shim so the driver-side module imports ray/torch-free.

Net −38 lines; the runtime layer covers both async paths in one place instead of one path plus per-trainer copies.

Also included: a one-line fix to AsyncARTrainer.__init__ (_rollout_anchor_device = None). It is a pre-existing main bug, not part of the refactor — AsyncARTrainer calls BaseTrainer.__init__ directly to skip the parent's colocate placement, but inherits ARTrainer.evaluate(), which reads that attribute. Every eval on the async AR path died with AttributeError before the first training step. It is bundled here because the refactor's own GPU validation could not run eval without it.

Related Issue

N/A (tracked in Linear as LIN-631).

Test Plan

Static checks (branch @ 22b27122, based on main @ b1499015):

ruff format --check + ruff check (v0.15.1, the pinned pre-commit rev) over all 28 changed Python files — 28 files already formatted, All checks passed!.
python -m compileall over the same set — clean.

CPU-only harnesses pinning the new buffer/pool invariants (complete-or-nothing reap retry without double-insert, first-error deferral, KeyboardInterrupt passthrough, freshest-first staleness eviction, launch-time vs completion-time version stamping, quiesce-poll-before-bump ordering, n-sibling group assembly):

python -m pytest tests/ -q → 24 passed in 0.04s.
Per the tests-tree policy (test: remove tests directory #99/test: remove tests directory #267, restated in CLAUDE.md §5), these harnesses are run but not committed.

GPU end-to-end, 1x8 H20 (4 train + 4 rollout), models and datasets on pod-local disk:

Async AR — examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml, Qwen3-4B-Base on DAPO-Math, num_devices=8 batch_size=16 sampling.samples_per_prompt=8:
- 22 rollouts, 0 errors. Mean reward 0.1605 (rollouts 1-11) → 0.2884 (rollouts 12-22), a 1.80x lift with near-disjoint ranges (first-half max 0.250 vs second-half min 0.203).
- ratio held at 0.991–1.000 and |Δlogp| at 1.1–3.4e-02 on every step. With weight_sync_interval=1 (on-policy) a correct implementation must pin the PPO ratio at ~1.0, so this exercises AsyncBatchRolloutEngine's launch-time stamping and the quiesce-before-bump ordering across 22 consecutive sync cycles.
- Separately with eval_interval=10 +eval_num_prompts=8: the rollout_id=-1 baseline eval runs anchored=False and completes cleanly (verifies the bundled _rollout_anchor_device fix).
BAGEL async diffusion — examples/diffusion/bagel/bagel_vllmomni_async.yaml, BAGEL-7B-MoT x PickScore x LoRA, vLLM-Omni rollout on a disjoint slab:
- Reaches training steps with all 4 rollout engines healthy: 0 worker deaths, 0 DiffusionExecutor is closed, 0 tracebacks.
- Replay is exact: ratio=1.0000±0.0000. Because BAGEL's opaque KV contexts cannot cross the IPC boundary, the adapter ships prompts and the stage rebuilds contexts trainer-side; an exact ratio is direct evidence those rebuilt contexts match the ones the rollout worker used, and that the cross-slab fp32 trajectory transfer and worker-vs-trainside SDE math agree.
- Reward rises to convergence. Over 1000 rollouts (~38 h) mean reward per
  100-rollout block climbs 0.7730 -> 0.8886 and plateaus (first-50 vs last-50:
  1.207x). The curve is monotonic until it flattens at ~0.889 — a complete
  learning curve, not a truncated segment.
- The same recipe on main @ b1499015 (separate pod, identical config)
  independently converges to the same place: 0.7796 -> 0.8874, lift 1.193x,
  tracking the branch within ~0.002 at every block. The refactor reproduces
  main's learning dynamics, not merely its crash-free behaviour. Early
  rollouts additionally match rollout-for-rollout (mean |delta| 0.0011 over the
  first 10), pinning identical group ordering and version stamping.
- 0 worker deaths, 0 executor closures, 0 tracebacks on either run across ~38 h.
ALFWorld partial/async agentic — Not run; reason: no ALFWorld recipe/env staged on the validation pod. AsyncAgenticRolloutEngine remains covered only by the CPU harnesses above.

Compatibility / Risk

No recipe/config changes required; _target_ paths and knobs are unchanged (reap_before_launch was internal, never a recipe key). The repo's check-recipe-targets hook guards this.
Import paths changed with no aliases: unirl.rollout.async_runtime and the names it exported (AsyncRolloutScheduler, RayGenerationDispatcher, VersionedGroupBuffer, GenerationDispatcher, InflightGeneration) are gone; unirl.rollout.engine.base → unirl.rollout.engine.synchronous. All in-repo consumers are migrated; out-of-tree code importing these must update.
Behavior-preserving by design. The riskiest spots are the per-trainer _next_step loops (reviewed line-by-line against the deleted next_step for both orderings) and the agentic quiesce() fold (harness-pinned: groups completing during the quiesce carry the pre-sync version).
_pending_carried removal is a verified dead-state deletion (only ever assigned []).

Reviewer Notes

Suggested review order: handle.py seam (additive) → engine/asynchronous.py mechanisms → protocol + engines → async_ar/async_diffusion loops (check reap ordering) → agentic_partial/agentic_async (check quiesce-before-bump and tail-policy counter divergence).
AI-assisted (Claude Code); the full diff was reviewed and every listed command run by the submitter.
Duplicate-work check: rebased onto main @ b1499015 mid-review, which already carried #276 (PDEATHSIG for thread-launched workers) and #277 (rebuild replay KV contexts in eval()). Both had been independently re-fixed on this branch during GPU debugging; those two commits were dropped in favour of upstream's. No other overlapping open PRs.
Both of those upstream fixes were required to get the BAGEL async recipe running at all — worth knowing if you are validating that recipe on an older base.
Follow-ups deliberately out of scope: trainer-level dedup of the still-duplicated _next_step/_build_tasks loop pairs, making handle_fn delegate to the seam halves, and the _connect_separate/two-slab-build duplication.

Checklist

I reviewed the changed code and removed unrelated/generated artifacts.
I updated tests, docs, and configs where needed, or explained why not.

Update 2026-08-02 (maintainer)

Rebased in place onto main @ 00cf5df0 (crossing #226 TP/EP full-weight sync, #259 AR PPO/GAE, #174 qwen3-5) and extended with a surface-trim commit; the branch keeps the original four commits plus one maintainer commit.

Rebase conflict resolution (unirl/trainer/async_ar.py, the only conflicted file):

Import block: kept feat(qwen3-5): support text and single-image inputs in AR training #174's validate_qwen3_5_training_contract alongside the new AsyncBatchRolloutEngine import; the deleted async_runtime import is gone.
Ctor: kept both feat(weight-sync): support SGLang TP/EP full-weight sync #226's hoisted rollout_parsed dedicated-engine check and this PR's _rollout_anchor_device = None fix (the stale ar.py:NNN line refs in that comment are rewritten as symbol refs).
Semantics preserved, not "fixed": evaluate()'s weight sync without a driver-side version bump predates both branches (b1499015 already synced inline; feat(weight-sync): support SGLang TP/EP full-weight sync #226 only extracted the _prepare_rollout hook), so no bump was added. Follow-up note: with weight_sync_interval>1, an eval landing mid-interval skews staleness accounting by one sync — pre-existing on both sides.

Surface trim (refactor(rollout): drop the async engine layer's zero-consumer surface) — five deletions, each verified zero-consumer in-tree:

engine/__init__.py lazy re-export shim → docstring-only init (no package-level importer existed; emptiness preserves the ray/torch-free import property with zero machinery).
chunked_engine_generate deleted (zero callers since the initial release, both before and after the move into synchronous.py).
AsyncRolloutEngine Protocol deleted (zero annotations, zero isinstance; neither trainer family can be typed against it — both use engine-specific verbs). The contract description folded into the module docstring.
__all__ trimmed to the consumed surface (AsyncBatchRolloutEngine, AsyncAgenticRolloutEngine, root_of).
InflightPool(method=...) parameter removed (single construction site, never overridden).

Net vs main is now −157 lines (34 files, +751/−908); the Summary's original −38 predates the trim.

Post-rebase verification: ruff v0.15.1 format --check + check clean over the 28 changed Python files; compileall clean; import unirl.rollout.engine.asynchronous on a torch-less interpreter stays stdlib-only; repo-wide grep finds zero references to any deleted name; handle_fn and launch_nowait dispatch gates remain in parity (#226 did not touch the dispatch half). GPU revalidation of the async-AR ratio-pinning recipe on this base is queued before merge (the original GPU runs predate #226).

Copilot

Pull request overview

Refactors async rollout to use engine-level sync/async contracts under unirl/rollout/engine/, replacing the previous unirl/rollout/async_runtime.py scheduler/dispatcher layer and moving async buffering/dispatch mechanisms into a shared driver-side engine module.

Changes:

Introduces driver-side async engines/mechanisms (AsyncBatchRolloutEngine, AsyncAgenticRolloutEngine, VersionedBuffer, InflightPool) in unirl/rollout/engine/asynchronous.py and updates async trainers to use them.
Adds a non-blocking Handle.launch_nowait() + PendingHandleCall seam to decouple async dispatch from Handle internals previously accessed by RayGenerationDispatcher.
Renames/reshapes worker-side sync engine contracts (BaseSingleTurnRolloutEngine → SyncRolloutEngine), moves chunked_engine_generate into synchronous.py, and updates engine implementations/docs/imports accordingly.

Reviewed changes

Copilot reviewed 34 out of 34 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
unirl/trainer/README.md	Documents the new driver-side `AsyncRolloutEngine` protocol and concrete engines used by async trainers.
unirl/trainer/async_diffusion.py	Migrates diffusion async loop from `AsyncRolloutScheduler` to `AsyncBatchRolloutEngine` with explicit poll/launch policy in `_next_step`.
unirl/trainer/async_ar.py	Migrates AR async loop to `AsyncBatchRolloutEngine` and fixes missing `_rollout_anchor_device` init for inherited `evaluate()`.
unirl/trainer/agentic_partial.py	Switches colocate partial agentic trainer to `AsyncAgenticRolloutEngine` for group assembly/buffering and folds quiesce logic into engine.
unirl/trainer/agentic_async.py	Removes trainer-local group assembler/buffer, delegating to `AsyncAgenticRolloutEngine` for async agentic rollout bookkeeping.
unirl/rollout/README.md	Updates rollout architecture docs for `synchronous.py` and new driver-side async engines.
unirl/rollout/loop/README.md	Updates agent loop contract docs to reference `SyncRolloutEngine`.
unirl/rollout/loop/engine_port.py	Updates docstring to reference `SyncRolloutEngine` as the nominal runtime counterpart.
unirl/rollout/engine/vllm_omni/weight_sync.py	Updates references from `base.py` to `synchronous.py` in docs/comments.
unirl/rollout/engine/vllm_omni/engine.py	Updates base class to `SyncRolloutEngine` and sync-surface documentation references.
unirl/rollout/engine/vllm_omni/config.py	Switches `BaseEngineConfig` import to `unirl.rollout.engine.synchronous`.
unirl/rollout/engine/vllm_omni/init.py	Updates lazy-import commentary to match `synchronous.py` naming.
unirl/rollout/engine/trainside/engine.py	Updates base class to `SyncRolloutEngine`.
unirl/rollout/engine/trainside/config.py	Switches `BaseEngineConfig` import to `unirl.rollout.engine.synchronous`.
unirl/rollout/engine/trainside/init.py	Updates doc reference to `unirl.rollout.engine.synchronous.BaseRolloutEngine`.
unirl/rollout/engine/synchronous.py	Renames/reshapes sync engine base classes and moves `chunked_engine_generate` here.
unirl/rollout/engine/sglang/weight_sync.py	Updates references from `base.py` to `synchronous.py` in docs/comments.
unirl/rollout/engine/sglang/engine.py	Updates base class to `SyncRolloutEngine` and sync-surface documentation references.
unirl/rollout/engine/sglang/config.py	Switches `BaseEngineConfig` import to `unirl.rollout.engine.synchronous`.
unirl/rollout/engine/sglang_diffusion/weight_sync.py	Updates references from `base.py` to `synchronous.py` in docs/comments.
unirl/rollout/engine/sglang_diffusion/engine.py	Updates base class to `SyncRolloutEngine` and sync-surface documentation references.
unirl/rollout/engine/sglang_diffusion/config.py	Switches `BaseEngineConfig` import to `unirl.rollout.engine.synchronous`.
unirl/rollout/engine/fastvideo/engine.py	Updates base class to `SyncRolloutEngine`.
unirl/rollout/engine/fastvideo/config.py	Switches `BaseEngineConfig` import to `unirl.rollout.engine.synchronous`.
unirl/rollout/engine/composed/engine.py	Updates child engine type checks/annotations to `SyncRolloutEngine`.
unirl/rollout/engine/composed/config.py	Switches `BaseEngineConfig` import to `unirl.rollout.engine.synchronous`.
unirl/rollout/engine/asynchronous.py	Adds new driver-side async engine protocol, batch/agentic engines, and shared mechanisms.
unirl/rollout/engine/agentic/engine.py	Updates inner-engine contract/type checks to `SyncRolloutEngine` and removes the deprecated `drained()` probe.
unirl/rollout/engine/agentic/config.py	Switches `BaseEngineConfig` import to `unirl.rollout.engine.synchronous`.
unirl/rollout/engine/init.py	Replaces eager exports with lazy re-export shim to keep driver-side imports ray/torch-free.
unirl/rollout/async_runtime.py	Deletes the old async scheduler/dispatcher runtime layer.
unirl/distributed/group/handle.py	Adds `PendingHandleCall` and `Handle.launch_nowait()` non-blocking dispatch seam.
examples/diffusion/bagel/bagel_vllmomni_async.yaml	Updates comments to reflect trainer-side poll-before-launch policy (no scheduler flag).
.github/CODEOWNERS	Updates ownership mapping from removed `async_runtime.py` to new `engine/asynchronous.py`.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+        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


+        try:
+            dispatch_mode, dispatch_fn, _, execute_fn = self._method_configs[method_name]
+        except KeyError:
+            raise AttributeError(
+                f"{method_name!r} is not a @distributed method of {_owning_class(self.role_cls).__name__}"
+            ) from None
+
+        batch_size = infer_batch_size(args, kwargs)
+        if (
+            dispatch_mode in (Dispatch.DP_SCATTER, Dispatch.DP_SCATTER_HEAD)
+            and batch_size is not None
+            and batch_size % self.dp_size != 0
+        ):
+            raise ValueError(f"batch_size={batch_size} not divisible by dp_size={self.dp_size}")
+


@distributed

Handle.launch_nowait runs the dispatch -> localize -> execute half of a @distributed method call and returns a PendingHandleCall future (ready/wait/result) instead of blocking in ray.get; result() runs the rebind + collect half at most once (cached on success, retryable on failure). _bind_methods records each method's dispatch tuple so the seam works for any @distributed method. handle_fn itself is untouched; the async rollout path builds on this instead of reaching into the private _execute_all/_rebind_tree seams.

… contracts Replace the layered async_runtime (AsyncRolloutScheduler + GenerationDispatcher protocol + RayGenerationDispatcher + VersionedGroupBuffer) and the agentic trainers' duplicated producer plumbing (_GroupAssembler/_GroupBuffer/ _ingest_completed/_pump/_drain_buffer) with two driver-side engines built on two mechanism classes, all under unirl/rollout/engine/: - synchronous.py (was base.py): the worker-side sync contracts, unchanged in substance; BaseSingleTurnRolloutEngine is renamed SyncRolloutEngine and chunked_engine_generate moves in from the package __init__. - asynchronous.py (new): VersionedBuffer + InflightPool (ray-free mechanisms over Handle.launch_nowait), the AsyncRolloutEngine protocol (poll / drain_freshest / pop_evicted / quiesce + engine-owned weight_version), and its two concretes: AsyncBatchRolloutEngine (launch-time version stamps; used by AsyncARTrainer/AsyncDiffusionTrainer) and AsyncAgenticRolloutEngine (facade over the agentic rank-0 coordinator: [0] unwraps, PendingGroups sibling assembly, completion-time stamps, quiesce = abort + final poll; used by the partial/async agentic trainers). - engine/__init__.py becomes a lazy re-export shim so the driver-side module imports ray/torch-free. Trainers keep the policy (launch ceiling, reap-vs-launch order as statement order, quiesce points, tail carry/drop) as explicit loops; the reap_before_launch flag, dead drained()/pending_roots()/_pending_carried, and the trainer-side _weight_version counters are removed. Import paths change with no aliases: unirl.rollout.async_runtime and the names it exported (AsyncRolloutScheduler, RayGenerationDispatcher, VersionedGroupBuffer, GenerationDispatcher, InflightGeneration) are gone, and unirl.rollout.engine.base is now unirl.rollout.engine.synchronous. All in-repo consumers are migrated; out-of-tree code importing these must update. Verified on GPU (1x8 H20): the async AR recipe trains end-to-end with the PPO ratio pinned at 0.99-1.00 across every weight-sync cycle, and the BAGEL vllm-omni async recipe replays at ratio=1.0000+/-0.0000. CPU-only harnesses pinning the buffer/pool invariants were run but not committed, per the tests-tree policy (Tencent-Hunyuan#99/Tencent-Hunyuan#267); commands and results are in the PR Test Plan.

AsyncARTrainer.__init__ calls BaseTrainer.__init__ directly to avoid the colocate placement(fraction=1.0) block that ARTrainer.__init__ opens, and mirrors ARTrainer's scalar fields by hand. That hand-mirror covers ar.py:62-88 but stops short of ar.py:101, where _rollout_anchor_device is set — while the inherited ARTrainer.evaluate reads it at ar.py:394. Every eval on the async AR path therefore died with: AttributeError: 'AsyncARTrainer' object has no attribute '_rollout_anchor_device' reached from the rollout_id=-1 baseline eval (async_ar.py:309), i.e. before the first training step whenever eval_interval > 0. None is the correct value: the anchored branch is the colocate TP-anchor path, and the disaggregated layout is always SPMD. It keeps evaluate() on its non-anchored route (nullcontext instead of _anchored_rollout_session), where rollout.wake_up() is a no-op on the resident SGLang engine (early-returns while _is_offloaded is False). The other three anchored-path fields ARTrainer.__init__ sets (_enable_fsdp_offload, _anchored_backend_offloaded, _anchored_rollout_awake) are deliberately NOT mirrored: they are read only by the _ensure_anchored_* helpers and the anchored branch of evaluate, none of which async_ar reaches once _rollout_anchor_device is None. Sibling AsyncAgenticTrainer uses the same BaseTrainer.__init__-direct pattern and also omits the field, but pins eval_interval = 0 (AgenticTrainer.evaluate raises), so it never reaches the read and needs no change. Pre-existing on main, not a LIN-631 regression: ar.py is byte-identical between main and this branch, and the refactor does not touch the call site or the __init__ chain. Found by GPU e2e validation of the async AR recipe.

Five deletions, all surfaces LIN-631 itself introduced or relocated, each verified to have zero in-tree consumers: - engine/__init__.py: replace the lazy re-export shim with a docstring-only init. No package-level importer exists (py/md and yaml _target_ paths all verified); the ray/torch-free property comes from importing nothing, not from lazy machinery, and the branch already breaks out-of-tree imports with no aliases elsewhere. - chunked_engine_generate: zero callers on main since the initial release and zero after the move into synchronous.py; recoverable from git history if a consumer ever materializes. - AsyncRolloutEngine Protocol: zero annotations, zero isinstance checks, and neither trainer family can be typed against it (both use off-protocol verbs). The contract description moves into the module docstring. - __all__ trimmed to the consumed surface (the two engines + root_of); VersionedBuffer/InflightPool/PendingGroups/Complete stay module-internal. - InflightPool(method=...): the only construction site never overrides the default; hardcode "generate". ruff format --check + ruff check (v0.15.1) clean over the three .py files; python3 -c "import unirl.rollout.engine.asynchronous" on a torch-less interpreter confirms the driver-side chain still imports stdlib-only.

…dger Closes the weight_version accounting gap documented as a follow-up in #287: ARTrainer.evaluate() pushed weights via _prepare_rollout without advancing the driver-side counter, leaving the engine on unnumbered weights until the next interval boundary. The async trainer families also disagreed on eval policy: AR synced-without-bump, diffusion explicitly passed sync_weights=False, agentic has no eval. Structure: both driver-side engines replace bump_weight_version() with sync_weights(weight_sync) - one call that pushes and advances the ledger, so the pairing stops being call-site discipline. The batch engine also enforces the quiesce invariant (raises before pushing if any generation is in flight). All async sync sites (boundary, resume, agentic drive) route through it; bump_weight_version drops to zero callers and is removed. Policy: async eval becomes read-only. ARTrainer.evaluate() gains sync_weights: bool = True (mirroring DiffusionTrainer.evaluate); the async trainers override evaluate() with async-correct defaults (no push; diffusion also defaults sleep_after=False), so bare evaluate() calls are safe outside train() too. A pre-train explicit push raises a lifecycle RuntimeError instead of an incidental AttributeError. Deliberate behavior changes: (1) async-AR eval no longer pushes weights, so BOTH the eval series and the training rollouts launched between an eval and the next interval boundary change - they previously ran under eval-synced fresher weights; deployment cadence is now governed solely by weight_sync_interval. At interval=1, or when eval_interval is a multiple of the sync interval, training is point-identical to main. (2) The scored eval policy is 1..interval optimizer steps old (exactly 1 at interval=1). (3) The resume-time push advances the ledger 0->1 (metric offset only - eviction math is relative; the launch ceiling is computed from rollout_id).

…ontrollers Naming and placement follow the control-plane role; loop ownership follows whose contract dictates the sequence: - engine/asynchronous.py -> controller.py at the rollout package root. AsyncBatchRolloutEngine -> BatchRolloutController; VersionedBuffer / InflightPool unchanged. The controller stays POLICY-FREE (#287's boundary): batch step loops remain in the trainers as visible statement order (AR launches before reaping; diffusion polls first), built on the verbs plus a new launch_until(gen_id_limit, max_inflight, build) window top-up whose parameters are all controller-native. The one shared batch-policy artifact is launch_ceiling(), a module-level formula the controller never calls — the on-policy invariant both trainers previously carried as duplicated arithmetic now exists exactly once. - The agentic driver side moves family-local: engine/agentic/controller.py with AgenticRolloutController (was AsyncAgenticRolloutEngine), root_of, PendingGroups. Its collect() pump DOES own its loop: that sequence is dictated by the coordinator's engine contract (finalize_if_drained atomicity vs submit's worker-buffer reset — misordering double-pulls), not by training policy, so it lives engine-side, parameterized by poll_interval_s / refill / on_evicted / context; the per-trainer _next_batch/_collect_until/_drain_buffer copies are deleted. - engine/agentic/__init__.py becomes docstring-only (its two re-exports had zero package-level importers; recipes name full dotpaths) so importing the agentic controller stays ray/torch-free like controller.py. Rule recorded in the docstrings: sequence dictated by an engine contract lives engine-side; sequence dictated by training policy lives trainer-side; shared invariants become named functions, not shared loops. Import paths change with no aliases (#287 precedent); in-repo consumers migrated.

Three duplications collapse, with zero renames and loops kept visible in the trainers (#287's mechanism/policy boundary unchanged): - launch_ceiling(): the on-policy launch-clamp arithmetic both batch trainers carried as duplicated inline math becomes one module-level function the engine classes never call. Glue may be copied; the load-bearing invariant may not. - AsyncBatchRolloutEngine.launch_until(gen_id_limit, max_inflight, build): the two-line window top-up both trainers repeated; both parameters are engine-native. The trainers' _next_step loops stay in place as visible statement order (AR launches then polls; diffusion polls first). - AsyncAgenticRolloutEngine.collect(): the producer pump both agentic trainers carried verbatim (poll -> drain -> finalize_if_drained -> backoff -> refill, plus the _drain_buffer eviction sweep). Unlike the batch loops this sequence is dictated by the coordinator's engine contract (finalize atomicity vs submit's worker-buffer reset; misordering double-pulls), not by training policy, so it lives with the engine, parameterized by poll_interval_s / refill / on_evicted / context. Trainer copies deleted. No import paths, class names, files, recipes, or knobs change.

…dger Closes the weight_version accounting gap documented as a follow-up in #287: ARTrainer.evaluate() pushed weights via _prepare_rollout without advancing the driver-side counter, leaving the engine on unnumbered weights until the next interval boundary. The async trainer families also disagreed on eval policy: AR synced-without-bump, diffusion explicitly passed sync_weights=False, agentic has no eval. Structure: both driver-side engines replace bump_weight_version() with sync_weights(weight_sync) - one call that pushes and advances the ledger, so the pairing stops being call-site discipline. The batch engine also enforces the quiesce invariant (raises before pushing if any generation is in flight). All async sync sites (boundary, resume, agentic drive) route through it; bump_weight_version drops to zero callers and is removed. Policy: async eval becomes read-only. ARTrainer.evaluate() gains sync_weights: bool = True (mirroring DiffusionTrainer.evaluate); the async trainers override evaluate() with async-correct defaults (no push; diffusion also defaults sleep_after=False), so bare evaluate() calls are safe outside train() too. A pre-train explicit push raises a lifecycle RuntimeError instead of an incidental AttributeError. Deliberate behavior changes: (1) async-AR eval no longer pushes weights, so BOTH the eval series and the training rollouts launched between an eval and the next interval boundary change - they previously ran under eval-synced fresher weights; deployment cadence is now governed solely by weight_sync_interval. At interval=1, or when eval_interval is a multiple of the sync interval, training is point-identical to main. (2) The scored eval policy is 1..interval optimizer steps old (exactly 1 at interval=1). (3) The resume-time push advances the ledger 0->1 (metric offset only - eviction math is relative; the launch ceiling is computed from rollout_id).

* fix(rollout): route async weight pushes through the engine version ledger Closes the weight_version accounting gap documented as a follow-up in #287: ARTrainer.evaluate() pushed weights via _prepare_rollout without advancing the driver-side counter, leaving the engine on unnumbered weights until the next interval boundary. The async trainer families also disagreed on eval policy: AR synced-without-bump, diffusion explicitly passed sync_weights=False, agentic has no eval. Structure: both driver-side engines replace bump_weight_version() with sync_weights(weight_sync) - one call that pushes and advances the ledger, so the pairing stops being call-site discipline. The batch engine also enforces the quiesce invariant (raises before pushing if any generation is in flight). All async sync sites (boundary, resume, agentic drive) route through it; bump_weight_version drops to zero callers and is removed. Policy: async eval becomes read-only. ARTrainer.evaluate() gains sync_weights: bool = True (mirroring DiffusionTrainer.evaluate); the async trainers override evaluate() with async-correct defaults (no push; diffusion also defaults sleep_after=False), so bare evaluate() calls are safe outside train() too. A pre-train explicit push raises a lifecycle RuntimeError instead of an incidental AttributeError. Deliberate behavior changes: (1) async-AR eval no longer pushes weights, so BOTH the eval series and the training rollouts launched between an eval and the next interval boundary change - they previously ran under eval-synced fresher weights; deployment cadence is now governed solely by weight_sync_interval. At interval=1, or when eval_interval is a multiple of the sync interval, training is point-identical to main. (2) The scored eval policy is 1..interval optimizer steps old (exactly 1 at interval=1). (3) The resume-time push advances the ledger 0->1 (metric offset only - eviction math is relative; the launch ceiling is computed from rollout_id). * fix(rollout): enforce decode-idle on agentic sync_weights; log eval/weight_version * fix(trainer): deterministic diffusion eval sends a pure-ODE request eval_eta=0 previously rode along with the training-resolved sde_indices — a contradictory request the central kernel silently degrades to ODE but BAGEL's worker-resident scheduler refuses (RuntimeError at the first gated step). Clearing the gate at eval_eta<=0 makes the request say what eval means; SD3-family trajectories are unchanged (they already ran ODE). * fix(rollout): log each ledger push at the engine The actor-side [LoRA-SYNC] lines never reach the driver log, leaving weight deployments invisible in stdout. sync_weights is now the single push path, so one driver-side INFO line covers every async push. --------- Co-authored-by: Celve <celve03@gmail.com> Co-authored-by: CjhHa1 <cjh18671720497@outlook.com>

github-actions Bot added the wip Draft / work in progress label Jul 31, 2026

celve mentioned this pull request Jul 31, 2026

refactor(rollout): engine-level sync/async contracts for async rollout celve/unirl#4

Closed

2 tasks

celve force-pushed the LIN-631/main branch 2 times, most recently from 762b04e to 22b2712 Compare July 31, 2026 09:13

celve marked this pull request as ready for review July 31, 2026 09:43

Copilot AI review requested due to automatic review settings July 31, 2026 09:43

celve requested review from CjhHa1, Ideny42, Jayce-Ping, Zcchill, haonan3, leviking98z-rgb and zzhuoxin1508 as code owners July 31, 2026 09:43

github-actions Bot added need review Ready and waiting for review and removed wip Draft / work in progress labels Jul 31, 2026

Copilot started reviewing on behalf of celve July 31, 2026 09:44 View session

Copilot AI reviewed Jul 31, 2026

View reviewed changes

This was referenced Jul 31, 2026

refactor(trainer): separate lifecycle and loop programs leviking98z-rgb/UniRL#25

Draft

refactor(rollout): split lifecycle and receiver capabilities leviking98z-rgb/UniRL#26

Draft

haonan3 force-pushed the LIN-631/main branch from 22b2712 to bfa496f Compare August 2, 2026 04:45

celve and others added 5 commits August 2, 2026 12:57

docs(rollout/engine): trim base.py comments to one-line essentials

f5d24da

haonan3 force-pushed the LIN-631/main branch from 47de5f5 to 525a418 Compare August 2, 2026 04:58

haonan3 approved these changes Aug 2, 2026

View reviewed changes

haonan3 merged commit fada67b into Tencent-Hunyuan:main Aug 2, 2026
5 checks passed

github-actions Bot removed the need review Ready and waiting for review label Aug 2, 2026

haonan3 mentioned this pull request Aug 2, 2026

fix(rollout): guard agentic sync and deterministic diffusion eval #294

Merged

4 tasks

haonan3 mentioned this pull request Aug 2, 2026

refactor(rollout): dedup async driver plumbing; extract ToolAgentHarness #295

Merged

2 tasks

haonan3 mentioned this pull request Aug 2, 2026

refactor(distributed): make handle_fn and launch_nowait share the call halves #297

Closed

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor(rollout): engine-level sync/async contracts for async rollout - #287

refactor(rollout): engine-level sync/async contracts for async rollout#287
haonan3 merged 5 commits into
Tencent-Hunyuan:mainfrom
celve:LIN-631/main

celve commented Jul 31, 2026 •

edited by haonan3

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

celve commented Jul 31, 2026 • edited by haonan3 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Related Issue

Test Plan

Compatibility / Risk

Reviewer Notes

Checklist

Update 2026-08-02 (maintainer)

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

celve commented Jul 31, 2026 •

edited by haonan3

Loading