From 2a6847171516b5977f4e9e835f427f5e731d73f3 Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Thu, 17 Sep 2026 01:09:23 +0800 Subject: [PATCH] [Bugfix] Reject failed diffusion benchmark warmups (#7092) Signed-off-by: hyw Signed-off-by: HuangYuwei Signed-off-by: Alex Brooks Co-authored-by: Alex Brooks --- .../diffusion/diffusion_benchmark_serving.py | 12 +- .../test_diffusion_benchmark_warmups.py | 128 ++++++++++++++++++ 2 files changed, 139 insertions(+), 1 deletion(-) create mode 100644 tests/benchmarks/test_diffusion_benchmark_warmups.py diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 1e5aaf29c..6bf64ec1d 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -1087,6 +1087,12 @@ async def _run_warmups( warmup_tasks = [asyncio.create_task(limited_warmup_request_func(req)) for req in warmup_requests] warmup_outputs = await asyncio.gather(*warmup_tasks) + failed_count = sum(not output.success for output in warmup_outputs) + if failed_count: + raise RuntimeError( + f"{failed_count}/{len(warmup_outputs)} warmup requests failed; " + "check server logs before running the benchmark." + ) return list(zip(warmup_requests, warmup_outputs)) @@ -1455,7 +1461,7 @@ async def benchmark(args): print(f"Metrics saved to {args.output_file}") -if __name__ == "__main__": +def build_parser(): parser = argparse.ArgumentParser(description="Benchmark serving for diffusion models.") parser.add_argument( "--base-url", @@ -1637,6 +1643,10 @@ if __name__ == "__main__": action="store_true", help="Request stage duration metrics from endpoints that support return_stage_metrics.", ) + return parser + +if __name__ == "__main__": + parser = build_parser() args = parser.parse_args() asyncio.run(benchmark(args)) diff --git a/tests/benchmarks/test_diffusion_benchmark_warmups.py b/tests/benchmarks/test_diffusion_benchmark_warmups.py new file mode 100644 index 000000000..9d72189b7 --- /dev/null +++ b/tests/benchmarks/test_diffusion_benchmark_warmups.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib +import sys + +import pytest + +from benchmarks.diffusion import backends +from benchmarks.diffusion.backends import RequestFuncInput, RequestFuncOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu, pytest.mark.asyncio] + +ENDPOINT = "/v1/chat/completions" +TEST_MODEL = "test-model" + + +@pytest.fixture +def diffusion_benchmark(monkeypatch): + # The standalone script imports its sibling backend module by name. + monkeypatch.setitem(sys.modules, "backends", backends) + return importlib.import_module("benchmarks.diffusion.diffusion_benchmark_serving") + + +@pytest.fixture +def warmup_args(diffusion_benchmark): + from benchmarks.diffusion.diffusion_benchmark_serving import build_parser + + parser = build_parser() + return parser.parse_args( + [ + f"--endpoint={ENDPOINT}", + "--dataset=random", + "--task=t2i", + "--warmup-requests=2", + "--warmup-num-inference-steps=4", + ] + ) + + +@pytest.fixture +def input_requests(): + return [ + RequestFuncInput( + prompt=prompt, + api_url=f"http://test.local{ENDPOINT}", + model=TEST_MODEL, + num_inference_steps=20, + ) + for prompt in ("first", "second") + ] + + +async def test_disabled_warmup_sends_no_requests(diffusion_benchmark, warmup_args, input_requests): + warmup_args.warmup_requests = 0 + + async def unexpected_request(*args): + pytest.fail("disabled warmup must not send requests") + + pairs = await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, unexpected_request) + + assert pairs == [] + + +async def test_successful_warmups_preserve_request_output_pairs(diffusion_benchmark, warmup_args, input_requests): + outputs = [RequestFuncOutput(success=True, latency=0.5), RequestFuncOutput(success=True, latency=1.0)] + pending_outputs = iter(outputs) + + async def successful_request(req, session, pbar): + return next(pending_outputs) + + pairs = await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, successful_request) + + assert [request.prompt for request, _ in pairs] == ["first", "second"] + assert [request.num_inference_steps for request, _ in pairs] == [4, 4] + assert [output for _, output in pairs] == outputs + assert [request.num_inference_steps for request in input_requests] == [20, 20] + + +@pytest.mark.parametrize("failed_count", [1, 2], ids=["partial-failure", "all-failed"]) +@pytest.mark.parametrize( + "error", + ["HTTP 503: private prompt rejected", "Cannot connect to https://test.local/?token=private"], + ids=["http-error", "transport-error"], +) +async def test_failed_warmups_abort_without_exposing_response_errors( + diffusion_benchmark, warmup_args, input_requests, failed_count, error +): + outputs = iter( + [RequestFuncOutput(success=False, error=error) for _ in range(failed_count)] + + [RequestFuncOutput(success=True) for _ in range(2 - failed_count)] + ) + + async def request(req, session, pbar): + return next(outputs) + + with pytest.raises(RuntimeError, match=f"{failed_count}/2 warmup requests failed") as exc_info: + await diffusion_benchmark._run_warmups(input_requests, warmup_args, None, request) + + assert error not in str(exc_info.value) + assert "server logs" in str(exc_info.value) + + +async def test_benchmark_does_not_start_measurement_after_failed_warmup( + monkeypatch, mocker, diffusion_benchmark, warmup_args, input_requests +): + sent = [] + + async def failed_request(req, session, pbar): + sent.append(req) + return RequestFuncOutput(success=False, error="HTTP 503: unavailable") + + monkeypatch.setitem( + diffusion_benchmark.backends_function_mapping["2i"], + warmup_args.endpoint, + (failed_request, warmup_args.endpoint), + ) + dataset = mocker.Mock(spec=diffusion_benchmark.RandomDataset) + dataset.get_requests.return_value = input_requests + monkeypatch.setattr(diffusion_benchmark, "RandomDataset", mocker.Mock(return_value=dataset)) + timer = mocker.Mock(spec=diffusion_benchmark.time) + timer.perf_counter.side_effect = AssertionError("measurement started after failed warmup") + monkeypatch.setattr(diffusion_benchmark, "time", timer) + + with pytest.raises(RuntimeError, match="2/2 warmup requests failed"): + await diffusion_benchmark.benchmark(warmup_args) + + assert len(sent) == warmup_args.warmup_requests