Give each rollout cell a health checker driven by its own state

miles/ray/rollout/inference_controller.py

@@ -21,6 +21,7 @@
     requires_lock,
     with_lock,
 )
+from miles.utils.ft_utils.health_checker import ActivenessTracker
 from miles.utils.misc import SimpleTicker
 from miles.utils.workers.worker_provider.base import BaseWorkerProvider, CellInfo, StopWatchFn
 from miles.utils.workers.worker_provider.ray import RayWorkerProvider
@@ -43,14 +44,19 @@ def __init__(self, args):
         self.rollout_id = -1
         self.eval_fleet: EvalFleet | None = None
         self._watcher_disposers: list[StopWatchFn] = []
+        self._health_checker_activeness = ActivenessTracker(active=True)
         self._ticker: SimpleTicker | None = None
 
     @lock_exempt
     async def init(self) -> None:
         if self.args.debug_train_only:
             return
 
-        self.servers = await create_rollout_servers(self.args, context_lock=self.context_lock)
+        self.servers = await create_rollout_servers(
+            self.args,
+            context_lock=self.context_lock,
+            global_health_checker_activeness=self._health_checker_activeness.get,
+        )
         if self.args.eval_num_gpus > 0:
             self.eval_fleet = EvalFleet(self.args, srv=self.servers["eval"])
 
@@ -90,6 +96,9 @@ async def dispose(self):
             await disposer()
         self._watcher_disposers = []
 
+        for srv in self.servers.values():
+            await srv.dispose()
+
     # -------------------------- offload/onload -----------------------------
 
     # TODO may parallelly execute offload/onload across services
@@ -267,25 +276,24 @@ async def _reconcile(self, cell_id: str, observed: CellInfo | None) -> None:
 
     @requires_lock
     async def _health_monitoring_pause(self) -> None:
-        self._assert_rollout_fault_tolerance_is_unsupported()
+        self._health_checker_activeness.bump_active(False)
+        await asyncio.gather(
+            *[
+                cell.cancel_inflight_health_probe()
+                for srv in self.servers.values()
+                for cell in srv.server_cells.values()
+            ]
+        )
 
     @requires_lock
     async def _health_monitoring_resume(self) -> None:
-        self._assert_rollout_fault_tolerance_is_unsupported()
+        self._health_checker_activeness.bump_active(True)
 
     @property
     @requires_lock
     def _rollout_ft_enabled(self) -> bool:
         return self.args.use_fault_tolerance and "rollout" in self.args.ft_components
 
-    @requires_lock
-    def _assert_rollout_fault_tolerance_is_unsupported(self) -> None:
-        if not self.args.debug_train_only and self._rollout_ft_enabled:
-            raise NotImplementedError(
-                "rollout fault tolerance is being rebuilt; health monitoring must pause before "
-                "get_updatable_engines snapshots the engines"
-            )
-
     @property
     @requires_lock
     def _server(self) -> RolloutServer | None:

miles/ray/rollout/rollout_server.py

@@ -1,6 +1,7 @@
 import asyncio
 import dataclasses
 import logging
+from collections.abc import Callable
 from typing import Any
 
 from miles.backends.sglang_utils.sglang_api_client import SGLangApiClient
@@ -9,6 +10,7 @@
 from miles.ray.rollout.router_manager import wait_router_ready
 from miles.ray.rollout.server_cell import ServerCell, ServerCellMetadata
 from miles.utils.context_lock import ContextLock, enforce_lock_discipline, lock_exempt, requires_lock
+from miles.utils.ft_utils.health_checker import ActiveAndEpoch
 from miles.utils.retry_utils import retry_until_deadline
 
 logger = logging.getLogger(__name__)
@@ -17,7 +19,9 @@
 WAIT_CELLS_MAX_DELAY_SECONDS = 5.0
 
 
