[CI/Build] Fix Qwen3-Omni Talker benchmark workload (#8057)

Signed-off-by: psv666 <2693925048@qq.com>

tests/benchmarks/patch/test_patch.py

@@ -2045,3 +2045,79 @@ def test_get_samples_forwards_upstream_multimodal_backends_kwarg(mocker: MockerF
     # No upstream kwargs: unchanged legacy delegate call.
     assert patch.get_samples(args, sentinel) == ["delegated"]
     assert calls[-1] == (args, sentinel, {})
+
+
+@pytest.mark.asyncio
+async def test_benchmark_preserves_stage_metrics_request_order_and_missing_snapshots(monkeypatch):
+    """Persist compact formal-request snapshots in input order, excluding warmups and retaining gaps."""
+    second_finished = asyncio.Event()
+    completion_order = []
+    # Request 1 carries the empty dict that openai-chat-omni initializes before any SSE merge.
+    snapshots = [
+        {"1": {"num_tokens_out": 1536, "finish_reason": "length", "vllm_itls_ms": [8.0, 9.0]}},
+        {},
+        {"1": {"num_tokens_out": 486, "finish_reason": "stop"}, "2": {"audio_frames": 24000, "audio_duration_s": 1.0}},
+    ]
+    expected = [
+        {"1": {"num_tokens_out": 1536, "finish_reason": "length"}},
+        None,
+        {"1": {"num_tokens_out": 486, "finish_reason": "stop"}, "2": {"audio_frames": 24000, "audio_duration_s": 1.0}},
+    ]
+
+    async def request_func(request_func_input, session, pbar=None):
+        request_id = request_func_input.request_id
+        if not request_id:
+            return MixRequestFuncOutput(success=True, stage_metrics={"1": {"num_tokens_out": 999}})
+        index = int(request_id)
+        if index == 0:
+            await second_finished.wait()
+        elif index == 1:
+            second_finished.set()
+        completion_order.append(index)
+        return MixRequestFuncOutput(
+            success=index != 1,
+            stage_metrics=snapshots[index],
+            prompt_len=1,
+            output_tokens=900,
+            ttft=0.01,
+            text_latency=0.1,
+            latency=0.1,
+        )
+
+    monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "test-stage-metrics", request_func)
+    result = await patch.benchmark(
+        task_type=patch.TaskType.GENERATION,
+        endpoint_type="test-stage-metrics",
+        api_url="http://unused/v1/chat/completions",
+        base_url="http://unused",
+        model_id="test-model",
+        model_name="test-model",
+        tokenizer=None,
+        input_requests=[
+            patch.SampleRequest(prompt="hello", prompt_len=1, expected_output_len=900, request_id=str(i))
+            for i in range(3)
+        ],
+        logprobs=None,
+        request_rate=float("inf"),
+        burstiness=1.0,
+        disable_tqdm=True,
+        num_warmups=2,
+        profile=False,
+        selected_percentile_metrics=[],
+        selected_percentiles=[],
+        ignore_eos=True,
+        goodput_config_dict={},
+        max_concurrency=3,
+        lora_modules=None,
+        extra_headers=None,
+        extra_body={"return_stage_metrics": True},
+        ready_check_timeout_sec=0,
+    )
+    assert completion_order.index(1) < completion_order.index(0)
+    assert result["request_stage_metrics"] == expected
+
+
+@pytest.mark.parametrize("snapshot", [None, {}, "not-a-dict"])
+def test_compact_request_stage_metrics_drops_empty_snapshots(snapshot):
+    """Empty chat-omni snapshots must not make every result persist request_stage_metrics."""
+    assert patch._compact_request_stage_metrics(snapshot) is None

tests/dfx/perf/scripts/run_benchmark.py

@@ -332,6 +332,36 @@ def _resolve_num_warmups(params: dict[str, Any], *, default: int) -> int:
     return value
 
 
+def _assert_fixed_stage_workloads(result: dict[str, Any], params: dict[str, Any], num_prompts: int) -> None:
+    """Reject successful responses that did less work than a fixed-length benchmark requested."""
+    extra_body = params.get("extra_body") or params.get("extra-body") or {}
+    fixed_lengths = {
+        str(stage): sampling["max_tokens"]
+        for stage, sampling in enumerate(extra_body.get("sampling_params_list", []))
+        if isinstance(sampling, dict)
+        and isinstance(sampling.get("max_tokens"), int)
+        and sampling["max_tokens"] > 0
+        and sampling.get("min_tokens") == sampling["max_tokens"]
+    }
+    if not fixed_lengths:
+        return
+
+    snapshots = result.get("request_stage_metrics")
+    assert isinstance(snapshots, list) and len(snapshots) == num_prompts, (
+        "Fixed stage workload: missing per-request stage metrics"
+    )
+    for request_index, snapshot in enumerate(snapshots):
+        for stage, expected in fixed_lengths.items():
+            metrics = snapshot.get(stage) if isinstance(snapshot, dict) else None
+            actual = metrics.get("num_tokens_out") if isinstance(metrics, dict) else None
+            assert actual == expected, (
+                f"Fixed stage workload: request {request_index}, stage {stage}: "
+                f"expected {expected} output tokens, got {actual!r}"
+            )
+    stages = ", ".join(f"stage {stage}={expected}" for stage, expected in fixed_lengths.items())
+    print(f"Fixed stage workload OK: {num_prompts} requests, num_tokens_out {stages}")
+
+
 def _is_finite_number(value: object) -> bool:
     return isinstance(value, int | float) and not isinstance(value, bool) and math.isfinite(value)
 
@@ -448,6 +478,7 @@ def _maybe_assert_aggregate_ia_qtf1(params: dict[str, object], acc_summary: dict
 
 def assert_result(result, params, num_prompt) -> None:
     assert result["completed"] == num_prompt, "Request failures exist"
+    _assert_fixed_stage_workloads(result, params, num_prompt)
     if params.get("dataset_name") == "omniinteract":
         summary = result.get("omniinteract")
         assert isinstance(summary, dict), "OmniInteract summary is missing"

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -1,6 +1,7 @@
 [
   {
     "test_name": "test_qwen3_omni",
+    "description": "random row: fixed workload (Thinker 900 / Talker 1536 tokens, min_tokens == max_tokens) for comparable audio RTF; H100 baseline pending recalibration, see #8100",
     "mark": [
       {
         "hardware_marks": {
@@ -43,17 +44,28 @@
         "random_output_len": 900,
         "ignore_eos": true,
         "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "H100": {
-            "mean_audio_rtf": [
-              0.169,
-              0.2041,
-              0.2382,
-              0.3177,
-              0.4595
-            ]
-          }
-        }
+        "seed": 0,
+        "random_range_ratio": 0.0,
+        "extra_body": {
+          "return_stage_metrics": true,
+          "sampling_params_list": [
+            {
+              "temperature": 0.0,
+              "min_tokens": 900,
+              "max_tokens": 900,
+              "ignore_eos": true
+            },
+            {
+              "temperature": 0.9,
+              "top_k": 50,
+              "repetition_penalty": 1.05,
+              "seed": 0,
+              "min_tokens": 1536,
+              "max_tokens": 1536
+            }
+          ]
+        },
+        "baseline": {}
       },
       {
         "dataset_name": "random-mm",

tests/dfx/perf/tests/test_runner_metadata.py

@@ -979,3 +979,58 @@ def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples
             {"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
             1,
         )
+
+
+# Mirrors the Thinker/Talker lengths in test_qwen3_omni_no_async_chunk.json.
+_FIXED_SAMPLING = {
+    "sampling_params_list": [{"min_tokens": 900, "max_tokens": 900}, {"min_tokens": 1536, "max_tokens": 1536}]
+}
+_FIXED_OK = {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 1536}}
+
+
+@pytest.mark.parametrize("extra_body_key", ["extra_body", "extra-body"])
+def test_fixed_stage_workload_accepts_exact_lengths(extra_body_key, capsys):
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    assert_result(
+        {"completed": 2, "failed": 0, "request_stage_metrics": [_FIXED_OK, _FIXED_OK]},
+        {extra_body_key: _FIXED_SAMPLING},
+        2,
+    )
+    assert "Fixed stage workload OK: 2 requests, num_tokens_out stage 0=900, stage 1=1536" in capsys.readouterr().out
+
+
+@pytest.mark.parametrize(
+    "snapshots",
+    [
+        None,
+        [],
+        [_FIXED_OK],
+        [_FIXED_OK, None],
+        [_FIXED_OK, {}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {}}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 486}}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 900}, "1": {"num_tokens_out": 1537}}],
+        [_FIXED_OK, {"1": {"num_tokens_out": 1536}}],
+        [_FIXED_OK, {"0": {"num_tokens_out": 512}, "1": {"num_tokens_out": 1536}}],
+    ],
+)
+def test_fixed_stage_workload_rejects_missing_or_wrong_lengths(snapshots):
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    with pytest.raises(AssertionError, match="Fixed stage workload"):
+        assert_result({"completed": 2, "request_stage_metrics": snapshots}, {"extra_body": _FIXED_SAMPLING}, 2)
+
+
+def test_fixed_stage_workload_rejects_request_failures():
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    with pytest.raises(AssertionError, match="Request failures"):
+        assert_result({"completed": 1, "failed": 1}, {}, 2)
+
+
+@pytest.mark.parametrize("extra_body", [{}, {"sampling_params_list": [{"max_tokens": 900}]}])
+def test_variable_stage_workload_does_not_require_stage_metrics(extra_body):
+    from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+    assert_result({"completed": 2}, {"extra_body": extra_body}, 2)

