Fix serving benchmark post-warmup cache flush race (#33663)

python/sglang/benchmark/serving.py

@@ -971,19 +971,27 @@ _BACKEND_API_PATHS = {
 
 _EMBEDDING_BACKENDS = frozenset(("sglang-embedding", "vllm-embedding"))
 
+_DEFAULT_SGLANG_FLUSH_CACHE_TIMEOUT = 60.0
 
-def flush_server_cache(base_url: str, backend: str) -> None:
+
+def flush_server_cache(
+    base_url: str,
+    backend: str,
+    flush_cache_timeout: float = _DEFAULT_SGLANG_FLUSH_CACHE_TIMEOUT,
+) -> None:
     """Flush an engine's prefix cache after benchmark warmup."""
-    cache_endpoint = (
-        "/reset_prefix_cache" if backend.startswith("vllm") else "/flush_cache"
-    )
-    # Pass timeout so the server waits for idle instead of failing immediately
-    params = {"timeout": 10.0} if not backend.startswith("vllm") else {}
-    response = requests.post(
-        base_url + cache_endpoint,
-        headers=get_auth_headers(),
-        params=params,
-    )
+    if backend.startswith("vllm"):
+        response = requests.post(
+            base_url + "/reset_prefix_cache", headers=get_auth_headers()
+        )
+    elif backend.startswith("sglang"):
+        response = requests.post(
+            base_url + "/flush_cache",
+            headers=get_auth_headers(),
+            params={"timeout": flush_cache_timeout},
+        )
+    else:
+        response = requests.post(base_url + "/flush_cache", headers=get_auth_headers())
     response.raise_for_status()
 
 
@@ -1343,6 +1351,7 @@ async def benchmark(
     profile: bool,
     pd_separated: bool = False,
     flush_cache: bool = False,
+    flush_cache_timeout: float = _DEFAULT_SGLANG_FLUSH_CACHE_TIMEOUT,
     warmup_requests: int = 1,
     use_trace_timestamps: bool = False,
     mooncake_slowdown_factor=1.0,
@@ -1452,7 +1461,7 @@ async def benchmark(
         "sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")
     ) or flush_cache
     if should_flush_cache:
-        flush_server_cache(base_url, backend)
+        flush_server_cache(base_url, backend, flush_cache_timeout)
 
     time.sleep(1.0)
 
@@ -2099,6 +2108,8 @@ def run_benchmark(args_: argparse.Namespace):
     # compatible with SimpleNamespace
     if not hasattr(args, "flush_cache"):
         args.flush_cache = False
+    if not hasattr(args, "flush_cache_timeout"):
+        args.flush_cache_timeout = _DEFAULT_SGLANG_FLUSH_CACHE_TIMEOUT
 
     # Prepare LoRA arguments
     lora_request_distribution = (
@@ -2129,6 +2140,7 @@ def run_benchmark(args_: argparse.Namespace):
             profile=args.profile,
             pd_separated=args.pd_separated,
             flush_cache=args.flush_cache,
+            flush_cache_timeout=args.flush_cache_timeout,
             warmup_requests=args.warmup_requests,
             use_trace_timestamps=args.use_trace_timestamps,
             mooncake_slowdown_factor=args.mooncake_slowdown_factor,
@@ -2577,6 +2589,12 @@ def cli_main():
         action="store_true",
         help="Flush the cache before running the benchmark",
     )
+    parser.add_argument(
+        "--flush-cache-timeout",
+        type=_finite_positive_float,
+        default=_DEFAULT_SGLANG_FLUSH_CACHE_TIMEOUT,
+        help="Maximum seconds to wait for an SGLang server to become idle before flushing the cache",
+    )
     parser.add_argument(
         "--warmup-requests",
         type=int,