Move MOPD teacher scoring out of TQReplayBuffer.commit() into its own stage

[terrykong]

Context

#3768 added text-only MOPD to the SingleController path. Teacher scoring is wired in as a post-write hook on the replay buffer: commit() writes the rollout rows to TransferQueue, awaits an enricher, and only then marks the slot trainable (replay_buffer.py:844-856).

The port is correct and the posted 5-step comparison matches the legacy path. This issue is about two follow-ups that turn out to have the same fix. Neither blocks #3768.

---

Problem 1 — an algorithm-specific hook now lives in shared infrastructure

set_post_write_enricher, _post_write_enricher, and PostWriteEnrichmentError are all new in #3768; there was no pre-existing extension point on main. TQReplayBuffer is used by every SC algorithm, and it gained a callback whose signature — (KVBatchMeta, PromptGroupRecord) -> Awaitable[KVBatchMeta] — is shaped around MOPD's needs and has exactly one implementation.

The sharper form: commit() used to be a bounded write (pack, put_samples, flip a flag). It now awaits a GPU forward pass on another node. That is a surprising property for a buffer method.

Problem 2 — the teacher is called at prompt-group granularity

The enricher runs once per prompt group, so with the shipped recipe:

num_prompts_per_step	8
num_generations_per_prompt	4 → 4 rows per teacher call
teacher: 1 node, 8 GPUs, TP2	DP = 4 → 1 sequence per rank per forward
groups per step	8, all serialized on one teacher's lock

Measured in #3768: 16.44 s of lock wait in the first bucket, 0.90–1.23 s steady state per 4-row call. A group parked in the enricher also still holds its generation permit (single_controller.py:570), so teacher latency directly reduces rollout concurrency.

Why these are one fix

Batching across groups is impossible inside commit() — a single commit only sees its own group. A separate stage can accumulate. So moving the teacher out of the buffer is exactly what unlocks batching.

---

Recommended: a teacher pump

Teacher scoring becomes its own long-lived pump, _teacher_pump, enabled only when MOPD is on. The buffer keeps one small generic primitive: a group can be committed not yet ready, and marked ready later.

This follows the pump pattern SC already uses. The actor is not a two-pump producer/consumer — it runs _rollout_pump, _train_pump, and _stall_watchdog_pump unconditionally, plus _gen_fleet_probe_pump conditionally on a feature flag (single_controller.py:322-334). A conditional _teacher_pump is the same shape as the fleet probe: created only when the feature is on, skipped entirely otherwise.

  _rollout_pump ──► commit(defer_ready=True) ──► pump.submit()   [returns immediately]
                          │                          │             generation permit released here
                          ▼                          ▼
                    rows in TQ                  batch by teacher
                    ready = False                    │
                                                     ▼
                                          one teacher call for N groups
                                                     │
                                                     ▼
                                          buffer.mark_ready(each group)
                                                     │
                                                     ▼
                                              _train_pump picks it up

1. Buffer: deferred readiness (the only change to shared code)

No callback field, no MOPD-shaped signature, no new exception type.

class TQReplayBuffer:
    async def commit(self, group_id, record, *, start_weight_version,
                     end_weight_version, defer_ready: bool = False) -> KVBatchMeta:
        ...
        await self._call_dp("put_samples", ...)
        meta = KVBatchMeta(...)
        idx = self._group_ids.index(group_id)
        self.meta_list[idx] = meta
        self.end_weight_list[idx] = end_weight_version
        self.ready_list[idx] = not defer_ready      # <-- the whole change
        return meta

    def mark_ready(self, group_id: str, meta: KVBatchMeta) -> None:
        """Make a deferred group visible to the sampler."""
        idx = self._group_ids.index(group_id)
        self.meta_list[idx] = meta
        self.ready_list[idx] = True

    async def discard(self, group_id: str) -> None:
        """Clear a deferred group's rows and drop its slot."""

defer_ready is generic: any later scorer (reward model, verifier) wants the same primitive. "Post-write enricher" is a description of MOPD.

