fix: align evals with applied policies

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

@@ -471,9 +471,6 @@ def auto_setup_weight_broadcast(self):
                 "PEFT-shaped directory on disk (LoRAModel.from_local_checkpoint) - in-memory transports "
                 "have no disk artifact to load from."
             )
-        # The final version v{max_steps} is broadcast iff something consumes it:
-        # training never samples from it, but a configured final eval measures it.
-        broadcast_final = self.orchestrator.eval is not None
         if self.weight_broadcast.type in ("nccl", "nixl"):
             inference_world_size = (
                 self.inference.vllm.data_parallel_size * self.inference.vllm.tensor_parallel_size
@@ -485,7 +482,6 @@ def auto_setup_weight_broadcast(self):
                 port=self.weight_broadcast.port,
                 timeout=self.weight_broadcast.timeout,
                 inference_world_size=inference_world_size,
-                broadcast_final=broadcast_final,
             )
             if self.weight_broadcast.type == "nccl":
                 transport_config = dict(
@@ -501,10 +497,10 @@ def auto_setup_weight_broadcast(self):
             self.orchestrator.weight_broadcast = orchestrator_config_type(**common_config, **transport_config)
         elif self.weight_broadcast.type == "filesystem":
             self.trainer.weight_broadcast = TrainerFileSystemWeightBroadcastConfig(
-                timeout=self.weight_broadcast.timeout, broadcast_final=broadcast_final
+                timeout=self.weight_broadcast.timeout
             )
             self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig(
-                timeout=self.weight_broadcast.timeout, broadcast_final=broadcast_final
+                timeout=self.weight_broadcast.timeout
             )
         if self.inference is not None:
             self.inference.weight_broadcast = InferenceWeightBroadcastConfig(type=self.weight_broadcast.type)

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

@@ -42,12 +42,6 @@ class BaseWeightBroadcastConfig(BaseConfig):
     """Timeout in seconds for the broadcast handshake and transfer. The trainer
     fails the run when no consumer acknowledges an offered version in time."""
 
-    broadcast_final: bool = True
-    """Internal - stamped by the RL/SFT launcher, never set by users: whether the
-    trainer broadcasts the final version v{max_steps}. True iff something consumes
-    it (a configured final eval) - training itself never samples from the final
-    version."""
-
 
 class RunConfig(BaseConfig):
     name: str | None = None

src/prime_rl/orchestrator/dispatcher.py

@@ -210,6 +210,8 @@ def __init__(
         # *why* — the orchestrator toggles this based on step / policy lead.
         self.dispatch_allowed = asyncio.Event()
         self.dispatch_allowed.set()
+        self.policy_update_pending = False
+        self.scheduling_lock = asyncio.Lock()
 
         self.stopped = asyncio.Event()
         self.task: asyncio.Task | None = None
@@ -371,6 +373,12 @@ async def on_version_pending(self, step: int) -> None:
         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``)."""
+        self.policy_update_pending = True
+        # Wait for a scheduling call that started before the pending update.
+        # No rollout can cross the inference weight swap after this barrier.
+        async with self.scheduling_lock:
+            pass
+
         if self.train_envs is None or self.progress is None:
             return
         min_version = min_fresh_version(self.progress.step, self.max_off_policy_steps)
@@ -392,7 +400,8 @@ async def on_version_pending(self, step: int) -> None:
             )
 
     async def on_new_version(self, step: int) -> None:
-        """No-op: the dispatcher drains in ``on_version_pending`` (pre-pause)."""
+        """Resume rollout scheduling after inference applies the new policy."""
+        self.policy_update_pending = False
 
     async def fill_inflight(self) -> None:
         """Schedule new rollouts up to ``max_inflight``, honoring
@@ -401,28 +410,33 @@ async def fill_inflight(self) -> None:
         respects it. When ``PREFER_EVAL``'s source exhausts we flip back to
         ``PREFER_TRAIN`` so the eval tail drains alongside fresh train."""
         while True:
+            if self.policy_update_pending:
+                return
             if self.available_permits <= 0 or self.admission_budget() <= 0:
                 return
 
-            if self.mode == DispatcherMode.PREFER_EVAL:
-                # PREFER_EVAL is only entered when the orchestrator triggers
-                # eval, which requires ``eval_source`` to be configured
-                assert self.eval_source is not None
-                if not self.eval_has_work:
-                    # Eval source + all eval groups fully dispatched. Flip
-                    # to PREFER_TRAIN so any remaining permits go to train
-                    # while the in-flight eval tail completes naturally
-                    self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained")
-                    continue
-                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")
-                if not scheduled:
+            async with self.scheduling_lock:
+                if self.policy_update_pending:
                     return
+                if self.mode == DispatcherMode.PREFER_EVAL:
+                    # PREFER_EVAL is only entered when the orchestrator triggers
+                    # eval, which requires ``eval_source`` to be configured
+                    assert self.eval_source is not None
+                    if not self.eval_has_work:
+                        # Eval source + all eval groups fully dispatched. Flip
+                        # to PREFER_TRAIN so any remaining permits go to train
+                        # while the in-flight eval tail completes naturally
+                        self.switch_mode(DispatcherMode.PREFER_TRAIN, reason="the eval queue drained")
+                        continue
+                    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")
+                    if not scheduled:
+                        return
 
     def switch_mode(self, new_mode: DispatcherMode, *, reason: str) -> None:
         if new_mode == self.mode:

src/prime_rl/orchestrator/eval_source.py

@@ -1,7 +1,7 @@
 """EvalSource: trigger-driven, finite-per-epoch pull of eval examples.
 
-The orchestrator pokes ``trigger(step)`` after each ship + once at
-startup; the dispatcher pulls via ``next_task()`` until
+The policy watcher calls ``trigger(step)`` after each applied policy,
+including startup. The dispatcher pulls via ``next_task()`` until
 ``bool(source) == False``. Constructed only when eval is configured."""
 
 from __future__ import annotations
@@ -37,8 +37,8 @@ def __init__(
 
         self.queue: deque[TaskRequest] = deque()
 
-        # On resume we skip the startup eval; on fresh start the first
-        # trigger fires every env (subject to ``skip_first_step``)
+        # A fresh run evaluates the base policy. Resumed runs apply interval
+        # rules to the loaded checkpoint and later policies.
         self.first_trigger = not is_resumed
 
     def trigger(self, step: int, *, force: bool = False) -> list[str]: