diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index ca75f5b43..128dcfe32 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -59,6 +59,7 @@ from sglang.srt.utils.common import ( maybe_reindex_device_id, ) from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter +from sglang.srt.utils.watchdog import Watchdog from sglang.utils import TypeBasedDispatcher, get_exception_traceback logger = logging.getLogger(__name__) @@ -173,6 +174,12 @@ class DataParallelController: self.init_dispatcher() + self.watchdog = Watchdog.create( + debug_name="DataParallelController", + watchdog_timeout=server_args.soft_watchdog_timeout, + soft=True, + ) + def send_to_all_workers(self, obj): for worker in self.workers: worker.send_pyobj(obj) @@ -519,6 +526,7 @@ class DataParallelController: def event_loop(self): while True: while True: + self.watchdog.feed() try: recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK) except zmq.ZMQError: diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 9a207b554..3ff06fe63 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -41,6 +41,7 @@ from sglang.srt.utils import ( kill_itself_when_parent_died, ) from sglang.srt.utils.hf_transformers_utils import get_tokenizer +from sglang.srt.utils.watchdog import Watchdog from sglang.utils import ( TypeBasedDispatcher, find_printable_text, @@ -111,13 +112,21 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): ] ) + self.watchdog = Watchdog.create( + debug_name="DetokenizerManager", + watchdog_timeout=server_args.soft_watchdog_timeout, + soft=True, + ) + def event_loop(self): """The event loop that handles requests""" while True: - recv_obj = self.recv_from_scheduler.recv_pyobj() + with self.watchdog.disable(): + recv_obj = self.recv_from_scheduler.recv_pyobj() output = self._request_dispatcher(recv_obj) if output is not None: self.send_to_tokenizer.send_pyobj(output) + self.watchdog.feed() def trim_matched_stop( self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5e48d95fd..189a32bb5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -147,7 +147,7 @@ from sglang.srt.managers.scheduler_profiler_mixin import SchedulerProfilerMixin from sglang.srt.managers.scheduler_recv_skipper import SchedulerRecvSkipper from sglang.srt.managers.scheduler_runtime_checker_mixin import ( SchedulerRuntimeCheckerMixin, - SchedulerWatchdog, + create_scheduler_watchdog, ) from sglang.srt.managers.scheduler_update_weights_mixin import ( SchedulerUpdateWeightsMixin, @@ -797,11 +797,13 @@ class Scheduler( def init_watch_dog_memory_saver_input_blocker(self): # Start watchdog thread - self.watchdog = SchedulerWatchdog( + self.watchdog = create_scheduler_watchdog( self, watchdog_timeout=self.server_args.watchdog_timeout ) if (x := self.server_args.soft_watchdog_timeout) is not None: - self.soft_watchdog = SchedulerWatchdog(self, watchdog_timeout=x, soft=True) + self.soft_watchdog = create_scheduler_watchdog( + self, watchdog_timeout=x, soft=True + ) # Init memory saver, profiler and metric stats self.memory_saver_adapter = TorchMemorySaverAdapter.create( diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index e00db8945..718f25d8b 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -1,15 +1,10 @@ from __future__ import annotations import logging -import signal -import sys -import threading import time import warnings from typing import TYPE_CHECKING -import psutil - from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ScheduleBatch @@ -18,9 +13,9 @@ from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache from sglang.srt.utils.common import ( ceil_align, disable_request_logging, - pyspy_dump_schedulers, raise_error_or_warn, ) +from sglang.srt.utils.watchdog import WatchdogRaw if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler @@ -341,62 +336,31 @@ class SchedulerRuntimeCheckerMixin: self.maybe_sleep_on_idle() -class SchedulerWatchdog: - """A watch dog thread that will try to kill the server itself if one forward batch takes too long.""" +def create_scheduler_watchdog( + scheduler: Scheduler, watchdog_timeout: float, soft: bool = False +) -> WatchdogRaw: + def dump_info() -> str: + if disable_request_logging(): + return "" + if scheduler.is_hybrid_swa: + _, info_msg = scheduler._check_hybrid_memory() + elif scheduler.is_hybrid_ssm and isinstance( + scheduler.tree_cache, MambaRadixCache + ): + _, info_msg = scheduler._check_mamba_memory() + else: + _, info_msg = scheduler._check_radix_cache_memory() + return ( + f"{scheduler.cur_batch.batch_size()=}\n" + f"{scheduler.cur_batch.reqs=}\n" + f"{info_msg}" + ) - def __init__( - self, scheduler: Scheduler, watchdog_timeout: float, soft: bool = False - ): - self.scheduler = scheduler - self.soft = soft - - self.watchdog_timeout = watchdog_timeout - t = threading.Thread(target=self._watchdog_thread, daemon=True) - t.start() - self.parent_process = psutil.Process().parent() - - def _watchdog_thread(self): - while True: - self._watchdog_once() - - def _watchdog_once(self): - watchdog_last_forward_ct = 0 - watchdog_last_time = time.perf_counter() - - while True: - current = time.perf_counter() - if self.scheduler.cur_batch is not None: - if watchdog_last_forward_ct == self.scheduler.forward_ct: - if current > watchdog_last_time + self.watchdog_timeout: - break - else: - watchdog_last_forward_ct = self.scheduler.forward_ct - watchdog_last_time = current - time.sleep(self.watchdog_timeout // 2) - - if not disable_request_logging(): - # TODO extract this duplicated logic w/ another place - # Print batch size and memory pool info to check whether there are de-sync issues. - if self.scheduler.is_hybrid_swa: - _, info_msg = self.scheduler._check_hybrid_memory() - elif self.scheduler.is_hybrid_ssm and isinstance( - self.scheduler.tree_cache, MambaRadixCache - ): - _, info_msg = self.scheduler._check_mamba_memory() - else: - _, info_msg = self.scheduler._check_radix_cache_memory() - logger.error( - f"{self.scheduler.cur_batch.batch_size()=}\n" - f"{self.scheduler.cur_batch.reqs=}\n" - f"{info_msg}" - ) - - pyspy_dump_schedulers() - logger.error(f"Watchdog timeout ({self.watchdog_timeout=}, {self.soft=})") - print(file=sys.stderr, flush=True) - print(file=sys.stdout, flush=True) - - if not self.soft: - # Wait for some time so that the parent process can print the error. - time.sleep(5) - self.parent_process.send_signal(signal.SIGQUIT) + return WatchdogRaw( + debug_name="Scheduler", + get_counter=lambda: scheduler.forward_ct, + is_active=lambda: scheduler.cur_batch is not None, + watchdog_timeout=watchdog_timeout, + soft=soft, + dump_info=dump_info, + ) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 073584044..41ef9822f 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -111,6 +111,7 @@ from sglang.srt.utils.hf_transformers_utils import ( get_tokenizer, get_tokenizer_from_processor, ) +from sglang.srt.utils.watchdog import Watchdog from sglang.utils import TypeBasedDispatcher, get_exception_traceback asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) @@ -404,6 +405,12 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi ) self.init_communicators(server_args) + self.watchdog = Watchdog.create( + debug_name="TokenizerManager", + watchdog_timeout=server_args.soft_watchdog_timeout, + soft=True, + ) + async def generate_request( self, obj: Union[GenerateReqInput, EmbeddingReqInput], @@ -1499,9 +1506,11 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi async def handle_loop(self): """The event loop that handles requests""" while True: - recv_obj = await self.recv_from_detokenizer.recv_pyobj() + with self.watchdog.disable(): + recv_obj = await self.recv_from_detokenizer.recv_pyobj() self._result_dispatcher(recv_obj) self.last_receive_tstamp = time.time() + self.watchdog.feed() def _add_metric_if_present( self, diff --git a/python/sglang/srt/utils/watchdog.py b/python/sglang/srt/utils/watchdog.py new file mode 100644 index 000000000..2bbdcffbe --- /dev/null +++ b/python/sglang/srt/utils/watchdog.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import logging +import signal +import sys +import threading +import time +from contextlib import contextmanager +from typing import Callable, Optional + +import psutil + +from sglang.srt.utils.common import pyspy_dump_schedulers + +logger = logging.getLogger(__name__) + + +class Watchdog: + @staticmethod + def create( + debug_name: str, + watchdog_timeout: Optional[float], + soft: bool = False, + ) -> Watchdog: + if watchdog_timeout is None: + return _WatchdogNoop() + return _WatchdogReal( + debug_name=debug_name, + watchdog_timeout=watchdog_timeout, + soft=soft, + ) + + def feed(self): + pass + + @contextmanager + def disable(self): + yield + + +class _WatchdogReal(Watchdog): + def __init__( + self, + debug_name: str, + watchdog_timeout: float, + soft: bool = False, + ): + self._counter = 0 + self._active = True + self._raw = WatchdogRaw( + debug_name=debug_name, + get_counter=lambda: self._counter, + is_active=lambda: self._active, + watchdog_timeout=watchdog_timeout, + soft=soft, + ) + logger.info(f"Watchdog {self._raw.debug_name} initialized.") + + def feed(self): + self._counter += 1 + + @contextmanager + def disable(self): + assert self._active + self._active = False + try: + yield + finally: + assert not self._active + self._active = True + + +class _WatchdogNoop(Watchdog): + pass + + +class WatchdogRaw: + def __init__( + self, + debug_name: str, + get_counter: Callable[[], int], + is_active: Callable[[], bool], + watchdog_timeout: float, + soft: bool = False, + dump_info: Optional[Callable[[], str]] = None, + ): + self.debug_name = debug_name + self.get_counter = get_counter + self.is_active = is_active + self.watchdog_timeout = watchdog_timeout + self.soft = soft + self.dump_info = dump_info + + self.parent_process = psutil.Process().parent() + t = threading.Thread(target=self._watchdog_thread, daemon=True) + t.start() + + def _watchdog_thread(self): + try: + while True: + self._watchdog_once() + except Exception as e: + logger.error( + f"{self.debug_name} watchdog thread crashed: {e}", exc_info=True + ) + + def _watchdog_once(self): + watchdog_last_counter = 0 + watchdog_last_time = time.perf_counter() + + while True: + current = time.perf_counter() + if self.is_active(): + current_counter = self.get_counter() + if watchdog_last_counter == current_counter: + if current > watchdog_last_time + self.watchdog_timeout: + break + else: + watchdog_last_counter = current_counter + watchdog_last_time = current + time.sleep(self.watchdog_timeout / 2) + + if self.dump_info is not None and (info_msg := self.dump_info()): + logger.error(f"{self.debug_name} debug info:\n{info_msg}") + + pyspy_dump_schedulers() + logger.error( + f"{self.debug_name} watchdog timeout " + f"({self.watchdog_timeout=}, {self.soft=})" + ) + print(file=sys.stderr, flush=True) + print(file=sys.stdout, flush=True) + + if not self.soft: + # Wait for some time so that the parent process can print the error. + time.sleep(5) + self.parent_process.send_signal(signal.SIGQUIT)