From 8a1637a479073ac27d30f6a43ff3955831e5c9b0 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 6 Aug 2026 23:27:10 +0800 Subject: [PATCH] Fix serving benchmark post-warmup cache flush race (#33663) --- python/sglang/benchmark/serving.py | 42 +++++-- .../bench_fn/test_benchmark_datasets_api.py | 104 +++++++++++++++++- 2 files changed, 128 insertions(+), 18 deletions(-) diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index afb8b5753..fc82ff8cd 100644 --- a/python/sglang/benchmark/serving.py +++ b/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, diff --git a/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py index d1df61578..8937ea29a 100644 --- a/test/registered/bench_fn/test_benchmark_datasets_api.py +++ b/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