Api add flush cache timeout (#21413)

Signed-off-by: root <wenjun7j@gmail.com>

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",
+    "```"
    ]
   },
   {

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",

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

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:

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

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