[HiCache] Check in-flight async ops in is_fully_idle() before attach/detach (#20746)
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user