Route all trainer worker calls through the worker handle instead of Ray handles

miles/dashboard/hooks.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import threading
 import time
@@ -321,7 +322,7 @@ def register_router(args) -> None:
         _warner.warn("dashboard router registration failed; engine metrics will be missing")
 
 
-def register_engines(servers) -> None:
+async def register_engines(servers) -> None:
     """Called at the top of every InferenceController.prepare_rollout(): pushes an engine
     topology snapshot whenever the set of engine actors changed (startup,
     fault-tolerance recovery). Steady state costs one worker-manager round trip
@@ -342,7 +343,7 @@ def register_engines(servers) -> None:
         )
         if fingerprint == _engines_fingerprint:
             return
-        engines = _compute_engine_infos(cells, worker_infos_per_cell)
+        engines = await _compute_engine_infos(cells, worker_infos_per_cell)
         handle.update_topology.remote(TopologySnapshot(ts=time.time(), engines=engines))
         _engines_fingerprint = fingerprint
     except Exception:
@@ -389,13 +390,11 @@ def _collect_worker_infos(cells) -> list[list]:
     return _ray_get(futures)
 
 
-def _compute_engine_infos(cells, worker_infos_per_cell) -> list[EngineInfo]:
-    uuid_refs = [
-        info.actor_handle._get_gpu_uuids.remote(info.gpu_ids)
-        for worker_infos in worker_infos_per_cell
-        for info in worker_infos
-    ]
-    probed_uuids = iter(_ray_get(uuid_refs))
+async def _compute_engine_infos(cells, worker_infos_per_cell) -> list[EngineInfo]:
+    flat_infos = [info for worker_infos in worker_infos_per_cell for info in worker_infos]
+    probed_uuids = iter(
+        await asyncio.gather(*[info.handle._get_gpu_uuids(gpu_ids=info.gpu_ids) for info in flat_infos])
+    )
 
     engines = []
     for engine_rank, (cell, worker_infos) in enumerate(zip(cells, worker_infos_per_cell, strict=True)):

miles/ray/rollout/inference_controller.py

@@ -78,7 +78,7 @@ async def init(self) -> None:
     @with_lock
     async def prepare_rollout(self, rollout_id):
         await self._health_monitoring_resume()
-        dashboard_hooks.register_engines(self.servers)
+        await dashboard_hooks.register_engines(self.servers)
 
     @with_lock
     async def prepare_eval(self):

miles/ray/train/cell.py

@@ -2,8 +2,6 @@
 import logging
 import time
 
-import ray
-
 from miles.ray.specs.train import MASTER_PORT_NAME
 from miles.ray.train.cell_monitor import compute_cell_status
 from miles.ray.train.cell_state import (
@@ -18,12 +16,14 @@
 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.worker_handle import BaseWorkerHandle, WorkerUnreachableError
 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
+CONFIRM_DEAD_TIMEOUT_S = 120.0
 
 
 class RayTrainCell:
@@ -54,9 +54,7 @@ def __init__(
         self._master_addr: HostAndPort = worker_infos[0].self_addrs[MASTER_PORT_NAME]
 
         # NOTE: do *NOT* directly modify `self._state`, but instead use `self._change_state`
-        self._state: CellState = StateAllocatedUninitialized(
-            actor_handles=[info.actor_handle for info in worker_infos]
-        )
+        self._state: CellState = StateAllocatedUninitialized(worker_handles=[info.handle for info in worker_infos])
 
     # ------------------------ API ------------------------
 
@@ -94,7 +92,7 @@ async def train(
         attempt: int,
         external_data: list | None = None,
     ) -> list:
-        if external_data is not None and len(external_data) != len(self._get_actor_handles()):
+        if external_data is not None and len(external_data) != len(self._get_worker_handles()):
             raise NonRetryableError("external_data must contain one payload per train worker")
 
         return await self._execute_raw(
@@ -141,14 +139,14 @@ async def prepare_indep_dp_mode_healing(
     # ------------------------ state transition ------------------------
 
     async def _kill_workers_and_confirm_dead(self) -> None:
-        handles = self._get_actor_handles() if self.is_allocated else []
+        handles = self._get_worker_handles() if self.is_allocated else []
 
         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])
+        await asyncio.gather(*[handle.wait_dead(timeout=CONFIRM_DEAD_TIMEOUT_S) for handle in handles])
         log_structured(
             logger.info,
             tag="ft",
@@ -162,14 +160,14 @@ def _mark_as_alive(self, indep_dp_info: IndepDPInfo) -> None:
         self._change_state(
             "_mark_as_alive",
             StateAllocatedUninitialized,
-            StateAllocatedAlive(actor_handles=self._state.actor_handles, indep_dp_info=indep_dp_info),
+            StateAllocatedAlive(worker_handles=self._state.worker_handles, indep_dp_info=indep_dp_info),
         )
 
     def _update_indep_dp_info(self, indep_dp_info: IndepDPInfo) -> None:
         self._change_state(
             "_update_indep_dp_info",
             StateAllocatedAlive,
-            StateAllocatedAlive(actor_handles=self._state.actor_handles, indep_dp_info=indep_dp_info),
+            StateAllocatedAlive(worker_handles=self._state.worker_handles, indep_dp_info=indep_dp_info),
         )
 
     def _mark_as_errored(self) -> None:
@@ -180,7 +178,7 @@ def _mark_as_errored(self) -> None:
         self._change_state(
             "_mark_as_errored",
             (StateAllocatedUninitialized, StateAllocatedAlive, StateAllocatedErrored),
-            StateAllocatedErrored(actor_handles=self._state.actor_handles, indep_dp_info=indep_dp_info),
+            StateAllocatedErrored(worker_handles=self._state.worker_handles, indep_dp_info=indep_dp_info),
         )
 
     def _change_state(
@@ -225,14 +223,14 @@ async def _execute_raw(
         compute_kwargs,
         kill_on_failure: bool = True,
     ) -> list:
-        handles = self._get_actor_handles()
+        handles = self._get_worker_handles()
         log_structured(
             logger.info, tag="ft", op="execute", phase="start", cell=self.cell_id, fn=fn_name, n_actors=len(handles)
         )
         start = time.monotonic()
         try:
             result = await asyncio.gather(
-                *[getattr(actor, fn_name).remote(**compute_kwargs(i)) for i, actor in enumerate(handles)]
+                *[getattr(handle, fn_name)(**compute_kwargs(i)) for i, handle in enumerate(handles)]
             )
             log_structured(
                 logger.info,
@@ -291,40 +289,17 @@ def indep_dp_info(self) -> IndepDPInfo | None:
         assert isinstance(self._state, (StateAllocatedAlive, StateAllocatedErrored))
         return self._state.indep_dp_info
 
-    def _get_actor_handles(self) -> list[ray.actor.ActorHandle]:
+    def _get_worker_handles(self) -> list[BaseWorkerHandle]:
         assert isinstance(
             self._state, StateAllocatedBase
         ), f"Cell {self.cell_id} is not allocated (state={type(self._state).__name__})"
-        return self._state.actor_handles
+        return self._state.worker_handles
 
 
-async def _kill_worker(handle: ray.actor.ActorHandle) -> None:
+async def _kill_worker(handle: BaseWorkerHandle) -> None:
     try:
-        await asyncio.wait_for(handle.kill_self.remote(), timeout=KILL_RPC_TIMEOUT_S)
-    except (ray.exceptions.RayActorError, ray.exceptions.RayTaskError):
+        await asyncio.wait_for(handle.kill_self(), timeout=KILL_RPC_TIMEOUT_S)
+    except WorkerUnreachableError:
         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
-
-    async def _probe() -> None:
-        await handle.__ray_ready__.remote()
-
-    deadline = time.monotonic() + CONFIRM_DEAD_TIMEOUT_S
-    while True:
-        try:
-            await asyncio.wait_for(_probe(), timeout=CONFIRM_DEAD_PROBE_INTERVAL_S)
-        except (ray.exceptions.RayActorError, ray.exceptions.RayTaskError):
-            return
-        except (TimeoutError, asyncio.TimeoutError):
-            pass
-
-        if time.monotonic() >= deadline:
-            logger.error("Timed out after %.0fs confirming actor death; proceeding anyway", CONFIRM_DEAD_TIMEOUT_S)
-            return
-
-        await asyncio.sleep(CONFIRM_DEAD_PROBE_INTERVAL_S)

miles/ray/train/cell_monitor.py

@@ -24,7 +24,7 @@ async def _check() -> None:
         # Cell health is liveness, not training progress: the heartbeat RPC runs on
         # a dedicated concurrency group and returns even while the training thread is
         # blocked in a (legitimately waiting) cross-cell collective. A returned result
-        # proves the process is alive; an RayActorError or RPC timeout proves it is not.
+        # proves the process is alive; a WorkerUnreachableError or RPC timeout proves it is not.
         if not cell.is_alive:
             return
 

miles/ray/train/cell_state.py

@@ -1,16 +1,15 @@
-import ray
-
 from pydantic import BaseModel, ConfigDict
 
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
+from miles.utils.workers.worker_handle import BaseWorkerHandle
 
 
 class StateBase(BaseModel):
     model_config = ConfigDict(extra="forbid", frozen=True, arbitrary_types_allowed=True)
 
 
 class StateAllocatedBase(StateBase):
-    actor_handles: list[ray.actor.ActorHandle]
+    worker_handles: list[BaseWorkerHandle]
 
 
 class StateAllocatedUninitialized(StateAllocatedBase):

miles/utils/workers/ray_worker_manager.py

@@ -15,6 +15,7 @@
 from miles.utils.workers.addr_allocator import PortAllocator
 from miles.utils.workers.command_actor import CommandActor
 from miles.utils.workers.naming import compute_cell_id, compute_worker_name
+from miles.utils.workers.ray_worker_handle import RayWorkerHandle
 from miles.utils.workers.worker_info import WorkerInfo
 from miles.utils.workers.worker_provider.base import CellInfo
 from miles.utils.workers.worker_spec import (
@@ -98,7 +99,7 @@ def get_worker_infos(self, cell_id: str) -> list[WorkerInfo]:
                 generation=actor.generation,
                 self_addrs=actor.self_addrs,
                 gpu_ids=actor.gpu_ids,
-                actor_handle=actor.actor_handle,
+                handle=RayWorkerHandle(actor.actor_handle),
             )
             for actor in (cell.actors if cell.actors is not None else [])
         ]

miles/utils/workers/worker_info.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-
-import ray.actor
+from pydantic import ConfigDict
 
+from miles.utils.pydantic_utils import StrictBaseModel
+from miles.utils.workers.worker_handle import BaseWorkerHandle
 from miles.utils.workers.worker_spec import NamedHostAndPorts
 
 
-@dataclass(kw_only=True)
-class WorkerInfo:
+class WorkerInfo(StrictBaseModel):
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
     name: str
     generation: int
     self_addrs: NamedHostAndPorts
     gpu_ids: list[int]
-    actor_handle: ray.actor.ActorHandle
+    handle: BaseWorkerHandle