[Refactor] Centralize metrics related helpers to metrics/ (#5168)

tests/engine/test_orchestrator.py

@@ -30,7 +30,6 @@ from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
     _build_terminal_empty_output,
-    _infer_stage_audio_sample_rate,
 )
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.experimental.fullduplex.engine.duplex_control_plane import DuplexControlPlane
@@ -201,12 +200,12 @@ class FakeStageClient:
 def test_terminal_empty_audio_output_uses_stage_sample_rate() -> None:
     final_stage = FakeStageClient(final_output=True, final_output_type="audio")
     final_stage.sample_rate = 44100
-    final_pool = SimpleNamespace(stage_client=final_stage, _stage_vllm_config=None)
+    final_pool = StagePool(0, [final_stage])
 
     terminal_output = _build_terminal_empty_output(
         "req-1",
         final_output_type="audio",
-        audio_sample_rate=_infer_stage_audio_sample_rate(final_pool),
+        audio_sample_rate=final_pool._infer_audio_sample_rate(),
     )
 
     assert terminal_output.outputs[0].multimodal_output["sr"] == 44100

tests/metrics/test_metrics_utils.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.metrics.utils import (
+    count_audio_chunk_frames,
+    count_audio_frames,
+    count_image_pixels,
+    count_tokens_from_outputs,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("audio_chunk", "expected"),
+    [
+        (SimpleNamespace(shape=(2, 320)), 320),
+        ([0.0, 0.0, 0.0], 3),
+        (1.0, 1),
+    ],
+)
+def test_count_audio_chunk_frames(audio_chunk: object, expected: int) -> None:
+    assert count_audio_chunk_frames(audio_chunk) == expected
+
+
+def test_count_audio_frames_sums_audio_chunks() -> None:
+    mm_out = {
+        "audio": [
+            SimpleNamespace(shape=(1, 120)),
+            SimpleNamespace(shape=(2, 80)),
+        ]
+    }
+
+    assert count_audio_frames(mm_out) == 200
+
+
+def test_count_audio_frames_supports_model_outputs() -> None:
+    mm_out = {
+        "model_outputs": [
+            SimpleNamespace(shape=(1, 120)),
+            SimpleNamespace(shape=(2, 80)),
+        ]
+    }
+
+    assert count_audio_frames(mm_out) == 200
+
+
+@pytest.mark.parametrize("mm_out", [{}, {"audio": None}, {"model_outputs": None}])
+def test_count_audio_frames_returns_zero_without_audio(mm_out: dict[str, object]) -> None:
+    assert count_audio_frames(mm_out) == 0
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (None, 0),
+        (object(), 0),
+        (SimpleNamespace(size=(10, 20)), 200),
+        (SimpleNamespace(shape=(2, 3, 4, 5)), 40),
+        (SimpleNamespace(shape=(3, 4, 5)), 20),
+        (SimpleNamespace(shape=(6, 7, 3)), 42),
+        (SimpleNamespace(shape=(6, 7)), 42),
+    ],
+)
+def test_count_image_pixels(value: object, expected: int) -> None:
+    assert count_image_pixels(value) == expected
+
+
+def test_count_image_pixels_sums_nested_values() -> None:
+    images = [
+        SimpleNamespace(size=(10, 20)),
+        (
+            SimpleNamespace(shape=(3, 4, 5)),
+            SimpleNamespace(shape=(6, 7, 3)),
+        ),
+    ]
+
+    assert count_image_pixels(images) == 262
+
+
+def test_count_tokens_from_outputs() -> None:
+    engine_outputs = [
+        SimpleNamespace(
+            outputs=[
+                SimpleNamespace(cumulative_token_ids=[1, 2, 3], token_ids=[3]),
+                SimpleNamespace(cumulative_token_ids=None, token_ids=[4, 5]),
+                SimpleNamespace(cumulative_token_ids=None, token_ids=None),
+            ]
+        ),
+        SimpleNamespace(outputs=None),
+        SimpleNamespace(
+            outputs=[
+                SimpleNamespace(cumulative_token_ids=[6], token_ids=[6]),
+            ]
+        ),
+    ]
+
+    assert count_tokens_from_outputs(engine_outputs) == 6

vllm_omni/benchmarks/patch/patch.py

@@ -48,6 +48,7 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
 from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDataset
 from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
 from vllm_omni.metrics import definitions as defs
+from vllm_omni.metrics.utils import coerce_positive_int_scalar
 
 logger = init_logger(__name__)
 
