fix(train): hand off to the engine before an off-cadence final eval

tests/fast/test_train.py

@@ -123,6 +123,29 @@ async def test_the_last_rollout_is_always_evaluated_even_off_cadence(self, monke
 
         assert [event for event in events if event.startswith("eval:")] == ["eval:0", "eval:1", "eval:2"]
 
+    async def test_an_off_cadence_final_eval_follows_the_last_handoff(self, monkeypatch: pytest.MonkeyPatch):
+        """The last rollout skips its handoff only when no final eval follows; an off-cadence
+        final eval must still see the engine onloaded with the trained weights."""
+        events: list[str] = []
+        args = _make_args(num_rollout=3, eval_interval=2, offload_rollout=True)
+        _install_driver_fakes(monkeypatch, args, events)
+
+        await train_driver.train(args)
+
+        final = events[events.index("actor_train:2") : events.index("eval:2")]
+        assert final.index("onload_weights") < final.index("update_weights:2") < final.index("onload_kv")
+
+    async def test_the_last_rollout_skips_the_handoff_without_an_eval(self, monkeypatch: pytest.MonkeyPatch):
+        events: list[str] = []
+        args = _make_args(num_rollout=3, offload_rollout=True)
+        _install_driver_fakes(monkeypatch, args, events)
+
+        await train_driver.train(args)
+
+        assert "update_weights:1" in events
+        assert "update_weights:2" not in events
+        assert "onload_weights" not in events[events.index("actor_train:2") :]
+
 
 class TestWeightEqualityCheck:
     async def test_weight_equality_check_is_routed_to_the_inference_controller(self, monkeypatch: pytest.MonkeyPatch):

train.py

@@ -146,9 +146,9 @@ async def save_training_model(model):
             if external_save:
                 os.remove(args.save_trigger_sentinel)
 
-        if rollout_id + 1 < args.num_rollout or should_run_periodic_action(
-            rollout_id, args.eval_interval, num_rollout_per_epoch
-        ):
+        # One predicate for both blocks: the handoff below exists to feed this eval on the last rollout.
+        eval_due = should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch, args.num_rollout)
+        if rollout_id + 1 < args.num_rollout or eval_due:
             if args.colocate_memory_peak_device == "gpu":
                 await actor_model.clear_memory()
                 if lora_rollout_enabled(args):
@@ -163,7 +163,7 @@ async def save_training_model(model):
             if args.offload_rollout:
                 await inference_controller.onload_kv()
 
-        if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch, args.num_rollout):
+        if eval_due:
             await inference_controller.prepare_eval()
             await eval_dispatcher.dispatch(rollout_id, force=rollout_id == args.num_rollout - 1)