[Bugfix][Bench] Save per-request stage metrics in detailed results

docs/cli/bench/serve.md

@@ -68,8 +68,33 @@ Print per-stage benchmark metrics for --omni serving when stage metrics are retu
 Specify to save benchmark results to a json file
 
 - `--save-detailed`
-"When saving the results, whether to include per request "
-        "information such as response, error, ttfs, tpots, etc."
+Include per-request information, such as responses, errors, and latency metrics, when saving results with `--save-result`.
+
+For `daily-omni`, `openai-chat-omni`, `/v1/videos`, and the image backends
+(`openai-image-edits-omni`, `/v1/images/generations`, `/v1/images/edits`), this
+also requests server-side stage snapshots by setting
+`extra_body.return_stage_metrics=true` unless explicitly overridden. This is
+not a local-JSON-only option: snapshots can increase response/SSE payload size
+and measurement overhead. Stage 0 timings can also supply token intervals
+when client-side intervals are unavailable. Keep the same options in
+performance comparisons. Image backends already request these metrics
+without `--save-detailed`; unsupported backends are unchanged.
+
+Benchmark results keep two distinct, request-aligned fields:
+
+- `request_stage_metrics`: the existing compact workload snapshots containing
+    token counts, finish reasons, and audio frames/duration when available.
+    This field does not depend on `--save-detailed`, omits per-token timing
+    arrays, uses `null` for missing snapshots, and is omitted when no compact
+    snapshots are available.
+- `stage_metrics`: full returned stage snapshots, including any per-token
+    timing arrays, only with `--save-detailed`. Missing snapshots use `{}`.
+    This field supplements rather than replaces `request_stage_metrics`.
+
+Both arrays follow the formal input request order, retain entries for failed
+requests, and exclude warmup requests. Saving snapshots does not create
+metrics that the server did not return. An explicit
+`extra_body.return_stage_metrics=false` is preserved.
 
 - `--result-dir`
  "Specify directory to save benchmark json results."

tests/benchmarks/patch/test_patch.py

@@ -120,6 +120,7 @@ class MockResponse:
 
     def __init__(self, status, chunks, delay_between_chunks=0):
         self.status = status
+        self.headers: dict[str, str] = {}
         self.reason = "OK" if status == 200 else "Error"
         self._chunks = chunks
         self._delay = delay_between_chunks
@@ -138,6 +139,111 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         pass
 
 
+def test_request_stage_metrics_preserve_request_alignment():
+    first = MixRequestFuncOutput(
+        stage_metrics={
+            "0": {
+                "stage_id": 0,
+                "stage_gen_time_ms": 12.5,
+            }
+        }
+    )
+    missing = MixRequestFuncOutput(stage_metrics=None)
+    multiple = MixRequestFuncOutput(
+        stage_metrics={
+            "0": {"stage_id": 0},
+            "1": {"stage_id": 1},
+        }
+    )
+
+    assert patch._request_stage_metrics([first, missing, multiple]) == [
+        {
+            "0": {
+                "stage_id": 0,
+                "stage_gen_time_ms": 12.5,
+            }
+        },
+        {},
+        {
+            "0": {"stage_id": 0},
+            "1": {"stage_id": 1},
+        },
+    ]
+
+
+@pytest.mark.parametrize(
+    "backend",
+    [
+        "daily-omni",
+        "openai-chat-omni",
+        "openai-image-edits-omni",
+        "/v1/images/generations",
+        "/v1/images/edits",
+        "/v1/videos",
+    ],
+)
+def test_save_detailed_requests_stage_metrics(backend):
+    args = SimpleNamespace(
+        print_stage=False,
+        save_detailed=True,
+        backend=backend,
+        extra_body={},
+    )
+
+    assert patch.should_request_stage_metrics(args) is True
+
+
+def test_save_detailed_does_not_request_stage_metrics_for_unsupported_backend():
+    args = SimpleNamespace(
+        print_stage=False,
+        save_detailed=True,
+        backend="openai",
+        percentile_metrics=(),
+        extra_body={},
+    )
+
+    assert patch.should_request_stage_metrics(args) is False
+
+
+def test_save_detailed_preserves_explicit_stage_metrics_opt_out():
+    extra_body = {"return_stage_metrics": False, "temperature": 0.0}
+    args = SimpleNamespace(
+        print_stage=False,
+        save_detailed=True,
+        backend="openai-chat-omni",
+        extra_body=extra_body,
+    )
+
+    actual = patch.maybe_enable_stage_metrics(extra_body, enabled=patch.should_request_stage_metrics(args))
+
+    assert actual == {"return_stage_metrics": False, "temperature": 0.0}
+    assert actual is not extra_body
+    assert extra_body == {"return_stage_metrics": False, "temperature": 0.0}
+
+
+def test_detailed_stage_metrics_are_opt_in():
+    outputs = [
+        MixRequestFuncOutput(stage_metrics={"0": {"stage_id": 0}}),
+        MixRequestFuncOutput(stage_metrics=None),
+    ]
+
+    result: dict[str, object] = {}
+    patch.set_save_detailed(False)
+    patch._add_detailed_stage_metrics(result, outputs)
+    assert "stage_metrics" not in result
+
+    patch.set_save_detailed(True)
+    try:
+        patch._add_detailed_stage_metrics(result, outputs)
+    finally:
+        patch.set_save_detailed(False)
+
+    assert result["stage_metrics"] == [
+        {"0": {"stage_id": 0}},
+        {},
+    ]
+
+
 @pytest.mark.asyncio
 async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch):
     class FakeRealtimeClient:
@@ -2091,7 +2197,8 @@ def fake_get_samples_old(args, tokenizer, **kwargs):
 
 
 @pytest.mark.asyncio
-async def test_benchmark_preserves_stage_metrics_request_order_and_missing_snapshots(monkeypatch):
+@pytest.mark.parametrize("save_detailed", [False, True])
+async def test_benchmark_preserves_stage_metrics_request_order_and_missing_snapshots(monkeypatch, save_detailed):
     """Persist compact formal-request snapshots in input order, excluding warmups and retaining gaps."""
     second_finished = asyncio.Event()
     completion_order = []
@@ -2128,6 +2235,7 @@ async def request_func(request_func_input, session, pbar=None):
         )
 
     monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "test-stage-metrics", request_func)
+    monkeypatch.setattr(patch, "_SAVE_DETAILED", save_detailed)
     result = await patch.benchmark(
         task_type=patch.TaskType.GENERATION,
         endpoint_type="test-stage-metrics",
@@ -2158,6 +2266,13 @@ async def request_func(request_func_input, session, pbar=None):
     )
     assert completion_order.index(1) < completion_order.index(0)
     assert result["request_stage_metrics"] == expected
+    if save_detailed:
+        assert result["stage_metrics"] == snapshots
+        serialized = json.loads(json.dumps(result))
+        assert serialized["request_stage_metrics"] == expected
+        assert serialized["stage_metrics"] == snapshots
+    else:
+        assert "stage_metrics" not in result
 
 
 @pytest.mark.parametrize("snapshot", [None, {}, "not-a-dict"])

tests/benchmarks/test_serve_cli.py

@@ -22,7 +22,8 @@
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
 
 
-def test_it2i_dataset_uses_upstream_warmups(monkeypatch, tmp_path):
+@pytest.mark.parametrize("save_detailed", [False, True])
+def test_it2i_dataset_uses_upstream_warmups(monkeypatch, tmp_path, save_detailed):
     """Exercise real dataset/benchmark code; replace only the HTTP backend."""
     from vllm_omni.benchmarks import serve
     from vllm_omni.benchmarks.patch import patch
@@ -58,13 +59,21 @@ def test_it2i_dataset_uses_upstream_warmups(monkeypatch, tmp_path):
             "--extra-body",
             '{"seed":42,"guidance_scale":2.5,"bot_task":"think_recaption"}',
             "--save-result",
+            *(["--save-detailed"] if save_detailed else []),
+            "--omni-request-timeout-s",
+            "42",
+            "--videomme-save-eval-items",
             "--result-dir",
             str(tmp_path),
             "--result-filename",
             "result.json",
         ]
     )
     requests = []
+    snapshot = {"0": {"num_tokens_out": 1, "finish_reason": "stop", "vllm_itls_ms": [5.0]}}
+    monkeypatch.setattr(patch, "_SAVE_DETAILED", not save_detailed)
+    monkeypatch.setattr(patch, "_REQUEST_TIMEOUT_OVERRIDE_S", None)
+    monkeypatch.delenv("VIDEOMME_SAVE_EVAL_ITEMS", raising=False)
 
     async def request_func(request_func_input, session, pbar=None):
         requests.append(request_func_input)
@@ -77,10 +86,14 @@ async def request_func(request_func_input, session, pbar=None):
             image_count=1,
             start_time=time.perf_counter(),
             latency=99.0 if len(requests) <= 2 else 0.1,
+            stage_metrics=snapshot,
         )
 
     monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "/v1/images/edits", request_func)
     result = serve.main(args)
+    assert patch._SAVE_DETAILED is save_detailed
+    assert patch._omni_request_timeout_s() == 42
+    assert os.environ["VIDEOMME_SAVE_EVAL_ITEMS"] == "1"
     assert len(requests) == 10
     for request in requests:
         assert request.prompt == sample["prompt"]
@@ -90,6 +103,11 @@ async def request_func(request_func_input, session, pbar=None):
     saved = json.loads((tmp_path / "result.json").read_text())
     assert saved["completed"] == 8
     assert saved["mean_e2el_ms"] == pytest.approx(100.0)  # Excludes both 99-second warmups.
+    assert saved["request_stage_metrics"] == [{"0": {"num_tokens_out": 1, "finish_reason": "stop"}}] * 8
+    if save_detailed:
+        assert saved["stage_metrics"] == [snapshot] * 8
+    else:
+        assert "stage_metrics" not in saved
 
 
 @pytest.mark.parametrize(

vllm_omni/benchmarks/patch/patch.py

@@ -117,7 +117,16 @@
         "openai-image-edits-omni",
     }
 )
+_STAGE_METRICS_BACKENDS = frozenset(
+    {
+        "daily-omni",
+        "openai-chat-omni",
+        "/v1/videos",
+        *_IMAGE_STAGE_METRICS_BACKENDS,
+    }
+)
 _PRINT_STAGE = False
+_SAVE_DETAILED = False
 
 
 def maybe_enable_stage_metrics(extra_body: dict[str, Any] | None, *, enabled: bool) -> dict[str, Any] | None:
@@ -135,6 +144,8 @@ def should_request_stage_metrics(args: Any) -> bool:
         return True
 
     backend = getattr(args, "backend", None)
+    if getattr(args, "save_detailed", False) and backend in _STAGE_METRICS_BACKENDS:
+        return True
     if backend in _IMAGE_STAGE_METRICS_BACKENDS:
         return True
 
@@ -160,6 +171,12 @@ def set_print_stage(enabled: bool) -> None:
     _PRINT_STAGE = bool(enabled)
 
 
+def set_save_detailed(enabled: bool) -> None:
+    """Set whether this benchmark run retains per-request detail."""
+    global _SAVE_DETAILED
+    _SAVE_DETAILED = bool(enabled)
+
+
 def _audio_continuity_threshold_s() -> float:
     """Return the per-request underrun budget (s).
 
@@ -969,6 +986,20 @@ class MixRequestFuncOutput(RequestFuncOutput):
     duplex_session_metrics: dict[str, object] | None = None
 
 
+def _request_stage_metrics(outputs: Iterable[object]) -> list[dict[str, dict]]:
+    """Return one stage snapshot per request while preserving alignment."""
+    return [dict(getattr(output, "stage_metrics", None) or {}) for output in outputs]
+
+
+def _add_detailed_stage_metrics(
+    result: dict[str, Any],
+    outputs: Iterable[object],
+) -> None:
+    """Attach request-scoped stage snapshots only to detailed results."""
+    if _SAVE_DETAILED:
+        result["stage_metrics"] = _request_stage_metrics(outputs)
+
+
 _IMAGE_EDITS_EXTRA_BODY_FORM_FIELDS = (
     "negative_prompt",
     "num_inference_steps",
@@ -3692,6 +3723,7 @@ def measured_ttft(output: RequestFuncOutput) -> float | None:
     if any(request_stage_metrics):
         result["request_stage_metrics"] = request_stage_metrics
 
+    _add_detailed_stage_metrics(result, outputs)
     # 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

vllm_omni/benchmarks/serve.py

@@ -21,6 +21,7 @@
     maybe_enable_stage_metrics,
     set_print_stage,
     set_request_timeout_s,
+    set_save_detailed,
     should_request_stage_metrics,
 )
 
@@ -70,6 +71,7 @@ def main(args: argparse.Namespace) -> dict[str, Any]:
         set_request_timeout_s(request_timeout_s)
     _use_endpoint_backend_when_implicit(args)
     set_print_stage(getattr(args, "print_stage", False))
+    set_save_detailed(getattr(args, "save_detailed", False))
     args.extra_body = maybe_enable_stage_metrics(
         getattr(args, "extra_body", None),
         enabled=should_request_stage_metrics(args),