2. Rollout side: submit and get out of the way

# RolloutManager.generate_and_push
meta = await self._tq_buffer.commit(
    group_id, record,
    start_weight_version=..., end_weight_version=...,
    defer_ready=self._teacher_stage is not None,
)
if self._teacher_stage is not None:
    self._teacher_stage.submit(group_id, meta, record)   # non-blocking
return meta
# rollout task ends here; sem.release() runs while the teacher is still working

3. The pump

class TeacherScoringPump:
    """Batches committed groups by teacher, scores them, marks them ready."""

    def __init__(self, *, buffer, dp_client, teacher_worker_groups,
                 alias_to_group_alias, opd_cfg, max_batch_groups, flush_after_s):
        self._buffer = buffer
        self._dp_client = dp_client
        self._teachers = teacher_worker_groups
        self._alias_to_group_alias = alias_to_group_alias
        self._opd_cfg = opd_cfg
        self._pending: dict[str, list[_Pending]] = defaultdict(list)
        self._in_flight: set[str] = set()          # one batch per teacher at a time
        self._max_batch_groups = max_batch_groups  # MUST be < max_inflight_prompts
        self._flush_after_s = flush_after_s
        self._wakeup = asyncio.Event()
        self._metrics = TeacherMetrics()           # counters, drained by _train_pump

    def submit(self, group_id, meta, record) -> None:
        """Called by the rollout task. Never blocks, never awaits."""
        group_alias = resolve_teacher(record, self._alias_to_group_alias, self._opd_cfg)
        self._pending[group_alias].append(_Pending(group_id, meta))
        if len(self._pending[group_alias]) >= self._max_batch_groups:
            self._wakeup.set()

    async def run(self) -> None:
        """The pump body, started beside _rollout_pump and _train_pump."""
        while not self._closed:
            await self._wait_for_work()             # wakeup event OR flush timer
            for group_alias in list(self._pending):
                if group_alias in self._in_flight:
                    continue                        # keeps this teacher's calls ordered
                batch = self._take(group_alias)     # pops up to max_batch_groups
                if batch:
                    self._in_flight.add(group_alias)
                    asyncio.create_task(self._score(group_alias, batch))
        await self._flush_all()                     # drain backlog before returning

    async def _score(self, group_alias, batch) -> None:
        # N groups -> ONE teacher call. 8 groups x 4 rows = 32 rows over DP 4 = 8 per rank.
        merged = batch[0].meta.concat(*[g.meta for g in batch[1:]])
        try:
            # A thread cannot be cancelled: shield and drain it, so rollback never
            # clears rows while teacher ranks are still reading them.
            await asyncio.shield(asyncio.to_thread(
                score_teacher_batch,                # module-level, needs no object
                self._dp_client, self._teachers[group_alias], merged,
            ))
        except Exception:
            for g in batch:
                await self._buffer.discard(g.group_id)
            raise
        finally:
            self._in_flight.discard(group_alias)
        for g in batch:
            self._buffer.mark_ready(
                g.group_id, g.meta.with_fields(["teacher_reference_logprobs"])
            )

_in_flight replaces the per-teacher asyncio.Lock. The requirement is unchanged — calls to one physical teacher must stay ordered so its collectives line up — but the pump is the only scheduler, so the guard belongs there rather than being enforced twice.

4. Controller wiring

Mirrors the existing conditional-pump pattern, right next to it:

rollout_task  = asyncio.create_task(self._rollout_pump())
train_task    = asyncio.create_task(self._train_pump())
watchdog_task = asyncio.create_task(self._stall_watchdog_pump())
tasks = [rollout_task, train_task, watchdog_task]

# Only with MOPD on, same as the fleet probe below: a pump for a feature
# this run does not use would just be an idle timer.
teacher_task = (
    asyncio.create_task(self._teacher_pump.run())
    if self._teacher_pump is not None
    else None
)
if teacher_task is not None:
    tasks.append(teacher_task)

5. TQTeacherLogprobCoordinator goes away

