Make ServerGroup callers cell-aware and drop the all_engines property

miles/dashboard/hooks.py

@@ -362,18 +362,15 @@ def report_data_buffer(length: int | None) -> None:
 
 
 def _alive_engine_chunks(servers) -> list[list]:
-    """Multi-node engines occupy ``nodes_per_engine`` consecutive entries of
-    ``group.all_engines``; only the first (master) owns the router-visible
-    URL. Chunks with any dead member are skipped until recovery completes."""
+    """A multi-node engine's cell holds ``nodes_per_engine`` entries; only the
+    first (master) owns the router-visible URL. Cells with any dead member are
+    skipped until recovery completes."""
     chunks = []
     for server in servers.values():
         for group in server.server_groups:
-            stride = group.nodes_per_engine
-            engines = group.all_engines
-            for i in range(0, len(engines), stride):
-                chunk = engines[i : i + stride]
-                if all(engine.is_allocated and engine.is_alive for engine in chunk):
-                    chunks.append(chunk)
+            for cell in group.cells:
+                if all(engine.is_allocated and engine.is_alive for engine in cell.engines):
+                    chunks.append(cell.engines)
     return chunks
 
 

miles/ray/rollout/rollout_server.py

@@ -299,7 +299,7 @@ async def wait_all_engines_alive(self, timeout: float = 600):
         # picture of init/recovery upper bounds across model sizes
         sleep_time = 2
         for _ in range(int(timeout // sleep_time)):
-            if all(e.is_alive for g in self.server_groups for e in g.all_engines):
+            if all(e.is_alive for g in self.server_groups for cell in g.cells for e in cell.engines):
                 return
             await asyncio.sleep(sleep_time)
             logger.info("wait_all_engines_alive looping...")

miles/ray/rollout/server_cell.py

@@ -12,6 +12,10 @@ class ServerCell:
     engines: list[ServerEngine]
 
 
+def flatten_cells(cells: list[ServerCell]) -> list[ServerEngine]:
+    return [engine for cell in cells for engine in cell.engines]
+
+
 class CellIndexer(NamedTuple):
     srv_key: str
     group_index: int
@@ -21,23 +25,22 @@ class CellIndexer(NamedTuple):
 def get_cell_indexer_of_id_map(servers: dict[str, "RolloutServer"]) -> list[CellIndexer]:
     """Flatten ``servers`` into a list whose position is the cell id.
 
-    A cell is one node-0 engine; ``engine_indices`` covers its ``nodes_per_engine``
-    underlying entries in ``group.all_engines``. Order is sorted by ``srv_key``, so
-    cell ids are stable across calls when the topology is unchanged.
+    ``engine_indices`` covers the cell's entries in the group's flat engine
+    list. Order is sorted by ``srv_key``, so cell ids are stable across calls
+    when the topology is unchanged.
     """
     result: list[CellIndexer] = []
     for srv_key in sorted(servers):
         srv = servers[srv_key]
         for group_index, group in enumerate(srv.server_groups):
-            assert len(group.all_engines) == len(group.engines) * group.nodes_per_engine
-            for local_index in range(len(group.engines)):
+            engine_offset = 0
+            for cell in group.cells:
                 result.append(
                     CellIndexer(
                         srv_key=srv_key,
                         group_index=group_index,
-                        engine_indices=list(
-                            range(local_index * group.nodes_per_engine, (local_index + 1) * group.nodes_per_engine)
-                        ),
+                        engine_indices=list(range(engine_offset, engine_offset + len(cell.engines))),
                     )
                 )
+                engine_offset += len(cell.engines)
     return result

miles/ray/rollout/server_group.py

@@ -15,7 +15,7 @@
     allocate_rollout_engine_addr_and_ports_external,
     allocate_rollout_engine_addr_and_ports_normal,
 )
-from miles.ray.rollout.server_cell import ServerCell
+from miles.ray.rollout.server_cell import ServerCell, flatten_cells
 from miles.ray.rollout.server_engine import AddrInfo, ServerEngine
 from miles.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST
 from miles.utils import async_utils, dumper_utils
@@ -38,7 +38,7 @@ class ServerGroup:
     pg: Any  # (placement_group, reordered_bundle_indices, reordered_gpu_ids)
     cells: list[ServerCell]
     num_gpus_per_engine: int
-    # NOTE: this may have risk when recovering engines parallelly; may use source of truth (all_engines) later
+    # NOTE: this may have risk when recovering engines parallelly; may use source of truth (cells) later
     has_new_engines: bool
     worker_type: str = "regular"  # "regular", "prefill", or "decode"
     rank_offset: int = 0
@@ -52,21 +52,17 @@ class ServerGroup:
 
     def __post_init__(self):
         assert (
-            not self.all_engines or self.rank_offset % self.nodes_per_engine == 0
+            not self.cells or self.rank_offset % self.nodes_per_engine == 0
         ), f"{self.rank_offset=} must be a multiple of {self.nodes_per_engine=}"
 
-    @property
-    def all_engines(self) -> list[ServerEngine]:
-        return [engine for cell in self.cells for engine in cell.engines]
-
     @property
     def nodes_per_engine(self):
         return max(1, self.num_gpus_per_engine // self.args.num_gpus_per_node)
 
     @property
     def engines(self) -> list[ServerEngine]:
         """Node-0 engines only (for multi-node serving)."""
-        return self.all_engines[:: self.nodes_per_engine]
+        return [cell.engines[0] for cell in self.cells]
 
     def start_engines(
         self, port_cursors: PortCursors, start_indices: list[int] | None = None
@@ -76,7 +72,7 @@ def start_engines(
         Mutates ``port_cursors`` in place to advance past any newly assigned ports.
         Returns ``(init_handles, new_engine_indices)`` where *init_handles* is a list
         of Ray ObjectRefs (one per newly created engine) and *new_engine_indices* is
-        the list of indices into ``self.all_engines`` that were just allocated.
+        the list of indices into the group's flat engine list that were just allocated.
         """
         assert not ({"host", "port"} & set(self.sglang_overrides)), (
             f"sglang_overrides must not override host/port ({self.sglang_overrides=}): the rollout process derives "
@@ -93,12 +89,14 @@ def start_engines(
 
         RolloutRayActor = ray.remote(SGLangEngine)
 
+        all_engines = flatten_cells(self.cells)
+
         new_engines = []
         new_engine_indices = []
-        for i in range(len(self.all_engines)):
+        for i in range(len(all_engines)):
             if (start_indices is not None) and (i not in start_indices):
                 continue
-            if self.all_engines[i].is_allocated:
+            if all_engines[i].is_allocated:
                 continue
 
             global_rank = self.rank_offset + i
@@ -153,7 +151,7 @@ def start_engines(
 
             new_engines.append((global_rank, rollout_engine))
             new_engine_indices.append(i)
-            self.all_engines[i].mark_allocated_uninitialized(rollout_engine)
+            all_engines[i].mark_allocated_uninitialized(rollout_engine)
 
         curr_num_new_engines = len(new_engines)
         self.has_new_engines |= curr_num_new_engines > 0
@@ -179,7 +177,7 @@ def start_engines(
 
         for index, _ in new_engines:
             engine_addr_and_ports = addr_and_ports[index]
-            self.all_engines[index - self.rank_offset].set_addressing(
+            all_engines[index - self.rank_offset].set_addressing(
                 AddrInfo(
                     server_url=build_server_url(
                         host=engine_addr_and_ports["host"], port=engine_addr_and_ports["port"]
@@ -220,10 +218,11 @@ async def unregister_workers(self, engine_indices: list[int]) -> None:
         )
 
     def _primary_engines_of(self, engine_indices: list[int]) -> list[ServerEngine]:
+        all_engines = flatten_cells(self.cells)
         return [
-            self.all_engines[index]
+            all_engines[index]
             for index in engine_indices
-            if index % self.nodes_per_engine == 0 and self.all_engines[index].is_allocated
+            if index % self.nodes_per_engine == 0 and all_engines[index].is_allocated
         ]
 
     @property
@@ -241,8 +240,9 @@ def stop_engines(self, engine_indices: list[int]):
             async_utils.run(asyncio.wait_for(self.unregister_workers(engine_indices), timeout=_SHUTDOWN_TIMEOUT))
         except Exception as e:
             logger.warning(f"Unregistering {engine_indices=} from the router failed, tearing down anyway (e: {e})")
+        all_engines = flatten_cells(self.cells)
         for i in engine_indices:
-            engine = self.all_engines[i]
+            engine = all_engines[i]
             if engine.is_allocated:
                 logger.info(f"Shutting down and killing engine at index {i}")
                 try:
@@ -256,12 +256,13 @@ def stop_engines(self, engine_indices: list[int]):
                     logger.warning(f"Fail to kill engine at index {i} (e: {e})")
             else:
                 logger.info(f"Engine at index {i} is already None")
-            self.all_engines[i].mark_stopped()
+            all_engines[i].mark_stopped()
 
     async def recover(self, port_cursors: PortCursors, filter_indices: list[int] | None = None):
+        all_engines = flatten_cells(self.cells)
         if filter_indices is None:
-            filter_indices = [i for i, engine in enumerate(self.all_engines) if not engine.is_allocated]
-        start_indices = [idx for idx in filter_indices if not self.all_engines[idx].is_allocated]
+            filter_indices = [i for i, engine in enumerate(all_engines) if not engine.is_allocated]
+        start_indices = [idx for idx in filter_indices if not all_engines[idx].is_allocated]
 
         handles, new_engine_indices = self.start_engines(port_cursors, start_indices=start_indices)
         await asyncio.gather(*handles)
@@ -273,7 +274,7 @@ async def recover(self, port_cursors: PortCursors, filter_indices: list[int] | N
             start_indices
         ), "curr_num_new_engines does not match start_indices length"
         if self.needs_offload and start_indices:
-            new_primary_engines = [self.all_engines[i] for i in start_indices if i % self.nodes_per_engine == 0]
+            new_primary_engines = [all_engines[i] for i in start_indices if i % self.nodes_per_engine == 0]
             release_handles.extend(engine.api_client.release_memory_occupation() for engine in new_primary_engines)
             if self.update_weights or self.model_path:
                 all_resume_engines.extend(new_primary_engines)
@@ -292,8 +293,9 @@ async def recover(self, port_cursors: PortCursors, filter_indices: list[int] | N
         await self.register_workers(new_engine_indices)
 
     def mark_alive(self, engine_indices: list[int]):
+        all_engines = flatten_cells(self.cells)
         for engine_index in engine_indices:
-            self.all_engines[engine_index].mark_alive()
+            all_engines[engine_index].mark_alive()
 
     async def offload(self, tags: list[str] | None = None):
         if not self.needs_offload:

tests/fast/dashboard/test_hooks.py

@@ -6,6 +6,7 @@
 from miles.dashboard import backend, hooks
 from miles.dashboard.hooks import BATCH_MAX_EVENTS, BATCH_MAX_SECONDS, _Identity
 from miles.dashboard.store import Role
+from miles.ray.rollout.server_cell import ServerCell
 from miles.utils.timer import Timer
 
 
@@ -152,7 +153,9 @@ def __init__(self, info, alive=True):
 
 class FakeGroup:
     def __init__(self, engines, nodes_per_engine=1):
-        self.all_engines = engines
+        self.cells = [
+            ServerCell(engines=engines[i : i + nodes_per_engine]) for i in range(0, len(engines), nodes_per_engine)
+        ]
         self.nodes_per_engine = nodes_per_engine
 
 

tests/fast/ray/rollout/real_ray/test_fault_tolerance.py

@@ -10,6 +10,7 @@
 from tests.fast.ray.rollout.conftest import chunk_engines_into_cells, make_args
 
 from miles.ray.rollout.addr_allocator import PortCursors
+from miles.ray.rollout.server_cell import flatten_cells
 from miles.ray.rollout.server_engine import ServerEngine
 from miles.ray.rollout.server_group import ServerGroup
 
@@ -44,7 +45,7 @@ def _start(group: ServerGroup) -> None:
 
 
 def _kill_all(group: ServerGroup) -> None:
-    for e in group.all_engines:
+    for e in flatten_cells(group.cells):
         if e.is_allocated:
             try:
                 ray.kill(e.actor_handle)
@@ -68,22 +69,22 @@ async def test_recover_creates_new_actor_after_real_kill(
         group = _build_group(pg_tuple=pg, num_engines=2)
         _start(group)
 
-        original_handles = [e.actor_handle for e in group.all_engines]
+        original_handles = [e.actor_handle for e in flatten_cells(group.cells)]
         # Real fault: kill engine 0 + mark its slot stopped (production code's
         # health monitor would do this; here we simulate it directly).
         ray.kill(original_handles[0])
-        group.all_engines[0].mark_stopped()
+        flatten_cells(group.cells)[0].mark_stopped()
 
         try:
             await group.recover(port_cursors=PortCursors.empty(), filter_indices=[0])
             # New actor for slot 0
-            assert group.all_engines[0].is_allocated
-            assert group.all_engines[0].actor_handle is not original_handles[0]
-            calls = ray.get(group.all_engines[0].actor_handle.get_calls.remote())
+            assert flatten_cells(group.cells)[0].is_allocated
+            assert flatten_cells(group.cells)[0].actor_handle is not original_handles[0]
+            calls = ray.get(flatten_cells(group.cells)[0].actor_handle.get_calls.remote())
             assert "init" in [c[0] for c in calls]
 
             # Slot 1 untouched, still the same actor
-            assert group.all_engines[1].actor_handle is original_handles[1]
+            assert flatten_cells(group.cells)[1].actor_handle is original_handles[1]
         finally:
             _kill_all(group)
 
@@ -99,17 +100,17 @@ async def test_recover_default_filter_picks_all_dead_slots(
         group = _build_group(pg_tuple=pg, num_engines=3)
         _start(group)
 
-        old = [e.actor_handle for e in group.all_engines]
+        old = [e.actor_handle for e in flatten_cells(group.cells)]
         for i in (0, 2):
             ray.kill(old[i])
-            group.all_engines[i].mark_stopped()
+            flatten_cells(group.cells)[i].mark_stopped()
 
         try:
             await group.recover(port_cursors=PortCursors.empty())
             for i in (0, 2):
-                assert group.all_engines[i].is_allocated
-                assert group.all_engines[i].actor_handle is not old[i]
-            assert group.all_engines[1].actor_handle is old[1]
+                assert flatten_cells(group.cells)[i].is_allocated
+                assert flatten_cells(group.cells)[i].actor_handle is not old[i]
+            assert flatten_cells(group.cells)[1].actor_handle is old[1]
         finally:
             _kill_all(group)
 
@@ -137,15 +138,15 @@ async def remove_worker(self, **kwargs):
         group = _build_group(pg_tuple=pg, num_engines=1)
         group.router_ip, group.router_port = "10.0.0.9", 9000
         _start(group)
-        ray.kill(group.all_engines[0].actor_handle)
-        group.all_engines[0].mark_stopped()
+        ray.kill(flatten_cells(group.cells)[0].actor_handle)
+        flatten_cells(group.cells)[0].mark_stopped()
 
         try:
             with patch.object(ServerGroup, "_router_api_client", property(lambda self: _Recorder())):
                 await group.recover(port_cursors=PortCursors.empty(), filter_indices=[0])
 
-            assert [event["worker_url"] for event in events] == [group.all_engines[0].addr_info.server_url]
-            assert group.all_engines[0].is_alive
+            assert [event["worker_url"] for event in events] == [flatten_cells(group.cells)[0].addr_info.server_url]
+            assert flatten_cells(group.cells)[0].is_alive
         finally:
             _kill_all(group)
 
@@ -161,14 +162,14 @@ async def test_recover_with_offload_calls_release_then_resume(
         pg = placement_group_factory(2)
         group = _build_group(pg_tuple=pg, num_engines=2, needs_offload=True, update_weights=True)
         _start(group)
-        old = [e.actor_handle for e in group.all_engines]
+        old = [e.actor_handle for e in flatten_cells(group.cells)]
 
         ray.kill(old[0])
-        group.all_engines[0].mark_stopped()
+        flatten_cells(group.cells)[0].mark_stopped()
 
         try:
             await group.recover(port_cursors=PortCursors.empty(), filter_indices=[0])
-            calls = ray.get(group.all_engines[0].actor_handle.get_calls.remote())
+            calls = ray.get(flatten_cells(group.cells)[0].actor_handle.get_calls.remote())
             assert "init" in [c[0] for c in calls]
 
             server = mock_engine_http_servers.for_rank(0)
@@ -223,18 +224,18 @@ async def test_two_groups_recover_in_parallel_completes_without_deadlock(
 
         # Kill one engine in each group
         for g in (a, b):
-            old = g.all_engines[0].actor_handle
+            old = flatten_cells(g.cells)[0].actor_handle
             ray.kill(old)
-            g.all_engines[0].mark_stopped()
+            flatten_cells(g.cells)[0].mark_stopped()
 
         try:
             # Real concurrent recover via asyncio.gather
             await asyncio.gather(
                 a.recover(port_cursors=PortCursors.empty(), filter_indices=[0]),
                 b.recover(port_cursors=PortCursors.empty(), filter_indices=[0]),
             )
-            assert a.all_engines[0].is_allocated
-            assert b.all_engines[0].is_allocated
+            assert flatten_cells(a.cells)[0].is_allocated
+            assert flatten_cells(b.cells)[0].is_allocated
         finally:
             _kill_all(a)
             _kill_all(b)
@@ -258,7 +259,7 @@ async def test_simulate_crash_then_health_check_still_returns(
         pg = placement_group_factory(1)
         group = _build_group(pg_tuple=pg, num_engines=1)
         _start(group)
-        actor = group.all_engines[0].actor_handle
+        actor = flatten_cells(group.cells)[0].actor_handle
 
         try:
             ray.get(actor.simulate_crash.remote())