diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index ca39d5923cc..6546eb7ddf3 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -1,7 +1,6 @@ import logging import multiprocessing import random -import threading import time from pathlib import Path from typing import List, Union @@ -13,6 +12,7 @@ from miles.backends.sglang_utils.sglang_engine import SGLangEngine from miles.ray.rollout_data_source import RolloutDataSourceWithBuffer +from miles.utils.health_monitor import RolloutHealthMonitor from miles.utils.http_utils import find_available_port, get_host_info, init_http_client from miles.utils.metric_checker import MetricChecker from miles.utils.misc import load_function @@ -63,13 +63,7 @@ def __init__(self, args, pg, wandb_run_id): self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() self._metric_checker = MetricChecker.maybe_create(args) - - # fault tolerance - self._health_monitor_thread = None - self._health_monitor_stop_event = None - self._health_check_interval = args.rollout_health_check_interval - self._health_check_timeout = args.rollout_health_check_timeout - self._health_check_first_wait = args.rollout_health_check_first_wait + self._health_monitor = RolloutHealthMonitor(self, args) def dispose(self): if self._metric_checker is not None: @@ -83,7 +77,7 @@ def get_num_rollout_per_epoch(self): return len(self.data_source.dataset) // self.args.rollout_batch_size def generate(self, rollout_id): - monitor_started = self._start_health_monitor() + monitor_started = self._health_monitor.start() start_time = time.time() try: data = self._get_rollout_data(rollout_id=rollout_id) @@ -93,7 +87,7 @@ def generate(self, rollout_id): return Box(ray.put(data)) finally: if monitor_started: - self._stop_health_monitor() + self._health_monitor.stop() self.num_new_engines = init_rollout_engines(self.args, self.pg, self.all_rollout_engines) self.rollout_engines = self.all_rollout_engines[:: self.nodes_per_engine] @@ -119,68 +113,6 @@ def offload(self): def onload(self, tags: List[str] = None): return [engine.resume_memory_occupation.remote(tags=tags) for engine in self.rollout_engines] - def _start_health_monitor(self) -> bool: - if not self.rollout_engines: - return False - - assert self._health_monitor_thread is None, "Health monitor thread is already running." - - self._health_monitor_stop_event = threading.Event() - self._health_monitor_thread = threading.Thread( - target=self._health_monitor_loop, - name="RolloutHealthMonitor", - daemon=True, - ) - self._health_monitor_thread.start() - return True - - def _stop_health_monitor(self) -> None: - if not self._health_monitor_thread: - return - - assert self._health_monitor_stop_event is not None - self._health_monitor_stop_event.set() - timeout = self._health_check_timeout + self._health_check_interval + 5 - self._health_monitor_thread.join(timeout=timeout) - if self._health_monitor_thread.is_alive(): - logging.warning("Rollout health monitor thread did not terminate within %.1fs", timeout) - - self._health_monitor_thread = None - self._health_monitor_stop_event = None - - def _health_monitor_loop(self) -> None: - assert self._health_monitor_stop_event is not None - # TODO: need to be waiting for the large moe to be ready. this is hacky. - if self._health_monitor_stop_event.wait(self._health_check_first_wait): - return - while not self._health_monitor_stop_event.is_set(): - self._run_health_checks() - if self._health_monitor_stop_event.wait(self._health_check_interval): - break - - def _run_health_checks(self) -> None: - for rollout_engine_id, engine in enumerate(self.rollout_engines): - if self._health_monitor_stop_event is not None and self._health_monitor_stop_event.is_set(): - break - self._check_engine_health(rollout_engine_id, engine) - - def _check_engine_health(self, rollout_engine_id, engine) -> None: - if engine is None: - return - - try: - ray.get(engine.health_generate.remote(timeout=self._health_check_timeout)) - except Exception as e: - print(f"Health check timed out for rollout engine {rollout_engine_id} (ray timeout). Killing actor.") - for i in range(rollout_engine_id * self.nodes_per_engine, (rollout_engine_id + 1) * self.nodes_per_engine): - engine = self.all_rollout_engines[i] - try: - ray.kill(engine) - except Exception: - pass - self.all_rollout_engines[i] = None - self.rollout_engines[rollout_engine_id] = None - def _get_rollout_data(self, rollout_id): if self.args.load_debug_rollout_data: data = torch.load( diff --git a/miles/utils/health_monitor.py b/miles/utils/health_monitor.py new file mode 100644 index 00000000000..9cc5297a480 --- /dev/null +++ b/miles/utils/health_monitor.py @@ -0,0 +1,84 @@ +import logging +import threading + +import ray + + +class RolloutHealthMonitor: + def __init__(self, rollout_manager, args): + # TODO may remove this dependency after refactoring + self._rollout_manager = rollout_manager + + self._thread = None + self._stop_event = None + 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 + + def start(self) -> bool: + if not self._rollout_manager.rollout_engines: + return False + + assert self._thread is None, "Health monitor thread is already running." + + self._stop_event = threading.Event() + self._thread = threading.Thread( + target=self._health_monitor_loop, + name="RolloutHealthMonitor", + daemon=True, + ) + self._thread.start() + return True + + def stop(self) -> None: + if not self._thread: + return + + assert self._stop_event is not None + self._stop_event.set() + 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) + + self._thread = None + self._stop_event = None + + def _health_monitor_loop(self) -> None: + assert self._stop_event is not None + # TODO: need to be waiting for the large moe to be ready. this is hacky. + if self._stop_event.wait(self._check_first_wait): + return + while not self._stop_event.is_set(): + self._run_health_checks() + if self._stop_event.wait(self._check_interval): + break + + def _run_health_checks(self) -> None: + for rollout_engine_id, engine in enumerate(self._rollout_manager.rollout_engines): + if self._stop_event is not None and self._stop_event.is_set(): + break + self._check_engine_health(rollout_engine_id, engine) + + def _check_engine_health(self, rollout_engine_id, engine) -> None: + if engine is None: + return + + try: + ray.get(engine.health_generate.remote(timeout=self._check_timeout)) + except Exception as e: + print(f"Health check timed out for rollout engine {rollout_engine_id} (ray timeout). Killing actor.") + self._kill_engine(rollout_engine_id=rollout_engine_id) + + def _kill_engine(self, rollout_engine_id: int): + for i in range( + rollout_engine_id * self._rollout_manager.nodes_per_engine, + (rollout_engine_id + 1) * self._rollout_manager.nodes_per_engine, + ): + engine = self._rollout_manager.all_rollout_engines[i] + try: + ray.kill(engine) + except Exception as e: + print(f"Fail to kill engine and skip (e: {e})") + self._rollout_manager.all_rollout_engines[i] = None + self._rollout_manager.rollout_engines[rollout_engine_id] = None