Support soft watchdog for tokenizer/detokenizer/dp-controller processes (#15607)

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:

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:
+            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

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(

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:
+def create_scheduler_watchdog(
-    """A watch dog thread that will try to kill the server itself if one forward batch takes too long."""
+    scheduler: Scheduler, watchdog_timeout: float, soft: bool = False
-
+) -> WatchdogRaw:
-    def __init__(
+    def dump_info() -> str:
-        self, scheduler: Scheduler, watchdog_timeout: float, soft: bool = False
+        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
         ):
-        self.scheduler = scheduler
+            _, info_msg = scheduler._check_mamba_memory()
-        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
+            _, info_msg = scheduler._check_radix_cache_memory()
-                    watchdog_last_time = current
+        return (
-            time.sleep(self.watchdog_timeout // 2)
+            f"{scheduler.cur_batch.batch_size()=}\n"
-
+            f"{scheduler.cur_batch.reqs=}\n"
-        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()
+    return WatchdogRaw(
-        logger.error(f"Watchdog timeout ({self.watchdog_timeout=}, {self.soft=})")
+        debug_name="Scheduler",
-        print(file=sys.stderr, flush=True)
+        get_counter=lambda: scheduler.forward_ct,
-        print(file=sys.stdout, flush=True)
+        is_active=lambda: scheduler.cur_batch is not None,
-
+        watchdog_timeout=watchdog_timeout,
-        if not self.soft:
+        soft=soft,
-            # Wait for some time so that the parent process can print the error.
+        dump_info=dump_info,
-            time.sleep(5)
+    )
-            self.parent_process.send_signal(signal.SIGQUIT)

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:
+            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,

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)