diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 5c1bedf5ee..e957e369f4 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/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..." diff --git a/tests/unit/orchestrator/test_orchestrator_shutdown.py b/tests/unit/orchestrator/test_orchestrator_shutdown.py new file mode 100644 index 0000000000..940394ed50 --- /dev/null +++ b/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()