The refactor removes a class rather than adding one. Today the coordinator is correctly an object: it owns the per-teacher locks and a stateful protocol — acquire, to_thread, shield, drain-on-cancel, temp-row cleanup — that nothing around it takes part in. Closed state, its own lifecycle, low reach-back.

The pump takes the lifecycle. What would be left does not justify an object:

Member today	Where it goes
_resolve_teacher	module-level resolve_teacher(record, alias_map, opd_cfg) — pure
_enrich_sync	module-level score_teacher_batch(dp_client, teacher, meta) — DP padding, temp rows, cleanup, the blocking call. Pure apart from its arguments; temporary_sample_ids is a local.
enrich()	dissolves: batching and scheduling → pump, to_thread/shield → pump, metrics → pump
_teacher_locks	pump's _in_flight guard
drain_metrics()	pump; _train_pump calls it in the same place it does today

The deciding fact: the coordinator's four constructor arguments — dp_client, teacher_worker_groups, alias_to_group_alias, opd_cfg — are exactly what the pump needs for its own job. Keeping both means two objects holding the same handles with one delegating to the other. And the two remaining behaviours have no mutable state at all, which is the shape that wants module-level functions, not a class whose constructor only memoizes its arguments.

This also makes the intricate part easier to test than it is now: score_teacher_batch is a plain call, where today reaching the same padding and cleanup paths means constructing a coordinator and driving asyncio.run(...).

Untouched by this refactor: resolve_reference_aliases, get_teacher_routing_metrics, reserve_teacher_clusters, create_teacher_worker_groups, and the OPD config schemas — all shared with the legacy path.

Small cleanup that comes along: _resolve_teacher currently rebuilds dict(opd_cfg.get("teacher_model_by_agent_name", {})) on every call; hoist it to the pump's __init__.

---

What this buys

• The buffer learns one generic concept instead of carrying an MOPD-shaped callback.
• commit() goes back to being a bounded write.
• The generation permit is released at commit, so teacher latency stops throttling rollouts.
• One teacher call per N groups: 8 sequences per rank instead of 1, in the shipped recipe.
• _contains_post_write_enrichment_error and the no-retry special case in rollout_manager (:64-72, :1303-1307) can be deleted — the rollout genuinely succeeded, so there is no teacher failure for it to classify.

---

Alternative considered — keep two pumps, batch inside the rollout path

Instead of a new pump, the rollout tasks themselves coalesce: each commits with defer_ready=True, adds itself to a per-teacher accumulator, and awaits a shared future that resolves when its batch is scored. No new long-lived task.

Why this is the weaker design.

1. It does not fix the back-pressure. A rollout task parked on its batch future still holds its generation permit (single_controller.py:570). Batching shortens the wait but the coupling stays: teacher latency keeps throttling generation. The pump version releases the permit at commit, so generation keeps running while teachers work — which is more pipelining, not less.
2. The inline variant is lumpy. If instead the batch-completing task does the call inline and resolves the others, most tasks return fast but one unlucky task pays for the whole batch while holding a permit — and if that task is cancelled mid-flight, the whole batch stalls.
3. It collapses into the pump anyway. A partial batch at end-of-epoch or shutdown needs a flush timer. If every rollout task is parked awaiting its future, nothing is left to fire that timer — so you add a background task to drive it, and now you have a third task with none of the benefits.
4. Worse failure surface. Rollout tasks waiting on each other's arrival, bounded by max_inflight_prompts, is a deadlock shape. The pump owns the batch and the timer in one place, where the bound is a local invariant instead of an emergent property of how many rollouts happen to be in flight.

What it would win. One less long-lived task to supervise, cancel and drain, and the producer/consumer story stays literally two-sided.

Why that is not worth much here. SC is already not two pumps — _stall_watchdog_pump runs unconditionally and _gen_fleet_probe_pump runs conditionally (single_controller.py:322-334). A conditional teacher pump adds a case to an existing pattern rather than introducing one.

Design points that need a decision