@@ -524,16 +525,8 @@ def _add_image_edit_extra_body_to_form(form: aiohttp.FormData, extra_body: dict[
             form.add_field(key, str(value))
 
 
-def _coerce_positive_int(value: Any) -> int | None:
-    try:
-        coerced = int(value)
-    except (TypeError, ValueError):
-        return None
-    return coerced if coerced > 0 else None
-
-
 def _extract_output_tokens_from_metrics(metrics: dict[str, Any]) -> int | None:
-    top_level_tokens = _coerce_positive_int(metrics.get(defs.NUM_TOKENS_OUT))
+    top_level_tokens = coerce_positive_int_scalar(metrics.get(defs.NUM_TOKENS_OUT))
     if top_level_tokens is not None:
         return top_level_tokens
 
@@ -545,7 +538,7 @@ def _extract_output_tokens_from_metrics(metrics: dict[str, Any]) -> int | None:
     for info in stage_snapshot.values():
         if not isinstance(info, dict):
             continue
-        num_tokens_out = _coerce_positive_int(info.get(defs.NUM_TOKENS_OUT))
+        num_tokens_out = coerce_positive_int_scalar(info.get(defs.NUM_TOKENS_OUT))
         if num_tokens_out is None:
             continue
         if info.get("final_output_type") == "text" or info.get("output_unit_type") == "token":
@@ -556,9 +549,9 @@ def _extract_output_tokens_from_metrics(metrics: dict[str, Any]) -> int | None:
 
 def _apply_usage_to_output(output: MixRequestFuncOutput, usage: dict[str, Any]) -> int | None:
     """Apply OpenAI ``usage`` fields to the benchmark output."""
-    if (pt := _coerce_positive_int(usage.get("prompt_tokens"))) is not None:
+    if (pt := coerce_positive_int_scalar(usage.get("prompt_tokens"))) is not None:
         output.prompt_len = pt
-    completion_tokens = _coerce_positive_int(usage.get("completion_tokens"))
+    completion_tokens = coerce_positive_int_scalar(usage.get("completion_tokens"))
     if completion_tokens is not None:
         output.output_tokens = max(int(output.output_tokens or 0), completion_tokens)
     return completion_tokens
@@ -901,12 +894,22 @@ async def async_request_openai_chat_omni_completions(
                     audio_frames = 0
                     if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None:
                         channels, sample_width, frame_rate = wav_audio_params
-                        frame_width = sample_width * channels
-                        if frame_width > 0:
-                            audio_frames = len(wav_pcm_buffer) // frame_width
-                            audio_duration_sec = audio_frames / frame_rate
+                        audio_frames = defs.compute_audio_frames(
+                            len(wav_pcm_buffer),
+                            sample_width=sample_width,
+                            channels=channels,
+                        )
+                        if audio_frames > 0 and frame_rate > 0:
+                            audio_duration_sec = audio_frames / float(frame_rate)
                         else:
-                            logger.warning("Audio frame width is zero")
+                            logger.warning(
+                                "Unable to derive audio frames/duration from wav pcm "
+                                "(pcm_nbytes=%d, sample_width=%d, channels=%d, sample_rate=%d)",
+                                len(wav_pcm_buffer),
+                                sample_width,
+                                channels,
+                                frame_rate,
+                            )
                     elif audio_bytes_buffer:
                         try:
                             from vllm.multimodal.audio import get_audio_duration
@@ -925,10 +928,8 @@ async def async_request_openai_chat_omni_completions(
                         output.audio_duration = audio_duration_sec
                         output.audio_frames = audio_frames
                         audio_duration = output.audio_duration
-                        if audio_duration > 0:
-                            output.audio_rtf = audio_generate_time / output.audio_duration
-                        else:
-                            output.audio_rtf = 0
+                        output.audio_rtf = defs.compute_audio_rtf(audio_generate_time, audio_duration)
+                        if audio_duration <= 0:
                             logger.warning("Audio duration is zero")
                         if _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False):
                             try:
@@ -1161,7 +1162,7 @@ async def async_request_openai_audio_speech(
 
     # PCM format: 16-bit signed; sample_rate/channels are model-dependent.
     sample_rate, channels = defs.stream_pcm_format_from_env()
-    sample_width = 2  # 16-bit = 2 bytes
+    sample_width = defs.DEFAULT_AUDIO_SAMPLE_WIDTH
 
     st = time.perf_counter()
     output.start_time = st
@@ -1190,13 +1191,17 @@ async def async_request_openai_audio_speech(
                 end_time = time.perf_counter()
                 output.latency = end_time - st
 
-                total_samples = total_pcm_bytes // (sample_width * channels)
-                output.audio_duration = total_samples / sample_rate
+                total_samples = defs.compute_audio_frames(
+                    total_pcm_bytes,
+                    sample_width=sample_width,
+                    channels=channels,
+                )
                 output.audio_frames = total_samples
-                if output.audio_duration > 0:
-                    output.audio_rtf = output.latency / output.audio_duration
-                else:
-                    output.audio_rtf = 0
+                output.audio_duration = (
+                    float(total_samples) / float(sample_rate) if total_samples > 0 and sample_rate > 0 else 0.0
+                )
+                output.audio_rtf = defs.compute_audio_rtf(output.latency, output.audio_duration)
+                if output.audio_duration <= 0:
                     logger.warning("Audio duration is zero")
 
                 continuity = compute_continuity_stats(

vllm_omni/engine/orchestrator.py

@@ -104,37 +104,6 @@ def _build_terminal_empty_output(
     )
 
 
-def _coerce_int_scalar(value: Any) -> int:
-    if value is None:
-        return 0
-    if isinstance(value, (list, tuple)):
-        for item in value:
-            coerced = _coerce_int_scalar(item)
-            if coerced > 0:
-                return coerced
-        return 0
-    if hasattr(value, "item"):
-        value = value.item()
-    try:
-        coerced = int(value)
-    except (TypeError, ValueError):
-        return 0
-    return coerced if coerced > 0 else 0
-
-
-def _infer_stage_audio_sample_rate(stage_pool: StagePool, default: int = 24000) -> int:
-    """Infer the final audio stage sample rate from stage metadata when possible."""
-    sample_rate_attrs = ("audio_sample_rate", "sample_rate", "sampling_rate", "output_sample_rate", "sr")
-    stage_client = getattr(stage_pool, "stage_client", None)
-    stage_config = getattr(stage_pool, "_stage_vllm_config", None)
-    for source in (stage_client, stage_config):
-        for attr in sample_rate_attrs:
-            sample_rate = _coerce_int_scalar(getattr(source, attr, None))
-            if sample_rate > 0:
-                return sample_rate
-    return default
-
-
 def build_engine_core_request_from_tokens(
     request_id: str,
     prompt: dict[str, Any],
@@ -1909,7 +1878,7 @@ class Orchestrator:
             terminal_output = _build_terminal_empty_output(
                 req_id,
                 final_output_type=final_output_type,
-                audio_sample_rate=_infer_stage_audio_sample_rate(final_pool),
+                audio_sample_rate=final_pool._infer_audio_sample_rate(),
             )
             submit_ts = _time.time()
             req_state.stage_submit_ts[final_stage_id] = submit_ts

vllm_omni/engine/stage_pool.py

@@ -25,10 +25,18 @@ from vllm_omni.engine.stage_client import (
     StagePoolLLMClient,
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.metrics import (
+    count_audio_frames,
+    count_image_pixels,
+    count_tokens_from_outputs,
+)
 from vllm_omni.metrics import definitions as defs
 from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
 from vllm_omni.metrics.stats import StageStats
-from vllm_omni.metrics.utils import count_tokens_from_outputs
+from vllm_omni.metrics.utils import (
+    coerce_positive_int_scalar,
+    iter_mm_outputs,
+)
 
 if TYPE_CHECKING:
     from vllm_omni.engine.orchestrator import OrchestratorRequestState
@@ -619,7 +627,7 @@ class StagePool:
             else 0.0
         )
         image_pixels = self._count_image_pixels(request_outputs) if output_unit_type == "image" else 0
-        num_inference_steps = self._coerce_int_scalar(getattr(sampling_params, "num_inference_steps", None))
+        num_inference_steps = coerce_positive_int_scalar(getattr(sampling_params, "num_inference_steps", None)) or 0
         denoise_step_latency_ms = (
             defs.compute_denoise_step_latency(stage_gen_time_ms, num_inference_steps)
             if output_unit_type == "image"
@@ -741,78 +749,40 @@ 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):
+            for mm_output in iter_mm_outputs(ro):
                 if isinstance(mm_output, Mapping) and mm_output.get("audio") is not None:
                     return True
         return False
 
-    def _collect_audio_metrics(self, request_outputs: list[Any]) -> tuple[int, int, float]:
+    def _collect_audio_metrics(
+        self,
+        request_outputs: list[Any],
+        *,
+        use_default_sample_rate: bool = False,
+    ) -> tuple[int, int, float]:
         total_frames = 0
         sample_rate = 0
         for ro in request_outputs:
-            for mm_output in self._iter_multimodal_outputs(ro):
-                if not isinstance(mm_output, Mapping):
-                    continue
+            for mm_output in iter_mm_outputs(ro):
                 if sample_rate <= 0:
-                    sample_rate = self._infer_audio_sample_rate(mm_output)
-                audio_output = mm_output.get("audio")
-                if audio_output is None:
-                    continue
-                items = audio_output if isinstance(audio_output, list) else [audio_output]
-                for item in items:
-                    total_frames += self._count_audio_frames(item)
+                    sample_rate = self._infer_audio_sample_rate(mm_output, use_default=use_default_sample_rate)
+                total_frames += count_audio_frames(mm_output)
         duration_s = (float(total_frames) / float(sample_rate)) if total_frames > 0 and sample_rate > 0 else 0.0
         return int(total_frames), int(sample_rate), duration_s
 
-    @staticmethod
-    def _count_audio_frames(audio_output: Any) -> int:
-        shape = getattr(audio_output, "shape", None)
-        if shape is not None and len(shape) > 0:
-            # Audio chunks are concatenated on dim=-1 in the output processor,
-            # so the frame/sample axis is the last dim (e.g. [channels, frames]).
-            # Keep this aligned with serving_chat.py: audio tensors are consumed
-            # as (T,), (C, T), or (B, C, T). Flattening would corrupt
-            # multi-channel audio.
-            return int(shape[-1])
-        try:
-            return len(audio_output)
-        except TypeError:
-            return 1
-
-    def _infer_audio_sample_rate(self, mm_output: dict[str, Any]) -> int:
-        for key in ("audio_sample_rate", "sample_rate", "sampling_rate", "sr"):
-            rate = self._coerce_int_scalar(mm_output.get(key))
-            if rate > 0:
-                return rate
-        for attr in ("audio_sample_rate", "sample_rate", "sampling_rate", "output_sample_rate"):
-            rate = self._coerce_int_scalar(getattr(self.stage_client, attr, None))
-            if rate > 0:
-                return rate
-            rate = self._coerce_int_scalar(getattr(self._stage_vllm_config, attr, None))
-            if rate > 0:
-                return rate
-        return 0
-
-    @classmethod
-    def _coerce_int_scalar(cls, value: Any) -> int:
-        if value is None:
-            return 0
-        if isinstance(value, (list, tuple)):
-            for item in value:
-                coerced = cls._coerce_int_scalar(item)
-                if coerced > 0:
-                    return coerced
-            return 0
-        item = getattr(value, "item", None)
-        if callable(item):
-            try:
-                return int(item())
-            except Exception:
-                return 0
-        try:
-            return int(value)
-        except (TypeError, ValueError):
-            return 0
+    def _infer_audio_sample_rate(
+        self,
+        mm_output: dict[str, Any] | None = None,
+        *,
+        use_default: bool = True,
+    ) -> int:
+        sources: list[Any] = []
+        if mm_output is not None:
+            sources.append(mm_output)
+        sources.extend([self.stage_client, self._stage_vllm_config])
+        if use_default:
+            return defs.resolve_audio_sample_rate(sources)
+        return defs.resolve_audio_sample_rate_or_none(sources) or 0
 
     def _has_image_output(self, request_outputs: list[Any]) -> bool:
         for ro in request_outputs:
@@ -834,7 +804,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):
+        for mm_output in iter_mm_outputs(request_output):
             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"))
@@ -843,43 +813,17 @@ class StagePool:
     def _count_image_pixels(self, request_outputs: list[Any]) -> int:
         total_pixels = 0
         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):
+            total_pixels += count_image_pixels(getattr(ro, "images", None))
+            for mm_output in iter_mm_outputs(ro):
                 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"))
+                    total_pixels += count_image_pixels(mm_output.get("image"))
+                    total_pixels += count_image_pixels(mm_output.get("images"))
         return total_pixels
 
-    @classmethod
-    def _count_image_value_pixels(cls, value: Any) -> int:
-        if value is None:
-            return 0
-        if isinstance(value, (list, tuple)):
-            return sum(cls._count_image_value_pixels(item) for item in value)
-
-        size = getattr(value, "size", None)
-        if isinstance(size, tuple) and len(size) >= 2:
-            try:
-                return int(size[0]) * int(size[1])
-            except (TypeError, ValueError):
-                return 0
-
-        shape = getattr(value, "shape", None)
-        if shape is None or len(shape) < 2:
-            return 0
-        dims = [int(dim) for dim in shape]
-        if len(dims) >= 4:
-            return dims[0] * dims[-2] * dims[-1]
-        if len(dims) == 3 and dims[0] in (1, 3, 4):
-            return dims[1] * dims[2]
-        if len(dims) == 3 and dims[-1] in (1, 3, 4):
-            return dims[0] * dims[1]
-        return dims[-2] * dims[-1]
-
     def _count_videos(self, request_output: Any) -> int:
         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):
+        for mm_output in iter_mm_outputs(request_output):
             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"))
@@ -904,7 +848,7 @@ class StagePool:
         custom_output = getattr(request_output, "_custom_output", None)
         if isinstance(custom_output, dict) and custom_output:
             return True
-        for mm_output in self._iter_multimodal_outputs(request_output):
+        for mm_output in iter_mm_outputs(request_output):
             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 []:
@@ -967,23 +911,15 @@ class StagePool:
             self._output_timestamps_by_request.setdefault(rid, []).append(output_ts)
             if self._has_non_empty_output(request_output):
                 self._non_empty_first_output_timestamps_by_request.setdefault(rid, output_ts)
-            audio_frames, audio_sample_rate, _ = self._collect_audio_metrics([request_output])
+            audio_frames, audio_sample_rate, _ = self._collect_audio_metrics(
+                [request_output],
+                use_default_sample_rate=False,
+            )
             if audio_frames > 0:
                 self._audio_frames_by_request[rid] = self._audio_frames_by_request.get(rid, 0) + audio_frames
             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: object) -> list[Mapping[str, Any]]:
-        multimodal_outputs: list[Mapping[str, Any]] = []
-        outer_mm = getattr(request_output, "multimodal_output", None)
-        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, Mapping) and inner_mm:
-                multimodal_outputs.append(inner_mm)
-        return multimodal_outputs
-
     # ---- Stage-local admission ----
 
     async def submit_initial(

vllm_omni/metrics/__init__.py

@@ -1,6 +1,11 @@
 from .prometheus import OmniPrometheusMetrics, OmniRequestCounter
 from .stats import OrchestratorAggregator, StageRequestStats, StageStats
-from .utils import count_tokens_from_outputs
+from .utils import (
+    count_audio_chunk_frames,
+    count_audio_frames,
+    count_image_pixels,
+    count_tokens_from_outputs,
+)
 
 __all__ = [
     "OmniPrometheusMetrics",
@@ -8,5 +13,8 @@ __all__ = [
     "OrchestratorAggregator",
     "StageStats",
     "StageRequestStats",
+    "count_audio_chunk_frames",
+    "count_audio_frames",
+    "count_image_pixels",
     "count_tokens_from_outputs",
 ]

vllm_omni/metrics/definitions.py

@@ -13,7 +13,9 @@ time-bearing metrics use the ``_s`` suffix (values in seconds), counters use
 
 import logging
 import os
-from typing import Any
+from collections.abc import Mapping, Sequence
+
+from vllm_omni.metrics.utils import resolve_int_by_sequential_keys
 
 # vllm_omni: namespace for omni-specific Prometheus families, distinct from
 # the upstream vllm:* families.
@@ -243,8 +245,20 @@ BYTES_BUCKETS = (
 
 
 # ============================================================================
-# Audio-continuity defaults
+# Audio defaults (shared by server-side observe and bench-side calculation)
 # ============================================================================
+# Most common across vllm-omni talker variants (cosyvoice3, omnivoice,
+# qwen3_tts, mimo_audio). voxcpm2 uses 48000, stable_audio 44100,
+# ming_flash 16000 — these models surface a rate at runtime on one of
+# ``_SAMPLE_RATE_KEYS``, so DEFAULT_AUDIO_SAMPLE_RATE only applies when no
+# usable positive rate is found (and as the bench PCM default when env is unset).
+# Streamed OpenAI ``response_format=pcm`` is PCM_16 (s16le); sample width is fixed.
+DEFAULT_AUDIO_SAMPLE_RATE = 24000
+DEFAULT_AUDIO_CHANNELS = 1
+DEFAULT_AUDIO_SAMPLE_WIDTH = 2  # PCM s16le bytes per sample
+AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
+AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
+
 # Default underrun threshold — kept aligned with the bench-side default and
 # the commonly-cited "audible gap" threshold for streaming TTS.
 AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S = 0.1
@@ -253,8 +267,8 @@ AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S = 0.1
 # ============================================================================
 # Formula helpers (shared by server-side observe and bench-side calculation)
 # ============================================================================
-def compute_audio_rtf(stage_gen_time_s: float, audio_duration_s: float) -> float:
-    """RTF = stage_gen_time / audio_content_duration.
+def compute_audio_rtf(audio_generation_latency_s: float, audio_duration_s: float) -> float:
+    """RTF = audio_generation_latency / audio_content_duration.
 
     SLO red line < 1 — must generate faster than content plays back to stream.
     Returns 0.0 when audio_duration_s is non-positive (caller decides whether
@@ -262,7 +276,30 @@ def compute_audio_rtf(stage_gen_time_s: float, audio_duration_s: float) -> float
     """
     if audio_duration_s <= 0:
         return 0.0
-    return stage_gen_time_s / audio_duration_s
+    return audio_generation_latency_s / audio_duration_s
+
+
+def compute_audio_frames(
+    pcm_nbytes: int,
+    *,
+    sample_width: int = DEFAULT_AUDIO_SAMPLE_WIDTH,
+    channels: int = DEFAULT_AUDIO_CHANNELS,
+) -> int:
+    """Convert a PCM byte length to sample/frame count.
+
+    ``pcm_nbytes`` is the raw byte size of interleaved PCM (e.g.
+    ``total_pcm_bytes`` or ``len(wav_pcm_buffer)``). Defaults match the
+    streamed s16le mono contract
+    (``DEFAULT_AUDIO_SAMPLE_WIDTH``, ``DEFAULT_AUDIO_CHANNELS``);
+    WAV callers should pass header values.
+
+    Returns 0 when the byte count or frame width (``sample_width * channels``)
+    is non-positive.
+    """
+    frame_width = sample_width * channels
+    if pcm_nbytes <= 0 or frame_width <= 0:
+        return 0
+    return pcm_nbytes // frame_width
 
 
 def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int) -> float:
@@ -278,40 +315,45 @@ def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int
 # ============================================================================
 # Audio sample-rate resolution
 # ============================================================================
-# Most common across vllm-omni talker variants (cosyvoice3, omnivoice,
-# qwen3_tts, mimo_audio). voxcpm2 uses 48000, stable_audio 44100,
-# ming_flash 16000 — these models populate multimodal_output["audio_sample_rate"]
-# at runtime so this default only kicks in when the field is missing.
-DEFAULT_AUDIO_SAMPLE_RATE = 24000
-DEFAULT_AUDIO_CHANNELS = 1
-AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
-AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
-
 _SAMPLE_RATE_KEYS = ("output_sample_rate", "audio_sample_rate", "sample_rate", "sampling_rate", "sr")
 
 
-def resolve_audio_sample_rate(source: dict[str, Any] | Any | None) -> int:
-    """Extract audio sample_rate from a dict or config object, with fallbacks.
+def resolve_audio_sample_rate_or_none(
+    source: Sequence[Mapping[str, object] | object | None] | Mapping[str, object] | object | None,
+) -> int | None:
+    """Extract an explicitly configured audio sample rate, or ``None`` when absent.
 
-    Tries the same key chain as serving_chat.py's audio response path so
-    /metrics audio_duration_s = audio_frames / sample_rate stays consistent
-    with what the OpenAI streaming endpoint reports back to clients. Also
-    accepts config objects that expose the same values as attributes.
-    Returns DEFAULT_AUDIO_SAMPLE_RATE when no usable value is present.
+    ``source`` may be:
+
+    - a Mapping (e.g. ``multimodal_output`` / config dict)
+    - an object exposing the same fields as attributes
+    - a Sequence of the above (``str`` / ``bytes`` excluded); the first
+      positive rate found across items wins
     """
-    if not source:
-        return DEFAULT_AUDIO_SAMPLE_RATE
-    for key in _SAMPLE_RATE_KEYS:
-        raw = source.get(key) if isinstance(source, dict) else getattr(source, key, None)
-        if raw is None:
-            continue
-        try:
-            value = int(raw)
-        except (TypeError, ValueError):
-            continue
-        if value > 0:
-            return value
-    return DEFAULT_AUDIO_SAMPLE_RATE
+    if isinstance(source, Sequence) and not isinstance(source, (str, bytes, bytearray)):
+        for item in source:
+            rate = resolve_int_by_sequential_keys(item, _SAMPLE_RATE_KEYS)
+            if rate is not None:
+                return rate
+        return None
+    return resolve_int_by_sequential_keys(source, _SAMPLE_RATE_KEYS)
+
+
+def resolve_audio_sample_rate(
+    source: Sequence[Mapping[str, object] | object | None] | Mapping[str, object] | object | None,
+    default: int = DEFAULT_AUDIO_SAMPLE_RATE,
+) -> int:
+    """Extract audio sample rate from a dict, config object, or list thereof.
+
+    Delegates to :func:`resolve_audio_sample_rate_or_none` (same key chain and
+    coerce rules). Returns ``default`` (``DEFAULT_AUDIO_SAMPLE_RATE`` unless
+    overridden) when no usable positive rate is present.
+
+    Used so /metrics ``audio_duration_s = audio_frames / sample_rate`` stays
+    consistent with rates surfaced on multimodal outputs and stage configs.
+    """
+    rate = resolve_audio_sample_rate_or_none(source)
+    return rate if rate is not None else default
 
 
 def stream_pcm_format_from_env(

vllm_omni/metrics/modality.py

@@ -15,9 +15,9 @@ Contents:
 - ``observe_audio_streaming_finalize``: emits ``audio_underrun_s`` +
   ``audio_continuity_ok_total`` at SSE close using accumulated per-chunk
   arrival timestamps.
-- ``_extract_mm_output`` / ``_count_audio_frames``: shape-tolerant helpers
-  for the heterogeneous multimodal_output payloads emitted by different
-  audio pipelines.
+- ``extract_mm_output`` / ``count_audio_frames`` (from ``metrics.utils``):
+  shape-tolerant helpers for the heterogeneous multimodal_output payloads
+  emitted by different audio pipelines.
 """
 
 from __future__ import annotations
@@ -27,6 +27,7 @@ from typing import Any
 from prometheus_client import Counter, Histogram
 
 from vllm_omni.metrics import definitions as defs
+from vllm_omni.metrics.utils import count_audio_frames, extract_mm_output
 
 _stage_labels = list(defs.STAGE_LABELS)
 
@@ -139,53 +140,6 @@ class OmniModalityMetrics:
         ).inc()
 
 
-def _extract_mm_output(engine_outputs: Any) -> dict[str, Any]:
-    """Return the multimodal_output dict regardless of where it's nested.
-
-    Three shapes seen in the wild:
-      * ``engine_outputs.multimodal_output`` — synthesized on OmniRequestOutput
-        for some pipelines (often empty for AR audio)
-      * ``engine_outputs.outputs[0].multimodal_output`` — vllm CompletionOutput
-        nesting (where actual qwen3-omni audio data lives)
-      * neither — returns ``{}``
-    """
-    mm = getattr(engine_outputs, "multimodal_output", None)
-    if isinstance(mm, dict) and mm:
-        return mm
-    outs = getattr(engine_outputs, "outputs", None)
-    if outs:
-        nested = getattr(outs[0], "multimodal_output", None)
-        if isinstance(nested, dict):
-            return nested
-    return {}
-
-
-def _count_audio_frames(mm_out: dict[str, Any]) -> int:
-    """Sum the per-tensor sample count of audio chunks in mm_out["audio"].
-
-    Returns the total number of audio frames (samples) across all chunks.
-    For multi-dim tensors (e.g. shape [channels, samples]) the last axis is
-    treated as the sample dim; for 1-D tensors the only axis is the sample
-    dim; scalars count as 1.
-    """
-    audio_chunks = mm_out.get("audio") if isinstance(mm_out, dict) else None
-    if audio_chunks is None:
-        return 0
-    chunks = audio_chunks if isinstance(audio_chunks, list) else [audio_chunks]
-    n = 0
-    for t in chunks:
-        try:
-            ndim = getattr(t, "ndim", 0)
-            shape = getattr(t, "shape", None)
-            if ndim == 0 or shape is None or len(shape) == 0:
-                n += 1
-            else:
-                n += int(shape[-1])
-        except Exception:
-            continue
-    return n
-
-
 def observe_modality_at_finalize(
     mod_metrics: OmniModalityMetrics,
     *,
@@ -213,7 +167,7 @@ def observe_modality_at_finalize(
     stage_label = str(stage_id)
     replica_label = str(replica_id)
     gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0
-    mm_out = _extract_mm_output(engine_outputs)
+    mm_out = extract_mm_output(engine_outputs)
 
     sample_rate = defs.resolve_audio_sample_rate(mm_out)
     # `stage_metrics.audio_generated_frames` is the legacy per-chunk
@@ -222,7 +176,7 @@ def observe_modality_at_finalize(
     # field lookup in place in case the accumulator gets re-wired upstream.
     n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
     if n_frames == 0:
-        n_frames = _count_audio_frames(mm_out)
+        n_frames = count_audio_frames(mm_out)
     mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
     duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
     if duration_s > 0:

vllm_omni/metrics/utils.py

@@ -1,10 +1,217 @@
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
 from dataclasses import fields
 from typing import Any
 
 from prettytable import PrettyTable
 
 
+def coerce_positive_int_scalar(value: object) -> int | None:
+    """Coerce a value to a positive int without importing tensor libs.
+
+    Meant to pull positive integers such as sample rate out of the shapes
+    that show up across omni stage outputs / configs:
+
+    - plain ``int``: ``44100``
+    - ``torch.Tensor`` / numpy scalar: ``tensor(44100)`` (via ``.item()``)
+    - ``list`` / ``tuple`` wrap: ``[tensor(44100)]`` or ``[44100]``
+    - Mapping field values: ``{"sr": tensor(44100)}`` → coerce the field
+      value (key lookup is done by the caller, e.g.
+      :func:`resolve_int_by_sequential_keys`)
+
+    Returns ``None`` when the value is missing, unparsable, or not positive.
+    """
+    if value is None:
+        return None
+    if isinstance(value, (list, tuple)):
+        for item in value:
+            coerced = coerce_positive_int_scalar(item)
+            if coerced is not None:
+                return coerced
+        return None
+    item = getattr(value, "item", None)
+    if callable(item):
+        try:
+            value = item()
+        except Exception:
+            return None
+    try:
+        value = int(value)
+    except (TypeError, ValueError):
+        return None
+    return value if value > 0 else None
+
+
+def resolve_int_by_sequential_keys(
+    source: Mapping[str, object] | object | None,
+    keys: Sequence[str],
+) -> int | None:
+    """Return the first positive int found by trying ``keys`` in order.
+
+    For each key, looks up ``source[key]`` when ``source`` is a Mapping, otherwise
+    ``getattr(source, key, None)``. Values are coerced via :func:`coerce_positive_int_scalar`.
+    Returns ``None`` when ``source`` is empty/missing or no key yields a usable int.
+    """
+    if not source:
+        return None
+    for key in keys:
+        raw = source.get(key) if isinstance(source, Mapping) else getattr(source, key, None)
+        value = coerce_positive_int_scalar(raw)
+        if value is not None:
+            return value
+    return None
+
+
+def extract_mm_output(mm_source: object) -> Mapping[str, Any]:
+    """Return the first non-empty multimodal_output Mapping on ``mm_source``.
+
+    Lookup order:
+      1. ``mm_source.multimodal_output`` — top-level attribute / property
+      2. ``mm_source.outputs[0].multimodal_output`` — CompletionOutput nesting
+         (typical AR audio path)
+
+    Accepts both plain ``dict`` and ``MultimodalPayload`` (a ``Mapping``).
+    Returns ``{}`` when neither location yields a non-empty Mapping.
+
+    Args:
+        mm_source: Duck-typed container with optional ``multimodal_output`` and/or
+            ``outputs`` (e.g. vLLM ``RequestOutput`` or ``OmniRequestOutput``).
+
+    Returns:
+        The first non-empty multimodal Mapping found, or an empty ``dict``.
+    """
+    mm = getattr(mm_source, "multimodal_output", None)
+    if isinstance(mm, Mapping) and mm:
+        return mm
+    outs = getattr(mm_source, "outputs", None)
+    if outs:
+        nested = getattr(outs[0], "multimodal_output", None)
+        if isinstance(nested, Mapping) and nested:
+            return nested
+    return {}
+
+
+def iter_mm_outputs(mm_source: object) -> list[Mapping[str, Any]]:
+    """Collect all non-empty multimodal_output Mappings from ``mm_source``.
+
+    Lookup order:
+      1. ``mm_source.multimodal_output`` — top-level attribute / property
+      2. each ``mm_source.outputs[i].multimodal_output`` — every CompletionOutput
+         entry that carries a non-empty multimodal Mapping
+
+    Accepts both plain ``dict`` and ``MultimodalPayload`` (a ``Mapping``).
+    Used by stage-pool metrics aggregation that needs to visit every mm payload.
+
+    Args:
+        mm_source: Duck-typed container with optional ``multimodal_output`` and/or
+            ``outputs`` (e.g. vLLM ``RequestOutput`` or ``OmniRequestOutput``).
+
+    Returns:
+        A list of non-empty multimodal Mappings in discovery order. Empty when
+        none are present.
+    """
+    multimodal_outputs: list[Mapping[str, Any]] = []
+    outer_mm = getattr(mm_source, "multimodal_output", None)
+    if isinstance(outer_mm, Mapping) and outer_mm:
+        multimodal_outputs.append(outer_mm)
+    for output in getattr(mm_source, "outputs", None) or []:
+        inner_mm = getattr(output, "multimodal_output", None)
+        if isinstance(inner_mm, Mapping) and inner_mm:
+            multimodal_outputs.append(inner_mm)
+    return multimodal_outputs
+
+
+def count_audio_chunk_frames(audio_chunk: object) -> int:
+    """Count frames (samples) in one audio tensor / chunk.
+
+    Audio chunks are concatenated on dim=-1 in the output processor, so the
+    frame/sample axis is the last dim (e.g. ``[channels, frames]``). Keep this
+    aligned with serving_chat.py: audio tensors are consumed as ``(T,)``,
+    ``(C, T)``, or ``(B, C, T)``. Flattening would corrupt multi-channel audio.
+
+    Args:
+        audio_chunk: A single audio tensor/array-like, or a scalar-like value.
+
+    Returns:
+        Frame count for this chunk. Uses ``shape[-1]`` when shaped; otherwise
+        ``len(...)``; scalars / unlenable values count as ``1``.
+    """
+    shape = getattr(audio_chunk, "shape", None)
+    if shape is not None and len(shape) > 0:
+        return int(shape[-1])
+    try:
+        return len(audio_chunk)  # type: ignore[arg-type]
+    except TypeError:
+        return 1
+
+
+def count_audio_frames(mm_out: Mapping[str, Any]) -> int:
+    """Sum frame counts over all audio chunks in ``mm_out["audio"]`` or with other related keys.
+
+    For multi-dim tensors (e.g. shape ``[channels, samples]``) the last axis is
+    the sample dim; for 1-D tensors the only axis is the sample dim; scalars
+    count as 1. Missing or empty ``audio`` yields ``0``.
+
+    Args:
+        mm_out: A multimodal_output Mapping (plain ``dict`` or
+            ``MultimodalPayload``) that may contain an ``audio`` or related key whose value
+            is one chunk or a list of chunks.
+
+    Returns:
+        Total audio frames (samples) across all chunks.
+    """
+    audio_chunks = None
+    if isinstance(mm_out, Mapping):
+        for key in ("audio", "model_outputs"):
+            audio_chunks = mm_out.get(key)
+            if audio_chunks is not None:
+                break
+    if audio_chunks is None:
+        return 0
+    chunks = audio_chunks if isinstance(audio_chunks, list) else [audio_chunks]
+    return sum(count_audio_chunk_frames(chunk) for chunk in chunks)
+
+
+def count_image_pixels(value: object) -> int:
+    """Count pixels in one image value, or sum over a nested list/tuple.
+
+    Accepts PIL-like objects (``size=(W, H)``), tensors / arrays with a
+    ``shape`` attribute, and nested ``list`` / ``tuple`` containers.
+
+    Shape heuristics (aligned with StagePool image metrics):
+
+    - ``ndim >= 4`` (e.g. ``BCHW``): ``B * H * W`` via
+      ``dims[0] * dims[-2] * dims[-1]``
+    - ``ndim == 3`` and ``dims[0] in (1, 3, 4)``: CHW → ``H * W``
+    - ``ndim == 3`` and ``dims[-1] in (1, 3, 4)``: HWC → ``H * W``
+    - otherwise: ``dims[-2] * dims[-1]``
+
+    Returns ``0`` when ``value`` is missing or cannot be interpreted.
+    """
+    if value is None:
+        return 0
+    if isinstance(value, (list, tuple)):
+        return sum(count_image_pixels(item) for item in value)
+
+    size = getattr(value, "size", None)
+    if isinstance(size, tuple) and len(size) >= 2:
+        try:
+            return int(size[0]) * int(size[1])
+        except (TypeError, ValueError):
+            return 0
+
+    shape = getattr(value, "shape", None)
+    if shape is None or len(shape) < 2:
+        return 0
+    dims = [int(dim) for dim in shape]
+    if len(dims) >= 4:
+        return dims[0] * dims[-2] * dims[-1]
+    if len(dims) == 3 and dims[0] in (1, 3, 4):
+        return dims[1] * dims[2]
+    if len(dims) == 3 and dims[-1] in (1, 3, 4):
+        return dims[0] * dims[1]
+    return dims[-2] * dims[-1]
+
+
 def _build_field_defs(
     cls: type,
     exclude: set[str],