Serve rollout cells over the api server

miles/ray/rollout/inference_controller.py

@@ -239,10 +239,6 @@ async def _tick_cells(self) -> None:
 
     @with_lock
     async def _reconcile(self, cell_id: str, observed: CellInfo | None) -> None:
-        observed_cell_meta: ServerCellMetadata | None = (
-            _compute_server_cell_meta_from_info(observed) if observed is not None else None
-        )
-
         actual_srv: RolloutServer | None = None
         actual_cell: ServerCell | None = None
         for srv in self.servers.values():

miles/utils/ft_utils/api_server/handles.py

@@ -2,8 +2,11 @@
 
 import abc
 
+import ray
+
+from miles.ray.rollout.server_cell import compute_pending_rollout_cell_status
 from miles.ray.train.group import RayTrainGroup
-from miles.utils.ft_utils.api_server.models import Cell, CellCondition, CellMetadata, CellSpec, CellStatus
+from miles.utils.ft_utils.api_server.models import Cell, CellMetadata, CellSpec
 from miles.utils.test_utils.fault_injector import FailureMode
 
 
@@ -79,9 +82,15 @@ async def inject_fault(self, *, mode: FailureMode, sub_index: int) -> None:
         actors[sub_index].inject_fault.remote(mode.value)
 
 
-# TODO the code will NOT work before implementing rollout ft
 class _RolloutCellHandle(_CellHandle):
-    def __init__(self, *, inference_controller: object, rollout_cell_id: str) -> None:
+    def __init__(
+        self,
+        *,
+        worker_manager: ray.actor.ActorHandle,
+        inference_controller: object,
+        rollout_cell_id: str,
+    ) -> None:
+        self._worker_manager = worker_manager
         self._inference_controller = inference_controller
         self._rollout_cell_id = rollout_cell_id
 
@@ -94,9 +103,8 @@ def cell_key(self) -> str:
         return self._rollout_cell_id
 
     async def get_cell(self) -> Cell:
-        phase = self._inference_controller.get_cell_phase(self._rollout_cell_id)
-        conditions_raw = self._inference_controller.get_cell_conditions(self._rollout_cell_id)
-        is_suspended = self._inference_controller.get_cell_is_suspended(self._rollout_cell_id)
+        statuses = self._inference_controller.get_cell_statuses()
+        status = statuses.get(self._rollout_cell_id) or compute_pending_rollout_cell_status()
         return Cell(
             metadata=CellMetadata(
                 name=self.cell_id,
@@ -105,15 +113,14 @@ async def get_cell(self) -> Cell:
                     "miles.io/cell-index": self.cell_key,
                 },
             ),
-            spec=CellSpec(suspend=is_suspended),
-            status=CellStatus(
-                phase=phase,
-                conditions=[CellCondition(**c) for c in conditions_raw],
-            ),
+            spec=CellSpec(suspend=status.phase == "Suspended"),
+            status=status,
         )
 
     async def suspend(self) -> None:
-        await self._inference_controller.stop_cell(self._rollout_cell_id)
+        await self._worker_manager.stop_cells.remote([self._rollout_cell_id])
+        self._inference_controller.notify_cell_suspended(self._rollout_cell_id)
 
     async def resume(self) -> None:
-        await self._inference_controller.start_cell(self._rollout_cell_id)
+        await self._worker_manager.start_cells.remote([self._rollout_cell_id])
+        self._inference_controller.notify_cell_resumed(self._rollout_cell_id)

miles/utils/ft_utils/api_server/server.py

@@ -4,14 +4,17 @@
 import logging
 import threading
 
+import ray
 import uvicorn
 from fastapi import FastAPI, Request
 from starlette.responses import JSONResponse
 
+from miles.ray.specs.inference import compute_engine_pool_ids
 from miles.ray.train.group import RayTrainGroup
 from miles.utils.ft_utils.api_server.handles import _ActorCellHandle, _CellHandle, _RolloutCellHandle
 from miles.utils.ft_utils.api_server.models import Cell, CellList, CellPatch, FaultInjection, K8sStatus, _OkResponse
 from miles.utils.ft_utils.api_server.registry import _CellRegistry
