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

docs/advanced_features/hicache_storage_runtime_attach_detach.md

@@ -35,11 +35,11 @@ The control path is:
 
 ## 2. Idle-state requirement (strict)
 
-The Scheduler uses a stricter `_is_idle_for_hicache_storage_op()`:
+The Scheduler uses `is_fully_idle()` which checks:
 
-- `_is_no_request()` is true (covers running/overlap/pp/disagg and other active states)
-- `waiting_queue` is empty
-- `grammar_queue` is empty (if the grammar backend is enabled)
+- No running batches (including chunked prefill, overlap, pipeline-parallel, and disaggregation paths)
+- No waiting requests in any queue (waiting, grammar, disagg bootstrap/prealloc/transfer/inflight)
+- No DLLM staging requests
 
 If the condition is not met, attach/detach returns an error like:
 

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
 

test/registered/rl/test_update_weights_from_tensor.py

@@ -249,44 +249,38 @@ class TestServerUpdateWeightsFromTensorNonBlocking(CustomTestCase):
         return ret
 
     def test_update_weights(self):
-        pause_generation_modes = ["in_place", "retract"]
-        for pause_generation_mode in pause_generation_modes:
-            num_requests = 32
-            with ThreadPoolExecutor(num_requests) as executor:
-                futures = [
-                    executor.submit(self.run_decode, 3000) for _ in range(num_requests)
-                ]
+        num_requests = 32
+        with ThreadPoolExecutor(num_requests) as executor:
+            futures = [
+                executor.submit(self.run_decode, 3000) for _ in range(num_requests)
+            ]
 
-                # ensure the decode has been started
-                time.sleep(2)
+            # ensure the decode has been started
+            time.sleep(2)
 
-                param_names = [
-                    f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 16)
-                ]
-                new_tensor = torch.full((16384, 2048), 1.5, device="cuda")
-                named_tensors = [(x, new_tensor) for x in param_names]
+            param_names = [f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 16)]
+            new_tensor = torch.full((16384, 2048), 1.5, device="cuda")
+            named_tensors = [(x, new_tensor) for x in param_names]
 
-                ret = self.pause_generation(pause_generation_mode)
-                ret = self.run_update_weights(
-                    named_tensors, flush_cache=pause_generation_mode == "retract"
+            # abort mode ensures server is totally idle before returning
+            ret = self.pause_generation("abort")
+            ret = self.run_update_weights(named_tensors, flush_cache=True)
+            self.assertTrue(ret["success"])
+            ret = self.continue_generation()
+
+            # requests were aborted by pause_generation("abort")
+            for future in as_completed(futures):
+                future.result()
+
+            for param_name in param_names[:3]:
+                response = requests.post(
+                    self.base_url + "/get_weights_by_name",
+                    json={"name": param_name},
                 )
-                self.assertTrue(ret["success"])
-                ret = self.continue_generation()
-
-                for future in as_completed(futures):
-                    self.assertNotEqual(
-                        future.result()["meta_info"]["finish_reason"]["type"], "abort"
-                    )
-
-                for param_name in param_names[:3]:
-                    response = requests.post(
-                        self.base_url + "/get_weights_by_name",
-                        json={"name": param_name},
-                    )
-                    actual_values = torch.tensor(response.json())[0, :5]
-                    assert torch.allclose(
-                        actual_values, torch.tensor([1.5] * 5), atol=0.002
-                    ), f"{actual_values=}"
+                actual_values = torch.tensor(response.json())[0, :5]
+                assert torch.allclose(
+                    actual_values, torch.tensor([1.5] * 5), atol=0.002
+                ), f"{actual_values=}"
 
 
 def _check_param(engine, param_name, expect_values):