diff --git a/miles/ray/train/cell.py b/miles/ray/train/cell.py index 8b0177d1c01..51343b07472 100644 --- a/miles/ray/train/cell.py +++ b/miles/ray/train/cell.py @@ -18,12 +18,13 @@ from miles.utils.ft_utils.indep_dp import IndepDPInfo from miles.utils.retry_utils import NonRetryableError from miles.utils.tracking_utils.structured_log import log_structured -from miles.utils.workers.ray_worker_manager import RayWorkerManager from miles.utils.workers.worker_provider.ray import RayWorkerProvider from miles.utils.workers.worker_spec import HostAndPort logger = logging.getLogger(__name__) +KILL_RPC_TIMEOUT_S = 10.0 + class RayTrainCell: def __init__( @@ -139,17 +140,14 @@ async def prepare_indep_dp_mode_healing( # ------------------------ state transition ------------------------ - async def stop(self) -> None: - await RayWorkerManager.get_handle().stop_cells.remote([self.cell_id]) - - async def _stop_and_confirm_dead(self) -> None: + async def _kill_workers_and_confirm_dead(self) -> None: handles = self._get_actor_handles() if self.is_allocated else [] - await self.stop() log_structured( logger.info, tag="ft", op="confirm_dead", phase="start", cell=self.cell_id, n_actors=len(handles) ) start = time.monotonic() + await asyncio.gather(*[_kill_worker(handle) for handle in handles]) await asyncio.gather(*[_confirm_actor_dead(handle) for handle in handles]) log_structured( logger.info, @@ -260,7 +258,7 @@ async def _execute_raw( ) if kill_on_failure: self._mark_as_errored() - await self._stop_and_confirm_dead() + await self._kill_workers_and_confirm_dead() raise # ------------------------ state and misc queries ------------------------ @@ -300,6 +298,15 @@ def _get_actor_handles(self) -> list[ray.actor.ActorHandle]: return self._state.actor_handles +async def _kill_worker(handle: ray.actor.ActorHandle) -> None: + try: + await asyncio.wait_for(handle.kill_self.remote(), timeout=KILL_RPC_TIMEOUT_S) + except (ray.exceptions.RayActorError, ray.exceptions.RayTaskError): + return + except (TimeoutError, asyncio.TimeoutError): + logger.warning("Timed out asking a worker to kill itself; falling back to the death confirmation probe") + + async def _confirm_actor_dead(handle: ray.actor.ActorHandle) -> None: CONFIRM_DEAD_TIMEOUT_S = 120.0 CONFIRM_DEAD_PROBE_INTERVAL_S = 1.0 diff --git a/miles/ray/train/group.py b/miles/ray/train/group.py index b76e7fa7709..63da19c9878 100644 --- a/miles/ray/train/group.py +++ b/miles/ray/train/group.py @@ -409,8 +409,8 @@ async def _gather_all_alive_and_catch(self, compute_coroutine, *, debug_name: st if not snapshot_alive_cells: raise NonRetryableError("No alive cells") # NOTE: no timeout here. If a cell hangs, the external FT controller - # detects stale heartbeat via cell_status(), calls cell.stop() to kill - # actors, which unblocks this gather with ActorDiedError. + # detects stale heartbeat via cell_status() and suspends the cell through + # the worker manager, which unblocks this gather with ActorDiedError. outputs = await asyncio.gather( *[compute_coroutine(cell) for cell in snapshot_alive_cells], return_exceptions=True, diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index d4c0950cd0f..c9b59b285f3 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -27,7 +27,7 @@ logger = logging.getLogger(__name__) -TRAINER_CONCURRENCY_GROUPS = {"heartbeat_status": 1, "default": 1, "fault_injector": 1} +TRAINER_CONCURRENCY_GROUPS = {"heartbeat_status": 1, "default": 1, "fault_injector": 1, "kill_self": 1} def get_local_gpu_id(): @@ -147,6 +147,10 @@ def get_heartbeat_status(self) -> HeartbeatStatus: def inject_fault(self, mode: str) -> None: _inject_fault(mode=mode) + @ray.method(concurrency_group="kill_self") + def kill_self(self) -> None: + os._exit(1) + def clear_memory(self): print_memory("before TrainRayActor.clear_memory") clear_memory() diff --git a/tests/fast/ray/train/dummy_actor.py b/tests/fast/ray/train/dummy_actor.py index 8c7ba032d82..9c847f83f01 100644 --- a/tests/fast/ray/train/dummy_actor.py +++ b/tests/fast/ray/train/dummy_actor.py @@ -3,6 +3,7 @@ Records all method calls so tests can verify what was dispatched. """ +import os from typing import Any import ray @@ -79,6 +80,10 @@ def save_model(self, *args: Any, **kwargs: Any) -> None: def update_weights(self) -> None: self._record("update_weights", (), {}) + def kill_self(self) -> None: + self._record("kill_self", (), {}) + os._exit(1) + def set_heartbeat_fail(self, fail: bool) -> None: self._heartbeat_fail = fail diff --git a/tests/fast/ray/train/test_cell.py b/tests/fast/ray/train/test_cell.py index 5a877c16bbc..999870bc2e6 100644 --- a/tests/fast/ray/train/test_cell.py +++ b/tests/fast/ray/train/test_cell.py @@ -29,29 +29,32 @@ def test_actor_handles_are_real_ray_actors(self): assert all(isinstance(h, ray.actor.ActorHandle) for h in handles) -class TestStop: - async def test_stop_asks_the_worker_manager_to_stop_the_cell(self): - """The manager owns the actors, so the cell only forwards the request.""" - cell = make_cell(actor_count=2) +class TestKillWorkers: + async def test_killing_reaches_every_worker(self): + """The dead workers must not linger in a cross-cell collective.""" + cell = make_alive_cell(0, alive_cell_indices=[0]) + handles = cell._get_actor_handles() - await cell.stop() + await cell._kill_workers_and_confirm_dead() - assert train_conftest.fake_worker_manager.stopped_cell_ids == [[cell.cell_id]] + for handle in handles: + with pytest.raises(ray.exceptions.RayActorError): + ray.get(handle.get_calls.remote()) - async def test_stop_from_alive_asks_the_worker_manager_too(self): - """An alive cell is torn down the same way; reconcile then drops the object.""" + async def test_killing_does_not_involve_the_worker_manager(self): + """The manager keeps reporting the cell alive so its errored status stays visible.""" cell = make_alive_cell(0, alive_cell_indices=[0]) - await cell.stop() + await cell._kill_workers_and_confirm_dead() - assert train_conftest.fake_worker_manager.stopped_cell_ids == [[cell.cell_id]] + assert train_conftest.fake_worker_manager.stopped_cell_ids == [] async def test_stop_kills_the_underlying_workers(self): """Stopping a cell really kills its workers, so every handle is confirmed dead and rejects new calls.""" cell = make_cell(actor_count=2) handles = cell._get_actor_handles() - await cell.stop() + await cell._kill_workers_and_confirm_dead() for handle in handles: await asyncio.wait_for(cell_module._confirm_actor_dead(handle), timeout=30.0) @@ -139,22 +142,26 @@ def test_transitions_uninitialized_to_errored_without_info(self): class TestErroredCellTeardown: - async def test_stop_from_errored_reaches_the_worker_manager(self): - """An errored cell is torn down through the manager like any other.""" + async def test_kill_from_errored_reaches_the_workers(self): + """An errored cell is torn down by killing its own workers.""" cell = make_alive_cell(0, alive_cell_indices=[0]) cell._mark_as_errored() assert cell.is_errored + handles = cell._get_actor_handles() - await cell.stop() + await cell._kill_workers_and_confirm_dead() - assert train_conftest.fake_worker_manager.stopped_cell_ids == [[cell.cell_id]] + for handle in handles: + with pytest.raises(ray.exceptions.RayActorError): + ray.get(handle.get_calls.remote()) async def test_the_replacement_cell_recovers_the_lifecycle(self): - """Errored → stop → reconcile builds a fresh cell on the new workers → alive.""" + """Errored → kill → heal restarts the cell → reconcile builds a fresh one → alive.""" cell = make_alive_cell(0, alive_cell_indices=[0]) cell._mark_as_errored() - await cell.stop() + await cell._kill_workers_and_confirm_dead() + train_conftest.fake_worker_manager._stop_cells([cell.cell_id]) replacement = make_cell(cell.cell_index) replacement._mark_as_alive(indep_dp_info=make_indep_dp_info(quorum_id=99)) @@ -175,9 +182,8 @@ async def test_dispatches_init_and_marks_alive(self): for handle in cell._get_actor_handles(): calls = ray.get(handle.get_calls.remote()) - assert len(calls) == 1 - assert calls[0][0] == "init" - kwargs = calls[0][2] + assert [name for name, _args, _kwargs in calls] == ["configure_master_addr_and_port", "init"] + kwargs = calls[1][2] assert kwargs["indep_dp_info"] == info assert kwargs["recv_ckpt_src_rank"] is None @@ -193,7 +199,9 @@ async def test_init_failure_leaves_cell_not_alive(self): await cell.init(indep_dp_info=make_indep_dp_info()) assert not cell.is_alive - assert train_conftest.fake_worker_manager.stopped_cell_ids == [[cell.cell_id]] + for handle in cell._get_actor_handles(): + with pytest.raises(ray.exceptions.RayActorError): + ray.get(handle.get_calls.remote()) class TestPrepareIndepDPModeAlive: @@ -292,8 +300,8 @@ def test_errored(self): class TestFullLifecycle: - async def test_full_stop_and_replacement_cycle(self): - """Full lifecycle: attach → alive → stop → reconcile replaces the object → alive again.""" + async def test_full_kill_and_replacement_cycle(self): + """Full lifecycle: attach → alive → kill → heal restarts → reconcile replaces the object → alive again.""" # Step 1: Create (attaches to the manager's workers) cell = make_cell(actor_count=2) assert cell.is_uninitialized and not cell.is_alive @@ -303,11 +311,11 @@ async def test_full_stop_and_replacement_cycle(self): cell._mark_as_alive(indep_dp_info=info_v1) assert cell.is_alive - # Step 3: Stop through the manager - await cell.stop() - assert train_conftest.fake_worker_manager.stopped_cell_ids == [[cell.cell_id]] + # Step 3: Kill the workers directly + await cell._kill_workers_and_confirm_dead() - # Step 4: The provider drops it and reconcile builds a fresh object on the new workers + # Step 4: The ft controller heals it and reconcile builds a fresh object on the new workers + train_conftest.fake_worker_manager._stop_cells([cell.cell_id]) cell = make_cell(actor_count=2) assert cell.is_uninitialized and not cell.is_alive diff --git a/tests/fast/ray/train/test_cell_restart.py b/tests/fast/ray/train/test_cell_restart.py index 84b6c1d0528..bc7629518d6 100644 --- a/tests/fast/ray/train/test_cell_restart.py +++ b/tests/fast/ray/train/test_cell_restart.py @@ -1,34 +1,80 @@ +import asyncio + import pytest +import ray from tests.fast.ray.train import conftest as train_conftest from tests.fast.ray.train.conftest import make_cell +from miles.ray.train import cell as cell_module + pytestmark = pytest.mark.asyncio -class TestCellRestartGoesThroughTheManager: - async def test_stopping_asks_the_manager_instead_of_killing_actors(self): - """The manager owns the actors, so a cell killing them behind its back desyncs it.""" +class _HangingKillSelfHandle: + """A worker handle whose kill_self never returns, so _kill_worker has to time it out.""" + + def __init__(self) -> None: + self.kill_self_call_count: int = 0 + self.wait_dead_call_count: int = 0 + + async def kill_self(self) -> None: + self.kill_self_call_count += 1 + await asyncio.Event().wait() + + async def wait_dead(self, *, timeout: float) -> None: + self.wait_dead_call_count += 1 + + +class TestCellKillAndRestart: + async def test_killing_a_failed_cell_reaches_the_workers_directly(self): + """Waiting for an external controller would leave the other cells hanging in NCCL.""" cell = make_cell(2) + handles = cell._get_actor_handles() - await cell.stop() + await cell._kill_workers_and_confirm_dead() - assert train_conftest.fake_worker_manager.stopped_cell_ids == [["trainer-actor-2"]] + assert train_conftest.fake_worker_manager.stopped_cell_ids == [] + for handle in handles: + with pytest.raises(ray.exceptions.RayActorError): + ray.get(handle.get_calls.remote()) async def test_a_replacement_cell_picks_up_the_fresh_actor_handles(self): """Reusing the dead handles would make every later call fail.""" cell = make_cell(0) old_handles = cell._get_actor_handles() - await cell.stop() + await cell._kill_workers_and_confirm_dead() + train_conftest.fake_worker_manager._stop_cells([cell.cell_id]) replacement = make_cell(0) assert replacement._get_actor_handles() != old_handles - async def test_stopping_twice_asks_the_manager_twice(self): - """The cell no longer tracks a stopped state, so idempotence is the manager's job.""" + async def test_killing_twice_is_harmless(self): + """Healing may tear down an already dead cell, which must not raise.""" cell = make_cell(0) - await cell.stop() + await cell._kill_workers_and_confirm_dead() + + await cell._kill_workers_and_confirm_dead() + + +class TestKillRpcTimeout: + async def test_a_kill_rpc_that_never_returns_gives_up_at_the_timeout(self, monkeypatch: pytest.MonkeyPatch): + """A worker whose kill_self RPC hangs forever must not block _kill_worker forever.""" + monkeypatch.setattr(cell_module, "KILL_RPC_TIMEOUT_S", 0.05) + handle = _HangingKillSelfHandle() + + await asyncio.wait_for(cell_module._kill_worker(handle), timeout=10.0) + + assert handle.kill_self_call_count == 1 + + async def test_a_hanging_kill_rpc_still_reaches_the_death_confirmation(self, monkeypatch: pytest.MonkeyPatch): + """When every kill_self RPC hangs, teardown must fall through to the death probe instead of stalling.""" + monkeypatch.setattr(cell_module, "KILL_RPC_TIMEOUT_S", 0.05) + cell = make_cell(2) + hanging_handles: list[_HangingKillSelfHandle] = [_HangingKillSelfHandle(), _HangingKillSelfHandle()] + monkeypatch.setattr(cell, "_get_worker_handles", lambda: hanging_handles) - await cell.stop() + await asyncio.wait_for(cell._kill_workers_and_confirm_dead(), timeout=10.0) - assert len(train_conftest.fake_worker_manager.stopped_cell_ids) == 2 + assert [handle.kill_self_call_count for handle in hanging_handles] == [1, 1] + assert [handle.wait_dead_call_count for handle in hanging_handles] == [1, 1] diff --git a/tests/fast/ray/train/test_group.py b/tests/fast/ray/train/test_group.py index 42d20dce1ee..f53b768a2b3 100644 --- a/tests/fast/ray/train/test_group.py +++ b/tests/fast/ray/train/test_group.py @@ -94,6 +94,16 @@ def _was_stopped(group: RayTrainGroup, cell_index: int) -> bool: return [f"{group._pool}-{cell_index}"] in train_conftest.fake_worker_manager.stopped_cell_ids +def _was_killed(group: RayTrainGroup, cell_index: int) -> bool: + for handle in _cell(group, cell_index)._get_actor_handles(): + try: + ray.get(handle.get_calls.remote()) + return False + except ray.exceptions.RayActorError: + pass + return True + + async def _init_controller(group: RayTrainGroup) -> None: """Call init and wait for all cells to become alive.""" await group.init() @@ -617,7 +627,7 @@ async def test_one_cell_failure_marks_errored_others_ok(self): # Step 3: Cell 1 is errored, others alive assert _cell(group, 0).is_alive - assert _was_stopped(group, 1) + assert _was_killed(group, 1) assert _cell(group, 2).is_alive # Step 4: Other cells received train call @@ -635,7 +645,7 @@ async def test_errored_cell_skipped_in_next_broadcast(self): ray.get(handle.set_fail_methods.remote(["train"])) await group._execute_all_alive_and_catch("train", rollout_id=0, rollout_data_ref="data") - assert _was_stopped(group, 0) + assert _was_killed(group, 0) # Step 2: Next broadcast only goes to cell 1 await group._execute_all_alive_and_catch("train", rollout_id=1, rollout_data_ref="data") @@ -659,7 +669,7 @@ async def test_first_cell_fails_retry_falls_back_to_next(self): await group.save_model(rollout_id=42) # Step 3: Cell 0 errored, cell 1 handled it - assert _was_stopped(group, 0) + assert _was_killed(group, 0) assert _cell(group, 1).is_alive for handle in _cell(group, 1)._get_actor_handles(): @@ -676,7 +686,7 @@ async def test_single_execute_first_alive_raises_on_failure(self): with pytest.raises(Exception): # noqa: B017 await group._execute_first_alive("save_model", rollout_id=42) - assert _was_stopped(group, 0) + assert _was_killed(group, 0) async def test_losing_the_last_cell_keeps_the_worker_error_as_the_cause(self): """Without the cause the driver traceback says nothing about why the last cell died.""" @@ -944,7 +954,7 @@ async def test_cell_errored_does_not_retry_when_others_normal(self): await group.train(rollout_id=0, rollout_data_pack=_DUMMY_DATA_PACK) # Step 3: Cell 1 errored, alive cells each got 1 train call (no retry) - assert _was_stopped(group, 1) + assert _was_killed(group, 1) for i in [0, 2]: assert _count_train_calls(group, i) == 1