fix(orchestrator): hold batch ship on the trainer's published version

.github/workflows/gpu_tests.yaml

@@ -121,3 +121,11 @@ jobs:
           PYTEST_OUTPUT_DIR: /tmp/outputs
           VLLM_WORKER_MULTIPROC_METHOD: spawn
         run: uv run pytest -vvs ${{ matrix.pytest_args }}
+      - name: Upload run logs on failure
+        if: failure()
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.test }}-logs
+          path: /tmp/outputs/**/*.log
+          retention-days: 7
+          if-no-files-found: ignore

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -539,8 +539,8 @@ class OrchestratorConfig(BaseConfig):
     max_steps: int | None = None
     """Maximum training steps. If None, runs indefinitely."""
 
-    max_off_policy_steps: int = Field(8, ge=0)
-    """Maximum policies allowed to generate a single rollout. Rollouts generated more than ``max_off_policy_steps`` ahead of training are discarded. Higher values yield better throughput at the cost of off-policy noise."""
+    max_off_policy_steps: int = Field(8, ge=1)
+    """Maximum policy versions a rollout may train behind its generation policy, queue time included (batch N trains on policy v{N-1}, so a rollout generated by v{k} shipping in batch N is (N-1)-k versions off-policy). Rollouts already past this bound — in flight or buffered awaiting a batch — are discarded. Higher values yield better throughput at the cost of off-policy noise. Must be at least 1: under NCCL weight broadcast the trainer never broadcasts the final policy versions, so the last batch necessarily trains one version off-policy."""
 
     bench: bool = False
     """Benchmark mode. Sets ``max_steps`` to 5 and disables W&B."""

src/prime_rl/orchestrator/dispatcher.py

