feat(orchestrator): redesign as async-pipelined signal-based orchestrator

[mikasenghaas]

Why

Three concrete wins over the previous prime_rl.orchestrator:

1. Train + eval rollouts overlap on every eval boundary. No more long tails on online evals. When policy.version reaches an eval interval, the dispatcher stops scheduling new train rollouts but lets the in-flight train tail drain naturally while eval queues up. Conversely, when the eval queue empties the dispatcher starts scheduling train again while the eval tail finishes. Both transitions overlap on the same shared semaphore — capacity stays saturated through the entire eval boundary instead of pausing the pipeline.

2. Shared concurrency limiter across train + eval. A single asyncio.Semaphore(max_inflight_rollouts) and a single AsyncLimiter(tasks_per_minute) govern both kinds of rollouts. The pipeline always runs at the rate the inference infra can sustain — no separate budgets to tune and no idle inference capacity while one kind drains.

3. Natural rollout / group / batch hierarchy amortizes the heavy per-batch post-processing. The old design did tokenization for an entire batch's worth of rollouts at once (asyncio.gather over interleave_rollout for all 128 rollouts) at ship time — a synchronization point that stalled the pipeline. The new design splits work by its natural scope:

   • rollout level — interleave_rollout / backfill_rollout_tokens happen as each rollout completes, overlapped with the dispatcher producing more rollouts.
   • group level — compute_advantages runs over one GRPO group at a time.
   • batch level — post_batch_filters + metrics build runs once per ship, on already-tokenized samples.

   By ship time everything is pre-computed; the ship step just serializes the batch and sends it.

Summary

• Replace the step-driven outer loop in prime_rl.orchestrator with an async-pipelined design driven by signal-based sinks. Clean deprecation — no flag, no parallel module, no compatibility shim.
• The atomic unit is a single Rollout. Sinks own grouping; three uniform processing levels (process_rollout / process_group / process_batch) drive both train and eval flows.
• The dispatcher schedules both train and eval rollouts through a single shared asyncio.Semaphore(max_inflight_rollouts) + shared AsyncLimiter(tasks_per_minute). Drain-switch overlap on each eval boundary recovers the cost of the previous "blocking eval" without cancelling in-flight work.
• Every dispatched rollout reaches the sink exactly once. Success, env-error, empty-trajectory, task-exception, off-policy-cancel — all flow through out_q as Rollouts; the dispatcher synthesizes minimal vf.RolloutOutput markers for the failure paths. The sink owns drop / partial-train policy (filters errored rollouts; drops the whole group when the env requires_group_scoring and any failed).
• Group / batch finalization is sink-derived by counting. The "last rollout in a group" is whichever straggler finishes last — not knowable at dispatch time — so no flags travel on the Rollout; the sink finalizes the (env, example_id) group on the group_size-th arrival.
• Adds pre_batch_filters / post_batch_filters config fields. filters is silently aliased to post_batch_filters (no warning); existing TOMLs keep parsing. Filters are train-only — the eval sink has no filter code by design.

Supersedes the prior brainstorming attempts in #2254 and #2336. Refs #2311.

How to compare against the previous design

Check out main and run the same config head-to-head. Existing configs/ and examples/ work as-is on both — no separate demo. Example with 2 GPUs (1 train + 1 inference):

# New design (this branch)
git checkout exp/orchestrator-v2
uv run rl @ configs/hendrycks_math/rl.toml \
  --output-dir outputs/orch_v2 --wandb.name orch_v2

# Previous design (main)
git checkout main
uv run rl @ configs/hendrycks_math/rl.toml \
  --output-dir outputs/prev --wandb.name prev

Signal model

The pipeline is driven by a small set of explicit signals — no exception-based control flow, no "is X ready?" magic. Each signal has a single producer and consumer:

Signal	Producer	Consumer	Triggers
Rollout (atomic)	RolloutDispatcher.out_q	sink router in Orchestrator.main_loop	next sink processing level
TrainSink.add(rollout) -> bool	train sink (group_size-th arrival + buffer full)	orchestrator main loop	process_one_step (ship)
EvalSink.add(rollout) -> EvalBatch | None	eval sink (num_examples × group_size-th arrival)	orchestrator main loop	log_eval_batch
Orchestrator.stopped: Event	train sink on max_steps, external stop()	orchestrator	drives shutdown
WeightWatcher.on_new_version(step)	watcher (broadcast dir poll)	dispatcher	off-policy cancel, eval trigger

