diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1f714baca..d5a946606 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -173,6 +173,7 @@ class Envs: SGLANG_DETECT_SLOW_RANK = EnvBool(False) SGLANG_TEST_STUCK_DETOKENIZER = EnvFloat(0) SGLANG_TEST_STUCK_DP_CONTROLLER = EnvFloat(0) + SGLANG_TEST_STUCK_SCHEDULER_INIT = EnvFloat(0) SGLANG_TEST_STUCK_TOKENIZER = EnvFloat(0) SGLANG_TEST_CRASH_AFTER_STREAM_OUTPUTS = EnvInt(0) IS_BLACKWELL = EnvBool(False) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 49088634a..35cb7a379 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -257,6 +257,9 @@ class Scheduler( pp_rank: int, dp_rank: Optional[int], ): + self.is_initializing = True + self.init_soft_watchdog(server_args) + # Parse args self.server_args = server_args self.tp_rank = tp_rank @@ -334,6 +337,9 @@ class Scheduler( # Launch a model worker and draft model worker if using speculative decoding self.init_model_worker() + if (t := envs.SGLANG_TEST_STUCK_SCHEDULER_INIT.get()) > 0: + time.sleep(t) + # Init cache and memory pool self.init_cache_with_memory_pool() @@ -367,6 +373,8 @@ class Scheduler( # Init request dispatcher self.init_request_dispatcher() + self.is_initializing = False + def init_model_config(self): self.model_config = ModelConfig.from_server_args(self.server_args) self.dllm_config = ( # For diffusion LLM @@ -793,15 +801,17 @@ class Scheduler( ) / envs.SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS.get() self.new_token_ratio = self.init_new_token_ratio + def init_soft_watchdog(self, server_args: ServerArgs): + if (x := server_args.soft_watchdog_timeout) is not None: + self.soft_watchdog = create_scheduler_watchdog( + self, watchdog_timeout=x, soft=True + ) + def init_watch_dog_memory_saver_input_blocker(self): # Start watchdog thread 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 = 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 60360574d..ac11e4ff8 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -337,7 +337,7 @@ def create_scheduler_watchdog( scheduler: Scheduler, watchdog_timeout: float, soft: bool = False ) -> WatchdogRaw: def dump_info() -> str: - if disable_request_logging(): + if scheduler.is_initializing or disable_request_logging(): return "" if scheduler.is_hybrid_swa: _, info_msg = scheduler._check_hybrid_memory() @@ -355,8 +355,9 @@ def create_scheduler_watchdog( return WatchdogRaw( debug_name="Scheduler", - get_counter=lambda: scheduler.forward_ct, - is_active=lambda: scheduler.cur_batch is not None, + get_counter=lambda: getattr(scheduler, "forward_ct", 0), + is_active=lambda: scheduler.is_initializing + or getattr(scheduler, "cur_batch", None) is not None, watchdog_timeout=watchdog_timeout, soft=soft, dump_info=dump_info, diff --git a/test/registered/debug_utils/test_soft_watchdog.py b/test/registered/debug_utils/test_soft_watchdog.py index a18d26d63..dfd523adf 100644 --- a/test/registered/debug_utils/test_soft_watchdog.py +++ b/test/registered/debug_utils/test_soft_watchdog.py @@ -73,5 +73,10 @@ class TestSoftWatchdogTokenizer(BaseTestSoftWatchdog, CustomTestCase): expected_message = "TokenizerManager watchdog timeout" +class TestSoftWatchdogSchedulerInit(BaseTestSoftWatchdog, CustomTestCase): + env_override = lambda: envs.SGLANG_TEST_STUCK_SCHEDULER_INIT.override(30) + expected_message = "Scheduler watchdog timeout" + + if __name__ == "__main__": unittest.main()