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

docs/mint.json

@@ -9,10 +9,11 @@
                 "training",
                 "inference",
                 "scaling",
+                "staleness",
                 "algorithms",
                 "advanced",
                 "development"
             ]
         }
     ]
-}
+}

docs/staleness.md

@@ -0,0 +1,80 @@
+# Staleness and pacing
+
+How the async RL pipeline bounds off-policy staleness, when dispatch starts and
+stops, and why shutdown is race-free. Two counters govern everything:
+
+- `N` — the batch the orchestrator is collecting (`progress.step`, 1-indexed).
+- `v` — the trainer's policy version (0-indexed). `v` advances when the trainer
+  *publishes* a version; inference applies it immediately and pauses generation
+  during the update, so published and applied versions are indistinguishable to
+  every consumer.
+
+A rollout is stamped with the version it was generated from (`v_gen`). If it is
+consumed in batch `N`, it trains `lag = (N-1) - v_gen` versions behind — batch
+`N` produces `v{N}` from `v{N-1}`, so `lag = 0` is fully on-policy.
+
+## The three rules
+
+Each rule bounds `N - v` (or a rollout's `lag`) at one lifecycle point. They are
+independent: each one prevents a failure mode the others cannot.
+
+| Rule | Bound | Enforced at | Failure mode it prevents |
+|---|---|---|---|
+| **START** | dispatch pauses while `(N-1) - v > TARGET_LAG` | dispatch gate | generating data that is born stale |
+| **ADVANCE** | batch `N` ships only once `v ≥ N-1-TARGET_LAG` | ship hold | the batch counter outrunning the trainer |
+| **DIE** | rollouts that straddle more than `max_off_policy_steps` weight updates in flight are cancelled | weight-update hook | long-running stragglers aging without bound |
+
+With `TARGET_LAG = 1`, START and ADVANCE together keep consumption lag at ≤ 2
+in steady state (≤ 3 with batch-boundary spill); DIE only fires on rollouts
+whose own duration spans many trainer steps (agentic long tails).
+
+## The two worlds
+
+Which rule binds depends on the ratio of trainer step time `T` to per-batch
+generation time `G`. The rules never change — only which one saturates:
+
+**World 1 — generation-bound (`T < G`, the trainer waits).** `v` catches up
+after every ship, so `(N-1) - v ≤ 1` always: START and ADVANCE never bind and
+DIE stays quiet. The pipeline is paced by generation; the trainer idles between
+batches. All staleness rules are inert insurance.
+
+**World 2 — trainer-bound (`T > G`, inference races).** Batches fill faster
+than versions arrive. ADVANCE binds every step (the ship hold paces `N` to
+`v`); START closes dispatch between version publishes (inference alternates
+generate/idle, idle fraction `1 - G/T`). Without ADVANCE this world is where
+the orchestrator finishes all its batches from buffered rollouts, exits, and
+strands the trainer (see below).
+
+The crossover is smooth: the same predicates are evaluated everywhere, and no
+configuration switch distinguishes the worlds.
+
+## Shutdown correctness
+
+For in-memory weight transports (NCCL, NIXL) the trainer *blocks* inside each
+broadcast until the orchestrator's weight watcher completes the matching apply.
+Two guarantees make teardown race-free:
+
+1. **ADVANCE is the liveness guarantee.** The orchestrator cannot exit before
+   shipping its final batch `M`, which requires `v ≥ M-2` — and `v{M-2}` is
+   exactly the last version whose broadcast needs a live watcher (the trainer
+   skips the final `TARGET_LAG + 1` in-memory broadcasts because their receiver
+   is torn down; this is also why the hold never waits for an unpublishable
+   version).
+2. **The watcher drains before dying.** `WeightWatcher.stop()` waits for an
+   in-flight apply to complete before cancelling, so a shutdown that races the
+   last handshake finishes it instead of stranding the trainer. The
+   orchestrator's global teardown budget bounds this wait.
+
+No other component participates in shutdown ordering: the trainer only ever
+waits on the watcher, and the watcher outlives every wait.
+
+## Who starts and stops dispatch
+
+Nobody signals — dispatch is *pulled*. The dispatcher re-checks the gate before
+each scheduling decision, so stop/resume is a stateless predicate over
+`(N, v)`, re-evaluated on its own loop. The single push-style signal in the
+system is the version-advance event that wakes a held ship. Throughput is
+identical to an unpaced orchestrator in both worlds: in World 1 no rule ever
+binds, and in World 2 the trainer is the bottleneck regardless — the hold only
+changes *where* a finished batch waits (in the orchestrator instead of the
+trainer's queue), never when the trainer gets to consume it.

src/prime_rl/orchestrator/orchestrator.py

@@ -176,6 +176,8 @@ def __init__(self, config: OrchestratorConfig) -> None:
         self.eval_triggered_at = {}
         self.consecutive_empty_batches = 0
         self.gate_closed_at = None
+        # Pulsed by the version hooks so a held ship can re-check ``policy.version``
+        self.version_advanced = asyncio.Event()
         self.wait_for_policy_time = 0.0
         self.component_tasks = []
 
@@ -601,6 +603,31 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
                 f"({n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config"
             )
 
+        # ADVANCE rule (docs/staleness.md): batch ``step`` ships only once the
+        # trainer has published policy v{step-1-TARGET_LAG}. On fast envs the sink
+        # otherwise fills batches from buffered rollouts faster than the trainer
+        # consumes them; the orchestrator would finish all its steps, tear down the
+        # weight watcher, and strand the trainer inside its next in-memory broadcast
+        # handshake (the watcher serves the receive side). Holding here paces the
+        # batch counter to the trainer and keeps the watcher alive for every
+        # broadcast the trainer will make. The requirement is always satisfiable:
+        # the trainer skips only the final TARGET_LAG+1 in-memory broadcasts, so
+        # v{max_steps-1-TARGET_LAG} is the last one published — exactly what the
+        # final batch needs. Bench runs have no trainer, so they ship freely.
+        required_version = step - 1 - TARGET_LAG
+        if not config.bench and self.policy.version < required_version:
+            get_logger().info(
+                f"Holding batch {step} until the trainer publishes policy v{required_version} "
+                f"(currently v{self.policy.version})"
+            )
+            hold_start = time.perf_counter()
+            while True:
+                self.version_advanced.clear()
+                if self.policy.version >= required_version:
+                    break
+                await self.version_advanced.wait()
+            self.wait_for_policy_time += time.perf_counter() - hold_start
+
         # The effective (clean, trained-on) subset lands in the per-step ``effective`` trace file
         # at ship time; the full arrival window already streamed into ``all`` on arrival.
         # to_record drops the per-node training tensors — they're for training, not the rollout
@@ -904,11 +931,14 @@ def update_dispatch_gate(self) -> None:
             gate.set()
 
     async def on_version_pending(self, step: int) -> None:
+        """``VersionObserver`` hook, fired at publish confirmation (pre-apply):
+        ``policy.version`` already carries the new version, so wake a held ship."""
         if self.model_express is not None:
             await asyncio.to_thread(self.model_express.set_status, p2p_pb2.SOURCE_STATUS_READY)
+        self.version_advanced.set()
 
     async def on_new_version(self, step: int) -> None:
-        """``VersionObserver`` hook: the watcher just advanced ``policy.version``;
+        """``VersionObserver`` hook: the weight update completed;
         re-evaluate the dispatch gate (may resume if the trainer caught up)."""
         if self.model_express is not None:
             await asyncio.to_thread(self.model_express.set_status, p2p_pb2.SOURCE_STATUS_INITIALIZING)

src/prime_rl/orchestrator/watcher.py

@@ -64,6 +64,12 @@ async def start(self) -> None:
 
     async def stop(self) -> None:
         self.stopped.set()
+        # Let an in-flight apply finish before dying: the trainer blocks inside
+        # its in-memory broadcast until the apply completes, so cancelling
+        # mid-apply would strand it. The orchestrator's global teardown budget
+        # bounds this wait.
+        async with self.update_lock:
+            pass
         if self.task is not None:
             await safe_cancel(self.task)
             self.task = None
@@ -103,6 +109,14 @@ async def apply_policy_update(self, next_step: int) -> None:
                     await wait_for_path(stable_marker)
             self.last_wait_for_ckpt_time = time.perf_counter() - t0
 
+            # Publish confirmed: the policy version advances here, before the
+            # apply — inference pauses during the update, so nothing can generate
+            # under the new number from the old weights, and a ship held on this
+            # version (see docs/staleness.md) releases without waiting out the
+            # inference weight reload.
+            self.ckpt_step = next_step
+            self.policy.version = next_step
+
             # Drain off-policy rollouts BEFORE pausing the inference engines.
             # Aborting a rollout triggers vLLM's KV-connector cleanup (NIXL's
             # ``_reqs_not_processed``), which is only propagated to the workers
@@ -128,8 +142,6 @@ async def apply_policy_update(self, next_step: int) -> None:
             self.update_count += 1
             get_logger().debug(f"Updated weights to step {next_step} in {format_time(self.last_update_weights_time)}")
 
-            self.ckpt_step = next_step
-            self.policy.version = next_step
             if self.lora_name is not None:
                 self.inference.update_model_name(self.lora_name)
                 self.policy.model_name = self.lora_name

src/prime_rl/trainer/rl/token_export.py

@@ -10,6 +10,7 @@
 
 from prime_rl.configs.trainer import DefaultLossConfig, TrainerConfig
 from prime_rl.trainer.rl.loss import compute_importance_ratio_and_mismatch_kl
+from prime_rl.utils.logger import get_logger
 
 SCHEMA_VERSION = 1
 
@@ -93,7 +94,12 @@ def mark_stable(self, ready_run_ids: set[str] | None = None) -> None:
         ready_run_ids = ready_run_ids or set()
         for run_id in [rid for rid in self._pending_stable_dirs if rid is None or rid in ready_run_ids]:
             for stable_dir in self._pending_stable_dirs.pop(run_id):
-                (stable_dir / "STABLE").touch()
+                try:
+                    (stable_dir / "STABLE").touch()
+                except FileNotFoundError:
+                    # A multi-run run dir can be deleted while its tail steps are
+                    # still flushing (mirrors the broadcast paths' guard)
+                    get_logger().warning(f"Run dir for {run_id} is deleted, skipping token-export STABLE marker")
 
     def _export_dir(self, export_step: int, run_id: str | None) -> Path:
         if run_id is not None:

tests/conftest.py

@@ -78,7 +78,10 @@ def output_dir(tmp_path_factory: pytest.TempPathFactory) -> Generator[Path, None
     output_dir = Path(os.environ.get("PYTEST_OUTPUT_DIR", tmp_path_factory.mktemp("outputs")))
     output_dir.mkdir(parents=True, exist_ok=True)
     yield output_dir
-    shutil.rmtree(output_dir, ignore_errors=True)
+    # In CI the workspace is discarded after the job; keep the outputs so the
+    # failure-artifact upload step can collect the run logs.
+    if os.environ.get("USERNAME_CI") != "CI_RUNNER":
+        shutil.rmtree(output_dir, ignore_errors=True)
 
 
 @pytest.fixture(scope="session")

tests/integration/test_reverse_text_multi_run.py

@@ -20,6 +20,17 @@
 ORCHESTRATOR_NAMES = ["alpha", "beta", "gamma"]
 
 
+def remove_run_dir(run_dir: Path) -> None:
+    """Delete a run dir, tolerating the shared trainer still flushing its tail
+    (token exports recreate directories mid-rmtree -> 'Directory not empty')."""
+    for _ in range(10):
+        shutil.rmtree(run_dir, ignore_errors=True)
+        if not run_dir.exists():
+            return
+        time.sleep(1)
+    shutil.rmtree(run_dir)
+
+
 def wait_for_file(
     file_path: Path,
     timeout: int = 600,
@@ -305,7 +316,7 @@ def multi_run_result(
     shutil.copy(run_dir / "logs" / "orchestrator.log", log_dir / "beta_orchestrator.log")
     shutil.copytree(beta_ckpt_dir, tmp_path / "beta_ckpt_step_20")
     print(f"Copied {beta_ckpt_dir} to {tmp_path / 'beta_ckpt_step_20'}")
-    shutil.rmtree(run_dir)
+    remove_run_dir(run_dir)
 
     # =====================
     # Queue beta's resume
@@ -329,7 +340,7 @@ def multi_run_result(
         timeout=TIMEOUT,
     )
     shutil.copy(output_dir / "run_gamma" / "logs" / "orchestrator.log", log_dir / "gamma_orchestrator.log")
-    shutil.rmtree(output_dir / "run_gamma")
+    remove_run_dir(output_dir / "run_gamma")
 
     # ================================================
     # Wait for alpha_resume and beta_resume to finish