[Scheduler] Unify idle checks into is_fully_idle() and fix weight update test (#20296)

python/sglang/srt/managers/scheduler.py

@@ -1495,35 +1495,12 @@ class Scheduler(
         now = time.monotonic()
         self.session_controller.maybe_reap(now)
         for recv_req in recv_reqs:
-            # If it is a health check generation request and there are running requests, ignore it.
-            if is_health_check_generate_req(recv_req):
-                # Check if there are requests being processed
-                has_running_requests = (
-                    self.chunked_req is not None
-                    or self.dllm_manager.any_staging_reqs()
-                    or not self.running_batch.is_empty()
-                    or len(self.offload_tags) > 0
-                )
-
-                # In PD disaggregation mode, also check if health check would be blocked
-                # in special queues (bootstrap/prealloc) due to external factors
-                will_block_in_pd_queue = False
-                if self.disaggregation_mode == DisaggregationMode.PREFILL:
-                    # If bootstrap queue has backlog, health check will also be blocked there
-                    will_block_in_pd_queue = (
-                        len(self.disagg_prefill_bootstrap_queue.queue) > 0
-                        or len(self.disagg_prefill_inflight_queue) > 0
-                    )
-                elif self.disaggregation_mode == DisaggregationMode.DECODE:
-                    # If prealloc/transfer queue has backlog, health check will also be blocked there
-                    will_block_in_pd_queue = (
-                        len(self.disagg_decode_prealloc_queue.queue) > 0
-                        or len(self.disagg_decode_transfer_queue.queue) > 0
-                    )
-
-                if has_running_requests or will_block_in_pd_queue:
-                    self.return_health_check_ct += 1
-                    continue
+            # Skip health check when server is busy — ongoing requests already carry health info.
+            if is_health_check_generate_req(recv_req) and not self.is_fully_idle(
+                for_health_check=True
+            ):
+                self.return_health_check_ct += 1
+                continue
 
             output = self._request_dispatcher(recv_req)
             if output is not None:
@@ -2647,20 +2624,36 @@ class Scheduler(
             if_success = False
         return ClearHiCacheReqOutput(success=if_success)
 
-    def _is_idle_for_hicache_storage_op(self) -> bool:
-        """Stricter idle check for storage attach/detach.
+    def is_fully_idle(self, for_health_check=False) -> bool:
+        # Batch running status
+        idle = (
+            self.running_batch.is_empty()
+            and self.chunked_req is None
+            and not self.dllm_manager.any_staging_reqs()
+            and (self.last_batch is None or self.last_batch.is_empty())
+            and (self.cur_batch is None or self.cur_batch.is_empty())
+            and (not self.enable_overlap or len(self.result_queue) == 0)
+            and (self.pp_size == 1 or all(x.is_empty() for x in self.running_mbs))
+        )
 
-        We require:
-        - no running batches (including overlap/pp/disagg paths) via `_is_no_request()`
-        - no queued requests in scheduler queues (waiting/grammar/disagg queues)
-        """
-        if not self._is_no_request():
-            return False
-        if len(self.waiting_queue) != 0:
-            return False
-        if len(self.grammar_manager.grammar_queue) != 0:
-            return False
-        return True
+        # Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode)
+        idle &= len(self.waiting_queue) == 0
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            idle &= len(self.disagg_prefill_bootstrap_queue.queue) == 0
+        if self.disaggregation_mode == DisaggregationMode.DECODE:
+            idle &= (
+                len(self.disagg_decode_prealloc_queue.queue) == 0
+                and len(self.disagg_decode_transfer_queue.queue) == 0
+            )
+
+        if not for_health_check:
+            # Grammar queue and prefill inflight queue may not produce batch results
+            # instantly, but they still indicate the server is not fully idle.
+            idle &= len(self.grammar_manager.grammar_queue) == 0
+            if self.disaggregation_mode == DisaggregationMode.PREFILL:
+                idle &= len(self.disagg_prefill_inflight_queue) == 0
+
+        return idle
 
     def attach_hicache_storage_wrapped(
         self, recv_req: AttachHiCacheStorageReqInput
@@ -2670,7 +2663,7 @@ class Scheduler(
                 success=False, message="Hierarchical cache is not enabled."
             )
 
-        if not self._is_idle_for_hicache_storage_op():
+        if not self.is_fully_idle():
             return AttachHiCacheStorageReqOutput(
                 success=False,
                 message=(
@@ -2723,7 +2716,7 @@ class Scheduler(
                 success=False, message="Hierarchical cache is not enabled."
             )
 
-        if not self._is_idle_for_hicache_storage_op():
+        if not self.is_fully_idle():
             return DetachHiCacheStorageReqOutput(
                 success=False,
                 message=(
@@ -2790,29 +2783,9 @@ class Scheduler(
             message=msg,
         )
 
-    def _is_no_request(self):
-        no_request = (
-            self.running_batch.is_empty()
-            and (self.last_batch is None or self.last_batch.is_empty())
-            and (self.cur_batch is None or self.cur_batch.is_empty())
-            and (not self.enable_overlap or len(self.result_queue) == 0)
-            and (self.pp_size == 1 or all(x.is_empty() for x in self.running_mbs))
-        )
-        if self.disaggregation_mode == DisaggregationMode.PREFILL:
-            no_request &= (
-                len(self.disagg_prefill_bootstrap_queue.queue) == 0
-                and len(self.disagg_prefill_inflight_queue) == 0
-            )
-        if self.disaggregation_mode == DisaggregationMode.DECODE:
-            no_request &= (
-                len(self.disagg_decode_prealloc_queue.queue) == 0
-                and len(self.disagg_decode_transfer_queue.queue) == 0
-            )
-        return no_request
-
     def flush_cache(self):
         """Flush the memory pool and cache."""
-        if self._is_no_request():
+        if self.is_fully_idle():
             self.cur_batch = None
             self.last_batch = None
             self.tree_cache.reset()

python/sglang/srt/managers/scheduler_update_weights_mixin.py

@@ -125,8 +125,8 @@ class SchedulerUpdateWeightsMixin:
         self: Scheduler, recv_req: ReleaseMemoryOccupationReqInput
     ):
         assert (
-            self._is_no_request()
-        ), "release_memory_occupation should be called only when no ongoing request."
+            self.is_fully_idle()
+        ), "release_memory_occupation should be called only when server is idle."
 
         tags = recv_req.tags