From ca973cd23efb6a9076ea116434c04beabb49eb65 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 3 Aug 2026 11:49:28 +0800 Subject: [PATCH] Delete the pre-reconcile engine recovery recover_updatable_engines restarted dead engines in place before each weight update. Recovery now happens the other way round: the worker manager relaunches a cell, reconcile observes the new generation and re-adds it, and the trainer reconnects because the cell hashes changed. Its rollout_id bookkeeping goes with it, since nothing else read the field. --- miles/ray/actor_group.py | 10 +--- miles/ray/rollout/inference_controller.py | 6 -- .../real_ray/test_inference_controller.py | 59 ------------------- tests/fast/ray/test_actor_group_shared_ppo.py | 38 ------------ .../fast/ray/test_update_weights_ordering.py | 38 ++++++++++++ 5 files changed, 41 insertions(+), 110 deletions(-) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index c818355fba2..78fa4532980 100644 --- a/miles/ray/actor_group.py +++ b/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 diff --git a/miles/ray/rollout/inference_controller.py b/miles/ray/rollout/inference_controller.py index dd26adc2fd3..1e8ff9962b9 100644 --- a/miles/ray/rollout/inference_controller.py +++ b/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] diff --git a/tests/fast/ray/rollout/real_ray/test_inference_controller.py b/tests/fast/ray/rollout/real_ray/test_inference_controller.py index a489d9a0b1e..a6647dc7e45 100644 --- a/tests/fast/ray/rollout/real_ray/test_inference_controller.py +++ b/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( diff --git a/tests/fast/ray/test_actor_group_shared_ppo.py b/tests/fast/ray/test_actor_group_shared_ppo.py index 86ab3f9b9d3..259aec42066 100644 --- a/tests/fast/ray/test_actor_group_shared_ppo.py +++ b/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) diff --git a/tests/fast/ray/test_update_weights_ordering.py b/tests/fast/ray/test_update_weights_ordering.py index da0dce6f127..320b83010f3 100644 --- a/tests/fast/ray/test_update_weights_ordering.py +++ b/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