diff --git a/miles/ray/rollout/inference_controller.py b/miles/ray/rollout/inference_controller.py index 621e3f5dbf2..f05edf0de3e 100644 --- a/miles/ray/rollout/inference_controller.py +++ b/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: diff --git a/miles/ray/rollout/server_cell.py b/miles/ray/rollout/server_cell.py index 42f1aaa625c..6381af66095 100644 --- a/miles/ray/rollout/server_cell.py +++ b/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=}" diff --git a/miles/utils/misc.py b/miles/utils/misc.py index 766bf7eddc9..036b352954c 100644 --- a/miles/utils/misc.py +++ b/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() diff --git a/miles/utils/workers/worker_provider/ray.py b/miles/utils/workers/worker_provider/ray.py index cb820fbf9fc..241ba19e33c 100644 --- a/miles/utils/workers/worker_provider/ray.py +++ b/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 diff --git a/tests/fast/ray/rollout/test_inference_controller_tick.py b/tests/fast/ray/rollout/test_inference_controller_tick.py new file mode 100644 index 00000000000..f6b7be0bb6a --- /dev/null +++ b/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 diff --git a/tests/fast/utils/test_misc.py b/tests/fast/utils/test_misc.py index 555ea3f7d0f..bc7873e767d 100644 --- a/tests/fast/utils/test_misc.py +++ b/tests/fast/utils/test_misc.py @@ -1,11 +1,25 @@ +import asyncio +import json import logging import socket +import subprocess +import sys from contextlib import ExitStack +from dataclasses import dataclass import pytest +from miles.utils.env_report import ENV_REPORT_PREFIX from miles.utils.http_utils import MILES_HOST_IP_ENV, get_host_info -from miles.utils.misc import NodeProbeMixin, filter_keys, get_current_node_ip, get_free_port +from miles.utils.misc import ( + NodeProbeMixin, + SimpleTicker, + cancel_and_await_task, + filter_keys, + get_current_node_ip, + get_free_port, + get_gpu_uuids, +) class TestFilterKeys: @@ -42,6 +56,73 @@ def test_missing_key_raises_key_error_and_logs(self, caplog): assert any("filter_keys" in record.message for record in caplog.records) +@dataclass(frozen=True) +class _FakeGpuHandle: + index: int + + +@dataclass(frozen=True) +class _FakeNvmlUuid: + text: str + + def __str__(self) -> str: + return self.text + + +class _FakeNvml: + def __init__( + self, + *, + uuid_by_index: dict[int, str], + init_error: Exception | None = None, + uuid_error_indices: frozenset[int] = frozenset(), + ) -> None: + self._uuid_by_index = uuid_by_index + self._init_error = init_error + self._uuid_error_indices = uuid_error_indices + + def nvmlInit(self) -> None: + if self._init_error is not None: + raise self._init_error + + def nvmlDeviceGetHandleByIndex(self, index: int) -> _FakeGpuHandle: + return _FakeGpuHandle(index=index) + + def nvmlDeviceGetUUID(self, handle: _FakeGpuHandle) -> _FakeNvmlUuid: + if handle.index in self._uuid_error_indices: + raise RuntimeError(f"nvml uuid lookup failed for {handle.index}") + return _FakeNvmlUuid(text=self._uuid_by_index[handle.index]) + + +class TestGetGpuUuids: + def test_get_gpu_uuids_returns_requested_nvml_uuids_in_order(self, monkeypatch) -> None: + """Each requested gpu index is resolved through NVML, coerced to str, and answered in request order.""" + fake_nvml = _FakeNvml(uuid_by_index={0: "GPU-zero", 1: "GPU-one", 2: "GPU-two"}) + monkeypatch.setitem(sys.modules, "pynvml", fake_nvml) + + uuids = get_gpu_uuids([2, 0]) + + assert uuids == ["GPU-two", "GPU-zero"] + assert all(isinstance(uuid, str) for uuid in uuids) + + def test_get_gpu_uuids_returns_none_per_gpu_when_nvml_fails(self, monkeypatch) -> None: + """A failing NVML init is swallowed and answered with exactly one None per requested gpu.""" + fake_nvml = _FakeNvml(uuid_by_index={}, init_error=RuntimeError("nvml unavailable")) + monkeypatch.setitem(sys.modules, "pynvml", fake_nvml) + + assert get_gpu_uuids([0, 1, 3]) == [None, None, None] + + def test_get_gpu_uuids_returns_all_none_when_one_lookup_fails(self, monkeypatch) -> None: + """A single failing uuid lookup yields all-None rather than a partial or short list.""" + fake_nvml = _FakeNvml( + uuid_by_index={0: "GPU-zero", 1: "GPU-one"}, + uuid_error_indices=frozenset({1}), + ) + monkeypatch.setitem(sys.modules, "pynvml", fake_nvml) + + assert get_gpu_uuids([0, 1]) == [None, None] + + class TestNodeProbeMixin: def test_get_node_ip_returns_nonempty_string(self): """The node ip probe answers with a usable address string.""" @@ -69,3 +150,194 @@ def test_get_gpu_uuids_returns_one_entry_per_gpu(self): uuids = NodeProbeMixin._get_gpu_uuids([0, 1, 2]) assert len(uuids) == 3 assert all(uuid is None or isinstance(uuid, str) for uuid in uuids) + + def test_collect_env_report_forwards_probe_context(self, monkeypatch, capsys) -> None: + """Role, rank and the launcher's partial report all reach the printed env report.""" + + def _failing_pip_inspect(*args, **kwargs) -> subprocess.CompletedProcess: + return subprocess.CompletedProcess(args=["pip", "inspect"], returncode=1, stdout="", stderr="no pip") + + monkeypatch.setattr("miles.utils.env_report.subprocess.run", _failing_pip_inspect) + + NodeProbeMixin._collect_env_report(role="rollout", rank=7, partial_env_report='{"flavor": "probe"}') + + lines = [line for line in capsys.readouterr().out.splitlines() if line.startswith(ENV_REPORT_PREFIX)] + assert len(lines) == 1 + parsed = json.loads(lines[0].removeprefix(ENV_REPORT_PREFIX)) + assert parsed["role"] == "rollout" + assert parsed["rank"] == 7 + assert parsed["launcher_env_report"] == {"flavor": "probe"} + + +async def _append(calls: list[int]) -> None: + calls.append(1) + + +class TestSimpleTicker: + async def test_it_keeps_calling_its_function(self): + """The ticked work only makes progress while the loop keeps coming back.""" + calls: list[int] = [] + + ticker = SimpleTicker(lambda: _append(calls), interval_seconds=0.0) + await asyncio.sleep(0.02) + await ticker.dispose() + + assert len(calls) > 1 + + async def test_it_survives_a_failing_call(self): + """A raising sweep must not silently kill the loop for every later round.""" + calls: list[int] = [] + + async def _boom() -> None: + calls.append(1) + raise RuntimeError("tick exploded") + + ticker = SimpleTicker(_boom, interval_seconds=0.0) + await asyncio.sleep(0.02) + await ticker.dispose() + + assert len(calls) > 1 + + async def test_dispose_stops_the_loop(self): + """A surviving loop would keep working after its owner is gone.""" + calls: list[int] = [] + + ticker = SimpleTicker(lambda: _append(calls), interval_seconds=0.0) + await asyncio.sleep(0.02) + await ticker.dispose() + calls_after_dispose = len(calls) + await asyncio.sleep(0.02) + + assert len(calls) == calls_after_dispose + + async def test_a_task_that_died_of_its_own_error_reports_it_on_dispose(self): + """Cancelling a task that already failed must surface the failure, not hide it behind the cancellation.""" + + async def _explode() -> None: + raise RuntimeError("ticker died") + + task = asyncio.create_task(_explode()) + await asyncio.sleep(0) + + with pytest.raises(RuntimeError, match="ticker died"): + await cancel_and_await_task(task) + + async def test_dispose_does_not_swallow_the_callers_cancellation(self): + """A teardown that eats its caller's cancellation lets the shutdown path run on regardless.""" + ticker = SimpleTicker(lambda: _append([]), interval_seconds=1000.0) + finished: list[str] = [] + + async def _dispose() -> None: + await ticker.dispose() + finished.append("returned") + + task = asyncio.create_task(_dispose()) + await asyncio.sleep(0) + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert task.cancelled() and finished == [] + + async def test_disposing_twice_is_harmless(self): + """Teardown paths overlap, so a second dispose must not raise.""" + ticker = SimpleTicker(lambda: _append([]), interval_seconds=0.0) + + await ticker.dispose() + await ticker.dispose() + + +class TestCancelAndAwaitTask: + async def test_it_swallows_the_cancellation_of_the_task_it_cancelled(self): + """The cancellation the helper itself asked for is teardown noise, not something to raise at the caller.""" + task = asyncio.create_task(asyncio.Event().wait()) + await asyncio.sleep(0) + + await cancel_and_await_task(task) + + assert task.cancelled() + + async def test_it_returns_only_after_the_task_finished_unwinding(self): + """Returning while the cancelled task is still unwinding lets it act once more after its owner is gone.""" + unwound: list[str] = [] + + async def _slow_teardown() -> None: + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(0) + unwound.append("cleaned") + + task = asyncio.create_task(_slow_teardown()) + await asyncio.sleep(0) + + await cancel_and_await_task(task) + + assert task.done() and unwound == ["cleaned"] + + async def test_it_requests_teardown_before_propagating_the_callers_cancellation(self): + """Being cancelled mid-teardown must still leave the task cancelled while the caller's cancellation travels on.""" + observed_cancel = asyncio.Event() + release = asyncio.Event() + returned: list[str] = [] + + async def _absorbs_the_first_cancel() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + observed_cancel.set() + await release.wait() + + task = asyncio.create_task(_absorbs_the_first_cancel()) + await asyncio.sleep(0) + + async def _caller() -> None: + await cancel_and_await_task(task) + returned.append("returned") + + caller = asyncio.create_task(_caller()) + await asyncio.wait_for(observed_cancel.wait(), timeout=2.0) + caller.cancel() + await asyncio.gather(caller, return_exceptions=True) + release.set() + await task + + assert caller.cancelled() and returned == [] + + async def test_a_shutdown_deadline_still_fires_when_the_task_is_slow_to_unwind(self): + """Eating the deadline's cancellation would let a timed-out shutdown continue as if it had made its deadline.""" + release = asyncio.Event() + returned: list[str] = [] + + async def _absorbs_the_first_cancel() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await release.wait() + + task = asyncio.create_task(_absorbs_the_first_cancel()) + await asyncio.sleep(0) + + with pytest.raises(TimeoutError): + async with asyncio.timeout(0.01): + await cancel_and_await_task(task) + returned.append("returned") + + release.set() + await task + + assert returned == [] + + async def test_a_task_that_fails_while_unwinding_reports_its_error(self): + """A teardown that raises its own error must not be mistaken for the cancellation the helper asked for.""" + + async def _fails_on_teardown() -> None: + try: + await asyncio.Event().wait() + finally: + raise RuntimeError("teardown exploded") + + task = asyncio.create_task(_fails_on_teardown()) + await asyncio.sleep(0) + + with pytest.raises(RuntimeError, match="teardown exploded"): + await cancel_and_await_task(task) diff --git a/tests/fast/utils/workers/worker_provider/test_ray.py b/tests/fast/utils/workers/worker_provider/test_ray.py index 938c9c21076..71d30ec2cb7 100644 --- a/tests/fast/utils/workers/worker_provider/test_ray.py +++ b/tests/fast/utils/workers/worker_provider/test_ray.py @@ -152,6 +152,22 @@ async def __call__(self, cell_id: str, info: CellInfo | None) -> None: await super().__call__(cell_id, info) +class _StuckOnCancelReconciler(_RecordingReconciler): + def __init__(self) -> None: + super().__init__() + self.entered = asyncio.Event() + self.release = asyncio.Event() + + async def __call__(self, cell_id: str, info: CellInfo | None) -> None: + await super().__call__(cell_id, info) + self.entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + await self.release.wait() + raise + + class TestRayWorkerProviderWatchCellsInitialSync: async def test_every_initial_cell_is_reconciled_before_the_watch_is_established(self): """Callers may assume the pool is fully observed once watch_cells returns.""" @@ -284,3 +300,28 @@ async def test_stopping_ends_the_polling(self): await asyncio.sleep(0.02) assert len(handle.get_cell_infos.calls) == settled + + async def test_stopping_does_not_swallow_the_callers_cancellation(self): + """A stop that eats its caller's cancellation lets a timed-out shutdown run on past its teardown.""" + handle = _make_watching_handle({}, {"cell-a": _cell_info("cell-a")}) + provider = RayWorkerProvider( + worker_manager_handle=handle, pool_ids=["inference-engine-0-0"], poll_interval_seconds=0.001 + ) + reconciler = _StuckOnCancelReconciler() + returned: list[str] = [] + + stop = await provider.watch_cells(reconciler) + await asyncio.wait_for(reconciler.entered.wait(), timeout=2.0) + + async def _stopper() -> None: + await stop() + returned.append("returned") + + stopper = asyncio.create_task(_stopper()) + await asyncio.sleep(0) + stopper.cancel() + await asyncio.gather(stopper, return_exceptions=True) + reconciler.release.set() + await stop() + + assert stopper.cancelled() and returned == []