fix(orchestrator): coordinate final NCCL shutdown

src/prime_rl/orchestrator/orchestrator.py

@@ -470,6 +470,17 @@ async def main_loop(self) -> None:
         ``None``) from ``add()``; we just dispatch on the result."""
         while not self.stopped.is_set():
             if self.draining and self.dispatcher.is_idle:
+                # Keep the NCCL receiver alive through the trainer's last
+                # required broadcast. The final two policy versions are never
+                # consumed by inference and are intentionally not broadcast.
+                final_nccl_policy_version = (
+                    max(self.config.max_steps - 2, 0)
+                    if self.config.weight_broadcast.type == "nccl" and self.config.max_steps is not None
+                    else self.policy.version
+                )
+                if self.policy.version < final_nccl_policy_version:
+                    await asyncio.sleep(0.5)
+                    continue
                 get_logger().info("Pipeline drained, exiting main loop")
                 self.stopped.set()
                 break
@@ -634,6 +645,17 @@ async def finalize_train_batch(self, batch: TrainBatch) -> None:
 
         self.train_sink.reset_pre_filter_stats()
         self.maybe_trigger_eval(self.progress.step)
+
+        # The final real batch has been shipped. Enter drain mode here instead
+        # of relying on a synthetic batch beyond max_steps to trigger shutdown.
+        if config.max_steps is not None and step >= config.max_steps:
+            self.draining = True
+            self.dispatcher.disable_train_scheduling()
+            n_cancelled = await self.dispatcher.cancel_inflight_train_rollouts()
+            get_logger().info(
+                f"Draining pipeline (cancelled {n_cancelled} in-flight train rollout(s); "
+                f"any in-flight evals will complete)"
+            )
         trim_process_memory()
 
     def maybe_trigger_eval(self, step: int) -> None:
@@ -828,18 +850,9 @@ def update_dispatch_gate(self) -> None:
         are 1-indexed while policy versions stay 0-indexed, so the shipped-batch
         count is ``progress.step - 1``."""
         lead = (self.progress.step - 1) - self.policy.version
-        # The trainer skips the final NCCL weight broadcast (inference group is
-        # torn down), so policy.version never reaches the last step. Without this
-        # the gate deadlocks waiting for a version that will never be published.
-        # The last batch uses the penultimate policy anyway, so let it through.
-        building_final_batch_nccl = (
-            self.config.weight_broadcast.type == "nccl"
-            and self.config.max_steps is not None
-            and self.progress.step >= self.config.max_steps - 1
-        )
         gate = self.dispatcher.dispatch_allowed
         was_set = gate.is_set()
-        if lead > TARGET_LAG and not building_final_batch_nccl:
+        if lead > TARGET_LAG:
             if was_set:
                 get_logger().info(
                     "Pausing dispatcher to prevent orchestrator from racing from trainer. Waiting for new policy..."

tests/unit/orchestrator/test_orchestrator_shutdown.py

@@ -0,0 +1,68 @@
+import asyncio
+from types import SimpleNamespace
+
+from prime_rl.orchestrator.orchestrator import Orchestrator
+
+
+def make_draining_orchestrator(*, broadcast_type: str, max_steps: int | None, policy_version: int) -> Orchestrator:
+    orchestrator = Orchestrator.__new__(Orchestrator)
+    orchestrator.config = SimpleNamespace(
+        max_steps=max_steps,
+        weight_broadcast=SimpleNamespace(type=broadcast_type),
+    )
+    orchestrator.policy = SimpleNamespace(version=policy_version)
+    orchestrator.dispatcher = SimpleNamespace(is_idle=True)
+    orchestrator.draining = True
+    orchestrator.stopped = asyncio.Event()
+    return orchestrator
+
+
+def test_nccl_drain_waits_for_last_required_broadcast():
+    async def run() -> None:
+        orchestrator = make_draining_orchestrator(
+            broadcast_type="nccl",
+            max_steps=20,
+            policy_version=17,
+        )
+
+        main_loop = asyncio.create_task(orchestrator.main_loop())
+        await asyncio.sleep(0.05)
+        assert not main_loop.done()
+        assert not orchestrator.stopped.is_set()
+
+        orchestrator.policy.version = 18
+        await asyncio.wait_for(main_loop, timeout=1.0)
+        assert orchestrator.stopped.is_set()
+
+    asyncio.run(run())
+
+
+def test_non_nccl_drain_exits_without_waiting_for_policy_update():
+    async def run() -> None:
+        orchestrator = make_draining_orchestrator(
+            broadcast_type="filesystem",
+            max_steps=20,
+            policy_version=0,
+        )
+
+        await asyncio.wait_for(orchestrator.main_loop(), timeout=0.1)
+        assert orchestrator.stopped.is_set()
+
+    asyncio.run(run())
+
+
+def test_dispatch_gate_remains_active_for_final_batches():
+    orchestrator = Orchestrator.__new__(Orchestrator)
+    orchestrator.progress = SimpleNamespace(step=20)
+    orchestrator.policy = SimpleNamespace(version=17)
+    orchestrator.dispatcher = SimpleNamespace(dispatch_allowed=asyncio.Event())
+    orchestrator.dispatcher.dispatch_allowed.set()
+    orchestrator.gate_closed_at = None
+    orchestrator.wait_for_policy_time = 0.0
+
+    orchestrator.update_dispatch_gate()
+    assert not orchestrator.dispatcher.dispatch_allowed.is_set()
+
+    orchestrator.policy.version = 18
+    orchestrator.update_dispatch_gate()
+    assert orchestrator.dispatcher.dispatch_allowed.is_set()