minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-08-06 15:27:10 +00:00
committed by GitHub
parent 7195b8e4c7
commit 8a1637a479
2 changed files with 128 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/serving.py
+30 -12
View File
@@ -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,
test/registered/bench_fn/test_benchmark_datasets_api.py
+98 -6
View File
@@ -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