diff --git a/docs/advanced_features/hicache_storage_runtime_attach_detach.md b/docs/advanced_features/hicache_storage_runtime_attach_detach.md index f40e36cd0..deaef8c12 100644 --- a/docs/advanced_features/hicache_storage_runtime_attach_detach.md +++ b/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: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 10615e072..c7e90ab6f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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() diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py index ba038877b..abcda6794 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/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 diff --git a/test/registered/rl/test_update_weights_from_tensor.py b/test/registered/rl/test_update_weights_from_tensor.py index 900020762..7cb023071 100644 --- a/test/registered/rl/test_update_weights_from_tensor.py +++ b/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):