Simplify flush_cache: reject concurrent requests, remove client-side retry (#21490)
This commit is contained in:
@@ -1,7 +1,6 @@
|
||||
import os
|
||||
import random
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from typing import Dict
|
||||
|
||||
@@ -14,7 +13,6 @@ 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,
|
||||
)
|
||||
|
||||
@@ -116,8 +114,12 @@ class DisaggregationHiCacheBase(PDDisaggregationServerBase):
|
||||
self.send_request(self.gen_prompt(1), max_tokens=150)
|
||||
|
||||
# Flush device cache to force remote storage access.
|
||||
time.sleep(2)
|
||||
flush_cache_with_retry(self.prefill_url)
|
||||
res = requests.post(
|
||||
f"{self.prefill_url}/flush_cache",
|
||||
params={"timeout": 30},
|
||||
timeout=40,
|
||||
)
|
||||
res.raise_for_status()
|
||||
|
||||
|
||||
class TestDisaggregationPrefillWithHiCache(DisaggregationHiCacheBase):
|
||||
|
||||
@@ -18,7 +18,6 @@ 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,8 +179,13 @@ class TestPPWithHiCache(unittest.TestCase):
|
||||
|
||||
return False
|
||||
|
||||
def flush_cache(self) -> bool:
|
||||
return flush_cache_with_retry(self.base_url)
|
||||
def flush_cache(self):
|
||||
res = requests.post(
|
||||
f"{self.base_url}/flush_cache",
|
||||
params={"timeout": 30},
|
||||
timeout=40,
|
||||
)
|
||||
res.raise_for_status()
|
||||
|
||||
def test_eval_accuracy(self):
|
||||
args = SimpleNamespace(
|
||||
@@ -197,8 +201,7 @@ class TestPPWithHiCache(unittest.TestCase):
|
||||
metrics_initial = run_eval_few_shot_gsm8k(args)
|
||||
self.assertGreater(metrics_initial["accuracy"], 0.6)
|
||||
|
||||
self.assertTrue(self.flush_cache())
|
||||
time.sleep(2)
|
||||
self.flush_cache()
|
||||
|
||||
metrics_cached = run_eval_few_shot_gsm8k(args)
|
||||
self.assertGreater(metrics_cached["accuracy"], 0.6)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -107,8 +106,12 @@ class TestQwen35WithHiCache(CustomTestCase):
|
||||
)
|
||||
|
||||
print(f"flush cache")
|
||||
time.sleep(2)
|
||||
requests.post(f"{self.base_url}/flush_cache", timeout=10)
|
||||
res = requests.post(
|
||||
f"{self.base_url}/flush_cache",
|
||||
params={"timeout": 30},
|
||||
timeout=40,
|
||||
)
|
||||
res.raise_for_status()
|
||||
|
||||
second_metrics = self._run_gsm8k()
|
||||
print(f"second_metrics={second_metrics}")
|
||||
|
||||
@@ -26,7 +26,6 @@ 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,
|
||||
)
|
||||
@@ -170,9 +169,14 @@ class HiCacheStorageBaseMixin:
|
||||
meta = response_json.get("meta_info", {})
|
||||
return int(meta.get("cached_tokens", 0))
|
||||
|
||||
def flush_cache(self) -> bool:
|
||||
def flush_cache(self):
|
||||
"""Flush device cache to force remote storage access."""
|
||||
return flush_cache_with_retry(self.base_url)
|
||||
res = requests.post(
|
||||
f"{self.base_url}/flush_cache",
|
||||
params={"timeout": 30},
|
||||
timeout=40,
|
||||
)
|
||||
res.raise_for_status()
|
||||
|
||||
def gen_prompt(self, token_num: int) -> str:
|
||||
"""Generate a random prompt of specified token length using tokenizer vocabulary."""
|
||||
@@ -186,8 +190,7 @@ class HiCacheStorageBaseMixin:
|
||||
self.send_request(self.gen_prompt(1), max_tokens=150)
|
||||
|
||||
# Flush device cache to force remote storage access
|
||||
time.sleep(2)
|
||||
self.assertTrue(self.flush_cache(), "Cache flush should succeed")
|
||||
self.flush_cache()
|
||||
|
||||
def test_basic_backup_and_prefetch(self):
|
||||
"""Test storage and retrieval of large context through remote cache"""
|
||||
@@ -304,8 +307,7 @@ def run_eval_accuracy_test(test_instance, accuracy_threshold: float = 0.03):
|
||||
|
||||
# Flush cache to force remote storage access
|
||||
print("Phase 2: Flushing device cache...")
|
||||
test_instance.assertTrue(test_instance.flush_cache(), "Cache flush should succeed")
|
||||
time.sleep(2)
|
||||
test_instance.flush_cache()
|
||||
|
||||
# Second evaluation - should use remote cache
|
||||
print("Phase 3: Running second GSM8K evaluation using remote cache...")
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
import unittest
|
||||
from collections import deque
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.srt.managers.io_struct import FlushCacheReqInput, FlushCacheReqOutput
|
||||
from sglang.srt.managers.io_struct import FlushCacheReqInput
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -12,13 +11,14 @@ register_cpu_ci(est_time=2, suite="stage-a-cpu-only")
|
||||
class TestSchedulerFlushCache(unittest.TestCase):
|
||||
def _new_scheduler(self) -> Scheduler:
|
||||
scheduler = Scheduler.__new__(Scheduler)
|
||||
scheduler._pending_flush = deque()
|
||||
scheduler._pending_flush = None
|
||||
scheduler.send_to_tokenizer = MagicMock()
|
||||
scheduler.flush_cache = MagicMock(return_value=True)
|
||||
scheduler.is_fully_idle = MagicMock(return_value=False)
|
||||
return scheduler
|
||||
|
||||
def test_flush_cache_wrapped_non_positive_timeout_immediate(self):
|
||||
def test_immediate_flush_no_timeout(self):
|
||||
"""No timeout → flush immediately regardless of idle state."""
|
||||
scheduler = self._new_scheduler()
|
||||
scheduler.flush_cache.return_value = False
|
||||
|
||||
@@ -26,12 +26,11 @@ class TestSchedulerFlushCache(unittest.TestCase):
|
||||
scheduler, FlushCacheReqInput(timeout_s=None)
|
||||
)
|
||||
|
||||
self.assertIsInstance(output, FlushCacheReqOutput)
|
||||
self.assertFalse(output.success)
|
||||
scheduler.flush_cache.assert_called_once()
|
||||
self.assertEqual(len(scheduler._pending_flush), 0)
|
||||
|
||||
def test_flush_cache_wrapped_positive_timeout_idle_immediate(self):
|
||||
def test_immediate_flush_when_idle(self):
|
||||
"""Positive timeout but already idle → flush immediately."""
|
||||
scheduler = self._new_scheduler()
|
||||
scheduler.is_fully_idle.return_value = True
|
||||
|
||||
@@ -39,12 +38,11 @@ class TestSchedulerFlushCache(unittest.TestCase):
|
||||
scheduler, FlushCacheReqInput(timeout_s=5.0)
|
||||
)
|
||||
|
||||
self.assertIsInstance(output, FlushCacheReqOutput)
|
||||
self.assertTrue(output.success)
|
||||
scheduler.flush_cache.assert_called_once()
|
||||
self.assertEqual(len(scheduler._pending_flush), 0)
|
||||
|
||||
def test_flush_cache_wrapped_positive_timeout_busy_enqueues(self):
|
||||
def test_defers_when_busy(self):
|
||||
"""Positive timeout + busy → defers, returns None."""
|
||||
scheduler = self._new_scheduler()
|
||||
req = FlushCacheReqInput(timeout_s=3.0)
|
||||
|
||||
@@ -52,52 +50,60 @@ class TestSchedulerFlushCache(unittest.TestCase):
|
||||
output = Scheduler.flush_cache_wrapped(scheduler, req)
|
||||
|
||||
self.assertIsNone(output)
|
||||
self.assertEqual(len(scheduler._pending_flush), 1)
|
||||
pending_req, deadline = scheduler._pending_flush[0]
|
||||
pending_req, deadline = scheduler._pending_flush
|
||||
self.assertIs(pending_req, req)
|
||||
self.assertEqual(deadline, 13.0)
|
||||
|
||||
def test_rejects_when_already_pending(self):
|
||||
"""Any new request is rejected while another is pending."""
|
||||
scheduler = self._new_scheduler()
|
||||
scheduler._pending_flush = (FlushCacheReqInput(timeout_s=10.0), 999.0)
|
||||
|
||||
for timeout in [None, 5.0]:
|
||||
output = Scheduler.flush_cache_wrapped(
|
||||
scheduler, FlushCacheReqInput(timeout_s=timeout)
|
||||
)
|
||||
self.assertFalse(output.success)
|
||||
self.assertIn("already in progress", output.message)
|
||||
|
||||
scheduler.flush_cache.assert_not_called()
|
||||
|
||||
def test_check_pending_flush_idle_flushes_all(self):
|
||||
def test_pending_flush_completes_on_idle(self):
|
||||
scheduler = self._new_scheduler()
|
||||
scheduler.is_fully_idle.return_value = True
|
||||
|
||||
req1 = FlushCacheReqInput(timeout_s=1.0)
|
||||
req2 = FlushCacheReqInput(timeout_s=2.0)
|
||||
scheduler._pending_flush = deque([(req1, 111.0), (req2, 222.0)])
|
||||
req = FlushCacheReqInput(timeout_s=1.0)
|
||||
scheduler._pending_flush = (req, 111.0)
|
||||
|
||||
Scheduler._check_pending_flush(scheduler)
|
||||
|
||||
self.assertIsNone(scheduler._pending_flush)
|
||||
scheduler.flush_cache.assert_called_once()
|
||||
self.assertEqual(len(scheduler._pending_flush), 0)
|
||||
calls = scheduler.send_to_tokenizer.send_output.call_args_list
|
||||
self.assertEqual(len(calls), 2)
|
||||
self.assertTrue(calls[0].args[0].success)
|
||||
self.assertIs(calls[0].args[1], req1)
|
||||
self.assertTrue(calls[1].args[0].success)
|
||||
self.assertIs(calls[1].args[1], req2)
|
||||
out = scheduler.send_to_tokenizer.send_output.call_args.args[0]
|
||||
self.assertTrue(out.success)
|
||||
|
||||
def test_check_pending_flush_busy_expires_only_timed_out(self):
|
||||
def test_pending_flush_expires_on_timeout(self):
|
||||
scheduler = self._new_scheduler()
|
||||
scheduler.is_fully_idle.return_value = False
|
||||
|
||||
expired_req = FlushCacheReqInput(timeout_s=1.0)
|
||||
alive_req = FlushCacheReqInput(timeout_s=5.0)
|
||||
scheduler._pending_flush = deque([(expired_req, 99.0), (alive_req, 101.0)])
|
||||
req = FlushCacheReqInput(timeout_s=1.0)
|
||||
scheduler._pending_flush = (req, 99.0)
|
||||
|
||||
with patch("sglang.srt.managers.scheduler.time.monotonic", return_value=100.0):
|
||||
Scheduler._check_pending_flush(scheduler)
|
||||
|
||||
self.assertIsNone(scheduler._pending_flush)
|
||||
scheduler.flush_cache.assert_not_called()
|
||||
self.assertEqual(len(scheduler._pending_flush), 1)
|
||||
pending_req, deadline = scheduler._pending_flush[0]
|
||||
self.assertIs(pending_req, alive_req)
|
||||
self.assertEqual(deadline, 101.0)
|
||||
out = scheduler.send_to_tokenizer.send_output.call_args.args[0]
|
||||
self.assertFalse(out.success)
|
||||
|
||||
calls = scheduler.send_to_tokenizer.send_output.call_args_list
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertFalse(calls[0].args[0].success)
|
||||
self.assertIs(calls[0].args[1], expired_req)
|
||||
def test_pending_flush_survives_before_deadline(self):
|
||||
scheduler = self._new_scheduler()
|
||||
req = FlushCacheReqInput(timeout_s=5.0)
|
||||
scheduler._pending_flush = (req, 101.0)
|
||||
|
||||
with patch("sglang.srt.managers.scheduler.time.monotonic", return_value=100.0):
|
||||
Scheduler._check_pending_flush(scheduler)
|
||||
|
||||
self.assertIsNotNone(scheduler._pending_flush)
|
||||
scheduler.send_to_tokenizer.send_output.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user