Add a periodic tick sweep over the rollout cells

miles/ray/rollout/inference_controller.py

@@ -11,11 +11,15 @@
 from miles.ray.rollout.router_manager import wait_session_server_ready
 from miles.ray.rollout.server_cell import ServerCell, ServerCellMetadata
 from miles.ray.specs.inference import compute_engine_pool_ids
+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
 
 logger = logging.getLogger(__name__)
 
+TICK_INTERVAL_SECONDS = 5.0
+CELL_TICK_TIMEOUT_SECONDS = 120.0
+
 
 class InferenceController:
     def __init__(self, args):
@@ -24,6 +28,7 @@ def __init__(self, args):
         self.rollout_id = -1
         self.eval_fleet: EvalFleet | None = None
         self._watcher_disposers: list[StopWatchFn] = []
+        self._ticker: SimpleTicker | None = None
 
     async def init(self) -> None:
         if self.args.debug_train_only:
@@ -38,6 +43,7 @@ async def init(self) -> None:
             pool_ids=compute_engine_pool_ids(self.args)
         )  # TODO inject instance
         self._watcher_disposers.append(await provider.watch_cells(self._reconcile))
+        self._ticker = SimpleTicker(self._tick_cells, interval_seconds=TICK_INTERVAL_SECONDS)
 
         dashboard_hooks.register_router(self.args)
         await wait_session_server_ready(self.args)
@@ -55,6 +61,10 @@ async def prepare_eval(self):
         await self._health_monitoring_resume()
 
     async def dispose(self):
+        if (ticker := self._ticker) is not None:
+            self._ticker = None
+            await ticker.dispose()
+
         for disposer in self._watcher_disposers:
             await disposer()
         self._watcher_disposers = []
@@ -154,6 +164,18 @@ async def check_weights(
             action=action, allow_quant_error=allow_quant_error, selector=selector, skip_list=skip_list
         )
 
+    # -------------------------- tick -----------------------------
+
+    async def _tick_cells(self) -> None:
+        cells = [cell for srv in list(self.servers.values()) for cell in list(srv.server_cells.values())]
+        results = await asyncio.gather(
+            *[asyncio.wait_for(cell.tick(), timeout=CELL_TICK_TIMEOUT_SECONDS) for cell in cells],
+            return_exceptions=True,
+        )
+        for cell, result in zip(cells, results, strict=True):
+            if isinstance(result, BaseException):
+                logger.error(f"Ticking cell {cell.meta.cell_id} failed", exc_info=result)
+
     # -------------------------- reconcile -----------------------------
 
     async def _reconcile(self, cell_id: str, observed: CellInfo | None) -> None:

miles/ray/rollout/server_cell.py

@@ -87,6 +87,9 @@ async def add(self) -> None:
         if not self.meta.update_weights or self.args.debug_rollout_only:
             await self.mark_weights_ready()
 
+    async def tick(self) -> None:
+        pass
+
     async def mark_weights_ready(self):
         assert isinstance(self._state, StatePendingWeights), f"{self._state=}"
 

miles/utils/misc.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from collections.abc import Sequence
+from collections.abc import Awaitable, Callable, Sequence
 from typing import Any
 
 import ray
