[HiCache] Check in-flight async ops in is_fully_idle() before attach/detach (#20746)

This commit is contained in:
Liangsheng Yin
2026-03-17 17:28:26 -07:00
committed by GitHub
parent c5d2528bff
commit 4d3976b6c5
5 changed files with 39 additions and 15 deletions

View File

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

View File

@@ -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")

View File

@@ -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):

View File

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

View File

@@ -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."""