Move engine teardown into ServerCell.stop

miles/ray/rollout/server_cell.py

@@ -1,11 +1,18 @@
+import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, NamedTuple
 
+import ray
+
 from miles.ray.rollout.server_engine import ServerEngine
 
 if TYPE_CHECKING:
     from miles.ray.rollout.rollout_server import RolloutServer
 
+logger = logging.getLogger(__name__)
+
+SHUTDOWN_TIMEOUT = 30
+
 
 @dataclass
 class ServerCell:
@@ -15,6 +22,25 @@ class ServerCell:
     def primary_engine(self) -> ServerEngine:
         return self.engines[0]
 
+    def stop(self):
+        for local_index, engine in enumerate(self.engines):
+            if engine.is_allocated:
+                logger.info(f"Shutting down and killing engine at cell-local index {local_index}")
+                try:
+                    ray.get(engine.actor_handle.shutdown.remote(), timeout=SHUTDOWN_TIMEOUT)
+                except Exception as e:
+                    logger.warning(
+                        f"Graceful shutdown of engine at cell-local index {local_index} failed, killing anyway (e: {e})"
+                    )
+                try:
+                    ray.kill(engine.actor_handle)
+                    logger.info(f"Successfully killed engine at cell-local index {local_index}")
+                except Exception as e:
+                    logger.warning(f"Fail to kill engine at cell-local index {local_index} (e: {e})")
+            else:
+                logger.info(f"Engine at cell-local index {local_index} is already None")
+            self.engines[local_index].mark_stopped()
+
     async def offload(self, tags: list[str] | None):
         return await self.primary_engine.api_client.release_memory_occupation(tags=tags)
 

miles/ray/rollout/server_group.py

@@ -15,15 +15,13 @@
     allocate_rollout_engine_addr_and_ports_external,
     allocate_rollout_engine_addr_and_ports_normal,
 )
-from miles.ray.rollout.server_cell import ServerCell, flatten_cells
+from miles.ray.rollout.server_cell import SHUTDOWN_TIMEOUT, 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
 
 logger = logging.getLogger(__name__)
 
-_SHUTDOWN_TIMEOUT = 30
-
 
 @dataclasses.dataclass
 class ServerGroup:
@@ -237,26 +235,16 @@ def _router_api_client(self) -> SGLangRouterApiClient:
     def stop_engines(self, engine_indices: list[int]):
         logger.info(f"Killing server {engine_indices=}...")
         try:
-            async_utils.run(asyncio.wait_for(self.unregister_workers(engine_indices), timeout=_SHUTDOWN_TIMEOUT))
+            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 = all_engines[i]
-            if engine.is_allocated:
-                logger.info(f"Shutting down and killing engine at index {i}")
-                try:
-                    ray.get(engine.actor_handle.shutdown.remote(), timeout=_SHUTDOWN_TIMEOUT)
-                except Exception as e:
-                    logger.warning(f"Graceful shutdown of engine at index {i} failed, killing anyway (e: {e})")
-                try:
-                    ray.kill(engine.actor_handle)
-                    logger.info(f"Successfully killed engine at index {i}")
-                except Exception as e:
-                    logger.warning(f"Fail to kill engine at index {i} (e: {e})")
-            else:
-                logger.info(f"Engine at index {i} is already None")
-            all_engines[i].mark_stopped()
+        for cell_index in sorted({i // self.nodes_per_engine for i in engine_indices}):
+            cell = self.cells[cell_index]
+            cell_engine_indices = range(cell_index * self.nodes_per_engine, (cell_index + 1) * self.nodes_per_engine)
+            assert set(cell_engine_indices) <= set(
+                engine_indices
+            ), f"stop_engines must cover whole cells ({engine_indices=}, {cell_index=})"
+            cell.stop()
 
     async def recover(self, port_cursors: PortCursors, filter_indices: list[int] | None = None):
         all_engines = flatten_cells(self.cells)

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

@@ -6,7 +6,7 @@
 import ray
 from tests.fast.ray.rollout.conftest import chunk_engines_into_cells, make_args
 
-import miles.ray.rollout.server_group as server_group_module
+import miles.ray.rollout.server_cell as server_cell_module
 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
@@ -61,7 +61,7 @@ def test_a_failing_shutdown_still_kills_the_actor(self, patched_sglang_engine, p
 
     def test_a_hanging_shutdown_does_not_block_teardown(self, monkeypatch, ray_local_mode):
         """A wedged engine must not stall teardown forever, since teardown is how a wedged engine is reclaimed."""
-        monkeypatch.setattr(server_group_module, "_SHUTDOWN_TIMEOUT", 0.5)
+        monkeypatch.setattr(server_cell_module, "SHUTDOWN_TIMEOUT", 0.5)
         group = _build_group(pg_tuple=(None, [], []))
         actor_handle = _HangingEngine.remote()
         flatten_cells(group.cells)[0].mark_allocated_uninitialized(actor_handle)

tests/fast/ray/rollout/test_server_group_router_registration.py

@@ -12,6 +12,7 @@
 from miles.utils import async_utils
 
 _MODULE = "miles.ray.rollout.server_group"
+_CELL_MODULE = "miles.ray.rollout.server_cell"
 
 
 class _RecordingRouterApiClient:
@@ -145,7 +146,7 @@ def test_stop_engines_unregisters_before_killing_the_actor():
 
     with (
         _with_recording_client(group),
-        patch(f"{_MODULE}.ray") as ray_mock,
+        patch(f"{_CELL_MODULE}.ray") as ray_mock,
     ):
         ray_mock.get.side_effect = lambda *args, **kwargs: events.append(("shutdown", {}))
         ray_mock.kill.side_effect = lambda handle: events.append(("kill", {}))
@@ -166,7 +167,7 @@ async def _reject():
 
     with (
         _with_recording_client(group),
-        patch(f"{_MODULE}.ray") as ray_mock,
+        patch(f"{_CELL_MODULE}.ray") as ray_mock,
     ):
         ray_mock.get.side_effect = lambda *args, **kwargs: events.append(("shutdown", {}))
         ray_mock.kill.side_effect = lambda handle: events.append(("kill", {}))
@@ -187,8 +188,8 @@ async def _hang():
 
     with (
         _with_recording_client(group),
-        patch(f"{_MODULE}._SHUTDOWN_TIMEOUT", 0.1),
-        patch(f"{_MODULE}.ray") as ray_mock,
+        patch(f"{_MODULE}.SHUTDOWN_TIMEOUT", 0.1),
+        patch(f"{_CELL_MODULE}.ray") as ray_mock,
     ):
         ray_mock.kill.side_effect = lambda handle: events.append(("kill", {}))
         group.stop_engines(engine_indices=[0])
@@ -204,7 +205,7 @@ def test_use_miles_router_reaches_both_router_calls():
 
     with (
         _with_recording_client(group),
-        patch(f"{_MODULE}.ray"),
+        patch(f"{_CELL_MODULE}.ray"),
     ):
         async_utils.run(group.register_workers([0]))
         group.stop_engines(engine_indices=[0])