518 lines
20 KiB
Python
518 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""CosyVoice3 streaming e2e benchmark: concurrency x utterance-length matrix.
|
|
|
|
Launches the server once per (concurrency, step-size) combo, streams TTS
|
|
requests cycling through several utterance lengths, and reports
|
|
TTFA/E2E/RTF/per-chunk latency (mean, stdev, range) per cell. "Step size" is
|
|
``codec_chunk_frames``, the token-hop length the first emitted chunk waits
|
|
for; later chunks wait for a doubling hop capped at ``4 * codec_chunk_frames``
|
|
(see ``_token_hop_lengths``). Per-chunk latency is reconstructed from that
|
|
same ramp rather than from raw transport read boundaries, which don't align
|
|
to model chunks (see ``one_request``'s docstring).
|
|
|
|
Running with no flags reproduces the default matrix. Per-request correctness
|
|
is checked: non-silent, long enough, and (with a fixed ``--seed``)
|
|
byte-identical PCM across requests for the same text.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import signal
|
|
import statistics
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import numpy as np
|
|
|
|
REPO = Path(__file__).resolve().parents[2]
|
|
DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
|
|
|
|
DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
|
|
DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。"
|
|
|
|
DEFAULT_TEXTS = [
|
|
"Thank you so much for the birthday gift you sent me.",
|
|
"The weather is lovely today, with warm sunshine and a gentle breeze. "
|
|
"We decided to take a walk in the park to enjoy the spring flowers and the fresh air.",
|
|
"The weather is lovely today, with warm sunshine and a gentle breeze. We decided to take a walk in the "
|
|
"park to enjoy the spring flowers and the fresh air. Many people were out walking, children were "
|
|
"running and playing on the grass, and older folks sat on benches chatting with one another. We "
|
|
"strolled slowly along the path and came across a cherry tree in full bloom, its pink petals drifting "
|
|
"on the wind. It was so beautiful that we stopped to take some photos to remember the moment. Then we "
|
|
"continued on and reached a small lake, its water so clear you could see straight to the bottom, with "
|
|
"the blue sky and white clouds reflected on its surface. A few ducks paddled back and forth across the "
|
|
"water, occasionally letting out a soft quack. We sat on a bench by the lake for a while, enjoying the "
|
|
"peaceful afternoon. Time passed quickly, and before we knew it evening had arrived, with the setting "
|
|
"sun painting the sky in shades of gold. Reluctant to leave, we finally headed home, already planning "
|
|
"our next visit. On the way back we talked about everything we had seen and heard that day, feeling "
|
|
"wonderfully happy. It was truly an unforgettable weekend.",
|
|
]
|
|
|
|
# PCM16 mono @ 24 kHz (CosyVoice3 output rate)
|
|
SAMPLE_RATE = 24000
|
|
|
|
# Well below any real speech's RMS; catches all-zero/near-silent PCM that byte
|
|
# count and duration checks alone pass through.
|
|
MIN_RMS = 50.0
|
|
|
|
# Overlay: only codec_chunk_frames is overridden; must live next to cosyvoice3.yaml (relative base_config)
|
|
DEPLOY_TEMPLATE = """# Auto-generated by bench_cosyvoice3_concurrency_matrix.py -- do not edit.
|
|
base_config: cosyvoice3.yaml
|
|
|
|
connectors:
|
|
connector_of_shared_memory:
|
|
name: SharedMemoryConnector
|
|
extra:
|
|
codec_streaming: true
|
|
connector_get_sleep_s: 0.01
|
|
connector_get_max_wait_first_chunk: 3000
|
|
connector_get_max_wait: 300
|
|
codec_chunk_frames: {chunk_frames}
|
|
codec_max_chunk_frames: {max_chunk_frames}
|
|
codec_stream_scale_factor: {stream_scale_factor}
|
|
codec_left_context_frames: 25
|
|
codec_vocab_size: 6561
|
|
"""
|
|
|
|
# CosyVoice3Config constants (transformers_utils/configs/cosyvoice3.py): tokens
|
|
# are consumed at token_mel_ratio mel frames each, and HiFT emits
|
|
# samples_per_mel PCM16 samples per mel frame. Pinned here (not left to the
|
|
# server's defaults) since one_request's byte-boundary reconstruction below
|
|
# depends on knowing them exactly.
|
|
TOKEN_MEL_RATIO = 2
|
|
SAMPLES_PER_MEL = 480
|
|
STREAM_SCALE_FACTOR = 2
|
|
|
|
|
|
def _stats(values: list[float]) -> dict[str, float]:
|
|
"""Mean/median/stdev/p95/min/max/n for a list of per-request measurements."""
|
|
n = len(values)
|
|
if n == 0:
|
|
return {
|
|
"n": 0,
|
|
"mean": float("nan"),
|
|
"median": float("nan"),
|
|
"stdev": float("nan"),
|
|
"p95": float("nan"),
|
|
"min": float("nan"),
|
|
"max": float("nan"),
|
|
}
|
|
return {
|
|
"n": n,
|
|
"mean": statistics.mean(values),
|
|
"median": statistics.median(values),
|
|
"stdev": statistics.stdev(values) if n >= 2 else 0.0,
|
|
"p95": statistics.quantiles(values, n=20)[18] if n >= 2 else values[0],
|
|
"min": min(values),
|
|
"max": max(values),
|
|
}
|
|
|
|
|
|
def vllm_bin() -> str:
|
|
for c in (REPO / ".venv/bin/vllm", Path(sys.executable).parent / "vllm"):
|
|
if c.is_file():
|
|
return str(c)
|
|
return "vllm"
|
|
|
|
|
|
def wait_ready(url: str, proc: subprocess.Popen, log: Path, timeout: float) -> None:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if proc.poll() is not None:
|
|
tail = log.read_text(errors="replace").splitlines()[-40:] if log.exists() else []
|
|
raise RuntimeError("server exited during startup:\n " + "\n ".join(tail))
|
|
try:
|
|
if httpx.get(f"{url}/health", timeout=2.0).status_code == 200:
|
|
return
|
|
except Exception:
|
|
pass
|
|
time.sleep(2.0)
|
|
raise TimeoutError(f"server not ready within {timeout}s (log: {log})")
|
|
|
|
|
|
def _token_hop_lengths(chunk_frames: int, max_chunk_frames: int):
|
|
"""The token-hop-length ramp talker2code2wav_async_chunk uses to decide
|
|
when to fire the next flow-solve/emission: chunk_frames, then doubling
|
|
(STREAM_SCALE_FACTOR) each step, capped at max_chunk_frames. This -- not
|
|
codec_chunk_frames alone -- is what actually paces emitted audio chunks:
|
|
with the defaults here, hops are chunk_frames, 2x, 4x, 4x, 4x, ...
|
|
|
|
The very first emission is consistently one token short of chunk_frames
|
|
(confirmed against real speech.audio.delta event sizes) even though the
|
|
ramp's own progression is computed from the nominal chunk_frames, not the
|
|
actual first hop -- so only the first yield gets the -1 adjustment.
|
|
"""
|
|
yield chunk_frames - 1
|
|
hop = chunk_frames
|
|
while True:
|
|
hop = min(max_chunk_frames, max(chunk_frames, hop * STREAM_SCALE_FACTOR))
|
|
yield hop
|
|
|
|
|
|
def one_request(client, url, model, text, ref_audio, ref_text, chunk_frames, seed=None):
|
|
"""Run one streaming request, returning TTFA, E2E, per-chunk gaps, and audio bytes.
|
|
|
|
``response_format=pcm`` + ``stream_format=audio`` streams raw bytes with no
|
|
framing, so a client read boundary alone is a transport artifact, not a
|
|
model chunk boundary. Reconstructed instead by replicating the server's
|
|
own token-hop-length ramp (see ``_token_hop_lengths``) to compute each
|
|
real chunk's expected byte size ahead of time, and slicing the
|
|
accumulated byte stream at those offsets -- independent of how the
|
|
transport happens to fragment the underlying reads. The final chunk
|
|
(finalize, which releases extra held-back frames) is whatever remains
|
|
when the stream closes, not a predicted hop.
|
|
"""
|
|
payload = {
|
|
"model": model,
|
|
"input": text,
|
|
"response_format": "pcm",
|
|
"stream": True,
|
|
"stream_format": "audio",
|
|
"ref_audio": ref_audio,
|
|
"ref_text": ref_text,
|
|
}
|
|
if seed is not None:
|
|
payload["seed"] = seed
|
|
t0 = time.perf_counter()
|
|
ttfa = None
|
|
nbytes = 0
|
|
audio_bytes = bytearray()
|
|
chunk_gaps_ms = [] # inter-chunk arrival gaps (per model chunk), by step index
|
|
prev_t = None
|
|
step_idx = 0
|
|
|
|
bytes_per_token = TOKEN_MEL_RATIO * SAMPLES_PER_MEL * 2
|
|
hop_gen = _token_hop_lengths(chunk_frames, 4 * chunk_frames)
|
|
next_boundary = next(hop_gen) * bytes_per_token
|
|
last_crossed_bytes = 0
|
|
received = 0
|
|
last_now = t0
|
|
|
|
with client.stream("POST", f"{url}/v1/audio/speech", json=payload, headers={"Authorization": "Bearer EMPTY"}) as r:
|
|
if r.status_code != 200:
|
|
raise RuntimeError(f"{r.status_code}: {r.read()[:300]!r}")
|
|
for chunk in r.iter_bytes():
|
|
if not chunk:
|
|
continue
|
|
now = time.perf_counter()
|
|
if ttfa is None:
|
|
ttfa = now - t0
|
|
last_now = now
|
|
received += len(chunk)
|
|
nbytes += len(chunk)
|
|
audio_bytes += chunk
|
|
while received >= next_boundary:
|
|
if prev_t is not None:
|
|
chunk_gaps_ms.append((step_idx, (now - prev_t) * 1e3))
|
|
step_idx += 1
|
|
prev_t = now
|
|
last_crossed_bytes = next_boundary
|
|
next_boundary += next(hop_gen) * bytes_per_token
|
|
|
|
if received > last_crossed_bytes:
|
|
# Finalize chunk: releases extra held-back frames, so it doesn't
|
|
# match a predicted hop size -- its arrival is the stream closing.
|
|
if prev_t is not None:
|
|
chunk_gaps_ms.append((step_idx, (last_now - prev_t) * 1e3))
|
|
step_idx += 1
|
|
|
|
e2e = time.perf_counter() - t0
|
|
audio_s = nbytes / 2 / SAMPLE_RATE
|
|
return {
|
|
"ttfa_s": ttfa,
|
|
"e2e_s": e2e,
|
|
"audio_s": audio_s,
|
|
"rtf": e2e / audio_s if audio_s > 0 else float("nan"),
|
|
"chunk_gaps_ms": chunk_gaps_ms,
|
|
"nbytes": nbytes,
|
|
"audio_bytes": audio_bytes,
|
|
}
|
|
|
|
|
|
def check_correctness(r: dict, min_audio_s: float) -> list[str]:
|
|
"""Return a list of correctness violations (empty = pass)."""
|
|
problems = []
|
|
if r["nbytes"] <= 0:
|
|
problems.append("empty audio (0 bytes)")
|
|
if r["audio_s"] < min_audio_s:
|
|
problems.append(f"audio too short ({r['audio_s']:.2f}s < {min_audio_s:.2f}s)")
|
|
if r["nbytes"] > 0:
|
|
samples = np.frombuffer(bytes(r["audio_bytes"]), dtype=np.int16).astype(np.float64)
|
|
rms = float(np.sqrt(np.mean(samples**2)))
|
|
if rms < MIN_RMS:
|
|
problems.append(f"near-silent audio (rms={rms:.1f} < {MIN_RMS})")
|
|
return problems
|
|
|
|
|
|
def drive(url, args, n, conc, label, results, lock, idx, chunk_frames):
|
|
"""Run ``n`` requests at concurrency ``conc``, cycling through ``args.texts``."""
|
|
texts = args.texts
|
|
|
|
def worker():
|
|
with httpx.Client(timeout=600.0) as c:
|
|
while True:
|
|
with lock:
|
|
i = idx[0]
|
|
idx[0] += 1
|
|
if i >= n:
|
|
return
|
|
text = texts[i % len(texts)]
|
|
try:
|
|
r = one_request(
|
|
c, url, args.model, text, args.ref_audio, args.ref_text, chunk_frames, seed=args.seed
|
|
)
|
|
r["text"] = text
|
|
with lock:
|
|
results.append(r)
|
|
except Exception as e:
|
|
print(f" [{label}] request {i} failed: {e}", flush=True)
|
|
|
|
ts = [threading.Thread(target=worker) for _ in range(conc)]
|
|
t0 = time.perf_counter()
|
|
for t in ts:
|
|
t.start()
|
|
for t in ts:
|
|
t.join()
|
|
return time.perf_counter() - t0
|
|
|
|
|
|
def run_combo(args, conc: int, chunk_frames: int, out: Path) -> dict:
|
|
tag = f"c{conc}_chunk{chunk_frames}"
|
|
run_dir = out / tag
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
log = run_dir / "server.log"
|
|
|
|
deploy = REPO / "vllm_omni/deploy" / f"_autogen_steps_{tag}.yaml"
|
|
deploy.write_text(
|
|
DEPLOY_TEMPLATE.format(
|
|
chunk_frames=chunk_frames,
|
|
max_chunk_frames=4 * chunk_frames,
|
|
stream_scale_factor=STREAM_SCALE_FACTOR,
|
|
)
|
|
)
|
|
|
|
cmd = [
|
|
vllm_bin(),
|
|
"serve",
|
|
args.model,
|
|
"--omni",
|
|
"--deploy-config",
|
|
str(deploy),
|
|
"--host",
|
|
"127.0.0.1",
|
|
"--port",
|
|
str(args.port),
|
|
"--allowed-local-media-path",
|
|
str(REPO),
|
|
]
|
|
env = dict(os.environ)
|
|
if args.cuda_visible_devices:
|
|
env["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
|
|
if args.trt is not None:
|
|
env["COSYVOICE3_TRT"] = args.trt
|
|
|
|
print(f"\n=== concurrency={conc} step_size={chunk_frames} ===", flush=True)
|
|
print(f" launching server (log: {log})", flush=True)
|
|
with open(log, "w") as lf:
|
|
proc = subprocess.Popen(cmd, env=env, stdout=lf, stderr=subprocess.STDOUT, start_new_session=True)
|
|
url = f"http://127.0.0.1:{args.port}"
|
|
try:
|
|
wait_ready(url, proc, log, args.server_timeout)
|
|
print(" server ready", flush=True)
|
|
|
|
results: list[dict] = []
|
|
lock = threading.Lock()
|
|
idx = [0]
|
|
if args.warmup:
|
|
print(f" warmup ({args.warmup} req)", flush=True)
|
|
drive(url, args, args.warmup, min(conc, args.warmup), "warmup", results, lock, idx, chunk_frames)
|
|
results.clear()
|
|
idx[0] = 0
|
|
|
|
print(f" running {args.num_prompts} req @ concurrency {conc}", flush=True)
|
|
wall = drive(url, args, args.num_prompts, conc, "bench", results, lock, idx, chunk_frames)
|
|
|
|
ok = [r for r in results if r.get("ttfa_s") is not None]
|
|
problems = []
|
|
for r in ok:
|
|
problems += check_correctness(r, args.min_audio_s)
|
|
# With a fixed seed, same-text requests must produce byte-identical PCM (per text group)
|
|
if args.seed is not None:
|
|
by_text: dict[str, bytes] = {}
|
|
for r in ok:
|
|
t = r.get("text")
|
|
ref = by_text.get(t)
|
|
if ref is None:
|
|
by_text[t] = r.get("audio_bytes")
|
|
elif r.get("audio_bytes") != ref:
|
|
problems.append(f"byte-identity violation for text {t!r}: seeded requests produced different PCM")
|
|
n_ok = len(ok)
|
|
total_audio_s = sum(r["audio_s"] for r in ok)
|
|
all_gaps = [g for r in ok for g in r["chunk_gaps_ms"]] # (step_index, gap_ms) per request
|
|
gap_vals = [g for _, g in all_gaps]
|
|
|
|
# bucket gaps by chunk position within a request, to see any O(n) slope
|
|
by_idx: dict[int, list[float]] = {}
|
|
for step_index, g in all_gaps:
|
|
by_idx.setdefault(step_index, []).append(g)
|
|
step_latency_by_idx = {i: statistics.median(vals) for i, vals in sorted(by_idx.items())}
|
|
|
|
agg = {
|
|
"concurrency": conc,
|
|
"step_size": chunk_frames,
|
|
"ok": n_ok,
|
|
"errors": args.num_prompts - n_ok,
|
|
"correctness_problems": len(problems),
|
|
"correctness_sample": problems[:5],
|
|
"total_audio_s": total_audio_s,
|
|
"throughput_audio_s_per_s": total_audio_s / wall if wall > 0 else float("nan"),
|
|
"ttfa_s": _stats([r["ttfa_s"] for r in ok]),
|
|
"e2e_s": _stats([r["e2e_s"] for r in ok]),
|
|
"rtf": _stats([r["rtf"] for r in ok]),
|
|
"step_latency_ms": _stats(gap_vals),
|
|
"step_latency_ms_by_idx": step_latency_by_idx,
|
|
"num_steps": len(all_gaps),
|
|
# per-utterance-length breakdown; n per text = num_prompts / len(texts)
|
|
"by_text": {
|
|
t: {
|
|
"audio_s": _stats([r["audio_s"] for r in ok if r.get("text") == t]),
|
|
"e2e_s": _stats([r["e2e_s"] for r in ok if r.get("text") == t]),
|
|
"rtf": _stats([r["rtf"] for r in ok if r.get("text") == t]),
|
|
}
|
|
for t in sorted({r.get("text") for r in ok})
|
|
},
|
|
}
|
|
print(json.dumps(agg, indent=2, ensure_ascii=False), flush=True)
|
|
return agg
|
|
finally:
|
|
print(" shutting down server", flush=True)
|
|
try:
|
|
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
|
|
except Exception:
|
|
proc.terminate()
|
|
try:
|
|
proc.wait(timeout=60)
|
|
except subprocess.TimeoutExpired:
|
|
try:
|
|
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
|
except Exception:
|
|
proc.kill()
|
|
proc.wait(timeout=20)
|
|
deploy.unlink(missing_ok=True)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
p.add_argument("--model", default=DEFAULT_MODEL)
|
|
p.add_argument(
|
|
"--text",
|
|
dest="texts",
|
|
nargs="+",
|
|
default=DEFAULT_TEXTS,
|
|
help="One or more utterance texts. With several, requests cycle through "
|
|
"them so a single server launch covers all utterance lengths; results "
|
|
"are bucketed per text.",
|
|
)
|
|
p.add_argument("--ref-audio", default=DEFAULT_REF_AUDIO)
|
|
p.add_argument("--ref-text", default=DEFAULT_REF_TEXT)
|
|
p.add_argument("--concurrency", type=int, nargs="+", default=[1, 8])
|
|
p.add_argument(
|
|
"--step-size",
|
|
type=int,
|
|
nargs="+",
|
|
default=[15],
|
|
help="codec_chunk_frames values to sweep (the streaming step size).",
|
|
)
|
|
p.add_argument("--num-prompts", type=int, default=30)
|
|
p.add_argument("--warmup", type=int, default=2)
|
|
p.add_argument("--port", type=int, default=8091)
|
|
p.add_argument(
|
|
"--server-timeout",
|
|
type=float,
|
|
default=1800.0,
|
|
help="Startup budget; first run also downloads weights and builds the TRT engine.",
|
|
)
|
|
p.add_argument("--output-dir", default="./perf_steps")
|
|
p.add_argument("--cuda-visible-devices", default="0")
|
|
p.add_argument(
|
|
"--trt",
|
|
choices=["0", "1"],
|
|
default=None,
|
|
help="Pin COSYVOICE3_TRT. 0=force torch estimator, 1=allow TensorRT. Default: unset.",
|
|
)
|
|
p.add_argument(
|
|
"--min-audio-s",
|
|
type=float,
|
|
default=1.0,
|
|
help="Minimum expected audio seconds per request for the correctness check.",
|
|
)
|
|
p.add_argument(
|
|
"--seed",
|
|
type=int,
|
|
default=42,
|
|
help="Fixed seed for reproducible (byte-identical) audio. Pass --seed '' "
|
|
"equivalent (any non-int) to disable -- there is no unset sentinel here "
|
|
"by design, since the default sweep is meant to be exactly reproducible.",
|
|
)
|
|
return p.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
out = Path(args.output_dir).resolve()
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
|
|
all_results = []
|
|
failed = False
|
|
for conc in args.concurrency:
|
|
for chunk_frames in args.step_size:
|
|
try:
|
|
result = run_combo(args, conc, chunk_frames, out)
|
|
all_results.append(result)
|
|
if result["errors"] > 0 or result["correctness_problems"] > 0:
|
|
failed = True
|
|
except Exception as e:
|
|
print(f" combo c={conc} chunk={chunk_frames} FAILED: {e}", flush=True)
|
|
failed = True
|
|
|
|
out_path = out / "steps_results.json"
|
|
with open(out_path, "w", encoding="utf-8") as f:
|
|
json.dump({"model": args.model, "texts": args.texts, "runs": all_results}, f, indent=2, ensure_ascii=False)
|
|
print(f"\nResults written to {out_path}", flush=True)
|
|
|
|
print("\n=== SUMMARY ===")
|
|
header = (
|
|
f"{'conc':>4} {'step':>4} {'ok':>4} {'corr':>4} "
|
|
f"{'TTFA med(ms)':>12} {'step med(ms)':>12} {'step p95(ms)':>12} {'RTF mean+-sd':>14}"
|
|
)
|
|
print(header)
|
|
print("-" * len(header))
|
|
for r in all_results:
|
|
print(
|
|
f"{r['concurrency']:>4} {r['step_size']:>4} {r['ok']:>4} {r['correctness_problems']:>4} "
|
|
f"{r['ttfa_s']['median'] * 1e3:>12.1f} "
|
|
f"{r['step_latency_ms']['median']:>12.1f} {r['step_latency_ms']['p95']:>12.1f} "
|
|
f"{r['rtf']['mean']:>6.3f}+-{r['rtf']['stdev']:>5.3f}"
|
|
)
|
|
for t, stats in r.get("by_text", {}).items():
|
|
print(
|
|
f" text {t[:24]!r}: e2e {stats['e2e_s']['mean'] * 1e3:8.1f}"
|
|
f"+-{stats['e2e_s']['stdev'] * 1e3:6.1f} ms (n={stats['e2e_s']['n']}, "
|
|
f"audio {stats['audio_s']['mean']:.1f}s)"
|
|
)
|
|
|
|
if failed:
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|