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

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,35 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
                 f"({n_trainable / len(batch.rollouts):.1%}) — consider reviewing task difficulty / filter config"
             )
 
+        # Ship batch ``step`` only once the trainer has published v{step-1-TARGET_LAG}.
+        # Without this, fast envs fill batches from buffered rollouts, the
+        # orchestrator finishes early, and its teardown strands the trainer inside
+        # an in-memory broadcast handshake that needs the live weight watcher.
+        # Always satisfiable: the trainer skips only the final TARGET_LAG+1
+        # in-memory broadcasts. 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
+
+        # Stamp each rollout's true staleness: batch ``step`` trains on policy
+        # v{step-1}, so a rollout generated from v{k} is (step-1)-k versions
+        # off-policy — queue time included, unlike the dispatcher's in-flight
+        # counter, which only sees weight updates during generation. Frozen-
+        # sourced rollouts stay 0 (their sampler doesn't follow the policy).
+        for r in batch.rollouts:
+            if self.train_envs.get(r.env_name).sampler.samples_from_live_policy:
+                r.off_policy_steps = (step - 1) - r.policy_version
+
         # 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 +935,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,13 @@ 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 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 +141,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

@@ -1,5 +1,4 @@
 import os
-import shutil
 import signal
 import socket
 import subprocess
@@ -74,11 +73,11 @@ def branch_name() -> str:
 
 @pytest.fixture(scope="session")
 def output_dir(tmp_path_factory: pytest.TempPathFactory) -> Generator[Path, None, None]:
-    """Fixture for temporary output directory for tests with automatic cleanup"""
+    """Fixture for the tests' output directory. Never deleted — CI uploads run logs
+    from it on failure; runs start with ``--clean-output-dir`` for a fresh workspace."""
     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)
 
 
 @pytest.fixture(scope="session")

tests/integration/test_alphabet_sort.py

@@ -38,6 +38,7 @@ def rl_process(
         "rl",
         "@",
         "configs/ci/integration/alphabet_sort.toml",
+        "--clean-output-dir",
         "--wandb.project",
         wandb_project,
         "--wandb.name",

tests/integration/test_reverse_text_moe.py

@@ -29,6 +29,7 @@ def rl_process(
         "rl",
         "@",
         "configs/ci/integration/reverse-text-moe/start.toml",
+        "--clean-output-dir",
         "--trainer.model.impl",
         "custom",
         "--wandb.project",

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