-async def create_rollout_servers(args, context_lock: ContextLock) -> dict[str, "RolloutServer"]:
+async def create_rollout_servers(
+    args, context_lock: ContextLock, global_health_checker_activeness: Callable[[], ActiveAndEpoch]
+) -> dict[str, "RolloutServer"]:
     """Create rollout servers: one per model, each with its own router."""
     assert args.sglang_router_ip is None, (
         "external router mode was removed: miles always starts its own routers "
@@ -43,6 +47,7 @@ async def create_rollout_servers(args, context_lock: ContextLock) -> dict[str, "
             router_port=router_addr.port,
             model_name=model_cfg.name,
             update_weights=model_cfg.update_weights,
+            global_health_checker_activeness=global_health_checker_activeness,
             expected_num_cells=model_cfg.num_server_cells,
         )
 
@@ -66,6 +71,9 @@ class RolloutServer:
     router_port: int | None = None
     model_name: str = "default"
     update_weights: bool = True
+    global_health_checker_activeness: Callable[[], ActiveAndEpoch] = lock_exempt(
+        lambda: ActiveAndEpoch(active=True, epoch=0)
+    )
     expected_num_cells: int = 0
 
     @property
@@ -102,17 +110,27 @@ async def probe_and_mark_dead(self):
     async def add_cell(self, cell_meta: ServerCellMetadata):
         cell_id = cell_meta.cell_id
         assert cell_id not in self.server_cells
-        cell = ServerCell(args=self.args, router_api_client=self._router_api_client, meta=cell_meta)
+        cell = ServerCell(
+            args=self.args,
+            router_api_client=self._router_api_client,
+            meta=cell_meta,
+            global_health_checker_activeness=self.global_health_checker_activeness,
+        )
+        self.server_cells[cell_id] = cell
         if not (self.args.colocate and cell_meta.needs_offload):
             await cell.init()
-        self.server_cells[cell_id] = cell
 
     @requires_lock
     async def remove_cell(self, cell_id: str):
         logger.info(f"Killing server {cell_id=}...")
         await self.server_cells[cell_id].dispose()
         del self.server_cells[cell_id]
 
+    @requires_lock
+    async def dispose(self) -> None:
+        for cell_id in list(self.server_cells.keys()):
+            await self.remove_cell(cell_id)
+
     @requires_lock
     async def offload(self, tags: list[str] | None = None):
         return await asyncio.gather(
@@ -160,9 +178,11 @@ async def _check(remaining_seconds: float) -> None:
 
     @lock_exempt
     def _count_startable_cells(self) -> int:
-        if self.args.colocate:
-            return len(self.server_cells)
-        return sum(1 for cell in self.server_cells.values() if cell.is_pending_weights_or_serving)
+        return sum(
+            1
+            for cell in self.server_cells.values()
+            if (self.args.colocate and cell.meta.needs_offload) or cell.is_pending_weights_or_serving
+        )
 
     @property
     @requires_lock

miles/ray/rollout/server_cell.py

@@ -1,6 +1,7 @@
 import asyncio
 import dataclasses
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any, Literal
 
@@ -18,6 +19,13 @@
     StateServing,
     StateUninitialized,
 )
+from miles.utils.ft_utils.health_checker import (
+    ActiveAndEpoch,
+    BaseHealthChecker,
+    NoopHealthChecker,
+    SimpleHealthChecker,
+    SimpleHealthCheckerConfig,
+)
 from miles.utils.pydantic_utils import FrozenStrictBaseModel
 from miles.utils.workers.launch_gate import GATE_PORT_NAME, activate_launch_gate
 from miles.utils.workers.worker_provider.base import BaseWorkerProvider
@@ -46,8 +54,35 @@ class ServerCell:
     args: Any
     meta: ServerCellMetadata
     router_api_client: SGLangRouterApiClient
+    global_health_checker_activeness: Callable[[], ActiveAndEpoch] = lambda: ActiveAndEpoch(active=True, epoch=0)
+    _health_checker: BaseHealthChecker = dataclasses.field(init=False)
     _state: CellState = dataclasses.field(default_factory=StateUninitialized)
 
+    def __post_init__(self) -> None:
+        self._health_checker = create_rollout_cell_health_checker(
+            args=self.args,
+            name=f"rollout-cell-{self.meta.cell_id}",
+            get_api_client=lambda: self.api_client,
+            get_activeness=self._get_health_checker_active_and_epoch,
+        )
+        self._health_checker.start()
+
+    def _get_health_checker_active_and_epoch(self) -> ActiveAndEpoch:
+        controller_active_and_epoch = self.global_health_checker_activeness()
+        cell_active = isinstance(self._state, (StatePendingWeights, StateServing))
+        return ActiveAndEpoch(
+            active=cell_active and controller_active_and_epoch.active, epoch=controller_active_and_epoch.epoch
+        )
+
+    def __del__(self) -> None:
+        assert isinstance(self._state, StateDisposed), (
+            f"ServerCell {self.meta.cell_id} was garbage collected without dispose() ({self._state=}); "
+            "every cell must be disposed so its health checker task is stopped"
+        )
+
+    async def cancel_inflight_health_probe(self) -> None:
+        await self._health_checker.cancel_inflight_probe()
+
     @property
     def is_uninitialized(self) -> bool:
         return isinstance(self._state, StateUninitialized)
@@ -122,9 +157,7 @@ async def _tick_when_initializing(self) -> None:
 
     async def mark_weights_ready(self) -> None:
         assert isinstance(self._state, StatePendingWeights), f"{self._state=}"
-
         await self._register_with_router(addr_info=self._state.addr_info)
-
         self._mark_serving()
 
     async def _register_with_router(self, addr_info: CellAddrInfo) -> None:
@@ -136,6 +169,8 @@ async def _register_with_router(self, addr_info: CellAddrInfo) -> None:
         )
 
     async def dispose(self) -> None:
