Delete the pre-reconcile engine recovery

miles/ray/actor_group.py

@@ -7,6 +7,7 @@
 
 from ray.util.placement_group import PlacementGroup
 
+from miles.ray.rollout.inference_controller import update_weights_window
 from miles.ray.train.actor_factory import allocate_gpus_for_actor
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
 
@@ -123,13 +124,8 @@ async def update_weights(self, rollout_id: int | None = None):
         if self.args.debug_train_only or self.args.debug_rollout_only:
             return
 
-        if self.args.use_fault_tolerance and "rollout" in self.args.ft_components:
-            await self._inference_controller.recover_updatable_engines()
-
-        info = await self._inference_controller.start_update_weights()
-
-        await self._broadcast("update_weights", info=info)
-        await self._inference_controller.end_update_weights(snapshot_cell_id_to_hashes=info.snapshot_cell_id_to_hashes)
+        async with update_weights_window(self._inference_controller) as info:
+            await self._broadcast("update_weights", info=info)
 
     async def reconcile_adapters(self) -> None:
         """Multi-LoRA: reconcile loaded adapters with the controller's active set

miles/ray/rollout/inference_controller.py

@@ -42,7 +42,6 @@ def __init__(self, args):
         self.args = args
         self.context_lock = ContextLock("InferenceController")
         self.servers: dict[str, RolloutServer] = {}
-        self.rollout_id = -1
         self.eval_fleet: EvalFleet | None = None
         self._watcher_disposers: list[StopWatchFn] = []
         self._health_checker_activeness = ActivenessTracker(active=True)
@@ -77,7 +76,6 @@ async def init(self) -> None:
 
     @with_lock
     async def prepare_rollout(self, rollout_id):
-        self.rollout_id = rollout_id
         await self._health_monitoring_resume()
         if self.args.ci_test and self._rollout_ft_enabled and rollout_id >= 2:
             await self._try_ci_fault_injection()
@@ -201,10 +199,6 @@ async def _ensure_cells_ready(self) -> None:
             async with self.context_lock.with_released():
                 await asyncio.sleep(CELLS_READY_POLL_INTERVAL_SECONDS)
 
-    @with_lock
-    async def recover_updatable_engines(self) -> None:
-        raise NotImplementedError("new ft to be implemented")
-
     @requires_lock
     def _get_updatable_server(self) -> RolloutServer | None:
         updatable = [srv for srv in self.servers.values() if srv.update_weights]

tests/fast/ray/rollout/real_ray/test_inference_controller.py

@@ -430,65 +430,6 @@ async def test_check_weights_targets_only_updatable_model(
             assert "/weights_checker" not in paths, f"frozen engine {cell.addr_info.server_url} must not be checked"
 
 
-@pytest.mark.asyncio
-class TestRecoverUpdatableEngines:
-    async def test_skips_recovery_when_no_rollout_started(
-        self,
-        ray_local_mode,
-        placement_group_factory,
-        tmp_path,
-        patch_low_level,
-    ):
-        """``recover_updatable_engines`` is a no-op while ``rollout_id == -1``
-        (initial state) — the trainer hasn't issued a rollout yet, so even if
-        a slot looks dead the controller must not pre-emptively recover."""
-        args = _make_test_args(tmp_path, models=[("actor", True)])
-        pg = placement_group_factory(2)
-
-        controller = InferenceController(args, pg)
-        await controller.init()
-        await controller.get_updatable_engines()
-        actor0_before = _cells(controller)[0].primary_actor_handle
-
-        # Kill engine 0 directly + mark stopped (simulates a fault before any
-        # rollout). recover_updatable_engines must not bring it back yet.
-        ray.kill(actor0_before)
-        controller.servers["actor"].server_cells["actor-0"]._mark_stopped()
-
-        await controller.recover_updatable_engines()
-
-        # Slot 0 is still de-allocated; recovery skipped because rollout_id=-1.
-        assert not controller.servers["actor"].server_cells["actor-0"].is_allocated
-
-    async def test_recovers_dead_engine_after_rollout_started(
-        self,
-        ray_local_mode,
-        placement_group_factory,
-        tmp_path,
-        patch_low_level,
-    ):
-        """Once ``rollout_id`` advances past -1 (mid-training), a dead slot on
-        the updatable server is brought back by ``recover_updatable_engines``."""
-        args = _make_test_args(tmp_path, models=[("actor", True)])
-        pg = placement_group_factory(2)
-
-        controller = InferenceController(args, pg)
-        await controller.init()
-        await controller.get_updatable_engines()
-        actor0_before = _cells(controller)[0].primary_actor_handle
-
-        ray.kill(actor0_before)
-        controller.servers["actor"].server_cells["actor-0"]._mark_stopped()
-
-        await controller.prepare_rollout(0)
-        await controller.recover_updatable_engines()
-
-        slot0 = controller.servers["actor"].server_cells["actor-0"]
-        assert slot0.is_allocated
-        assert slot0.primary_actor_handle is not actor0_before
-        assert isinstance(ray.get(slot0.primary_actor_handle.get_calls.remote()), list)
-
-
 @pytest.mark.asyncio
 class TestRolloutFaultToleranceIsUnsupported:
     async def test_fault_injection_is_skipped_when_fault_tolerance_skips_rollout(

tests/fast/ray/test_actor_group_shared_ppo.py

@@ -17,17 +17,6 @@ def __init__(self, rank, calls):
         self.train = _RemoteTrain(rank, calls)
 
 
-class _AsyncCall:
-    def __init__(self, name, calls, result=None):
-        self.name = name
-        self.calls = calls
-        self.result = result
-
-    async def __call__(self, *args, **kwargs):
-        self.calls.append((self.name, args, kwargs))
-        return self.result
-
-
 async def test_train_routes_each_critic_payload_to_matching_actor_rank():
     from miles.ray.actor_group import RayTrainGroup
 
@@ -71,30 +60,3 @@ async def test_train_rejects_wrong_number_of_rank_payloads():
     with pytest.raises(ValueError, match="one payload per train worker"):
         await group.train(5, {"data_ref": "rollout"}, external_data=[{"values": []}])
 
-
-async def test_train_only_ft_does_not_recover_rollout_engines():
-    from types import SimpleNamespace
-    from unittest.mock import AsyncMock
-
-    from miles.ray.actor_group import RayTrainGroup
-
-    calls = []
-    info = SimpleNamespace(snapshot_cell_id_to_hashes={})
-    group = object.__new__(RayTrainGroup)
-    group.args = SimpleNamespace(
-        debug_train_only=False,
-        debug_rollout_only=False,
-        use_fault_tolerance=True,
-        ft_components=["train"],
-    )
-    group._inference_controller = SimpleNamespace(
-        recover_updatable_engines=_AsyncCall("recover", calls),
-        start_update_weights=_AsyncCall("start", calls, result=info),
-        end_update_weights=_AsyncCall("end", calls),
-    )
-    group._broadcast = AsyncMock()
-
-    await group.update_weights(rollout_id=1)
-
-    assert [name for name, _, _ in calls] == ["start", "end"]
-    group._broadcast.assert_awaited_once_with("update_weights", info=info)

tests/fast/ray/test_update_weights_ordering.py

@@ -186,6 +186,44 @@ async def test_v2_hands_end_update_weights_the_snapshot_start_returned():
     _assert_the_snapshot_is_handed_back_unchanged(group._inference_controller)
 
 
+@pytest.mark.asyncio
+async def test_v1_aborts_the_window_when_the_broadcast_raises():
+    """A failed weight transfer must close the lock window instead of leaving it open forever."""
+    order: list[str] = []
+    group = RayTrainGroup.__new__(RayTrainGroup)
+    group.args = Namespace(debug_train_only=False, debug_rollout_only=False, use_fault_tolerance=False)
+    group._inference_controller = _OrderRecordingInferenceController(order)
+    group._broadcast = AsyncMock(side_effect=RuntimeError("weight transfer died"))
+
+    with pytest.raises(RuntimeError, match="weight transfer died"):
+        await group.update_weights()
+
+    assert order == ["start_update_weights", "abort_update_weights"]
+
+
+@pytest.mark.asyncio
+async def test_v2_aborts_the_window_when_the_broadcast_raises(monkeypatch):
+    """Same abort requirement on the fault-tolerant trainer group, after its retries are exhausted."""
+    from miles.ray.train import group as train_group_module
+
+    async def _retry_once(fn, **kwargs):
+        await fn(0)
+
+    monkeypatch.setattr(train_group_module, "retry", _retry_once)
+
+    order: list[str] = []
+    group = train_group_module.RayTrainGroup.__new__(train_group_module.RayTrainGroup)
+    group.args = Namespace(debug_train_only=False, debug_rollout_only=False)
+    group._inference_controller = _OrderRecordingInferenceController(order)
+    group._execute_first_alive = AsyncMock(side_effect=RuntimeError("weight transfer died"))
+    group._maybe_log_inference_engine_weight_checksums = AsyncMock()
+
+    with pytest.raises(RuntimeError, match="weight transfer died"):
+        await group.update_weights()
+
+    assert order == ["start_update_weights", "abort_update_weights"]
+
+
 def test_fsdp_updater_flushes_only_after_every_engine_is_paused():
     """Every engine is paused before any engine is flushed."""
     from unittest.mock import patch