1. Deadlock bound. max_batch_groups must be < max_inflight_prompts, or a batch can wait on a group that will never be dispatched. A flush timer is required regardless.
2. Failure blast radius. One teacher error now discards N groups instead of 1. Decide: discard the batch, or fall back to per-group retry once.
3. Shutdown and epoch end. Pending unflushed groups must be drained, or the run hangs at TaskGroup exit with rows written and never marked ready.
4. Checkpointing. A written-but-unscored group is a state that does not exist today. Decide whether save flushes the stage first or treats those groups as uncommitted.
5. Sampler accounting. Deferred groups still occupy buffer capacity, which is correct — they are buffered rollouts. Confirm count_for_target_step and the samplers treat unready groups the way they already treat reserved-but-uncommitted ones.

---

---

Beyond this issue: SC will keep growing pumps, and the wiring does not scale

Adding a pump is cheap. Supervising one is not. The precedence ladder in run() (single_controller.py:336-353) is hand-written per pump — if probe_task in done: await probe_task, then watchdog, then rollout, then await train_task — with a comment at each step explaining why it sits there. A teacher pump means editing that ladder, and every future pump means editing it again.

The ordering is not arbitrary: it already encodes three kinds of pump, just implicitly.

• Monitor — _stall_watchdog_pump (:1271), _gen_fleet_probe_pump (:1359). Loops forever, so finishing at all means it raised. Surfaced first, because its diagnosis beats a downstream pump whose only symptom is "waiting".
• Producer — _rollout_pump (:481). May finish normally when the data runs out; a failure must propagate, but a normal end just lets downstream drain.
• Terminal — _train_pump (:879). Awaited last; its return ends the run.

A teacher pump is a fourth kind that does not exist yet: a stage — it should end once the producer has ended and its own backlog is drained, and its failure should propagate like a producer's.

A small protocol would let the ladder be written once:

class PumpKind(Enum):
    MONITOR  = auto()   # loops forever; finishing at all means it raised
    PRODUCER = auto()   # may finish normally; downstream then drains
    STAGE    = auto()   # ends when upstream ended and its backlog is empty
    TERMINAL = auto()   # awaited last; its return ends the run

class Pump(Protocol):
    name: str
    kind: PumpKind
    async def run(self) -> None: ...
    async def drain(self) -> None: ...   # PRODUCER/STAGE: finish backlog, then return

run() then builds pumps: list[Pump] from whatever the config enabled and hands it to one _supervise(pumps) that encodes the precedence once, instead of run() growing a branch per feature. MONITOR and the conditional-creation rule are already the pattern — this just names it.

Two notes on shape:

• Keep _rollout_pump named _rollout_pump. The role belongs on kind, not in the name. SFT would supply a different PRODUCER — rows straight from the dataloader, no generation — and rollout stays the honest name for the GRPO one. A generic rename buys nothing and churns logs and call sites.
• Readiness generalizes past one boolean. With more than one stage, ready_list[i] stops being expressive: a group is trainable when it has all the fields training needs, which the code already states in _advantage_input_fields() (:1869). Tracking the field set per group and deriving readiness from it supports 0, 1 or N stages with no per-stage buffer change — and defer_ready becomes just the one-stage case of it.

This is larger than the MOPD refactor and should not be bundled with it. It is written here because a teacher pump is the change that makes the cost visible; if anyone picks it up, split it into its own issue.

---

First concrete step — measure before refactoring

The batching case assumes the teacher forward is latency-bound. At 1.7B with 4 rows it almost certainly is; with a large teacher and long sequences it may already saturate, and then there is nothing to win.

Time one 32-row teacher call against eight 4-row calls on the same teacher. If 32 rows costs about what 4 rows costs, this refactor is worth doing. If it costs roughly 8x, close this issue — Problem 1 alone is a much smaller cleanup and would not justify a new stage.

Non-goal

Not a blocker on #3768 and not a request to revert. The hook is ~25 lines and reversible.

Related: #2625 (SC / Async RL cleanup tracking).