@@ -143,3 +143,28 @@ def filter_keys(d: dict[str, Any], interest_keys: Sequence[str]) -> dict[str, An
     except Exception:
         logger.error(f"filter_keys d.keys={list(d)} {interest_keys=}", exc_info=True)
         raise
+
+
+class SimpleTicker:
+    def __init__(self, fn: Callable[[], Awaitable[None]], *, interval_seconds: float):
+        self._fn = fn
+        self._interval_seconds = interval_seconds
+        self._task = asyncio.create_task(self._loop())
+
+    async def dispose(self) -> None:
+        await cancel_and_await_task(self._task)
+
+    async def _loop(self) -> None:
+        while True:
+            await asyncio.sleep(self._interval_seconds)
+            try:
+                await self._fn()
+            except Exception:
+                logger.exception(f"Ticking {self._fn} failed; retrying")
+
+
+async def cancel_and_await_task(task: asyncio.Task) -> None:
+    task.cancel()
+    await asyncio.wait([task])
+    if not task.cancelled():
+        task.result()

miles/utils/workers/worker_provider/ray.py

@@ -4,6 +4,7 @@
 
 import ray.actor
 
+from miles.utils.misc import cancel_and_await_task
 from miles.utils.workers.ray_worker_manager import RayWorkerManager
 from miles.utils.workers.worker_provider.base import BaseWorkerProvider, CellInfo, ReconcileFn, StopWatchFn
 from miles.utils.workers.worker_spec import NamedHostAndPorts
@@ -38,7 +39,7 @@ async def watch_cells(self, reconcile: ReconcileFn) -> StopWatchFn:
         # the initial sync must complete (and raise on failure) before the watch is considered established
         await self._poll_once(reconcile, seen_infos=seen_infos, pool_ids=pool_ids)
         task = asyncio.create_task(self._watch_loop(reconcile, seen_infos, pool_ids=pool_ids))
-        return partial(_cancel_and_await_task, task)
+        return partial(cancel_and_await_task, task)
 
     def _watched_pool_ids(self) -> list[str]:
         assert self._pool_ids is not None, "this provider was built without the pool_ids it is meant to observe"
@@ -69,10 +70,3 @@ async def _poll_once(
             else:
                 seen_infos[cell_id] = observed_info
 
-
-async def _cancel_and_await_task(task) -> None:
-    task.cancel()
-    try:
-        await task
-    except asyncio.CancelledError:
-        pass

tests/fast/ray/rollout/test_inference_controller_tick.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import asyncio
+
+from types import SimpleNamespace
+
+from miles.ray.rollout import inference_controller as inference_controller_module
+from miles.ray.rollout.inference_controller import InferenceController
+from miles.utils.misc import SimpleTicker
+
+
+class _RecordingCell:
+    def __init__(self, *, error: Exception | None = None, delay: float = 0.0, cell_id: str = "cell"):
+        self.tick_count = 0
+        self.finished_count = 0
+        self.meta = SimpleNamespace(cell_id=cell_id)
+        self._error = error
+        self._delay = delay
+
+    async def tick(self) -> None:
+        self.tick_count += 1
+        if self._error is not None:
+            raise self._error
+        if self._delay:
+            await asyncio.sleep(self._delay)
+        self.finished_count += 1
+
+
+class _StubServer:
+    def __init__(self, server_cells: dict):
+        self.server_cells = server_cells
+
+
+def _make_controller(servers: dict) -> InferenceController:
+    controller = InferenceController.__new__(InferenceController)
+    controller.servers = servers
+    controller._watcher_disposers = []
+    controller._ticker = None
+    return controller
+
+
+def _start_ticker(controller: InferenceController) -> None:
+    controller._ticker = SimpleTicker(controller._tick_cells, interval_seconds=0.0)
+
+
+class TestTickCells:
+    async def test_it_drives_every_cell_of_every_server(self):
+        """A cell only makes progress when ticked, so no server may be left out of the sweep."""
+        first, second, third = _RecordingCell(), _RecordingCell(), _RecordingCell()
+        controller = _make_controller(
+            {"default": _StubServer({"a": first, "b": second}), "frozen": _StubServer({"c": third})}
+        )
+
+        await controller._tick_cells()
+
+        assert [cell.tick_count for cell in (first, second, third)] == [1, 1, 1]
+
+    async def test_one_failing_cell_does_not_let_its_siblings_escape_the_sweep(self):
+        """A sweep that returns early would release the lock while sibling ticks still mutate state."""
+        broken = _RecordingCell(error=RuntimeError("cell exploded"), cell_id="broken")
+        slow = _RecordingCell(delay=0.02, cell_id="slow")
+        controller = _make_controller({"default": _StubServer({"a": broken, "b": slow})})
+
+        await controller._tick_cells()
+
+        assert slow.finished_count == 1
+
+    async def test_a_wedged_cell_cannot_stall_the_sweep_forever(self, monkeypatch):
+        """A hung engine holds the controller lock for as long as its tick runs, so it must be bounded."""
+        wedged = _RecordingCell(delay=60.0, cell_id="wedged")
+        healthy = _RecordingCell(cell_id="healthy")
+        controller = _make_controller({"default": _StubServer({"a": wedged, "b": healthy})})
+        monkeypatch.setattr(inference_controller_module, "CELL_TICK_TIMEOUT_SECONDS", 0.01)
+
+        await controller._tick_cells()
+
+        assert wedged.finished_count == 0
+        assert healthy.finished_count == 1
+
+    async def test_a_cell_added_after_the_loop_started_is_picked_up(self):
+        """Cells appear from reconcile long after startup, so the sweep must re-read the bookkeeping."""
+        srv = _StubServer({})
+        controller = _make_controller({"default": srv})
+
+        _start_ticker(controller)
+        await asyncio.sleep(0.01)
+        late = _RecordingCell()
+        srv.server_cells["late"] = late
+        await asyncio.sleep(0.02)
+        await controller.dispose()
+
+        assert late.tick_count > 0
+
+    async def test_the_sweep_keeps_running_after_one_cell_raises(self):
+        """One wedged engine must not stop every other cell from making progress."""
+        broken, healthy = _RecordingCell(error=RuntimeError("cell exploded")), _RecordingCell()
+        controller = _make_controller({"default": _StubServer({"a": broken, "b": healthy})})
+
+        _start_ticker(controller)
+        await asyncio.sleep(0.02)
+        await controller.dispose()
+
+        assert broken.tick_count > 1
+        assert healthy.tick_count > 1
+
+
+class TestControllerDisposal:
+    async def test_dispose_stops_the_ticker(self):
+        """A surviving loop would keep dialing engines after the controller is gone."""
+        cell = _RecordingCell()
+        controller = _make_controller({"default": _StubServer({"a": cell})})
+
+        _start_ticker(controller)
+        await asyncio.sleep(0.02)
+        await controller.dispose()
+        ticks_after_dispose = cell.tick_count
+        await asyncio.sleep(0.02)
+
+        assert cell.tick_count == ticks_after_dispose
+
+    async def test_dispose_without_a_running_ticker_is_harmless(self):
+        """debug_train_only never starts the ticker, and teardown still has to work."""
+        controller = _make_controller({})
+
+        await controller.dispose()
+
+        assert controller._ticker is None