diff --git a/miles/ray/rollout/server_cell.py b/miles/ray/rollout/server_cell.py index 0c88fdf8993..a67e1eea689 100644 --- a/miles/ray/rollout/server_cell.py +++ b/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) diff --git a/miles/ray/rollout/server_group.py b/miles/ray/rollout/server_group.py index 112910f7fa5..b5a3fe73fff 100644 --- a/miles/ray/rollout/server_group.py +++ b/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) diff --git a/tests/fast/ray/rollout/real_ray/test_server_group_teardown.py b/tests/fast/ray/rollout/real_ray/test_server_group_teardown.py index f7540b873ad..7c8b8322773 100644 --- a/tests/fast/ray/rollout/real_ray/test_server_group_teardown.py +++ b/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) diff --git a/tests/fast/ray/rollout/test_server_group_router_registration.py b/tests/fast/ray/rollout/test_server_group_router_registration.py index 7d04298cd4d..e217ffafaca 100644 --- a/tests/fast/ray/rollout/test_server_group_router_registration.py +++ b/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])