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 = (
+    if backend.startswith("vllm"):
-        "/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,
+            base_url + "/reset_prefix_cache", headers=get_auth_headers()
-        headers=get_auth_headers(),
-        params=params,
         )
+    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,

test/registered/bench_fn/test_benchmark_datasets_api.py

@@ -1,3 +1,4 @@
+import argparse
 import asyncio
 import base64
 import io
@@ -7,11 +8,15 @@ import random
 import subprocess
 import sys
 import tempfile
+import threading
 import unittest
 from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
+from urllib.parse import parse_qs, urlparse
 
 import numpy as np
 from PIL import Image
@@ -47,10 +52,12 @@ from sglang.benchmark.serving import (
     _BACKEND_API_PATHS,
     _EMBEDDING_BACKENDS,
     ASYNC_REQUEST_FUNCS,
+    _finite_positive_float,
     async_request_openai_embeddings,
     flush_server_cache,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=40, suite="base-a-test-cpu")
 register_cpu_ci(est_time=7, suite="base-c-test-cpu")
@@ -95,7 +102,7 @@ def create_lightweight_tokenizer() -> PreTrainedTokenizerFast:
     return hf_tokenizer
 
 
-class TestEmbeddingBenchmarkBackends(unittest.TestCase):
+class TestEmbeddingBenchmarkBackends(CustomTestCase):
     def test_vllm_embedding_reuses_the_openai_embedding_request_path(self):
         self.assertIn("vllm-embedding", _EMBEDDING_BACKENDS)
         self.assertIs(
@@ -103,7 +110,10 @@ class TestEmbeddingBenchmarkBackends(unittest.TestCase):
         )
         self.assertEqual(_BACKEND_API_PATHS["vllm-embedding"], "/v1/embeddings")
 
-    def test_embedding_cache_flush_uses_the_engine_specific_endpoint(self):
+
+class TestBenchmarkCacheFlush(CustomTestCase):
+    def test_cache_flush_uses_the_backend_specific_request(self):
+        """SGLang forwards its timeout without changing other backend requests."""
         with (
             patch("sglang.benchmark.serving.get_auth_headers", return_value={}),
             patch("sglang.benchmark.serving.requests.post") as post,
@@ -114,16 +124,76 @@ class TestEmbeddingBenchmarkBackends(unittest.TestCase):
             post.assert_called_once_with(
                 "http://127.0.0.1:8000/reset_prefix_cache",
                 headers={},
-                params={},
             )
             post.reset_mock()
 
-            flush_server_cache("http://127.0.0.1:30000", "sglang-embedding")
+            flush_server_cache("http://127.0.0.1:30000", "sglang")
             post.assert_called_once_with(
                 "http://127.0.0.1:30000/flush_cache",
                 headers={},
-                params={"timeout": 10.0},
+                params={"timeout": 60.0},
             )
+            post.reset_mock()
+
+            flush_server_cache("http://127.0.0.1:23333", "lmdeploy")
+            post.assert_called_once_with(
+                "http://127.0.0.1:23333/flush_cache", headers={}
+            )
+            post.reset_mock()
+
+            flush_server_cache(
+                "http://127.0.0.1:30000",
+                "sglang-embedding",
+                flush_cache_timeout=7.5,
+            )
+            post.assert_called_once_with(
+                "http://127.0.0.1:30000/flush_cache",
+                headers={},
+                params={"timeout": 7.5},
+            )
+
+    def test_sglang_cache_flush_waits_for_idle(self):
+        """A busy server can become idle before the benchmark's flush times out."""
+        request_received = threading.Event()
+        server_idle = threading.Event()
+
+        class DeferredFlushHandler(BaseHTTPRequestHandler):
+            def do_POST(self):
+                url = urlparse(self.path)
+                timeout = float(parse_qs(url.query).get("timeout", ["0"])[0])
+                if timeout <= 0:
+                    status = 400
+                    request_received.set()
+                else:
+                    request_received.set()
+                    status = 200 if server_idle.wait(timeout) else 400
+                self.send_response(status)
+                self.end_headers()
+
+            def log_message(self, format, *args):
+                pass
+
+        server = HTTPServer(("127.0.0.1", 0), DeferredFlushHandler)
+        server_thread = threading.Thread(target=server.serve_forever, daemon=True)
+        server_thread.start()
+        base_url = f"http://127.0.0.1:{server.server_port}"
+
+        try:
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                flush = executor.submit(
+                    flush_server_cache,
+                    base_url,
+                    "sglang",
+                    5.0,
+                )
+                self.assertTrue(request_received.wait(timeout=5))
+                self.assertFalse(flush.done())
+                server_idle.set()
+                flush.result(timeout=5)
+        finally:
+            server.shutdown()
+            server.server_close()
+            server_thread.join(timeout=5)
 
 
 class DummyProcessor:
@@ -218,7 +288,7 @@ def make_args(**overrides):
     return SimpleNamespace(**args)
 
 
-class TestBenchmarkDatasetsAPI(unittest.TestCase):
+class TestBenchmarkDatasetsAPI(CustomTestCase):
     def setUp(self):
         self.tokenizer = create_lightweight_tokenizer()
         self.processor = DummyProcessor(self.tokenizer)
@@ -1327,6 +1397,28 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase):
         self.assertNotEqual(bad_choice_res.returncode, 0)
         self.assertIn("invalid choice", (bad_choice_res.stderr + bad_choice_res.stdout))
 
+    def test_serving_benchmark_cli_rejects_invalid_flush_cache_timeout(self):
+        """Invalid timeouts fail before the benchmark contacts a server or hangs."""
+        res = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "sglang.benchmark.serving",
+                "--flush-cache-timeout",
+                "inf",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=90,
+        )
+        self.assertEqual(res.returncode, 2, res.stderr)
+        self.assertIn("expected a finite float > 0", res.stderr)
+
+        for value in ("1e999", "0", "-1"):
+            with self.subTest(value=value):
+                with self.assertRaises(argparse.ArgumentTypeError):
+                    _finite_positive_float(value)
+
     def test_bench_serving_cli_rejects_zipf_without_alpha_before_server(self):
         # Malformed CLI combinations (zipf with no alpha) must fail at
         # argparse time so users see the GSP-flag error directly, not a