refactor(rollout): dedup async driver plumbing; extract ToolAgentHarness - #295
Merged
Conversation
haonan3
force-pushed
the
refactor/rollout-controllers
branch
2 times, most recently
from
August 2, 2026 12:35
d194ece to
bccb3a0
Compare
4 tasks
2 tasks
…l halves Extract _launch_half (dispatch -> localize -> execute) and _collect_half (ray.get -> rebind -> collect) on Handle; handle_fn, launch_nowait, and PendingHandleCall.result now sequence the same two methods instead of carrying keep-in-parity copies. No behavior change: the grad-context preamble/postamble stays in handle_fn, launch_nowait stays grad_mode=False/call_id=None.
Two minimal driver-side fixes, no loop moves: - launch_ceiling(): the on-policy launch-clamp arithmetic both batch trainers carried as duplicated inline math becomes one module-level function (trainer POLICY by ownership — the engine classes never call it; hosted here as the trainers' one shared torch-free home). The trainers' _next_step loops are otherwise untouched. - AsyncAgenticRolloutEngine.submit() now enforces its documented precondition: a second submit while the prior drive is live raises instead of double-pulling the coordinator queue (_drive_live flag, reset by finalize_if_drained reporting done and by quiesce). Agentic trainer loops are untouched — byte-identical to main.
…worker-side plugin Step 1 of the harness split: the sequence a task runs by (turns, stop conditions, teardown) is dictated by TASK SEMANTICS — a third sequence owner next to training policy (trainer loops) and engine contracts — and now has its own home, inside the rollout worker so multi-turn env state and intermediates never cross to the driver. - unirl/rollout/harness/: protocol.py (RolloutHarness / HarnessContext / HarnessOutcome — named engines + a suspend probe; NO residency verbs and NO trainer vocabulary; 'suspended' documents honestly that resumability is the ENV's property — stateless envs resume, stateful envs pair with tail_policy: drop) and tool_agent.py (ToolAgentHarness — the generate -> env.step -> observe loop, ported statement-for-statement from AgenticRolloutEngine._run_one: resume-aware turn budget, turn-boundary suspension, env_reward capture, failed-with-partial-trace fault isolation, guaranteed env.close teardown). Harness modules import ray/torch-free. - AgenticRolloutEngine constructs the harness from its existing config fields (no recipe change); _run_one delegates with an EXHAUSTIVE status match inside the try — an unknown status or malformed outcome from a plugin is NaN-marked instead of slipping into training as completed or sinking the drain. The runtime keeps NaN-marking and tensor-side env-reward attach. - Deletes the dead AgentLoop prototype (loop/agent_loop.py) and its RolloutEnginePort (loop/engine_port.py) — zero live callers; the loop now exists exactly once, in the harness. loop/ keeps environments and tools; README/docstrings updated. No recipe/config/knob changes; queueing, concurrency, buffers, abort, and the coordinator protocol are untouched.
haonan3
force-pushed
the
refactor/rollout-controllers
branch
from
August 2, 2026 15:05
1229e6b to
df3002a
Compare
haonan3
marked this pull request as ready for review
August 2, 2026 17:40
haonan3
requested review from
CjhHa1,
Ideny42,
leviking98z-rgb and
zzhuoxin1508
as code owners
August 2, 2026 17:40
celve
approved these changes
Aug 3, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-up to #287, three self-contained commits, deliberately minimal against
main(13 files, +396/−226; the agentic trainers are byte-identical tomain):Commit 1 —
handle_fnshares its halves with the nowait seam._launch_half(dispatch→localize→execute) and_collect_half(ray.get→rebind→collect) extracted onHandle;handle_fn,launch_nowait, andPendingHandleCall.resultsequence the same two methods — the "keep the halves in parity" comment convention (#287's own listed follow-up) is deleted. Grad-context handling stays inhandle_fn; no behavior change.Commit 2 — single-copy launch clamp + enforced double-pull guard. Two minimal driver-side fixes, no loop moves:
launch_ceiling(): the on-policy launch-clamp arithmetic both batch trainers carried as duplicated inline math becomes one module-level function. Ownership is stated in its docstring: trainer POLICY (rollout_id/sync_interval/num_rolloutsvocabulary), never called by the engine classes, hosted inengine/asynchronous.pyonly as the trainers' one shared torch-free home. The batch trainers change by 2 lines each; their_next_steploops are otherwise untouched.AsyncAgenticRolloutEngine.submit()now enforces its documented precondition: a secondsubmitwhile the prior drive is live raises instead of double-pulling the coordinator queue (_drive_liveflag; reset whenfinalize_if_drainedreports done and byquiesce). Agentic trainer loops untouched.Commit 3 — ToolAgentHarness: task control flow as a worker-side plugin (harness split, step 1). The sequence a task runs by (turns, stop conditions, teardown) is dictated by task semantics — a third sequence owner next to training policy (trainer loops) and engine contracts — and now has its own home, inside the rollout worker, so multi-turn env state and intermediates never cross to the driver:
unirl/rollout/harness/:protocol.py—RolloutHarness/HarnessContext(named engines + a suspend probe; no residency verbs — engine wake/sleep is the hosting runtime's deployment concern — and no trainer vocabulary) /HarnessOutcome(completed|suspended|failed+ optionalenv_reward; thesuspendeddoc is honest: resumability is the ENV's property — stateless envs resume, stateful envs pair withtail_policy: drop; config-time validation of that pairing is deferred to the config-selected-harness step).tool_agent.py— thegenerate → env.step → observeloop ported statement-for-statement fromAgenticRolloutEngine._run_one(resume-aware turn budget, turn-boundary suspension, env-reward capture, failed-with-partial-trace fault isolation, guaranteedenv.closeteardown). Harness modules import ray/torch-free.AgenticRolloutEngineconstructs the harness from its existing config fields (no recipe change);_run_onedelegates with an EXHAUSTIVE status match inside thetry— an unknown status or malformed outcome from a plugin is NaN-marked instead of slipping into training as completed or sinking the drain.AgentLoopprototype (loop/agent_loop.py) and itsRolloutEnginePort(loop/engine_port.py) — zero live callers; the loop now exists exactly once.loop/keeps environments and tools.Next steps (separate PRs): config-selected harness on the runtime; a second, genuinely different multimodal harness to validate the boundary; then migrate
ComposedRolloutEngine's PE flow and delete it.Related Issue
N/A (follow-up to #287; harness direction per the team architecture discussion).
Test Plan
Static (branch @ HEAD, base
main@ 8d77a54):ruff format --check+ruff checkover all changed files — clean;python3 -m py_compile— clean;lint/check_recipe_targets.py→2434 recipe _target_ paths resolve(no recipe/config change anywhere);lint/check_experimental_boundaries.py→ ok; stale-reference sweeps (AgentLoop|RolloutEnginePort|engine_port, and the revertedcollect()/launch_until/poll_drivedrafts) → zero hits.CPU harnesses (uncommitted per tests-tree policy) — 20 passed:
engine.asynchronous; buffer freshest-first + staleness eviction +pop_evicted; pool complete-or-nothing retry + immediateKeyboardInterrupt;launch_ceilingvalues (rollout_id=3, interval=2, stale=1 → 6;stale=0, interval=1 → rollout_id+1;num_rolloutsclamp); the trainer-shaped launch-then-reap loop over the verbs; launch- vs completion-time version stamping; submit double-pull guard lifecycle (second submit raises while live;finalize_if_draineddone-report re-arms;quiescere-arms);quiescefold (abort + mandatory follow-up poll, pre-bump stamping);PendingGroupsassembly/discard.harness/*; runs until envdonewith last-reward-wins capture + teardown;max_turns+ resume-awareness; suspension exactly at the turn boundary (in-flight turn finishes; teardown still runs); env fault →failedwith partial trace preserved, partial reward dropped, teardown still runs; unknown engine name →KeyErrorlisting provided names; frozen outcome.The
_run_oneexhaustive-status net is not CPU-testable (engine imports torch); it is review-verified and exercised by the ALFWorld smoke below.GPU end-to-end: Not run; reason: CPU-only dev box. Before undraft: async AR (
examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml, expect ratio≈1) + BAGEL async diffusion (examples/diffusion/bagel/bagel_vllmomni_async.yaml, expect replayratio=1.0000, ~8s/rollout) + one ALFWorld partial/async agentic smoke (exercises the submit guard and the harness delegation: suspension/carry, env reward, NaN-failure path).Compatibility / Risk
AgentLoop,RolloutEnginePort) had zero live callers (re-exports removed fromloop/__init__; out-of-tree importers of those two names must update — refactor(rollout): engine-level sync/async contracts for async rollout #287 no-alias precedent).main; the batch trainers differ by 2 lines each (the ceiling call). The known pump-loop duplication across the agentic pair deliberately REMAINS (glue may be copied; only the invariant got a single copy) — revisiting it is out of scope by decision, not oversight._run_one→ harness port — statement-for-statement, with the failure-path contract preserved (partial trace kept, partial env reward dropped, close-on-all-paths) and pinned by the 7 CPU tests.engine/asynchronous.py+ trainers (mechanical rebase either direction; version semantics unchanged here); feat(rollout): dynamic async rollout scheduler #273/feat(rollout): per-lane dynamic dispatch for AR rollout (sync + async) #289 evolve the AR trainer loop, untouched here; refactor(trainer): extract the sync-loop template into BaseTrainer #282 disjoint.Reviewer Notes
handle.pyhalves (pure extraction) →asynchronous.py(+launch_ceiling, +guard) with the two 2-line trainer hunks →harness/protocol.py(the boundary and its rules) →harness/tool_agent.pyvs the deleted_run_onebody (side-by-side diff recommended) →engine.pydelegation.collect()withrefill/on_evictedcallbacks — rejected (trainer policy + callback inversion in the engine; the parameter-list vocabulary test is the guard) and reverted tomain's trainer loops with only the guard kept engine-side;launch_untilwas dropped (2 lines of glue don't justify a cross-layerbuildcallback);launch_ceilingstays inengine/asynchronous.pyover a trainer-side home deliberately (a torch-heavy trainer module would kill the invariant's CPU numeric tests; a one-function module is ceremony) — its docstring carries the ownership statement; a PR split of the handle commit was made and then folded back per lead decision (commit-level separation suffices for review).Checklist