[Bugfix] Reject failed diffusion benchmark warmups (#7092)

Signed-off-by: hyw <yuweih205@gmail.com>
Signed-off-by: HuangYuwei <yuweih205@gmail.com>
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Alex Brooks <albrooks@redhat.com>

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))

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