diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 68d9c2c58..4217e2156 100644 --- a/tests/engine/test_orchestrator.py +++ b/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 diff --git a/tests/metrics/test_metrics_utils.py b/tests/metrics/test_metrics_utils.py new file mode 100644 index 000000000..544c7819a --- /dev/null +++ b/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 diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index e4f11b1e7..00af6ee23 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/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( diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 9372fee28..e10ce8845 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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 diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 4c4d3b7ee..32a35ac8c 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/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( diff --git a/vllm_omni/metrics/__init__.py b/vllm_omni/metrics/__init__.py index 6814a5891..7b6a2979f 100644 --- a/vllm_omni/metrics/__init__.py +++ b/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", ] diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 9660e4f34..23a1364a4 100644 --- a/vllm_omni/metrics/definitions.py +++ b/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( diff --git a/vllm_omni/metrics/modality.py b/vllm_omni/metrics/modality.py index 73d46ee8b..fd455ad36 100644 --- a/vllm_omni/metrics/modality.py +++ b/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: diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py index 2062eb5ba..ef8f69ed1 100644 --- a/vllm_omni/metrics/utils.py +++ b/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],