From 2e65c27b29152cc99fac5589a89a79edbc975a81 Mon Sep 17 00:00:00 2001 From: SevenJ <166966490+Wenjun7J@users.noreply.github.com> Date: Fri, 27 Mar 2026 05:44:37 +0800 Subject: [PATCH] Api add flush cache timeout (#21413) Signed-off-by: root --- docs/basic_usage/native_api.ipynb | 10 +- python/sglang/srt/entrypoints/http_server.py | 4 +- python/sglang/srt/managers/io_struct.py | 2 +- python/sglang/srt/managers/scheduler.py | 48 +++++++- .../managers/tokenizer_communicator_mixin.py | 8 +- .../managers/test_scheduler_flush_cache.py | 104 ++++++++++++++++++ 6 files changed, 167 insertions(+), 9 deletions(-) create mode 100644 test/registered/unit/managers/test_scheduler_flush_cache.py diff --git a/docs/basic_usage/native_api.ipynb b/docs/basic_usage/native_api.ipynb index 05f4f3688..65ff5809e 100644 --- a/docs/basic_usage/native_api.ipynb +++ b/docs/basic_usage/native_api.ipynb @@ -185,7 +185,15 @@ "source": [ "## Flush Cache\n", "\n", - "Flush the radix cache. It will be automatically triggered when the model weights are updated by the `/update_weights` API." + "Flush the radix cache. It will be automatically triggered when the model weights are updated by the `/update_weights` API.\n", + "\n", + "Parameters:\n", + "- `timeout` (query, float, default `0`, unit: seconds): Wait time for idle state before flushing. `0` means fail fast if not idle. When HiCache async operations are in-flight, a non-zero timeout allows the server to wait until idle before flushing, avoiding unnecessary 400 errors.\n", + "\n", + "```bash\n", + "# With timeout (wait up to 30s for idle state)\n", + "curl -s -X POST \"http://127.0.0.1:30000/flush_cache?timeout=30\"\n", + "```" ] }, { diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index acf60ccfe..e7eb05f96 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -733,9 +733,9 @@ async def classify_request(obj: EmbeddingReqInput, request: Request): @app.api_route("/flush_cache", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) -async def flush_cache(): +async def flush_cache(timeout: float = Query(0.0, ge=0.0)): """Flush the radix cache.""" - ret = await _global_state.tokenizer_manager.flush_cache() + ret = await _global_state.tokenizer_manager.flush_cache(timeout_s=timeout) return Response( content="Cache flushed.\nPlease check backend logs for more details. " "(When there are running or waiting requests, the operation will not be performed.)\n", diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d9c97a0b4..468143be2 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1166,7 +1166,7 @@ class ClearHiCacheReqOutput(BaseReq): @dataclass class FlushCacheReqInput(BaseReq): - pass + timeout_s: Optional[float] = None @dataclass diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 703393c15..b050a72c9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -834,6 +834,7 @@ class Scheduler( self.last_batch: Optional[ScheduleBatch] = None self.forward_ct = 0 self.return_health_check_ipcs: Deque[Optional[str]] = deque() + self._pending_flush: Deque[Tuple[FlushCacheReqInput, float]] = deque() self.num_retracted_reqs: int = 0 self.num_paused_reqs: int = 0 self.session_controller = SessionController(self.tree_cache) @@ -1562,6 +1563,8 @@ class Scheduler( if self.recv_from_rpc is not None: self.recv_from_rpc.send_pyobj(output) + self._check_pending_flush() + def init_req_max_new_tokens(self, req): req.sampling_params.max_new_tokens = min( ( @@ -2781,9 +2784,48 @@ class Scheduler( ) ) - def flush_cache_wrapped(self, recv_req: FlushCacheReqInput): - success = self.flush_cache() - return FlushCacheReqOutput(success=success) + def _check_pending_flush(self): + if not self._pending_flush: + return + + if self.is_fully_idle(): + success = self.flush_cache() + while self._pending_flush: + pending_req, _ = self._pending_flush.popleft() + self.send_to_tokenizer.send_output( + FlushCacheReqOutput(success=success), pending_req + ) + return + + self._expire_timed_out_pending_flushes(time.monotonic()) + + def _expire_timed_out_pending_flushes(self, now: float): + remaining: Deque[Tuple[FlushCacheReqInput, float]] = deque() + while self._pending_flush: + pending_req, deadline = self._pending_flush.popleft() + if now >= deadline: + logging.warning( + "Deferred flush_cache timed out while waiting for idle state." + ) + self.send_to_tokenizer.send_output( + FlushCacheReqOutput(success=False), pending_req + ) + else: + remaining.append((pending_req, deadline)) + self._pending_flush = remaining + + def flush_cache_wrapped( + self, recv_req: FlushCacheReqInput + ) -> Optional[FlushCacheReqOutput]: + timeout_s = float(recv_req.timeout_s or 0.0) + if timeout_s <= 0.0: + return FlushCacheReqOutput(success=self.flush_cache()) + + if self.is_fully_idle(): + return FlushCacheReqOutput(success=self.flush_cache()) + + self._pending_flush.append((recv_req, time.monotonic() + timeout_s)) + return None def clear_hicache_storage_wrapped(self, recv_req: ClearHiCacheReqInput): if self.enable_hierarchical_cache: diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 3faf15cd3..087337684 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -352,9 +352,13 @@ class TokenizerCommunicatorMixin: ] ) - async def flush_cache(self: TokenizerManager) -> FlushCacheReqOutput: + async def flush_cache( + self: TokenizerManager, timeout_s: Optional[float] = None + ) -> FlushCacheReqOutput: self.auto_create_handle_loop() - return (await self.flush_cache_communicator(FlushCacheReqInput()))[0] + return ( + await self.flush_cache_communicator(FlushCacheReqInput(timeout_s=timeout_s)) + )[0] async def clear_hicache_storage(self: TokenizerManager) -> ClearHiCacheReqOutput: """Clear the hierarchical cache storage.""" diff --git a/test/registered/unit/managers/test_scheduler_flush_cache.py b/test/registered/unit/managers/test_scheduler_flush_cache.py new file mode 100644 index 000000000..133716db8 --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_flush_cache.py @@ -0,0 +1,104 @@ +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.scheduler import Scheduler +from sglang.test.ci.ci_register import register_cpu_ci + +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.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): + scheduler = self._new_scheduler() + scheduler.flush_cache.return_value = False + + output = Scheduler.flush_cache_wrapped( + 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): + scheduler = self._new_scheduler() + scheduler.is_fully_idle.return_value = True + + output = Scheduler.flush_cache_wrapped( + 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): + scheduler = self._new_scheduler() + req = FlushCacheReqInput(timeout_s=3.0) + + with patch("sglang.srt.managers.scheduler.time.monotonic", return_value=10.0): + output = Scheduler.flush_cache_wrapped(scheduler, req) + + self.assertIsNone(output) + self.assertEqual(len(scheduler._pending_flush), 1) + pending_req, deadline = scheduler._pending_flush[0] + self.assertIs(pending_req, req) + self.assertEqual(deadline, 13.0) + scheduler.flush_cache.assert_not_called() + + def test_check_pending_flush_idle_flushes_all(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)]) + + Scheduler._check_pending_flush(scheduler) + + 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) + + def test_check_pending_flush_busy_expires_only_timed_out(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)]) + + with patch("sglang.srt.managers.scheduler.time.monotonic", return_value=100.0): + Scheduler._check_pending_flush(scheduler) + + 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) + + 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) + + +if __name__ == "__main__": + unittest.main()