265 lines
9.9 KiB
Python
265 lines
9.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Client-local helpers for duplex JSON metric blocks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Mapping, Sequence
|
|
from typing import NamedTuple
|
|
|
|
__all__ = [
|
|
"STREAMING_OUTPUT_UNIT_TYPES",
|
|
"StageModalityFlags",
|
|
"distribution_summary",
|
|
"metric_mean",
|
|
"stage_modality_flags",
|
|
"summarize_stage_metrics",
|
|
]
|
|
|
|
# Intermediate streaming units that are not the pipeline's text or audio
|
|
# final output (e.g. MiniCPM Talker codec tokens).
|
|
STREAMING_OUTPUT_UNIT_TYPES = frozenset({"text", "stream", "audio"})
|
|
|
|
|
|
class StageModalityFlags(NamedTuple):
|
|
"""How duplex JSON classifies one engine stage."""
|
|
|
|
is_text_stage: bool
|
|
is_audio_stage: bool
|
|
is_image_stage: bool
|
|
is_video_stage: bool
|
|
is_internal_stream_stage: bool
|
|
|
|
|
|
def stage_modality_flags(
|
|
final_output_type: object,
|
|
output_unit_type: object,
|
|
) -> StageModalityFlags:
|
|
"""Classify a stage from wire ``final_output_type`` / ``output_unit_type``."""
|
|
final_type = final_output_type if isinstance(final_output_type, str) else ""
|
|
unit_type = output_unit_type if isinstance(output_unit_type, str) else ""
|
|
is_text_stage = final_type == "text" or unit_type == "text"
|
|
is_audio_stage = final_type == "audio" or unit_type == "audio"
|
|
is_video_stage = final_type in {"video", "videos"} or unit_type == "video"
|
|
# Video diffusion may still report output_unit_type="image" when frames are
|
|
# stored in ``images``; prefer video when final_output_type says so.
|
|
is_image_stage = (not is_video_stage) and (final_type in {"image", "images"} or unit_type == "image")
|
|
is_internal_stream_stage = unit_type in STREAMING_OUTPUT_UNIT_TYPES and not is_text_stage and not is_audio_stage
|
|
return StageModalityFlags(
|
|
is_text_stage=is_text_stage,
|
|
is_audio_stage=is_audio_stage,
|
|
is_image_stage=is_image_stage,
|
|
is_video_stage=is_video_stage,
|
|
is_internal_stream_stage=is_internal_stream_stage,
|
|
)
|
|
|
|
|
|
def _rounded_ms(value: float) -> float:
|
|
return round(float(value), 3)
|
|
|
|
|
|
def _finite_number(value: object, *, nonnegative: bool = False) -> float | None:
|
|
if not isinstance(value, int | float) or isinstance(value, bool) or not math.isfinite(float(value)):
|
|
return None
|
|
number = float(value)
|
|
if nonnegative and number < 0:
|
|
return None
|
|
return number
|
|
|
|
|
|
def _interval_summary(values: list[float]) -> dict[str, float | int]:
|
|
clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0)
|
|
if not clean:
|
|
return {"count": 0, "mean": 0.0, "p50": 0.0, "p95": 0.0, "max": 0.0}
|
|
|
|
def nearest_rank(percentile: float) -> float:
|
|
index = max(0, math.ceil(percentile * len(clean)) - 1)
|
|
return clean[min(index, len(clean) - 1)]
|
|
|
|
return {
|
|
"count": len(clean),
|
|
"mean": _rounded_ms(sum(clean) / len(clean)),
|
|
"p50": nearest_rank(0.50),
|
|
"p95": nearest_rank(0.95),
|
|
"max": clean[-1],
|
|
}
|
|
|
|
|
|
def distribution_summary(values: Sequence[float], *, digits: int = 3) -> dict[str, float | int] | None:
|
|
"""Summarize values as ``{count, mean, p50, p99}`` for duplex report fields."""
|
|
clean = sorted(float(value) for value in values if math.isfinite(float(value)))
|
|
if not clean:
|
|
return None
|
|
|
|
def nearest_rank(percentile: float) -> float:
|
|
index = max(0, math.ceil(percentile * len(clean)) - 1)
|
|
return clean[min(index, len(clean) - 1)]
|
|
|
|
return {
|
|
"count": len(clean),
|
|
"mean": round(sum(clean) / len(clean), digits),
|
|
"p50": round(nearest_rank(0.50), digits),
|
|
"p99": round(nearest_rank(0.99), digits),
|
|
}
|
|
|
|
|
|
def metric_mean(value: object) -> float | None:
|
|
"""Read a scalar mean, or the ``mean`` field of a distribution summary."""
|
|
if isinstance(value, Mapping):
|
|
nested = value.get("mean")
|
|
if isinstance(nested, int | float) and not isinstance(nested, bool) and math.isfinite(float(nested)):
|
|
return float(nested)
|
|
return None
|
|
if isinstance(value, int | float) and not isinstance(value, bool) and math.isfinite(float(value)):
|
|
return float(value)
|
|
return None
|
|
|
|
|
|
def _stage_id_sort_key(stage_id: str) -> tuple[int, int | str]:
|
|
try:
|
|
return (0, int(stage_id))
|
|
except ValueError:
|
|
return (1, stage_id)
|
|
|
|
|
|
def _stage_type_labels(stage_metrics: Mapping[str, object]) -> dict[str, object]:
|
|
labels: dict[str, object] = {}
|
|
for key in ("final_output_type", "output_unit_type"):
|
|
value = stage_metrics.get(key)
|
|
if isinstance(value, str) and value:
|
|
labels[key] = value
|
|
return labels
|
|
|
|
|
|
def _stage_ms_list(stage_metrics: Mapping[str, object], field_name: str) -> list[float]:
|
|
raw = stage_metrics.get(field_name)
|
|
return [float(value) for value in raw if isinstance(value, int | float)] if isinstance(raw, list) else []
|
|
|
|
|
|
def _has_legacy_vllm_token_metrics(stage_metrics: Mapping[str, object]) -> bool:
|
|
if stage_metrics.get("vllm_ttft_ms") or stage_metrics.get("vllm_tpot_ms"):
|
|
return True
|
|
raw_itls = stage_metrics.get("vllm_itls_ms")
|
|
return isinstance(raw_itls, list) and bool(raw_itls)
|
|
|
|
|
|
def _text_stage_engine_metrics_block(stage_metrics: Mapping[str, object]) -> dict[str, object]:
|
|
itls = _stage_ms_list(stage_metrics, "vllm_itls_ms")
|
|
return {
|
|
"source": "engine_stage_metrics",
|
|
**_stage_type_labels(stage_metrics),
|
|
"output_token_count": int(stage_metrics.get("num_tokens_out") or 0),
|
|
"ttft_ms": float(stage_metrics.get("vllm_ttft_ms") or 0.0),
|
|
"tpot_ms": float(stage_metrics.get("vllm_tpot_ms") or 0.0),
|
|
"itls_ms": itls,
|
|
"inter_token_interval_ms": _interval_summary(itls),
|
|
}
|
|
|
|
|
|
def _audio_stage_engine_metrics_block(stage_metrics: Mapping[str, object]) -> dict[str, object]:
|
|
block: dict[str, object] = {
|
|
"source": "engine_stage_metrics",
|
|
**_stage_type_labels(stage_metrics),
|
|
"output_unit_count": int(stage_metrics.get("output_unit_count") or 0),
|
|
"audio_generated_frames": int(stage_metrics.get("audio_generated_frames") or 0),
|
|
"audio_duration_s": float(stage_metrics.get("audio_duration_s") or 0.0),
|
|
}
|
|
serving_time = _finite_number(
|
|
stage_metrics.get("serving_time_to_first_output_ms"),
|
|
nonnegative=True,
|
|
)
|
|
if serving_time:
|
|
block["ttfp_ms"] = serving_time
|
|
return block
|
|
|
|
|
|
def _stream_stage_engine_metrics_block(stage_metrics: Mapping[str, object]) -> dict[str, object]:
|
|
icls = _stage_ms_list(stage_metrics, "inter_output_latencies_ms")
|
|
return {
|
|
"source": "engine_stage_metrics",
|
|
**_stage_type_labels(stage_metrics),
|
|
"output_unit_count": int(stage_metrics.get("output_unit_count") or 0),
|
|
"ttfc_ms": float(stage_metrics.get("serving_time_to_first_output_ms") or 0.0),
|
|
"tpop_ms": float(stage_metrics.get("time_per_output_unit_ms") or 0.0),
|
|
"icls_ms": icls,
|
|
"inter_chunk_interval_ms": _interval_summary(icls),
|
|
}
|
|
|
|
|
|
def _stage_engine_metrics_block(stage_metrics: Mapping[str, object]) -> dict[str, object]:
|
|
flags = stage_modality_flags(
|
|
stage_metrics.get("final_output_type"),
|
|
stage_metrics.get("output_unit_type"),
|
|
)
|
|
if flags.is_text_stage or (
|
|
not flags.is_audio_stage
|
|
and not flags.is_internal_stream_stage
|
|
and not flags.is_image_stage
|
|
and not flags.is_video_stage
|
|
and _has_legacy_vllm_token_metrics(stage_metrics)
|
|
):
|
|
return _text_stage_engine_metrics_block(stage_metrics)
|
|
if flags.is_audio_stage:
|
|
return _audio_stage_engine_metrics_block(stage_metrics)
|
|
if flags.is_internal_stream_stage:
|
|
return _stream_stage_engine_metrics_block(stage_metrics)
|
|
block: dict[str, object] = {"source": "engine_stage_metrics", **_stage_type_labels(stage_metrics)}
|
|
gen_time_ms = stage_metrics.get("stage_gen_time_ms")
|
|
if isinstance(gen_time_ms, int | float) and not isinstance(gen_time_ms, bool):
|
|
block["stage_gen_time_ms"] = float(gen_time_ms)
|
|
return block
|
|
|
|
|
|
def _finite_metric_values(
|
|
rows: Sequence[Mapping[str, object]],
|
|
metric: str,
|
|
*,
|
|
positive: bool = False,
|
|
) -> list[float]:
|
|
return [
|
|
float(row[metric])
|
|
for row in rows
|
|
if isinstance(row.get(metric), int | float)
|
|
and not isinstance(row.get(metric), bool)
|
|
and math.isfinite(float(row[metric]))
|
|
and (not positive or float(row[metric]) > 0)
|
|
]
|
|
|
|
|
|
def summarize_stage_metrics(
|
|
request_metrics: Sequence[Mapping[str, object]],
|
|
) -> dict[str, dict[str, object]] | None:
|
|
"""Roll per-response engine stage blocks into ``{count, mean, p50, p99}``.
|
|
|
|
Reads ``stages`` when present, otherwise ``stage0_tokens`` as stage ``"0"``.
|
|
Zero or missing ``tpot_ms`` / ``tpop_ms`` values are omitted, matching
|
|
session ``tpot_ms``.
|
|
"""
|
|
buckets: dict[str, list[Mapping[str, object]]] = {}
|
|
for request in request_metrics:
|
|
stages = request.get("stages")
|
|
if not isinstance(stages, dict):
|
|
stage0 = request.get("stage0_tokens")
|
|
stages = {"0": stage0} if isinstance(stage0, dict) else {}
|
|
for stage_id, stage_snapshot in stages.items():
|
|
if isinstance(stage_snapshot, dict):
|
|
buckets.setdefault(str(stage_id), []).append(stage_snapshot)
|
|
summary: dict[str, dict[str, object]] = {}
|
|
for stage_id in sorted(buckets, key=_stage_id_sort_key):
|
|
rows = buckets[stage_id]
|
|
stage_summary: dict[str, object] = {}
|
|
for metric, positive in (
|
|
("ttft_ms", False),
|
|
("tpot_ms", True),
|
|
("ttfc_ms", False),
|
|
("tpop_ms", True),
|
|
("ttfp_ms", False),
|
|
):
|
|
rolled = distribution_summary(_finite_metric_values(rows, metric, positive=positive))
|
|
if rolled is not None:
|
|
stage_summary[metric] = rolled
|
|
if stage_summary:
|
|
summary[stage_id] = stage_summary
|
|
return summary or None
|