vllm_omni/benchmarks/patch/patch.py

@@ -1203,6 +1203,27 @@ def _update_output_stage_metrics_from_payload(
         output.stage_metrics.update(stage_snapshot)
 
 
+# Per-request stage fields persisted in benchmark results. Full snapshots carry
+# per-token latency lists, which would inflate every chat-omni result file.
+_REQUEST_STAGE_METRIC_FIELDS = (
+    defs.NUM_TOKENS_OUT,
+    "finish_reason",
+    defs.AUDIO_FRAMES,
+    f"{defs.AUDIO_DURATION}_s",
+)
+
+
+def _compact_request_stage_metrics(snapshot: object) -> dict[str, dict] | None:
+    """Keep the stage fields used for workload checks; empty snapshots become None."""
+    if not isinstance(snapshot, dict) or not snapshot:
+        return None
+    return {
+        stage: {field: metrics[field] for field in _REQUEST_STAGE_METRIC_FIELDS if field in metrics}
+        for stage, metrics in snapshot.items()
+        if isinstance(metrics, dict)
+    }
+
+
 def _apply_chat_stage0_token_timings(output: MixRequestFuncOutput) -> bool:
     """Apply native Stage 0 timings from a chat response snapshot."""
     stage_metrics = output.stage_metrics
@@ -3635,6 +3656,14 @@ async def benchmark(
             "input_lens": [output.prompt_len for output in outputs],
             "errors": [output.error for output in outputs],
         }
+    # Preserve request order, including missing snapshots, so CI can verify
+    # fixed stage workloads without parsing logs or storing audio payloads.
+    request_stage_metrics = [
+        _compact_request_stage_metrics(getattr(output, "stage_metrics", None)) for output in outputs
+    ]
+    if any(request_stage_metrics):
+        result["request_stage_metrics"] = request_stage_metrics
+
     # Plain-vLLM backends (e.g. the vLLM-text perf config) return upstream
     # RequestFuncOutput objects without the Mix duplex fields; read them
     # tolerantly or the whole benchmark result is discarded ("fallback to