@@ -11,13 +11,17 @@
   schedule next. Transitions are level-triggered (driven by the eval
   source's emptiness), so in-flight rollouts of the opposite kind drain
   naturally on either side of an eval boundary.
-- ``on_version_pending`` (called by the watcher before the engines pause for
-  the weight update) bumps ``off_policy_steps`` on in-flight train rollouts and
-  drops groups past ``max_off_policy_steps``.
-  Eval rollouts are measurements for the policy version they started with,
-  so they are allowed to finish even if training advances. Train rollouts
-  sampled from a frozen model never age — their sampler doesn't change
-  with policy updates.
+- Train scheduling is demand-driven: the injected ``train_needed`` predicate
+  (owned by the orchestrator) says whether the batch being collected still
+  needs rollouts; ``fill_inflight`` re-evaluates it per scheduled group, so
+  dispatch stops the moment the batch is covered.
+- ``cancel_stale_train_groups`` (driven by the orchestrator's
+  ``on_version_pending``, before the engines pause for the weight update)
+  drops in-flight train groups whose generation policy fell below the
+  staleness cutoff. Eval rollouts are measurements for the policy version
+  they started with, so they are allowed to finish even if training
+  advances; train rollouts sampled from a frozen model never go stale —
+  their sampler doesn't change with policy updates.
   Cancellations surface as synthetic ``Cancelled`` markers so the sink's
   count-to-``group_size`` finalization still fires.
 """
@@ -29,7 +33,7 @@
 from collections import Counter, defaultdict
 from dataclasses import dataclass, field
 from enum import Enum, auto
-from typing import Literal
+from typing import Callable, Literal
 
 import verifiers.v1 as vf
 from aiolimiter import AsyncLimiter
@@ -117,8 +121,9 @@ class RolloutDispatcher:
     """``await dispatcher.start()`` runs the dispatch loop until ``stop()``.
     Pulls examples from ``TrainSource`` / ``EvalSource``, schedules
     rollouts under shared capacity, and emits ``Rollout``\\ s to
-    ``out_q``. The watcher drives ``on_version_pending`` for off-policy
-    cancellation; the orchestrator triggers eval epochs."""
+    ``out_q``. The orchestrator supplies ``train_needed``, drives
+    ``cancel_stale_train_groups`` on weight updates, and triggers eval
+    epochs."""
 
     def __init__(
         self,
@@ -131,7 +136,7 @@ def __init__(
         policy: Policy,
         max_inflight_rollouts: int,
         tasks_per_minute: float | None,
-        max_off_policy_steps: int,
+        train_needed: Callable[[], bool],
     ) -> None:
         self.policy = policy
         self.train_envs = train_envs
@@ -141,7 +146,7 @@ def __init__(
         self.policy_pool = policy_pool
         self.train_source = train_source
         self.eval_source = eval_source
-        self.max_off_policy_steps = max_off_policy_steps
+        self.train_needed = train_needed
 
         self.max_inflight = max_inflight_rollouts
         self.inflight_permits = 0
@@ -161,12 +166,9 @@ def __init__(
         # winds down without scheduling new train rollouts
         self.train_scheduling_disabled: bool = False
         self.metrics = DispatcherMetrics()
-
-        # Orchestrator-owned gate. When clear, ``fill_inflight`` returns
-        # without scheduling new groups. The dispatcher itself doesn't know
-        # *why* — the orchestrator toggles this based on step / policy lead.
-        self.dispatch_allowed = asyncio.Event()
-        self.dispatch_allowed.set()
+        # Last observed value of ``train_needed`` so pause/resume transitions
+        # log once instead of every scheduling pass
+        self._train_paused = False
 
         self.stopped = asyncio.Event()
         self.task: asyncio.Task | None = None
@@ -223,16 +225,6 @@ def disable_train_scheduling(self) -> None:
         triggered eval drain naturally."""
         self.train_scheduling_disabled = True
 
-    @property
-    def max_off_policy_level(self) -> int:
-        steps = [m.off_policy_steps for m in self.inflight.values() if m.kind == "train"]
-        return max(steps) if steps else 0
-
-    @property
-    def mean_off_policy_level(self) -> float:
-        steps = [m.off_policy_steps for m in self.inflight.values() if m.kind == "train"]
-        return sum(steps) / len(steps) if steps else 0.0
-
     # ── lifecycle ──────────────────────────────────────────────────────────
 
     async def start(self) -> None:
@@ -267,48 +259,49 @@ async def stop(self) -> None:
             await safe_cancel(self.task)
             self.task = None
 
-    async def on_version_pending(self, step: int) -> None:
-        """Bump off-policy counters and drop groups past
-        ``max_off_policy_steps`` (drop_group emits ``Cancelled`` markers so
-        the sink still finalizes the partial group). Eval rollouts are not
-        aged because they are tied to their start-time policy version.
-
-        Runs *before* the inference engines are paused for the weight update so
-        the resulting aborts are processed while the engine is still stepping —
-        otherwise the orphaned KV transfers crash the decode engine on resume
-        (see ``WeightWatcher.apply_policy_update``)."""
+    async def cancel_stale_train_groups(self, min_policy_version: int) -> int:
+        """Drop in-flight train groups generated by a policy older than
+        ``min_policy_version`` (``drop_group`` emits ``Cancelled`` markers so
+        the sink still finalizes the partial group). Eval rollouts are tied
+        to their start-time policy version and never dropped; frozen-sourced
+        train rollouts never go stale — their sampler doesn't change with
+        policy updates.
+
+        Must run *before* the inference engines are paused for a weight update
+        so the resulting aborts are processed while the engine is still
+        stepping — otherwise the orphaned KV transfers crash the decode engine
+        on resume (see ``WeightWatcher.apply_policy_update``). The
+        orchestrator drives this from its ``on_version_pending`` hook."""
         stale_groups: set[uuid.UUID] = set()
-        cancelled = 0
         for meta in self.inflight.values():
             if meta.kind != "train":
                 continue
-            # Frozen-sourced rollouts never go stale — their sampler doesn't
-            # change with policy updates.
             if not self.train_envs.get(meta.env_name).sampler.samples_from_live_policy:
                 continue
-            meta.off_policy_steps += 1
-            if meta.off_policy_steps > self.max_off_policy_steps:
+            if meta.policy_version < min_policy_version:
                 stale_groups.add(meta.group_id)
 
+        cancelled = 0
         for gid in stale_groups:
-            removed = await self.drop_group(gid)
-            cancelled += removed
+            cancelled += await self.drop_group(gid)
 
         if cancelled:
             get_logger().warning(
-                f"Cancelled {cancelled} train rollouts past max_off_policy_steps={self.max_off_policy_steps}. "
+                f"Cancelled {cancelled} in-flight train rollouts generated by policy < v{min_policy_version} — "
+                "they would exceed max_off_policy_steps by the time they train. "
                 "Consider increasing it to avoid this."
             )
-
-    async def on_new_version(self, step: int) -> None:
-        """No-op: the dispatcher drains in ``on_version_pending`` (pre-pause)."""
+        return cancelled
 
     async def fill_inflight(self) -> None:
         """Schedule new rollouts up to ``max_inflight``, honoring
-        ``self.mode``. Eval scheduling ignores the orchestrator's dispatch
-        gate (evals are version-pinned measurements); only train scheduling
-        respects it. When ``PREFER_EVAL``'s source exhausts we flip back to
-        ``PREFER_TRAIN`` so the eval tail drains alongside fresh train."""
+        ``self.mode``. Eval scheduling is unbounded by train demand (evals
+        are version-pinned measurements); train scheduling re-evaluates
+        ``train_needed`` per group, so it stops the moment the batch being
+        collected is covered and resumes as soon as demand reappears (next
+        batch, staleness drops, errored arrivals). When ``PREFER_EVAL``'s
+        source exhausts we flip back to ``PREFER_TRAIN`` so the eval tail
+        drains alongside fresh train."""
         while True:
             if self.available_permits <= 0:
                 return
@@ -326,10 +319,23 @@ async def fill_inflight(self) -> None:
                 scheduled = await self.try_schedule("eval")
                 if not scheduled:
                     return
-            else:  # PREFER_TRAIN — respects the orchestrator's dispatch gate
-                if not self.dispatch_allowed.is_set():
-                    return
-                scheduled = await self.try_schedule("train")
+            else:  # PREFER_TRAIN — fresh groups bounded by the orchestrator's train demand
+                needed = self.train_needed()
+                paused = not needed
+                if paused != self._train_paused:
+                    self._train_paused = paused
+                    get_logger().debug(
+                        "Pausing fresh train dispatch (batch covered or policy too far behind)"
+                        if paused
+                        else "Resuming train dispatch"
+                    )
+                # Demand only gates *opening* fresh groups. A group schedules its
+                # group_size rollouts one at a time, and the sink can only
+                # finalize it once all of them arrive — so cutting a group off
+                # mid-schedule would leave it unable to ever finalize while its
+                # arrived members still count as batch coverage, wedging the
+                # batch a few rollouts short with dispatch paused.
+                scheduled = await self.try_schedule("train", allow_fresh=needed)
                 if not scheduled:
                     return
 
@@ -340,11 +346,11 @@ def switch_mode(self, new_mode: DispatcherMode, *, reason: str) -> None:
         get_logger().info(f"Switching dispatcher mode to prefer {prefer} rollouts because {reason}")
         self.mode = new_mode
 
-    async def try_schedule(self, kind: RolloutKind) -> bool:
+    async def try_schedule(self, kind: RolloutKind, *, allow_fresh: bool = True) -> bool:
         """Schedule one rollout of ``kind``: prefer continuing an existing
-        group (keeps prefix-cache hits); otherwise open a fresh group from
-        the corresponding source. Returns False if nothing could be
-        scheduled."""
+        group (keeps prefix-cache hits); otherwise, when ``allow_fresh``, open
+        a fresh group from the corresponding source. Returns False if nothing
+        could be scheduled."""
         if kind == "train" and self.train_scheduling_disabled:
             return False
         envs = self.train_envs if kind == "train" else self.eval_envs
@@ -359,6 +365,8 @@ async def try_schedule(self, kind: RolloutKind) -> bool:
             if cost <= self.available_permits:
                 return await self.schedule_group_rollout(gid, group)
 
+        if not allow_fresh:
+            return False
         fresh = self.next_fresh_group(kind, envs)
         if fresh is None:
             return False
@@ -558,7 +566,6 @@ async def emit_episode(self, meta: InflightRollout, group: GroupState | None, ro
             rollout.env_name = meta.env_name
             rollout.group_id = meta.group_id
             rollout.policy_version = policy_version
-            rollout.off_policy_steps = meta.off_policy_steps
             if meta.kind == "eval":
                 assert eval_step is not None, "eval rollout missing eval_step"
                 rollout.eval_step = eval_step
@@ -688,12 +695,18 @@ async def cancel_inflight_train_rollouts(self) -> int:
 
     def gauges(self) -> dict[str, float]:
         """Instantaneous, read-only gauges sampled by the periodic logger."""
+        # Versions each in-flight live-sampled train rollout is behind the policy
+        lags = [
+            self.policy.version - m.policy_version
+            for m in self.inflight.values()
+            if m.kind == "train" and self.train_envs.get(m.env_name).sampler.samples_from_live_policy
+        ]
         return {
             "dispatcher/inflight_train": float(self.inflight_train_count),
             "dispatcher/inflight_eval": float(self.inflight_eval_count),
             "dispatcher/queued/eval": float(self.queued_eval_examples),
             "dispatcher/mode": float(self.mode == DispatcherMode.PREFER_EVAL),
             "dispatcher/groups_in_flight": float(len(self.groups)),
-            "dispatcher/off_policy_level_max": float(self.max_off_policy_level),
-            "dispatcher/off_policy_level_mean": self.mean_off_policy_level,
+            "dispatcher/off_policy_level_max": float(max(lags) if lags else 0),
+            "dispatcher/off_policy_level_mean": sum(lags) / len(lags) if lags else 0.0,
         }