+from miles.utils.workers.ray_worker_manager import RayWorkerManager
 
 logger = logging.getLogger(__name__)
 
@@ -21,6 +24,7 @@
 
 def start_api_server(
     *,
+    args,
     actor_model: RayTrainGroup,
     inference_controller: object,
     port: int,
@@ -33,12 +37,15 @@ def start_api_server(
             registry.register(_ActorCellHandle(group=actor_model, cell_index=i))
 
     if "rollout" in ft_components:
-        # TODO the code will NOT work before implementing rollout ft
-        for rollout_cell_id in inference_controller.list_cell_ids():
+        worker_manager = RayWorkerManager.get_handle()
+        engine_pool_ids = compute_engine_pool_ids(args)
+        summaries = ray.get(worker_manager.get_cell_infos.remote(pool_ids=engine_pool_ids))
+        for cell_id in sorted(summaries):
             registry.register(
                 _RolloutCellHandle(
+                    worker_manager=worker_manager,
                     inference_controller=inference_controller,
-                    rollout_cell_id=rollout_cell_id,
+                    rollout_cell_id=cell_id,
                 )
             )
 

miles/utils/workers/ray_worker_manager.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import logging
-from collections.abc import Awaitable, Callable
+from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
@@ -58,26 +58,14 @@ async def init(self, specs: list[BaseWorkerSpec], pgs: dict[str, PlacementGroupI
 
     async def start_cells(self, cell_ids: list[str]) -> None:
         cells = [cell for cell_id in cell_ids if (cell := self._find_cell(cell_id)).actors is None]
-        phases: list[Callable[[_CellManager], Awaitable[None]]] = [
-            lambda c: c.launch_actors(),
-            lambda c: c.alloc_ports(),
-            lambda c: c.post_setup(),
-        ]
-
-        errors: list[BaseException] = []
-        for phase in phases:
-            outcomes = await asyncio.gather(*[phase(c) for c in cells], return_exceptions=True)
-            failed = [c for c, outcome in zip(cells, outcomes, strict=True) if isinstance(outcome, BaseException)]
-            errors += [outcome for outcome in outcomes if isinstance(outcome, BaseException)]
-            if failed:
-                logger.warning(f"Rolling back cells that failed to start ({[c.cell_id for c in failed]}) ({errors=})")
-                await asyncio.gather(*[c.stop() for c in failed], return_exceptions=True)
-                cells = [
-                    c for c, outcome in zip(cells, outcomes, strict=True) if not isinstance(outcome, BaseException)
-                ]
-
-        if errors:
-            raise errors[0]
+        try:
+            await _gather_or_raise([c.launch_actors() for c in cells])
+            await _gather_or_raise([c.alloc_ports() for c in cells])
+            await _gather_or_raise([c.post_setup() for c in cells])
+        except Exception:
+            logger.error(f"Starting cells {[c.cell_id for c in cells]} failed, rolling back", exc_info=True)
+            await asyncio.gather(*[c.stop() for c in cells], return_exceptions=True)
+            raise
 
     async def stop_cells(self, cell_ids: list[str]) -> None:
         await asyncio.gather(*[self._find_cell(cell_id).stop() for cell_id in cell_ids])
@@ -328,3 +316,10 @@ async def post_setup(self) -> None:
         )
         launch_cmd = self.spec.launch_command(ctx)
         self.actor_handle.run.remote(cmd=launch_cmd, envs={})
+
+
+async def _gather_or_raise(coros: list[Coroutine[Any, Any, None]]) -> None:
+    results = await asyncio.gather(*coros, return_exceptions=True)
+    for result in results:
+        if isinstance(result, BaseException):
+            raise result

tests/fast/utils/api_server/conftest.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 import asyncio
+import dataclasses
+from collections.abc import Callable
 
 import httpx
 import pytest
 
 from miles.utils.ft_utils.api_server.models import Cell, CellCondition, CellMetadata, CellSpec, CellStatus
 from miles.utils.ft_utils.api_server.registry import _CellRegistry
 from miles.utils.ft_utils.api_server.server import _create_api_app
+from miles.utils.workers.worker_provider.base import CellInfo
 
 
 class MockHandle:
@@ -73,49 +76,71 @@ async def resume(self) -> None:
 
 
 class MockRemoteCall:
-    def __init__(self, return_value: object) -> None:
+    def __init__(self, return_value: object, effect: Callable[..., None] | None = None) -> None:
         self._return_value = return_value
+        self._effect = effect
         self.calls: list[tuple[tuple[object, ...], dict[str, object]]] = []
 
     def remote(self, *args: object, **kwargs: object) -> asyncio.Future[object]:
         self.calls.append((args, kwargs))
+        if self._effect is not None:
+            self._effect(*args, **kwargs)
         future: asyncio.Future[object] = asyncio.get_event_loop().create_future()
         future.set_result(self._return_value)
         return future
 
 
 class MockInferenceController:
-    """Plain object, not a Ray actor: the controller lives in the driver process."""
+    def __init__(self, statuses: dict[str, CellStatus] | None = None) -> None:
+        self._statuses = dict(statuses or {})
+        self.status_calls: int = 0
 
-    def __init__(
-        self,
-        phase: str = "Running",
-        conditions: list[dict[str, str | None]] | None = None,
-        is_suspended: bool = False,
-    ) -> None:
-        self._phase = phase
-        self._conditions = conditions or [
-            {"type": "Allocated", "status": "True"},
-            {"type": "Healthy", "status": "True"},
-        ]
-        self._is_suspended = is_suspended
-        self.stopped_cells: list[str] = []
-        self.started_cells: list[str] = []
+    def get_cell_statuses(self) -> dict[str, CellStatus]:
+        self.status_calls += 1
+        return dict(self._statuses)
 
-    def get_cell_phase(self, cell_id: str) -> str:
-        return self._phase
+    def observe_cell(self, cell_id: str, status: CellStatus) -> None:
+        self._statuses[cell_id] = status
 
-    def get_cell_conditions(self, cell_id: str) -> list[dict[str, str | None]]:
-        return self._conditions
 
-    def get_cell_is_suspended(self, cell_id: str) -> bool:
-        return self._is_suspended
+class MockWorkerManager:
+    def __init__(self, summaries: dict[str, CellInfo] | None = None) -> None:
+        self._summaries = dict(summaries or {})
+        self.stopped_cells: list[list[str]] = []
+        self.started_cells: list[list[str]] = []
+        self.cell_info_calls: list[dict[str, object]] = []
 
-    async def stop_cell(self, cell_id: str) -> None:
-        self.stopped_cells.append(cell_id)
+    @property
+    def get_cell_infos(self) -> MockRemoteCall:
+        return MockRemoteCall(dict(self._summaries), effect=lambda **kwargs: self.cell_info_calls.append(kwargs))
 
-    async def start_cell(self, cell_id: str) -> None:
-        self.started_cells.append(cell_id)
+    @property
+    def stop_cells(self) -> MockRemoteCall:
+        return MockRemoteCall(None, effect=lambda ids: self._record(self.stopped_cells, ids, suspended=True))
+
+    @property
+    def start_cells(self) -> MockRemoteCall:
+        return MockRemoteCall(None, effect=lambda ids: self._record(self.started_cells, ids, suspended=False))
+
+    def _record(self, log: list[list[str]], cell_ids: list[str], *, suspended: bool) -> None:
+        log.append(list(cell_ids))
+        for cell_id in cell_ids:
+            previous = self._summaries[cell_id]
+            self._summaries[cell_id] = dataclasses.replace(previous, alive=not suspended)
+
+
+def make_cell_summaries(*cell_ids: str, suspended: bool = False) -> dict[str, CellInfo]:
+    return {
+        cell_id: CellInfo(
+            cell_id=cell_id,
+            pool_id=cell_id.rsplit("-", 1)[0],
+            alive=not suspended,
+            worker_names=[] if suspended else [f"{cell_id}-0"],
+            workers_hash="pseudo-hash-0",
+            meta={"model_id": "default"},
+        )
+        for cell_id in cell_ids
+    }
 
 
 class MockRayTrainCell: