[Bugfix] Avoid invalid TPOT metrics (#6696)

Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

vllm_omni/benchmarks/patch/patch.py

@@ -1492,6 +1492,72 @@ def _nonnegative_number(value: object) -> bool:
     return isinstance(value, int | float) and not isinstance(value, bool) and np.isfinite(value) and value >= 0
 
 
+def _apply_stage0_token_timings(
+    output: MixRequestFuncOutput,
+    raw_stages: list[object],
+    *,
+    expected_output_tokens: int,
+) -> bool:
+    """Populate aggregate token timing from complete per-response Stage 0 metrics."""
+    output.itl = []
+    output.text_latency = output.ttft
+    output.tpot_measured = False
+
+    stages: list[tuple[int, dict[str, object]]] = []
+    for stage in raw_stages:
+        if not isinstance(stage, dict):
+            return False
+        token_count = stage.get("output_token_count")
+        if not isinstance(token_count, int) or isinstance(token_count, bool) or token_count < 0:
+            return False
+        stages.append((token_count, stage))
+
+    if not stages or sum(count for count, _ in stages) != expected_output_tokens:
+        return False
+
+    timed = [(count - 1, stage) for count, stage in stages if count > 1]
+    if not timed:
+        return False
+
+    measured_itls: list[float] = []
+    exact_itls = True
+    for intervals, stage in timed:
+        values = stage.get("itls_ms")
+        if (
+            not isinstance(values, list)
+            or len(values) != intervals
+            or not all(_nonnegative_number(value) for value in values)
+        ):
+            exact_itls = False
+            break
+        measured_itls.extend(float(value) / 1000.0 for value in values if isinstance(value, int | float))
+    if exact_itls:
+        if sum(measured_itls) > 0:
+            output.itl = measured_itls
+            output.text_latency = output.ttft + sum(output.itl)
+            output.tpot_measured = True
+            return True
+
+    weighted_tpot_ms = 0.0
+    for intervals, stage in timed:
+        tpot_ms = stage.get("tpot_ms")
+        if not _nonnegative_number(tpot_ms) or not isinstance(tpot_ms, int | float) or tpot_ms <= 0:
+            break
+        weighted_tpot_ms += float(tpot_ms) * intervals
+    else:
+        decode_intervals = sum(intervals for intervals, _ in timed)
+        weighted_tpot = weighted_tpot_ms / decode_intervals / 1000.0
+        # ``calculate_metrics`` consumes a request-shaped latency. Preserve
+        # the weighted per-response TPOT when projecting it onto the combined
+        # output token count; gaps and TTFTs between turns are intentionally
+        # excluded.
+        output.text_latency = output.ttft + weighted_tpot * (expected_output_tokens - 1)
+        output.tpot_measured = True
+        return True
+
+    return False
+
+
 async def _async_request_omniinteract(
     request_func_input: RequestFuncInput,
     *,
@@ -1538,35 +1604,12 @@ async def _async_request_omniinteract(
         output.ttft = float(session_metrics.get("mean_ttft_ms") or 0.0) / 1000.0
         output.audio_ttfp = float(session_metrics.get("mean_ttfp_ms") or 0.0) / 1000.0
         output.audio_rtf = float(session_metrics.get("mean_rtf") or 0.0)
-        stages: list[tuple[int, dict[str, object]]] = []
-        for request_metric in case_result.duplex_request_metrics:
-            stage0 = request_metric.get("stage0_tokens")
-            if not isinstance(stage0, dict):
-                continue
-            token_count = stage0.get("output_token_count")
-            if isinstance(token_count, int) and not isinstance(token_count, bool) and token_count >= 0:
-                stages.append((token_count, stage0))
-        timed = [(count - 1, stage) for count, stage in stages if count > 1]
-        complete = (
-            len(stages) == len(case_result.duplex_request_metrics)
-            and sum(count for count, _ in stages) == output.output_tokens
+        token_timing_measured = _apply_stage0_token_timings(
+            output,
+            [request_metric.get("stage0_tokens") for request_metric in case_result.duplex_request_metrics],
+            expected_output_tokens=output.output_tokens,
         )
-        exact = complete and all(
-            isinstance(values := stage.get("itls_ms"), list)
-            and len(values) == intervals
-            and all(_nonnegative_number(value) for value in values)
-            for intervals, stage in timed
-        )
-        output.itl = [float(value) / 1000.0 for _, stage in timed for value in stage["itls_ms"]] if exact else []
-        if output.itl:
-            output.text_latency = output.ttft + sum(output.itl)
-        elif complete and timed and all(_nonnegative_number(stage.get("tpot_ms")) for _, stage in timed):
-            intervals = sum(count for count, _ in timed)
-            weighted_tpot = sum(float(stage["tpot_ms"]) * count for count, stage in timed) / intervals / 1000.0
-            output.text_latency = output.ttft + weighted_tpot * (output.output_tokens - 1)
-        else:
-            output.text_latency = output.ttft
-            output.tpot_measured = False
+        if not token_timing_measured:
             if output.output_tokens > 1 or (output.output_tokens == 0 and output.generated_text):
                 logger.warning(
                     "OmniInteract session %s omitted complete engine token timing; standard TPOT/ITL are unavailable",
@@ -1610,6 +1653,9 @@ async def async_request_openai_realtime_duplex(
     speech_extra = getattr(request_func_input, "seed_tts_speech_extra", None)
     if not isinstance(speech_extra, dict):
         speech_extra = {}
+    else:
+        speech_extra = dict(speech_extra)
+    speech_extra["return_stage_metrics"] = True
     configured_turns = getattr(request_func_input, "seed_tts_turns", ())
     turn_prompts = [
         (
@@ -1740,6 +1786,16 @@ async def async_request_openai_realtime_duplex(
                 for timing in turn_timings
                 if isinstance((stage0 := timing.get("stage0_tokens")), dict)
             )
+            token_timing_measured = _apply_stage0_token_timings(
+                output,
+                [timing.get("stage0_tokens") for timing in turn_timings],
+                expected_output_tokens=output.output_tokens,
+            )
+            if not token_timing_measured and output.output_tokens > 1:
+                logger.warning(
+                    "Realtime TTS session %s omitted complete engine token timing; standard TPOT/ITL are unavailable",
+                    session_id,
+                )
             output.success = True
     except Exception:
         output.success = False