[Bugfix] clean per-stage audio RTF and align audio metrics (#4794)

vllm_omni/benchmarks/metrics/metrics.py

@@ -63,11 +63,9 @@ _STAGE_BENCHMARK_FIELDS = [
     (defs.VLLM_TPOTS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.VLLM_ITLS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.AUDIO_TTFPS, _FLOAT_LIST_TYPE, field(default_factory=list)),
-    (defs.AUDIO_RTFS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.AUDIO_DURATIONS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.AUDIO_FRAMES, _INT_LIST_TYPE, field(default_factory=list)),
     (defs.MISSING_AUDIO_DURATION_COUNT, int, field(default=0)),
-    (defs.MISSING_AUDIO_RTF_COUNT, int, field(default=0)),
     (defs.STAGE_GEN_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     (defs.POSTPROCESS_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
     ("output_unit_type", str, field(default="")),
@@ -443,12 +441,6 @@ def _print_audio_stage_metrics(
             getattr(sm, defs.AUDIO_TTFPS),
             selected_percentiles,
         )
-    if defs.AUDIO_RTF in selected_percentile_metrics:
-        audio_rtfs = getattr(sm, defs.AUDIO_RTFS)
-        if audio_rtfs:
-            _print_percentile_metric("Real Time Factor", "AUDIO_RTF", audio_rtfs, selected_percentiles, to_ms=False)
-        elif getattr(sm, defs.MISSING_AUDIO_RTF_COUNT) > 0:
-            print("{:<40} {:<10}".format("AUDIO_RTF skipped:", "missing stage-local audio duration/sample rate"))
     if defs.AUDIO_DURATION in selected_percentile_metrics:
         audio_durations = getattr(sm, defs.AUDIO_DURATIONS)
         if audio_durations:
@@ -643,11 +635,9 @@ def _build_stage_metrics_from_outputs(
             total_output = sum(int((info or {}).get(defs.NUM_TOKENS_OUT) or 0) for _, info in rows)
 
         audio_ttfps: list[float] = []
-        audio_rtfs: list[float] = []
         audio_durations: list[float] = []
         audio_frames: list[int] = []
         missing_audio_duration_count = 0
-        missing_audio_rtf_count = 0
         if is_audio_stage:
             for _, info in rows:
                 audio_ttfp_ms = float((info or {}).get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS) or 0.0)
@@ -665,12 +655,6 @@ def _build_stage_metrics_from_outputs(
                 else:
                     missing_audio_duration_count += 1
 
-                rtf = float((info or {}).get(defs.AUDIO_RTF) or 0.0)
-                if rtf > 0:
-                    audio_rtfs.append(rtf)
-                else:
-                    missing_audio_rtf_count += 1
-
         result.append(
             StageBenchmarkMetrics(
                 **{
@@ -685,11 +669,9 @@ def _build_stage_metrics_from_outputs(
                     defs.VLLM_TPOTS: vllm_tpots,
                     defs.VLLM_ITLS: vllm_itls,
                     defs.AUDIO_TTFPS: audio_ttfps,
-                    defs.AUDIO_RTFS: audio_rtfs,
                     defs.AUDIO_DURATIONS: audio_durations,
                     defs.AUDIO_FRAMES: audio_frames,
                     defs.MISSING_AUDIO_DURATION_COUNT: missing_audio_duration_count,
-                    defs.MISSING_AUDIO_RTF_COUNT: missing_audio_rtf_count,
                     defs.STAGE_GEN_TIMES_MS: stage_gen_times_ms,
                     defs.POSTPROCESS_TIMES_MS: postprocess_times_ms,
                     "output_unit_type": output_unit_type,

vllm_omni/engine/stage_pool.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import time as _time
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
 
@@ -597,7 +598,6 @@ class StagePool:
         batch_id = metrics.batch_seq
         metrics.agg_total_tokens += num_tokens_out
         metrics.agg_total_gen_time_ms += stage_gen_time_ms
-        audio_rtf = defs.compute_audio_rtf(stage_gen_time_ms / 1000.0, audio_duration_s)
 
         return StageRequestMetrics(
             num_tokens_in=num_tokens_in,
@@ -616,7 +616,6 @@ class StagePool:
             audio_generated_frames=audio_generated_frames,
             audio_sample_rate=audio_sample_rate,
             audio_duration_s=audio_duration_s,
-            audio_rtf=audio_rtf,
             image_pixels=image_pixels,
             denoise_step_latency_ms=denoise_step_latency_ms,
             output_unit_type=output_unit_type,
@@ -686,7 +685,7 @@ class StagePool:
     def _has_audio_output(self, request_outputs: list[Any]) -> bool:
         for ro in request_outputs:
             for mm_output in self._iter_multimodal_outputs(ro):
-                if isinstance(mm_output, dict) and mm_output.get("audio") is not None:
+                if isinstance(mm_output, Mapping) and mm_output.get("audio") is not None:
                     return True
         return False
 
@@ -695,7 +694,7 @@ class StagePool:
         sample_rate = 0
         for ro in request_outputs:
             for mm_output in self._iter_multimodal_outputs(ro):
-                if not isinstance(mm_output, dict):
+                if not isinstance(mm_output, Mapping):
                     continue
                 if sample_rate <= 0:
                     sample_rate = self._infer_audio_sample_rate(mm_output)
@@ -779,7 +778,7 @@ class StagePool:
     def _count_images(self, request_output: Any) -> int:
         total_images = self._count_value_units(getattr(request_output, "images", None))
         for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict):
+            if isinstance(mm_output, Mapping):
                 total_images += self._count_value_units(mm_output.get("image"))
                 total_images += self._count_value_units(mm_output.get("images"))
         return total_images
@@ -789,7 +788,7 @@ class StagePool:
         for ro in request_outputs:
             total_pixels += self._count_image_value_pixels(getattr(ro, "images", None))
             for mm_output in self._iter_multimodal_outputs(ro):
-                if isinstance(mm_output, dict):
+                if isinstance(mm_output, Mapping):
                     total_pixels += self._count_image_value_pixels(mm_output.get("image"))
                     total_pixels += self._count_image_value_pixels(mm_output.get("images"))
         return total_pixels
@@ -824,7 +823,7 @@ class StagePool:
         total_videos = self._count_video_units(getattr(request_output, "video", None))
         total_videos += self._count_video_units(getattr(request_output, "videos", None))
         for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict):
+            if isinstance(mm_output, Mapping):
                 total_videos += self._count_video_units(mm_output.get("video"))
                 total_videos += self._count_video_units(mm_output.get("videos"))
         return total_videos
@@ -849,7 +848,7 @@ class StagePool:
         if isinstance(custom_output, dict) and custom_output:
             return True
         for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict) and any(self._is_non_empty_value(value) for value in mm_output.values()):
+            if isinstance(mm_output, Mapping) and any(self._is_non_empty_value(value) for value in mm_output.values()):
                 return True
         for output in getattr(request_output, "outputs", None) or []:
             if getattr(output, "text", None):
@@ -917,14 +916,14 @@ class StagePool:
             if self._audio_sample_rate_by_request.get(rid, 0) <= 0 and audio_sample_rate > 0:
                 self._audio_sample_rate_by_request[rid] = audio_sample_rate
 
-    def _iter_multimodal_outputs(self, request_output: Any) -> list[dict[str, Any]]:
-        multimodal_outputs: list[dict[str, Any]] = []
+    def _iter_multimodal_outputs(self, request_output: object) -> list[Mapping[str, Any]]:
+        multimodal_outputs: list[Mapping[str, Any]] = []
         outer_mm = getattr(request_output, "multimodal_output", None)
-        if isinstance(outer_mm, dict) and outer_mm:
+        if isinstance(outer_mm, Mapping) and outer_mm:
             multimodal_outputs.append(outer_mm)
         for output in getattr(request_output, "outputs", None) or []:
             inner_mm = getattr(output, "multimodal_output", None)
-            if isinstance(inner_mm, dict) and inner_mm:
+            if isinstance(inner_mm, Mapping) and inner_mm:
                 multimodal_outputs.append(inner_mm)
         return multimodal_outputs
 

vllm_omni/entrypoints/omni_base.py

@@ -221,6 +221,13 @@ class OmniBase(PDDisaggregationMixin):
         """Expose engine stage configs for PD disaggregation detection and validation."""
         return self.engine.stage_configs
 
+    def _consumed_metric_message_ids(self, request_id: str) -> set[int]:
+        consumed_by_request = getattr(self, "_consumed_metric_messages", None)
+        if consumed_by_request is None:
+            consumed_by_request = {}
+            self._consumed_metric_messages = consumed_by_request
+        return consumed_by_request.setdefault(request_id, set())
+
     def _has_dead_stage(self) -> bool:
         for stage_client in self.engine.stage_clients:
             if getattr(stage_client, "_engine_dead", False):
@@ -310,7 +317,9 @@ class OmniBase(PDDisaggregationMixin):
             )
         finally:
             self.request_states.pop(request_id, None)
-            self._consumed_metric_messages.pop(request_id, None)
+            consumed_by_request = getattr(self, "_consumed_metric_messages", None)
+            if consumed_by_request is not None:
+                consumed_by_request.pop(request_id, None)
             # Republish gauges so any stale value left by the per-stage
             # publish in _process_single_result (which runs while the request
             # is still in self.request_states) is corrected after the pop.
@@ -397,7 +406,7 @@ class OmniBase(PDDisaggregationMixin):
             stage_meta = self.engine.get_stage_metadata(stage_id)
             output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type)
             msg_id = id(msg)
-            consumed = self._consumed_metric_messages.setdefault(req_id, set())
+            consumed = self._consumed_metric_message_ids(req_id)
             if msg_id not in consumed:
                 req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type)
                 submit_ts = msg.stage_submit_ts
@@ -505,7 +514,11 @@ class OmniBase(PDDisaggregationMixin):
         stage_meta = self.engine.get_stage_metadata(stage_id)
         output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
         if finished and _m is not None:
-            metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
+            msg_id = id(result)
+            consumed = self._consumed_metric_message_ids(req_id)
+            if msg_id not in consumed:
+                metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
+                consumed.add(msg_id)
 
         if not stage_meta.final_output:
             return None

vllm_omni/metrics/definitions.py

@@ -76,10 +76,8 @@ VLLM_TTFTS = "vllm_ttfts"
 VLLM_TPOTS = "vllm_tpots"
 VLLM_ITLS = "vllm_itls"
 AUDIO_TTFPS = "audio_ttfps"
-AUDIO_RTFS = "audio_rtfs"
 AUDIO_DURATIONS = "audio_durations"
 MISSING_AUDIO_DURATION_COUNT = "missing_audio_duration_count"
-MISSING_AUDIO_RTF_COUNT = "missing_audio_rtf_count"
 STAGE_GEN_TIME = "stage_gen_time"
 STAGE_GEN_TIME_MS = f"{STAGE_GEN_TIME}_ms"
 STAGE_GEN_TIMES_MS = f"{STAGE_GEN_TIME}s_ms"

vllm_omni/metrics/stats.py

@@ -48,7 +48,6 @@ class StageRequestStats:
     audio_generated_frames: int = 0
     audio_sample_rate: int = 0
     audio_duration_s: float = 0.0
-    audio_rtf: float = 0.0
     image_pixels: int = 0
     denoise_step_latency_ms: float = 0.0
     pipeline_timings: dict[str, float] | None = None
@@ -471,7 +470,6 @@ class OrchestratorAggregator:
                 defs.AUDIO_FRAMES: int(evt.audio_generated_frames),
                 defs.AUDIO_SAMPLE_RATE: int(evt.audio_sample_rate),
                 f"{defs.AUDIO_DURATION}_s": float(evt.audio_duration_s),
-                defs.AUDIO_RTF: float(evt.audio_rtf),
                 defs.IMAGE_PIXELS: int(evt.image_pixels),
                 defs.DENOISE_STEP_LATENCY_MS: float(evt.denoise_step_latency_ms),
                 "output_unit_type": evt.output_unit_type,
@@ -541,10 +539,6 @@ class OrchestratorAggregator:
             0.0,
         )
         current[defs.TIME_PER_OUTPUT_UNIT_MS] = remaining_ms / float(output_count - 1) if output_count > 1 else 0.0
-        duration_s = float(current.get(f"{defs.AUDIO_DURATION}_s", 0.0))
-        current[defs.AUDIO_RTF] = defs.compute_audio_rtf(
-            float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) / 1000.0, duration_s
-        )
         return current
 
     def _build_stage_metrics_snapshot(self, req_id: str) -> dict[str, dict[str, Any]]: