Skip to content

refactor(rollout): move the driver-side async engines to rollout/manager/ - #7

Closed
celve wants to merge 1 commit into
LIN-693/addressed-callsfrom
LIN-693/manager-queue
Closed

refactor(rollout): move the driver-side async engines to rollout/manager/#7
celve wants to merge 1 commit into
LIN-693/addressed-callsfrom
LIN-693/manager-queue

Conversation

@celve

@celve celve commented Aug 4, 2026

Copy link
Copy Markdown
Owner

Summary

AsyncBatchRolloutEngine and AsyncAgenticRolloutEngine hold no model and cannot
generate. Calling them engines made "engine" mean three different things at once:
the BaseRolloutEngine ABC, a concrete backend, and these driver-side buffer
managers. 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,
LrSchedulerConfig), the diffusion noise scheduler (flow_match_sde_scheduler.py,
patch_scheduler.py, patch_wan_scheduler.py), and SGLang's own scheduler
subprocesses (sglang/config.py:178). A fourth meaning would be worse than the
current name.

Adds manager/protocol.py recording the consumer surface both implementations
share, and splits the policy-free mechanisms (VersionedBuffer, PendingGroups,
root_of) out of the file that held both engines and all their machinery.

launch_ceiling stays with BatchManager, whose two trainers are its only
callers, so its ownership note ("trainer-side POLICY... hosted here only because
it is the two batch trainers' one shared torch-free home") remains accurate.

Pure move + rename. No behaviour change.

Related Issue

LIN-693 — first commit of the change that moves the agentic rollout task queue
from rank 0 of the rollout slab to the driver.

Test Plan

  • AST comparison against the pre-move engine/asynchronous.py proving
    VersionedBuffer, PendingGroups, InflightPool, _InflightJob, root_of
    and BatchManager byte-identical modulo the rename — all six report IDENTICAL.
  • python -m compileall unirl/ — passes.
  • python3 lint/check_recipe_targets.py — 2434 recipe _target_ paths resolve.
  • python3 lint/check_experimental_boundaries.py — ok.
  • grep sweep for residual engine.asynchronous / old class names outside the
    new package: only intentional historical notes in the moved docstrings.

Not run; reason: no Ray, torch or GPU in the authoring environment, and this
repo's CI is lint-only, so nothing here has been executed. The rename is proven
structurally rather than behaviourally.

Compatibility / Risk

Import paths change for four trainers (async_ar, async_diffusion,
agentic_async, agentic_partial), all updated in this PR. No recipe, config,
checkpoint or data-format change; no _target_ in any recipe points at the moved
module (confirmed by check_recipe_targets).

Anything on a branch importing unirl.rollout.engine.asynchronous will need the
one-line path update.

Reviewer Notes

Stacked on #6 (LIN-693/addressed-calls), which is its base — review that
first. This PR's diff against that base is the move alone.

Draft, because this is the first of five commits in the intended change. Still
to come on this branch: exposing run_trajectory/set_stopping on
AgenticRolloutEngine, giving AgenticManager the task queue and per-replica
placement, rewiring all three agentic trainers, and deleting the rank-0
coordinator. The move stands alone and is mergeable on its own merit, but merging
it in isolation buys only the naming.

The design and staged plan are written up outside this repo; the short version is
in the linked issue.

Prepared with AI assistance. The AST proof is what carries the "pure move" claim —
worth re-running rather than eyeballing the diff. Checked for overlapping open
PRs — there are none.

Checklist

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

…ger/

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.
@github-actions github-actions Bot added the wip label Aug 4, 2026
@celve

celve commented Aug 4, 2026

Copy link
Copy Markdown
Owner Author

Closing: opened against the fork by mistake. Re-opening against Tencent-Hunyuan/UniRL, with the @addressed primitive and the manager move combined into one PR.

@celve celve closed this Aug 4, 2026
@github-actions github-actions Bot removed the wip label Aug 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant