[Scheduler] Decouple maybe_send_health_check_signal from process_batch_result (#20227)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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:

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,

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:

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()
 

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()