diff --git a/miles/ray/rollout/inference_controller.py b/miles/ray/rollout/inference_controller.py index be236296712..2b8a6f251d2 100644 --- a/miles/ray/rollout/inference_controller.py +++ b/miles/ray/rollout/inference_controller.py @@ -1,4 +1,3 @@ -import asyncio import logging from dataclasses import dataclass @@ -12,7 +11,6 @@ from miles.ray.rollout.router_manager import start_session_server from miles.ray.rollout.server_cell import get_cell_indexer_of_id_map from miles.ray.utils import Lock -from miles.utils.health_monitor import RolloutHealthMonitor logger = logging.getLogger(__name__) @@ -33,33 +31,20 @@ def __init__(self, args, pg): self.rollout_id = -1 self.eval_fleet = EvalFleet(args, srv=self.servers["eval"]) if args.eval_num_gpus > 0 else None - # TODO will be replaced by full ft, thus temporarily leave it without modifications - self._health_monitors = [] - self._rollout_ft_enabled = self.args.use_fault_tolerance and "rollout" in self.args.ft_components - self._ci_fault_injection_pending = False - if not self.args.debug_train_only and self._rollout_ft_enabled: - for srv in self.servers.values(): - for group in srv.server_groups: - monitor = RolloutHealthMonitor(group, args) - monitor.start() - self._health_monitors.append(monitor) - self._ci_fault_injection_pending = self.args.ci_test - # -------------------------- rollout lifecycle hooks ----------------------------- async def prepare_rollout(self, rollout_id): self.rollout_id = rollout_id - self._health_monitoring_resume() + await self.health_monitoring_resume() if self.args.ci_test and self._rollout_ft_enabled and rollout_id >= 2: await self._try_ci_fault_injection() dashboard_hooks.register_engines(self.servers) async def prepare_eval(self): - self._health_monitoring_resume() + await self.health_monitoring_resume() async def dispose(self): - for monitor in self._health_monitors: - monitor.stop() + pass # -------------------------- offload/onload ----------------------------- @@ -176,12 +161,21 @@ async def check_weights( # -------------------------- utils ----------------------------- async def health_monitoring_pause(self) -> None: - for monitor in self._health_monitors: - monitor.pause() + self._assert_rollout_fault_tolerance_is_unsupported() - def _health_monitoring_resume(self) -> None: - for monitor in self._health_monitors: - monitor.resume() + async def health_monitoring_resume(self) -> None: + self._assert_rollout_fault_tolerance_is_unsupported() + + @property + def _rollout_ft_enabled(self) -> bool: + return self.args.use_fault_tolerance and "rollout" in self.args.ft_components + + def _assert_rollout_fault_tolerance_is_unsupported(self) -> None: + if not self.args.debug_train_only and self._rollout_ft_enabled: + raise NotImplementedError( + "rollout fault tolerance is being rebuilt; health monitoring must pause before " + "get_updatable_engines_and_lock snapshots the engines" + ) @property def _server(self) -> RolloutServer | None: @@ -190,31 +184,8 @@ def _server(self) -> RolloutServer | None: return None return next(iter(self.servers.values())) - # TODO will be replaced by full ft, thus temporarily leave it without modifications async def _try_ci_fault_injection(self): - """Try to inject fault during generate (when health monitor is running).""" - if not self._ci_fault_injection_pending: - return - - # Only inject fault once - self._ci_fault_injection_pending = False - - if ( - self._server - and self._server.server_groups[0].all_engines - and self._server.server_groups[0].all_engines[0].is_allocated - ): - logger.info("CI Fault Injection: Simulating crash on engine 0 during generate") - try: - # This will cause the ray actor to exit - self._server.server_groups[0].all_engines[0].actor_handle.simulate_crash.remote() - # Wait for health monitor to detect the crash and mark engine as None - # health_check_interval + health_check_timeout + buffer - wait_time = self.args.rollout_health_check_interval + self.args.rollout_health_check_timeout + 5 - logger.info(f"CI Fault Injection: Waiting {wait_time}s for health monitor to detect crash") - await asyncio.sleep(wait_time) - except Exception as e: - logger.warning(f"CI Fault Injection failed: {e}") + raise NotImplementedError("rollout fault injection is being rebuilt with rollout fault tolerance") @dataclass(frozen=True) diff --git a/miles/ray/rollout/server_group.py b/miles/ray/rollout/server_group.py index c73271ebae1..b44384746b8 100644 --- a/miles/ray/rollout/server_group.py +++ b/miles/ray/rollout/server_group.py @@ -169,17 +169,11 @@ def start_engines( ] return init_handles, new_engine_indices - # There are two callers, only one of them will exist in a running system - # 1. For new callers (InferenceController.stop_cell, main thread, async), - # deliberately make this function non-async here to avoid introducing two states - # like "stopping (but not stopped)" vs "stopped", since single-thread async code will not yield - # without an await point - # it has the drawback of freezing the whole async thread, which may be avoided later by - # moving `shutdown` mainly to local code - # 2. For legacy callers (RolloutHealthMonitor, another thread, sync) - # it is still unsafe to be called in another thread - # because engine may be observed as non-stopped while being shutdown, - # but that is same as the original code + # Called from InferenceController.stop_cell (main thread, async): deliberately non-async here + # to avoid introducing two states like "stopping (but not stopped)" vs "stopped", since + # single-thread async code will not yield without an await point + # it has the drawback of freezing the whole async thread, which may be avoided later by + # moving `shutdown` mainly to local code def stop_engines(self, engine_indices: list[int]): logger.info(f"Killing server {engine_indices=}...") for i in engine_indices: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 035cf4ea95d..36ec3326a65 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -3671,7 +3671,7 @@ def _maybe_apply_dumper_overrides(args) -> None: return if args.use_fault_tolerance: - logger.info("Dumper mode: disabling --use-fault-tolerance to suppress RolloutHealthMonitor heartbeats") + logger.info("Dumper mode: disabling --use-fault-tolerance to suppress fault tolerance heartbeats") args.use_fault_tolerance = False logger.info("Dumper mode: all heartbeat mechanisms disabled") diff --git a/miles/utils/ft_utils/health_checker.py b/miles/utils/ft_utils/health_checker.py index ad98116f022..805ff0ee1ca 100644 --- a/miles/utils/ft_utils/health_checker.py +++ b/miles/utils/ft_utils/health_checker.py @@ -83,8 +83,7 @@ def resume(self) -> None: ... class SimpleHealthChecker(BaseHealthChecker): """Periodic async health checker. Calls *check_fn*; reports result via *on_result*. - After each ``resume()``, waits ``first_wait`` seconds before the first check - (matching ``RolloutHealthMonitor._need_first_wait`` semantics). + After each ``resume()``, waits ``first_wait`` seconds before the first check. """ def __init__( diff --git a/miles/utils/health_monitor.py b/miles/utils/health_monitor.py deleted file mode 100644 index a85f46e923c..00000000000 --- a/miles/utils/health_monitor.py +++ /dev/null @@ -1,167 +0,0 @@ -import logging -import threading - -import ray - -from miles.ray.rollout.server_group import ServerGroup - -logger = logging.getLogger(__name__) - - -class RolloutHealthMonitor: - """Health monitor for rollout engines. - - The monitor runs continuously once started, but can be paused/resumed - based on whether the engines are offloaded (cannot health check when offloaded). - - Lifecycle: - - start(): Start the monitor thread (called once during initialization) - - pause(): Pause health checking (called when offloading engines) - - resume(): Resume health checking (called when onloading engines) - - stop(): Stop the monitor thread completely (called during dispose) - """ - - def __init__(self, server_group: ServerGroup, args): - self._server_group = server_group - - self._thread = None - self._stop_event = None - self._pause_event = None # When set, health checking is paused - self._check_interval = args.rollout_health_check_interval - self._check_timeout = args.rollout_health_check_timeout - self._check_first_wait = args.rollout_health_check_first_wait - self._need_first_wait = True # Need to wait after each resume - self._is_checking_enabled = False # Track if health checking should be active - - def start(self) -> bool: - """Start the health monitor thread. Called once during initialization. - - Returns: - True if the monitor was started, False if there are no engines to monitor. - """ - if not self._server_group.all_engines: - return False - - if self._thread is not None: - logger.warning("Health monitor thread is already running.") - return True - - logger.info("Starting RolloutHealthMonitor...") - self._stop_event = threading.Event() - self._pause_event = threading.Event() - self._pause_event.set() # Start in paused state until resume() is called - self._thread = threading.Thread( - target=self._health_monitor_loop, - name="RolloutHealthMonitor", - daemon=True, - ) - self._thread.start() - logger.info("RolloutHealthMonitor started (in paused state).") - return True - - def stop(self) -> None: - """Stop the health monitor thread completely. Called during dispose.""" - if not self._thread: - return - - logger.info("Stopping RolloutHealthMonitor...") - assert self._stop_event is not None - self._stop_event.set() - # Also clear pause to let the thread exit - if self._pause_event: - self._pause_event.clear() - timeout = self._check_timeout + self._check_interval + 5 - self._thread.join(timeout=timeout) - if self._thread.is_alive(): - logging.warning("Rollout health monitor thread did not terminate within %.1fs", timeout) - else: - logger.info("RolloutHealthMonitor stopped.") - - self._thread = None - self._stop_event = None - self._pause_event = None - self._is_checking_enabled = False - - def pause(self) -> None: - """Pause health checking. Called when engines are offloaded.""" - if self._pause_event is None: - return - logger.info("Pausing health monitor...") - self._pause_event.set() - self._is_checking_enabled = False - - def resume(self) -> None: - """Resume health checking. Called when engines are onloaded.""" - if self._pause_event is None: - return - logger.info("Resuming health monitor...") - self._need_first_wait = True # Need to wait after each resume - self._pause_event.clear() - self._is_checking_enabled = True - - def is_checking_enabled(self) -> bool: - """Return whether health checking is currently enabled (not paused).""" - return self._is_checking_enabled - - def _health_monitor_loop(self) -> None: - assert self._stop_event is not None - assert self._pause_event is not None - - while not self._stop_event.is_set(): - # Wait while paused - while self._pause_event.is_set() and not self._stop_event.is_set(): - self._stop_event.wait(timeout=0.5) - - if self._stop_event.is_set(): - break - - # Do first wait after each resume (for large MoE models to be ready) - if self._need_first_wait: - logger.info(f"Health monitor doing first wait after resume: {self._check_first_wait}s") - if self._stop_event.wait(self._check_first_wait): - logger.info("Health monitor stopped during first wait.") - break - if self._pause_event.is_set(): - # Got paused during first wait, skip this round and wait again next resume - logger.info("Health monitor paused during first wait, will wait again next resume.") - continue - self._need_first_wait = False - - # Run health checks - if not self._pause_event.is_set() and not self._stop_event.is_set(): - self._run_health_checks() - - # Wait for next check interval - if self._stop_event.wait(self._check_interval): - break - - def _run_health_checks(self) -> None: - for rollout_engine_id, engine in enumerate(self._server_group.engines): - if self._stop_event is not None and self._stop_event.is_set(): - break - if self._pause_event is not None and self._pause_event.is_set(): - break - self._check_engine_health(rollout_engine_id, engine) - - def _check_engine_health(self, rollout_engine_id, engine) -> None: - if not engine.is_allocated: - logger.info(f"Skipping health check for engine {rollout_engine_id} (None)") - return - - try: - ray.get(engine.actor_handle.health_generate.remote(timeout=self._check_timeout)) - except Exception as e: - logger.error( - f"Health check failed for rollout engine {rollout_engine_id} (ray timeout or error). Killing actor. Exception: {e}" - ) - nodes_per_engine = self._server_group.nodes_per_engine - self._server_group.stop_engines( - engine_indices=list( - range( - rollout_engine_id * nodes_per_engine, - (rollout_engine_id + 1) * nodes_per_engine, - ) - ) - ) - else: - logger.debug(f"Health check passed for rollout engine {rollout_engine_id}") diff --git a/tests/fast/ray/rollout/conftest.py b/tests/fast/ray/rollout/conftest.py index ff06d614bd8..060abda9451 100644 --- a/tests/fast/ray/rollout/conftest.py +++ b/tests/fast/ray/rollout/conftest.py @@ -94,6 +94,7 @@ def make_args(**overrides: Any) -> Namespace: # offload / fault tolerance offload_rollout=False, use_fault_tolerance=False, + ft_components=[], rollout_health_check_interval=10.0, rollout_health_check_timeout=30.0, # checkpoint / data source diff --git a/tests/fast/ray/rollout/real_ray/test_inference_controller.py b/tests/fast/ray/rollout/real_ray/test_inference_controller.py index 794614df5cd..48bc292734b 100644 --- a/tests/fast/ray/rollout/real_ray/test_inference_controller.py +++ b/tests/fast/ray/rollout/real_ray/test_inference_controller.py @@ -112,51 +112,6 @@ async def _assert_engine_dies(actor_handle, *, deadline_s: float = 15.0, poll_in @pytest.mark.asyncio class TestInferenceControllerInit: - @pytest.mark.parametrize( - ("ft_components", "expected_monitor_count", "expected_injection_pending"), - [ - (["train"], 0, False), - (["rollout"], 1, True), - ], - ) - async def test_rollout_ft_lifecycle_follows_selected_component( - self, - ray_local_mode, - placement_group_factory, - tmp_path, - patch_low_level, - monkeypatch, - ft_components, - expected_monitor_count, - expected_injection_pending, - ): - import miles.ray.rollout.inference_controller as ictl - - started_monitors = [] - - class FakeMonitor: - def __init__(self, group, args): - self.group = group - - def start(self): - started_monitors.append(self) - - def stop(self): - pass - - monkeypatch.setattr(ictl, "RolloutHealthMonitor", FakeMonitor) - args = _make_test_args(tmp_path, models=[("actor", True)]) - args.use_fault_tolerance = True - args.ft_components = ft_components - args.ci_test = True - pg = placement_group_factory(2) - - controller = _make_controller(args, pg) - - assert len(started_monitors) == expected_monitor_count - assert len(controller._health_monitors) == expected_monitor_count - assert controller._ci_fault_injection_pending is expected_injection_pending - async def test_init_creates_live_mock_engines_via_real_start_rollout_servers( self, ray_local_mode, @@ -493,3 +448,76 @@ async def test_recovers_dead_engine_after_rollout_started( assert slot0.is_allocated assert slot0.actor_handle is not actor0_before assert ray.get(slot0.actor_handle.health_generate.remote(timeout=1.0)) is True + + +@pytest.mark.asyncio +class TestRolloutFaultToleranceIsUnsupported: + async def test_health_monitoring_hooks_are_noops_without_fault_tolerance( + self, + ray_local_mode, + placement_group_factory, + tmp_path, + patch_low_level, + ): + """A plain run never asked for fault tolerance, so the hooks stay out of its way.""" + args = _make_test_args(tmp_path, models=[("actor", True)]) + pg = placement_group_factory(2) + + controller = _make_controller(args, pg) + + await controller.health_monitoring_pause() + await controller.health_monitoring_resume() + + async def test_health_monitoring_hooks_refuse_to_run_under_fault_tolerance( + self, + ray_local_mode, + placement_group_factory, + tmp_path, + patch_low_level, + ): + """Asking for fault tolerance must fail loudly, not run unmonitored.""" + args = _make_test_args(tmp_path, models=[("actor", True)]) + pg = placement_group_factory(2) + + controller = _make_controller(args, pg) + controller.args.use_fault_tolerance = True + controller.args.ft_components = ["rollout"] + + with pytest.raises(NotImplementedError): + await controller.health_monitoring_pause() + with pytest.raises(NotImplementedError): + await controller.health_monitoring_resume() + + async def test_health_monitoring_hooks_are_noops_when_fault_tolerance_skips_rollout( + self, + ray_local_mode, + placement_group_factory, + tmp_path, + patch_low_level, + ): + """Fault tolerance limited to training never monitored the engines, so nothing is lost.""" + args = _make_test_args(tmp_path, models=[("actor", True)]) + pg = placement_group_factory(2) + + controller = _make_controller(args, pg) + controller.args.use_fault_tolerance = True + controller.args.ft_components = ["train"] + + await controller.health_monitoring_pause() + await controller.health_monitoring_resume() + + async def test_fault_injection_refuses_to_run( + self, + ray_local_mode, + placement_group_factory, + tmp_path, + patch_low_level, + ): + """The injector depended on the deleted monitor to observe the crash.""" + args = _make_test_args(tmp_path, models=[("actor", True)]) + pg = placement_group_factory(2) + + controller = _make_controller(args, pg) + + with pytest.raises(NotImplementedError): + await controller._try_ci_fault_injection()