diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1e74d2f8e..938503567 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2674,8 +2674,8 @@ class Scheduler( idle &= len(self.waiting_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. + # Grammar queue and prefill inflight queue may not produce batch + # results instantly, but they still indicate the server is not idle. idle &= len(self.grammar_manager.grammar_queue) == 0 if self.disaggregation_mode == DisaggregationMode.PREFILL: idle &= len(self.disagg_prefill_inflight_queue) == 0 @@ -2685,6 +2685,16 @@ class Scheduler( idle &= len(self.disagg_decode_prealloc_queue.queue) == 0 idle &= len(self.disagg_decode_transfer_queue.queue) == 0 + # HiCache: in-flight async ops (GPU↔Host↔L3) must drain before + # destructive operations like attach/detach/flush_cache. + if self.enable_hierarchical_cache: + tc = self.tree_cache + idle &= len(tc.ongoing_write_through) == 0 + idle &= len(tc.ongoing_load_back) == 0 + if tc.enable_storage: + idle &= len(tc.ongoing_prefetch) == 0 + idle &= len(tc.ongoing_backup) == 0 + return idle def attach_hicache_storage_wrapped( diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 064626d67..32db3c07d 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -167,6 +167,25 @@ def download_image_with_retry(image_url: str, max_retries: int = 3) -> Image.Ima time.sleep(2**i) +def flush_cache_with_retry( + base_url: str, retries: int = 5, interval: float = 2.0 +) -> bool: + """Flush device cache with retry. + + The scheduler may still have in-flight HiCache async ops (GPU↔Host↔L3) + that prevent is_fully_idle() from returning True, so we retry. + """ + for _ in range(retries): + try: + response = requests.post(f"{base_url}/flush_cache", timeout=10) + if response.status_code == 200: + return True + except requests.RequestException: + pass + time.sleep(interval) + return False + + def is_in_ci(): """Return whether it is in CI runner.""" return get_bool_env_var("SGLANG_IS_IN_CI") diff --git a/test/manual/hicache/test_disaggregation_hicache.py b/test/manual/hicache/test_disaggregation_hicache.py index c95cccdf3..42ad6609f 100644 --- a/test/manual/hicache/test_disaggregation_hicache.py +++ b/test/manual/hicache/test_disaggregation_hicache.py @@ -14,6 +14,7 @@ from sglang.test.server_fixtures.disaggregation_fixture import ( from sglang.test.test_utils import ( DEFAULT_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + flush_cache_with_retry, popen_launch_pd_server, ) @@ -114,9 +115,9 @@ class DisaggregationHiCacheBase(PDDisaggregationServerBase): # Trigger offloading self.send_request(self.gen_prompt(1), max_tokens=150) - # Flush device cache to force remote storage access + # Flush device cache to force remote storage access. time.sleep(2) - requests.post(self.prefill_url + "/flush_cache") + flush_cache_with_retry(self.prefill_url) class TestDisaggregationPrefillWithHiCache(DisaggregationHiCacheBase): diff --git a/test/manual/hicache/test_pp_with_hicache.py b/test/manual/hicache/test_pp_with_hicache.py index d89935969..0ed3bdeac 100644 --- a/test/manual/hicache/test_pp_with_hicache.py +++ b/test/manual/hicache/test_pp_with_hicache.py @@ -18,6 +18,7 @@ from sglang.test.test_utils import ( DEFAULT_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, find_available_port, + flush_cache_with_retry, popen_launch_server, ) @@ -180,11 +181,7 @@ class TestPPWithHiCache(unittest.TestCase): return False def flush_cache(self) -> bool: - try: - response = requests.post(f"{self.base_url}/flush_cache", timeout=10) - return response.status_code == 200 - except requests.RequestException: - return False + return flush_cache_with_retry(self.base_url) def test_eval_accuracy(self): args = SimpleNamespace( diff --git a/test/registered/hicache/test_hicache_storage_file_backend.py b/test/registered/hicache/test_hicache_storage_file_backend.py index ed9fca934..acd28dbb3 100644 --- a/test/registered/hicache/test_hicache_storage_file_backend.py +++ b/test/registered/hicache/test_hicache_storage_file_backend.py @@ -26,6 +26,7 @@ from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, + flush_cache_with_retry, is_in_ci, popen_launch_server, ) @@ -168,12 +169,8 @@ class HiCacheStorageBaseMixin: return int(meta.get("cached_tokens", 0)) def flush_cache(self) -> bool: - """Flush device cache to force remote storage access""" - try: - response = requests.post(f"{self.base_url}/flush_cache", timeout=10) - return response.status_code == 200 - except requests.RequestException: - return False + """Flush device cache to force remote storage access.""" + return flush_cache_with_retry(self.base_url) def gen_prompt(self, token_num: int) -> str: """Generate a random prompt of specified token length using tokenizer vocabulary."""