From d267ee087bfc8f09f39b8ec99b1a884214806709 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Sat, 11 Oct 2025 17:04:23 +0800 Subject: [PATCH 1/6] more --- miles/ray/rollout.py | 75 ++-------------------------------- miles/utils/health_monitor.py | 76 +++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 71 deletions(-) create mode 100644 miles/utils/health_monitor.py diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index ca39d5923cc..f52d1eadfb4 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -15,6 +15,7 @@ from miles.ray.rollout_data_source import RolloutDataSourceWithBuffer 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.health_monitor import RolloutHealthMonitor from miles.utils.misc import load_function from miles.utils.ray_utils import Box from miles.utils.types import Sample @@ -63,13 +64,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(args) def dispose(self): if self._metric_checker is not None: @@ -83,7 +78,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 +88,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 +114,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..811c7a397eb --- /dev/null +++ b/miles/utils/health_monitor.py @@ -0,0 +1,76 @@ +import logging +import threading + +import ray + + +class RolloutHealthMonitor: + def __init__(self, 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 + + def start(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(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 From c9597de00711d707bfef83c8ddbd161b0ee22e97 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Sat, 11 Oct 2025 17:04:50 +0800 Subject: [PATCH 2/6] more --- miles/utils/health_monitor.py | 46 +++++++++++++++++------------------ 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/miles/utils/health_monitor.py b/miles/utils/health_monitor.py index 811c7a397eb..d6317d1c0fd 100644 --- a/miles/utils/health_monitor.py +++ b/miles/utils/health_monitor.py @@ -7,54 +7,54 @@ class RolloutHealthMonitor: def __init__(self, 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._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_engines: return False - assert self._health_monitor_thread is None, "Health monitor thread is already running." + assert self._thread is None, "Health monitor thread is already running." - self._health_monitor_stop_event = threading.Event() - self._health_monitor_thread = threading.Thread( + self._stop_event = threading.Event() + self._thread = threading.Thread( target=self._health_monitor_loop, name="RolloutHealthMonitor", daemon=True, ) - self._health_monitor_thread.start() + self._thread.start() return True def stop(self) -> None: - if not self._health_monitor_thread: + if not self._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(): + 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._health_monitor_thread = None - self._health_monitor_stop_event = None + self._thread = None + self._stop_event = None def _health_monitor_loop(self) -> None: - assert self._health_monitor_stop_event is not 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._health_monitor_stop_event.wait(self._health_check_first_wait): + if self._stop_event.wait(self._check_first_wait): return - while not self._health_monitor_stop_event.is_set(): + while not self._stop_event.is_set(): self._run_health_checks() - if self._health_monitor_stop_event.wait(self._health_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.rollout_engines): - if self._health_monitor_stop_event is not None and self._health_monitor_stop_event.is_set(): + if self._stop_event is not None and self._stop_event.is_set(): break self._check_engine_health(rollout_engine_id, engine) @@ -63,7 +63,7 @@ def _check_engine_health(self, rollout_engine_id, engine) -> None: return try: - ray.get(engine.health_generate.remote(timeout=self._health_check_timeout)) + 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.") for i in range(rollout_engine_id * self.nodes_per_engine, (rollout_engine_id + 1) * self.nodes_per_engine): From af0a35943b03f0e5d5d086455616f3fd67e2ccb5 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Sat, 11 Oct 2025 17:05:45 +0800 Subject: [PATCH 3/6] more --- miles/utils/health_monitor.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/miles/utils/health_monitor.py b/miles/utils/health_monitor.py index d6317d1c0fd..0023ad2c6bb 100644 --- a/miles/utils/health_monitor.py +++ b/miles/utils/health_monitor.py @@ -66,11 +66,14 @@ def _check_engine_health(self, rollout_engine_id, engine) -> None: 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.") - 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 + 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.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 From cfb32d250b4f9841092055abe4a71008181a1ea1 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Sat, 11 Oct 2025 17:06:47 +0800 Subject: [PATCH 4/6] more --- miles/ray/rollout.py | 2 +- miles/utils/health_monitor.py | 18 ++++++++++-------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index f52d1eadfb4..d27503b57e4 100644 --- a/miles/ray/rollout.py +++ b/miles/ray/rollout.py @@ -64,7 +64,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) - self._health_monitor = RolloutHealthMonitor(args) + self._health_monitor = RolloutHealthMonitor(self, args) def dispose(self): if self._metric_checker is not None: diff --git a/miles/utils/health_monitor.py b/miles/utils/health_monitor.py index 0023ad2c6bb..a8b21911318 100644 --- a/miles/utils/health_monitor.py +++ b/miles/utils/health_monitor.py @@ -5,8 +5,10 @@ class RolloutHealthMonitor: - def __init__(self, args): - # fault tolerance + 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 @@ -14,7 +16,7 @@ def __init__(self, args): self._check_first_wait = args.rollout_health_check_first_wait def start(self) -> bool: - if not self.rollout_engines: + if not self._rollout_manager.rollout_engines: return False assert self._thread is None, "Health monitor thread is already running." @@ -53,7 +55,7 @@ def _health_monitor_loop(self) -> None: break def _run_health_checks(self) -> None: - for rollout_engine_id, engine in enumerate(self.rollout_engines): + 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) @@ -69,11 +71,11 @@ def _check_engine_health(self, rollout_engine_id, engine) -> None: 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.nodes_per_engine, (rollout_engine_id + 1) * self.nodes_per_engine): - engine = self.all_rollout_engines[i] + 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: pass - self.all_rollout_engines[i] = None - self.rollout_engines[rollout_engine_id] = None + self._rollout_manager.all_rollout_engines[i] = None + self._rollout_manager.rollout_engines[rollout_engine_id] = None From d93a1a22cbe3d0a7d5214e53190f1f992abff3f5 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Sat, 11 Oct 2025 17:07:01 +0800 Subject: [PATCH 5/6] fmt --- miles/ray/rollout.py | 3 +-- miles/utils/health_monitor.py | 5 ++++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index d27503b57e4..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,9 +12,9 @@ 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.health_monitor import RolloutHealthMonitor from miles.utils.misc import load_function from miles.utils.ray_utils import Box from miles.utils.types import Sample diff --git a/miles/utils/health_monitor.py b/miles/utils/health_monitor.py index a8b21911318..73886aacc78 100644 --- a/miles/utils/health_monitor.py +++ b/miles/utils/health_monitor.py @@ -71,7 +71,10 @@ def _check_engine_health(self, rollout_engine_id, engine) -> None: 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): + 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) From 11fc39c3937920720ce3876461e78e4c5ac85a43 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Sat, 11 Oct 2025 17:07:58 +0800 Subject: [PATCH 6/6] fix --- miles/utils/health_monitor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miles/utils/health_monitor.py b/miles/utils/health_monitor.py index 73886aacc78..9cc5297a480 100644 --- a/miles/utils/health_monitor.py +++ b/miles/utils/health_monitor.py @@ -78,7 +78,7 @@ def _kill_engine(self, rollout_engine_id: int): engine = self._rollout_manager.all_rollout_engines[i] try: ray.kill(engine) - except Exception: - pass + 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