diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 958cff1cd..2a9d28ea3 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -980,6 +980,8 @@ class SchedulerDisaggregationDecodeMixin: # Update last_batch self.last_batch = batch + self.maybe_send_health_check_signal() + @torch.no_grad() def event_loop_overlap_disagg_decode(self: Scheduler): self.result_queue = deque() @@ -1017,6 +1019,8 @@ class SchedulerDisaggregationDecodeMixin: # Update last_batch self.last_batch = batch + self.maybe_send_health_check_signal() + def _run_batch_prebuilt( self: Scheduler, batch: ScheduleBatch ) -> GenerationBatchResult: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 381e92590..e1c4cadfb 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -381,6 +381,8 @@ class SchedulerDisaggregationPrefillMixin: # Update last_batch self.last_batch = batch + self.maybe_send_health_check_signal() + @torch.no_grad() def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self.result_queue = deque() @@ -421,6 +423,8 @@ class SchedulerDisaggregationPrefillMixin: # Update last_batch self.last_batch = batch + self.maybe_send_health_check_signal() + def process_batch_result_disagg_prefill( self: Scheduler, batch: ScheduleBatch, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 774036268..8d9d25345 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1259,6 +1259,8 @@ class Scheduler( # Update last_batch self.last_batch = batch + + self.maybe_send_health_check_signal() if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.self_check_during_busy() @@ -1314,6 +1316,8 @@ class Scheduler( # Update last_batch self.last_batch = batch + + self.maybe_send_health_check_signal() if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.self_check_during_busy() @@ -2617,7 +2621,6 @@ class Scheduler( self.log_batch_result_stats(batch, result) self._maybe_clear_mm_inputs(batch) - self.maybe_send_health_check_signal() def maybe_send_health_check_signal(self): if self.return_health_check_ct: diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index c39a8fc13..f255a4425 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -139,6 +139,8 @@ class SchedulerPPMixin: self.pp_outputs = next_pp_outputs + self.maybe_send_health_check_signal() + # When the server is idle, self-check and re-init some states if server_is_idle: self.self_check_during_idle() @@ -314,6 +316,8 @@ class SchedulerPPMixin: self.running_batch.batch_is_full = False + self.maybe_send_health_check_signal() + # When the server is idle, self-check and re-init some states if server_is_idle and len(self.disagg_prefill_inflight_queue) == 0: self.self_check_during_idle() @@ -504,6 +508,8 @@ class SchedulerPPMixin: if self.server_args.disaggregation_decode_enable_offload_kvcache: queue_size += len(self.decode_offload_manager.ongoing_offload) + self.maybe_send_health_check_signal() + if server_is_idle and queue_size == 0: self.self_check_during_idle() diff --git a/python/sglang/srt/multiplex/multiplexing_mixin.py b/python/sglang/srt/multiplex/multiplexing_mixin.py index 1e1e858ae..620b18974 100644 --- a/python/sglang/srt/multiplex/multiplexing_mixin.py +++ b/python/sglang/srt/multiplex/multiplexing_mixin.py @@ -219,3 +219,5 @@ class SchedulerMultiplexMixin: self.split_prefill_batch = None wait_prefill_kernel_done = False adjust_stream_group = True + + self.maybe_send_health_check_signal()