cancel all rollout eval (#1671)

Author: samsja

src/prime_rl/orchestrator/orchestrator.py

@@ -308,8 +308,10 @@ async def orchestrate(config: OrchestratorConfig):
             last_eval_step = ckpt_step
             logger.info(f"Running evals for checkpoint step {ckpt_step} (blocking, subprocess)")
 
-            # Pause weight updates during eval
+            # Pause weight updates during eval and cancel inflight rollouts
+            # this avoid doing eval across different checkpoints and avoid congestion
             scheduler.checkpoint_ready.clear()
+            scheduler.cancel_all_inflight_rollouts()
 
             await run_evals_subprocess(
                 client_config=config.client,

src/prime_rl/orchestrator/scheduler.py

@@ -152,6 +152,18 @@ async def stop(self):
             for worker in workers:
                 worker.stop()
 
+    def cancel_all_inflight_rollouts(self):
+        """Cancel all in-flight rollout requests.
+
+        Used when weights are updated to discard stale rollouts generated with old weights.
+        """
+        count = len(self.inflight_group_rollouts)
+        for future in list(self.inflight_group_rollouts.keys()):
+            if not future.done():
+                future.cancel()
+        self.inflight_group_rollouts.clear()
+        self.cancelled_rollouts_count += count
+
     async def schedule_group_rollout(self):
         """Asynchronously schedules a group rollout request."""
         example = self.buffer.sample_examples(n=1)[0]