Architecture

                                ┌──────────────────────────────────┐
                                │     WeightWatcher (task)         │
                                │ polls broadcasts/, on_new_version│
                                └────────────┬─────────────────────┘
                                             │ Policy.version, observers
                                             ▼
                                ┌──────────────────────────────────┐
                                │       RolloutDispatcher          │  shared:
                                │           (task)                 │   • Semaphore(max_inflight_rollouts)
                                │  PREFER_TRAIN ⇄ PREFER_EVAL      │   • AsyncLimiter(tasks_per_minute)
                                │  off-policy cancel               │   • Policy (read at dispatch)
                                │  emits every dispatched rollout  │
                                │  (success or error marker)       │
                                └────────────┬─────────────────────┘
                                             │ Rollout (raw.error optional)
                                             ▼
                                       asyncio.Queue (out_q)
                                             │
                          ┌──────────────────┴───────────────────┐
                          ▼                                      ▼
              ┌────────────────────────┐         ┌───────────────────────────┐
              │      TrainSink         │         │       EvalSink            │
              │ process_rollout: tok-  │         │ process_rollout: no-op    │
              │   enize (interleave)   │         │                           │
              │ process_group: filter  │         │ process_group: bucket     │
              │   errors, advantages,  │         │   per-example into env    │
              │   pre_filters          │         │ process_batch: build per- │
              │ process_batch: post_   │         │   env reward/pass@k       │
              │   filters + samples    │         │ finalizes on count        │
              │ finalizes on count     │         │   = num_examples × group  │
              │   = group_size         │         │                           │
              └───────────┬────────────┘         └───────────┬───────────────┘
                          │ TrainBatch                        │ EvalBatch
                          ▼                                   ▼
              ┌────────────────────────────────────────────────────────────────┐
              │   Orchestrator main_loop / process_one_step / log_eval_batch   │
              │  ckpt save, ship to trainer, monitor.log, heartbeat, step++    │
              └────────────────────────────────────────────────────────────────┘

   ┌─────────────────────────┐
   │   IntervalLogger (task) │   every log_loop_interval seconds:
   │                         │     dispatcher/{inflight_train, inflight_eval,
   │                         │                 sched_mode, available_permits, ...}
   │                         │     event_loop_lag/{p90, p99, max, ...}
   │                         │   written to wandb on the _timestamp axis
   └─────────────────────────┘

Four background tasks (dispatcher, watcher, log_loop, and the in-task main loop) under Orchestrator.start().

Module layout

File	Single responsibility	Construction deps
types.py	All dataclasses, type aliases, protocols (Policy, Progress, Rollout, RolloutMeta, GroupState, TrainBatch, EvalBatch, ProcessResult, Kind, SchedMode, VersionObserver).	—
watcher.py	Polls broadcasts/, advances Policy.version, notifies observers.	policy, student_inference, observers
dispatcher.py	Schedules rollouts; emits every dispatched rollout (with raw["error"] on failures). Off-policy cancel, eval triggers.	config, train_envs, eval_envs, policy, inference pools
train_sink.py	Three-level train sink. Per-rollout tokenize, per-group filter-errors + advantages + pre-filter, per-batch post-filter + samples. Sink-derived batch boundaries via group_size-counting.	config, tokenizer, renderer, train_envs, pre_filters, post_filters
eval_sink.py	Three-level eval sink. No filters. Sink-derived epoch boundary via num_examples × group_size-counting. Builds per-env reward/pass@k/completion-len.	eval_envs (for group_size / num_examples lookup)
metrics.py	MetricsBuilder.build(...) and the ProcessResult dataclass it consumes.	config
log_loop.py	IntervalLogger writes dispatcher gauges + event-loop lag to wandb on the _timestamp axis.	dispatcher, policy, interval
ckpt.py	Lean Progress(step, totals) checkpoint manager.	—
orchestrator.py	Owns everything. __init__ / setup / start / stop / shutdown / main_loop / process_one_step / log_eval_batch.	config

No back-pointers — every component takes only the deps it needs.

Overlap semantics

Scheduling priority transitions are level-triggered:

• PREFER_TRAIN → PREFER_EVAL when an eval-trigger fires (policy.version % env.interval == 0, gated by eval_base_model / skip_eval_on_resume). In-flight train drains naturally; eval fills the freed capacity.
• PREFER_EVAL → PREFER_TRAIN when the eval queue is empty. New train rollouts overlap the eval tail.

No cancellations. Both transitions overlap on the same semaphore.

Config surface

Existing TOMLs in configs/ and examples/ keep parsing without changes or warnings.

Additive in packages/prime-rl-configs/.../orchestrator.py:

• pre_batch_filters: list[FilterConfig] (default: gibberish, repetition, zero-advantage all in monitor mode).
• post_batch_filters: list[FilterConfig] with AliasChoices("post_batch_filters", "filters") (default unchanged from the prior filters: gibberish + repetition monitor, zero-advantage enforce). The old filters key still parses silently — no deprecation warning.
• OrchestratorExperimentalConfig.log_loop_interval: float = 5.0.

Parsed but inert (not used by the new orchestrator):

• [orchestrator.buffer] and every key inside it. Difficulty pools are replaced by pre_batch_filters; the block is silently ignored.
• eval.cancel_inflight_rollouts_on_eval. Eval transitions always use the drain-switch overlap.

Breaking

The Python module surface inside prime_rl.orchestrator changes (the orchestrator class, RolloutDispatcher, sinks, etc. all move to new files); external users importing those internals will need to update imports. The user-facing CLI (rl, orchestrator, sft) and TOML config schema are unchanged.

Dropped fields (parsed but inert):

• BufferConfig / difficulty pools (easy/hard pool tracking, fractions, hash_keys, online_difficulty_filtering) — replaced by pre_batch_filters.
• eval.cancel_inflight_rollouts_on_eval — drain-switch overlap is the only eval transition mode.

Filters on eval rollouts no longer apply (filters are train-only).