Kill failed trainer workers through their own RPC channel instead of the manager

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

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,

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()

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
 

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
 

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]

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