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

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)