Api add flush cache timeout (#21413)

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

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