vllm-project--vllm-omni
wehub-resource-sync
325 行
11 KiB
Python
325 行
11 KiB
Python
"""Plot universal TTS benchmark results.
|
||
|
||
Reads JSON files saved by ``bench_tts.py`` (via ``vllm bench serve --omni``)
|
||
and generates comparison bar charts grouped by task type.
|
||
|
||
Metrics plotted:
|
||
- AUDIO_TTFP (mean audio time-to-first-packet, ms)
|
||
- E2EL (mean end-to-end latency, ms)
|
||
- Audio RTF (mean real-time factor)
|
||
- Audio throughput (audio-seconds / wall-second)
|
||
|
||
Quality metrics (WER / SIM / UTMOS) are printed in a table when present.
|
||
|
||
Usage::
|
||
|
||
# Single run — one JSON per task, all in results/
|
||
python benchmarks/tts/plot_results.py \\
|
||
--results results/bench_tts_*.json \\
|
||
--output results/tts_benchmark.png
|
||
|
||
# Compare two runs (e.g. async_chunk on vs off)
|
||
python benchmarks/tts/plot_results.py \\
|
||
--results run_a/bench_tts_*.json \\
|
||
--results run_b/bench_tts_*.json \\
|
||
--labels "async_chunk_on" "async_chunk_off" \\
|
||
--output results/comparison.png
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import math
|
||
from pathlib import Path
|
||
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# JSON loading
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def load_run(paths: list[str]) -> list[dict]:
|
||
"""Load and merge all JSON files for one run into a flat list of records.
|
||
|
||
Each record is expected to have at least ``_concurrency`` (int) and
|
||
``_task`` (str) keys injected by ``bench_tts.py``. Records that come
|
||
from a file that contains a list are flattened.
|
||
"""
|
||
records: list[dict] = []
|
||
for p in paths:
|
||
raw = json.loads(Path(p).read_text(encoding="utf-8"))
|
||
if isinstance(raw, list):
|
||
records.extend(raw)
|
||
elif isinstance(raw, dict):
|
||
records.append(raw)
|
||
return records
|
||
|
||
|
||
def _get(record: dict, key: str) -> float:
|
||
v = record.get(key, float("nan"))
|
||
if v is None or (isinstance(v, float) and math.isnan(v)):
|
||
return float("nan")
|
||
try:
|
||
return float(v)
|
||
except (TypeError, ValueError):
|
||
return float("nan")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Plotting helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _bar_group(
|
||
ax: plt.Axes,
|
||
x: np.ndarray,
|
||
data_per_label: dict[str, list[float]],
|
||
width: float,
|
||
colors: list[str],
|
||
ylabel: str,
|
||
title: str,
|
||
concurrency_labels: list[str],
|
||
fmt: str = ".1f",
|
||
) -> None:
|
||
n = len(data_per_label)
|
||
offsets = np.linspace(-(n - 1) * width / 2, (n - 1) * width / 2, n) if n > 1 else [0.0]
|
||
|
||
for i, (label, values) in enumerate(data_per_label.items()):
|
||
plot_vals = [0.0 if math.isnan(v) else v for v in values]
|
||
bar = ax.bar(x + offsets[i], plot_vals, width, label=label, color=colors[i % len(colors)], alpha=0.85)
|
||
max_val = max((v for v in values if not math.isnan(v)), default=1.0)
|
||
for rect, val in zip(bar, values):
|
||
if not math.isnan(val) and val > 0:
|
||
ax.text(
|
||
rect.get_x() + rect.get_width() / 2,
|
||
rect.get_height() + max_val * 0.02,
|
||
f"{val:{fmt}}",
|
||
ha="center",
|
||
va="bottom",
|
||
fontsize=8,
|
||
fontweight="bold",
|
||
)
|
||
|
||
ax.set_xlabel("Concurrency", fontsize=11)
|
||
ax.set_ylabel(ylabel, fontsize=11)
|
||
ax.set_title(title, fontsize=12, fontweight="bold")
|
||
ax.set_xticks(x)
|
||
ax.set_xticklabels(concurrency_labels)
|
||
ax.legend(fontsize=9)
|
||
ax.grid(axis="y", alpha=0.3)
|
||
ax.set_axisbelow(True)
|
||
|
||
|
||
COLORS = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107", "#9C27B0"]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Comparison plot (multiple labels / runs)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def plot_comparison(
|
||
all_runs: list[list[dict]],
|
||
labels: list[str],
|
||
output_path: str,
|
||
task_filter: str | None = None,
|
||
title_prefix: str = "TTS",
|
||
) -> None:
|
||
"""One 2×2 subplot per task found in the data."""
|
||
# Determine tasks to plot
|
||
tasks: list[str] = []
|
||
for run in all_runs:
|
||
for r in run:
|
||
t = r.get("_task", "unknown")
|
||
if t not in tasks:
|
||
tasks.append(t)
|
||
if task_filter:
|
||
tasks = [t for t in tasks if t == task_filter]
|
||
|
||
n_tasks = len(tasks)
|
||
if n_tasks == 0:
|
||
print("[plot_results] No tasks found in data.")
|
||
return
|
||
|
||
fig, axes_grid = plt.subplots(n_tasks, 4, figsize=(18, 4.5 * n_tasks))
|
||
fig.suptitle(f"{title_prefix} Benchmark", fontsize=15, fontweight="bold")
|
||
|
||
# Ensure axes_grid is always 2D
|
||
if n_tasks == 1:
|
||
axes_grid = [axes_grid]
|
||
|
||
for row_idx, task in enumerate(tasks):
|
||
# Collect concurrencies across all runs for this task
|
||
all_concs: set[int] = set()
|
||
for run in all_runs:
|
||
for r in run:
|
||
if r.get("_task") == task:
|
||
c = r.get("_concurrency")
|
||
if c is not None:
|
||
all_concs.add(int(c))
|
||
concurrencies = sorted(all_concs)
|
||
x = np.arange(len(concurrencies))
|
||
conc_labels = [str(c) for c in concurrencies]
|
||
|
||
def _series(run: list[dict], metric_key: str) -> list[float]:
|
||
conc_map = {int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r}
|
||
return [_get(conc_map.get(c, {}), metric_key) for c in concurrencies]
|
||
|
||
metrics = [
|
||
("mean_audio_ttfp_ms", "TTFP (ms)", "Time-to-First-Packet", ".0f"),
|
||
("mean_e2el_ms", "E2E Latency (ms)", "End-to-End Latency", ".0f"),
|
||
("mean_audio_rtf", "RTF", "Real-Time Factor (RTF)", ".3f"),
|
||
("audio_throughput", "audio-s / wall-s", "Audio Throughput", ".2f"),
|
||
]
|
||
|
||
axes_row = axes_grid[row_idx]
|
||
for col_idx, (key, ylabel, subtitle, fmt) in enumerate(metrics):
|
||
data_per_label = {lbl: _series(run, key) for lbl, run in zip(labels, all_runs)}
|
||
_bar_group(
|
||
axes_row[col_idx],
|
||
x,
|
||
data_per_label,
|
||
width=0.3 if len(labels) > 1 else 0.5,
|
||
colors=COLORS,
|
||
ylabel=ylabel,
|
||
title=f"{task} — {subtitle}",
|
||
concurrency_labels=conc_labels,
|
||
fmt=fmt,
|
||
)
|
||
|
||
plt.tight_layout()
|
||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||
print(f"Plot saved to {output_path}")
|
||
plt.close()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Markdown comparison table
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def print_comparison_table(all_runs: list[list[dict]], labels: list[str]) -> None:
|
||
tasks: list[str] = []
|
||
for run in all_runs:
|
||
for r in run:
|
||
t = r.get("_task", "unknown")
|
||
if t not in tasks:
|
||
tasks.append(t)
|
||
|
||
perf_metrics = [
|
||
("TTFP (ms)", "mean_audio_ttfp_ms", ".1f"),
|
||
("E2E (ms)", "mean_e2el_ms", ".1f"),
|
||
("RTF", "mean_audio_rtf", ".3f"),
|
||
("Throughput (a-s/s)", "audio_throughput", ".2f"),
|
||
]
|
||
quality_metrics = [
|
||
("WER (%)", "seed_tts_mean_wer", ".1f"),
|
||
("SIM", "seed_tts_mean_sim", ".3f"),
|
||
("UTMOS", "seed_tts_mean_utmos", ".2f"),
|
||
]
|
||
|
||
for task in tasks:
|
||
all_concs: set[int] = set()
|
||
for run in all_runs:
|
||
for r in run:
|
||
if r.get("_task") == task:
|
||
c = r.get("_concurrency")
|
||
if c is not None:
|
||
all_concs.add(int(c))
|
||
concurrencies = sorted(all_concs)
|
||
|
||
print(f"\n## {task}\n")
|
||
col_header = "| Metric | Concurrency |" + "".join(f" {lbl} |" for lbl in labels)
|
||
sep = "| --- | --- |" + " --- |" * len(labels)
|
||
print(col_header)
|
||
print(sep)
|
||
|
||
for metric, key, fmt in perf_metrics + quality_metrics:
|
||
for c in concurrencies:
|
||
row = f"| {metric} | {c} |"
|
||
for run in all_runs:
|
||
conc_map = {
|
||
int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r
|
||
}
|
||
val = _get(conc_map.get(c, {}), key)
|
||
row += f" {val:{fmt}} |" if not math.isnan(val) else " n/a |"
|
||
print(row)
|
||
|
||
# Improvement column (2-run comparison only)
|
||
if len(all_runs) == 2:
|
||
print(f"\n### Improvement ({labels[0]} vs {labels[1]})\n")
|
||
print("| Metric | Concurrency | Change |")
|
||
print("| --- | --- | --- |")
|
||
for metric, key, _ in perf_metrics:
|
||
for c in concurrencies:
|
||
conc_map0 = {
|
||
int(r["_concurrency"]): r for r in all_runs[0] if r.get("_task") == task and "_concurrency" in r
|
||
}
|
||
conc_map1 = {
|
||
int(r["_concurrency"]): r for r in all_runs[1] if r.get("_task") == task and "_concurrency" in r
|
||
}
|
||
v0 = _get(conc_map0.get(c, {}), key)
|
||
v1 = _get(conc_map1.get(c, {}), key)
|
||
if not math.isnan(v0) and not math.isnan(v1) and v1 > 0:
|
||
pct = (v1 - v0) / v1 * 100
|
||
print(f"| {metric} | {c} | {pct:+.1f}% |")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description=__doc__,
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
)
|
||
parser.add_argument(
|
||
"--results",
|
||
type=str,
|
||
nargs="+",
|
||
action="append",
|
||
required=True,
|
||
metavar="FILE",
|
||
help="JSON result file(s) for one run. Repeat --results for multiple runs to compare.",
|
||
)
|
||
parser.add_argument(
|
||
"--labels",
|
||
type=str,
|
||
nargs="+",
|
||
default=None,
|
||
help="Label for each --results group (must match the number of --results groups).",
|
||
)
|
||
parser.add_argument("--output", type=str, default="results/tts_benchmark.png", help="Output image path.")
|
||
parser.add_argument("--title", type=str, default="TTS", help="Title prefix for the plot.")
|
||
parser.add_argument("--task", type=str, default=None, help="Filter to a single task (e.g. voice_clone).")
|
||
return parser.parse_args()
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
|
||
# args.results is a list-of-lists due to action="append"
|
||
all_runs: list[list[dict]] = [load_run(group) for group in args.results]
|
||
n_runs = len(all_runs)
|
||
|
||
labels: list[str]
|
||
if args.labels:
|
||
if len(args.labels) != n_runs:
|
||
raise SystemExit(f"--labels count ({len(args.labels)}) must match --results groups ({n_runs})")
|
||
labels = args.labels
|
||
else:
|
||
labels = [f"run{i + 1}" for i in range(n_runs)]
|
||
|
||
print_comparison_table(all_runs, labels)
|
||
plot_comparison(all_runs, labels, args.output, task_filter=args.task, title_prefix=args.title)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|