+        self._health_checker.stop()
+
         match self._state:
             case StateServing():
                 await self._unregister_from_router()
@@ -211,3 +246,21 @@ async def check_weights(self, action: str, allow_quant_error: bool, selector: st
 
 def compute_nodes_per_engine(*, num_gpus_per_engine: int, num_gpus_per_node: int) -> int:
     return max(1, num_gpus_per_engine // num_gpus_per_node)
+
+
+def create_rollout_cell_health_checker(
+    *,
+    args: Any,
+    name: str,
+    get_api_client: Callable[[], SGLangApiClient],
+    get_activeness: Callable[[], ActiveAndEpoch],
+) -> BaseHealthChecker:
+    if "rollout" not in args.ft_components:
+        return NoopHealthChecker()
+
+    config = SimpleHealthCheckerConfig.from_args(args, prefix="rollout_health_check")
+
+    async def _check() -> None:
+        await get_api_client().health_generate(timeout=config.timeout)
+
+    return SimpleHealthChecker(name=name, check_fn=_check, get_activeness=get_activeness, config=config)

miles/ray/train/cell_monitor.py

@@ -10,7 +10,7 @@
     StateStopped,
 )
 from miles.utils.ft_utils.api_server.models import CellCondition, CellStatus, TriState
-from miles.utils.ft_utils.health_checker import ActivenessState, SimpleHealthChecker, SimpleHealthCheckerConfig
+from miles.utils.ft_utils.health_checker import ActiveAndEpoch, SimpleHealthChecker, SimpleHealthCheckerConfig
 
 if TYPE_CHECKING:
     from miles.ray.train.cell import RayTrainCell
@@ -20,7 +20,7 @@ def create_trainer_cell_health_checker(
     *,
     cell: "RayTrainCell",
     config: SimpleHealthCheckerConfig,
-    get_activeness: Callable[[], ActivenessState],
+    get_activeness: Callable[[], ActiveAndEpoch],
 ) -> SimpleHealthChecker:
     async def _check() -> None:
         # Cell health is liveness, not training progress: the heartbeat RPC runs on

miles/utils/arguments.py

@@ -979,6 +979,7 @@ def add_fault_tolerance_arguments(parser):
                 interval_default=30.0,
                 timeout_default=30.0,
                 first_wait_default=0.0,
+                failure_threshold_default=1,
             )
             parser.add_argument(
                 "--api-server-port",