502 lines
17 KiB
Python
502 lines
17 KiB
Python
"""Shared benchmark infrastructure for Fish Speech serving benchmarks.
|
|
|
|
Provides common dataclasses, metrics computation, streaming HTTP client,
|
|
and result formatting used by model-specific benchmark scripts.
|
|
|
|
Model-specific scripts supply a ``create_payload_fn(prompt) -> dict``
|
|
callback and audio parameters; everything else is handled here.
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import time
|
|
from collections.abc import Callable
|
|
from dataclasses import asdict, dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import aiohttp
|
|
import numpy as np
|
|
from tqdm.asyncio import tqdm
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared test prompts (varying length for realistic workload)
|
|
# ---------------------------------------------------------------------------
|
|
PROMPTS = [
|
|
"Hello, welcome to the voice synthesis benchmark test.",
|
|
"She said she would be here by noon, but nobody showed up.",
|
|
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
|
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
|
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
|
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
|
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
|
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
|
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
|
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
|
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
|
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dataclasses
|
|
# ---------------------------------------------------------------------------
|
|
@dataclass
|
|
class RequestResult:
|
|
success: bool = False
|
|
ttfp: float = 0.0 # Time to first audio packet (seconds)
|
|
e2e: float = 0.0 # End-to-end latency (seconds)
|
|
audio_bytes: int = 0 # Total audio bytes received
|
|
audio_duration: float = 0.0 # Audio duration in seconds
|
|
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
|
|
prompt: str = ""
|
|
error: str = ""
|
|
|
|
|
|
@dataclass
|
|
class BenchmarkResult:
|
|
config_name: str = ""
|
|
concurrency: int = 0
|
|
num_prompts: int = 0
|
|
completed: int = 0
|
|
failed: int = 0
|
|
duration_s: float = 0.0
|
|
# TTFP stats (ms)
|
|
mean_ttfp_ms: float = 0.0
|
|
median_ttfp_ms: float = 0.0
|
|
std_ttfp_ms: float = 0.0
|
|
p90_ttfp_ms: float = 0.0
|
|
p95_ttfp_ms: float = 0.0
|
|
p99_ttfp_ms: float = 0.0
|
|
# E2E stats (ms)
|
|
mean_e2e_ms: float = 0.0
|
|
median_e2e_ms: float = 0.0
|
|
std_e2e_ms: float = 0.0
|
|
p90_e2e_ms: float = 0.0
|
|
p95_e2e_ms: float = 0.0
|
|
p99_e2e_ms: float = 0.0
|
|
# RTF stats
|
|
mean_rtf: float = 0.0
|
|
median_rtf: float = 0.0
|
|
std_rtf: float = 0.0
|
|
p99_rtf: float = 0.0
|
|
# Audio stats
|
|
mean_audio_duration_s: float = 0.0
|
|
total_audio_duration_s: float = 0.0
|
|
audio_throughput: float = 0.0 # audio_duration / wall_time
|
|
request_throughput: float = 0.0 # requests / second
|
|
# Per-request details
|
|
per_request: list = field(default_factory=list)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Audio helpers
|
|
# ---------------------------------------------------------------------------
|
|
def pcm_bytes_to_duration(
|
|
num_bytes: int,
|
|
sample_rate: int = 24000,
|
|
sample_width: int = 2,
|
|
) -> float:
|
|
"""Convert raw PCM byte count to duration in seconds."""
|
|
return num_bytes / sample_width / sample_rate
|
|
|
|
|
|
def _is_sse_response(response: aiohttp.ClientResponse) -> bool:
|
|
content_type = (response.headers.get("Content-Type") or "").lower()
|
|
return "text/event-stream" in content_type
|
|
|
|
|
|
async def _read_raw_audio_stream(
|
|
response: aiohttp.ClientResponse,
|
|
*,
|
|
start_time: float,
|
|
) -> tuple[int, float]:
|
|
first_audio_at = 0.0
|
|
total_bytes = 0
|
|
|
|
async for chunk in response.content.iter_any():
|
|
if chunk and first_audio_at <= 0:
|
|
first_audio_at = time.perf_counter() - start_time
|
|
total_bytes += len(chunk)
|
|
|
|
return total_bytes, first_audio_at
|
|
|
|
|
|
def _extract_sse_payload(raw_event: bytes) -> bytes | None:
|
|
data_lines: list[bytes] = []
|
|
for raw_line in raw_event.splitlines():
|
|
line = raw_line.rstrip(b"\r")
|
|
if line.startswith(b"data: "):
|
|
data_lines.append(line[6:])
|
|
elif line.startswith(b"data:"):
|
|
data_lines.append(line[5:].lstrip())
|
|
|
|
if not data_lines:
|
|
return None
|
|
return b"\n".join(data_lines).strip()
|
|
|
|
|
|
async def _read_sse_audio_stream(
|
|
response: aiohttp.ClientResponse,
|
|
*,
|
|
start_time: float,
|
|
) -> tuple[int, float]:
|
|
"""Decode SSE events and count raw audio bytes from base64 payloads."""
|
|
first_audio_at = 0.0
|
|
total_bytes = 0
|
|
pending = b""
|
|
|
|
async for chunk in response.content.iter_any():
|
|
if not chunk:
|
|
continue
|
|
pending += chunk
|
|
pending = pending.replace(b"\r\n", b"\n")
|
|
|
|
while b"\n\n" in pending:
|
|
raw_event, pending = pending.split(b"\n\n", 1)
|
|
payload_bytes = _extract_sse_payload(raw_event)
|
|
if payload_bytes is None:
|
|
continue
|
|
if payload_bytes == b"[DONE]":
|
|
return total_bytes, first_audio_at
|
|
|
|
try:
|
|
payload = json.loads(payload_bytes)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"Invalid SSE JSON payload: {exc}") from exc
|
|
|
|
audio = payload.get("audio")
|
|
if not isinstance(audio, dict):
|
|
continue
|
|
|
|
audio_b64 = audio.get("data")
|
|
if not audio_b64:
|
|
continue
|
|
|
|
try:
|
|
audio_bytes = base64.b64decode(audio_b64)
|
|
except Exception as exc:
|
|
raise ValueError(f"Invalid base64 audio chunk: {exc}") from exc
|
|
|
|
if audio_bytes and first_audio_at <= 0:
|
|
first_audio_at = time.perf_counter() - start_time
|
|
total_bytes += len(audio_bytes)
|
|
|
|
return total_bytes, first_audio_at
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Metrics
|
|
# ---------------------------------------------------------------------------
|
|
def compute_stats(
|
|
results: list[RequestResult],
|
|
wall_time: float,
|
|
) -> BenchmarkResult:
|
|
"""Compute aggregate statistics from per-request results."""
|
|
successful = [r for r in results if r.success]
|
|
failed = [r for r in results if not r.success]
|
|
|
|
bench = BenchmarkResult(
|
|
completed=len(successful),
|
|
failed=len(failed),
|
|
duration_s=wall_time,
|
|
)
|
|
|
|
if not successful:
|
|
return bench
|
|
|
|
ttfps = [r.ttfp * 1000 for r in successful]
|
|
e2es = [r.e2e * 1000 for r in successful]
|
|
rtfs = [r.rtf for r in successful]
|
|
audio_durs = [r.audio_duration for r in successful]
|
|
|
|
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
|
bench.median_ttfp_ms = float(np.median(ttfps))
|
|
bench.std_ttfp_ms = float(np.std(ttfps))
|
|
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
|
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
|
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
|
|
|
bench.mean_e2e_ms = float(np.mean(e2es))
|
|
bench.median_e2e_ms = float(np.median(e2es))
|
|
bench.std_e2e_ms = float(np.std(e2es))
|
|
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
|
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
|
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
|
|
|
bench.mean_rtf = float(np.mean(rtfs))
|
|
bench.median_rtf = float(np.median(rtfs))
|
|
bench.std_rtf = float(np.std(rtfs))
|
|
bench.p99_rtf = float(np.percentile(rtfs, 99))
|
|
|
|
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
|
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
|
bench.audio_throughput = bench.total_audio_duration_s / wall_time
|
|
bench.request_throughput = len(successful) / wall_time
|
|
|
|
bench.per_request = [
|
|
{
|
|
"ttfp_ms": r.ttfp * 1000,
|
|
"e2e_ms": r.e2e * 1000,
|
|
"rtf": r.rtf,
|
|
"audio_duration_s": r.audio_duration,
|
|
"prompt": r.prompt,
|
|
}
|
|
for r in successful
|
|
]
|
|
|
|
return bench
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Output formatting
|
|
# ---------------------------------------------------------------------------
|
|
def print_benchmark_results(bench: BenchmarkResult) -> None:
|
|
"""Print benchmark results in standardized format."""
|
|
W = 50
|
|
print("")
|
|
print(f"{'=' * W}")
|
|
print(f"{'Serving Benchmark Result':^{W}}")
|
|
print(f"{'=' * W}")
|
|
print(f"{'Successful requests:':<40}{bench.completed:<10}")
|
|
print(f"{'Failed requests:':<40}{bench.failed:<10}")
|
|
print(f"{'Maximum request concurrency:':<40}{bench.concurrency:<10}")
|
|
print(f"{'Benchmark duration (s):':<40}{bench.duration_s:<10.2f}")
|
|
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
|
|
print(f"{'-' * W}")
|
|
print(f"{'End-to-end Latency':^{W}}")
|
|
print(f"{'-' * W}")
|
|
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
|
|
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
|
|
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
|
|
print(f"{'=' * W}")
|
|
print(f"{'Audio Result':^{W}}")
|
|
print(f"{'=' * W}")
|
|
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
|
|
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
|
|
print(f"{'-' * W}")
|
|
print(f"{'Time to First Packet':^{W}}")
|
|
print(f"{'-' * W}")
|
|
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
|
|
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
|
|
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
|
|
print(f"{'-' * W}")
|
|
print(f"{'Real Time Factor':^{W}}")
|
|
print(f"{'-' * W}")
|
|
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
|
|
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
|
|
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
|
|
print(f"{'=' * W}")
|
|
print("")
|
|
|
|
|
|
def save_results(
|
|
all_results: list[dict],
|
|
result_dir: str,
|
|
config_name: str,
|
|
) -> Path:
|
|
"""Save benchmark results as JSON and return the file path."""
|
|
out = Path(result_dir)
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
result_file = out / f"bench_{config_name}_{timestamp}.json"
|
|
|
|
with open(result_file, "w") as f:
|
|
json.dump(all_results, f, indent=2)
|
|
print(f"Results saved to {result_file}")
|
|
return result_file
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Streaming HTTP client
|
|
# ---------------------------------------------------------------------------
|
|
async def send_streaming_request(
|
|
session: aiohttp.ClientSession,
|
|
api_url: str,
|
|
payload: dict,
|
|
sample_rate: int,
|
|
sample_width: int,
|
|
pbar: tqdm | None = None,
|
|
) -> RequestResult:
|
|
"""Send a streaming TTS request and measure latency metrics."""
|
|
result = RequestResult(prompt=payload.get("input", ""))
|
|
st = time.perf_counter()
|
|
|
|
try:
|
|
async with session.post(api_url, json=payload) as response:
|
|
if response.status != 200:
|
|
result.error = f"HTTP {response.status}: {await response.text()}"
|
|
else:
|
|
if _is_sse_response(response):
|
|
total_bytes, result.ttfp = await _read_sse_audio_stream(
|
|
response,
|
|
start_time=st,
|
|
)
|
|
else:
|
|
total_bytes, result.ttfp = await _read_raw_audio_stream(
|
|
response,
|
|
start_time=st,
|
|
)
|
|
|
|
result.e2e = time.perf_counter() - st
|
|
result.audio_bytes = total_bytes
|
|
result.audio_duration = pcm_bytes_to_duration(total_bytes, sample_rate, sample_width)
|
|
|
|
if total_bytes <= 0 or result.ttfp <= 0:
|
|
result.error = "HTTP 200 but no audio bytes were received"
|
|
else:
|
|
if result.audio_duration > 0:
|
|
result.rtf = result.e2e / result.audio_duration
|
|
result.success = True
|
|
|
|
except Exception as e:
|
|
result.error = str(e)
|
|
result.e2e = time.perf_counter() - st
|
|
|
|
finally:
|
|
if pbar:
|
|
pbar.update(1)
|
|
return result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Benchmark runner
|
|
# ---------------------------------------------------------------------------
|
|
async def run_benchmark(
|
|
host: str,
|
|
port: int,
|
|
num_prompts: int,
|
|
max_concurrency: int,
|
|
create_payload_fn: Callable[[str], dict],
|
|
sample_rate: int,
|
|
sample_width: int = 2,
|
|
num_warmups: int = 3,
|
|
request_timeout_s: float = 120.0,
|
|
) -> BenchmarkResult:
|
|
"""Run a TTS streaming benchmark at a given concurrency level.
|
|
|
|
Args:
|
|
create_payload_fn: Model-specific function that takes a prompt string
|
|
and returns the request JSON payload dict.
|
|
sample_rate: PCM sample rate for audio duration calculation.
|
|
sample_width: PCM sample width in bytes (default 2 for 16-bit).
|
|
"""
|
|
api_url = f"http://{host}:{port}/v1/audio/speech"
|
|
|
|
connector = aiohttp.TCPConnector(
|
|
limit=max_concurrency,
|
|
limit_per_host=max_concurrency,
|
|
keepalive_timeout=60,
|
|
)
|
|
session = aiohttp.ClientSession(
|
|
connector=connector,
|
|
timeout=aiohttp.ClientTimeout(
|
|
total=request_timeout_s,
|
|
connect=min(10.0, request_timeout_s),
|
|
sock_connect=min(10.0, request_timeout_s),
|
|
sock_read=request_timeout_s,
|
|
),
|
|
)
|
|
|
|
try:
|
|
# Warmup
|
|
if num_warmups > 0:
|
|
print(f" Warming up with {num_warmups} requests...")
|
|
warmup_tasks = [
|
|
send_streaming_request(
|
|
session,
|
|
api_url,
|
|
create_payload_fn(PROMPTS[i % len(PROMPTS)]),
|
|
sample_rate,
|
|
sample_width,
|
|
)
|
|
for i in range(num_warmups)
|
|
]
|
|
warmup_results = await asyncio.gather(*warmup_tasks)
|
|
warmup_ok = sum(1 for r in warmup_results if r.success)
|
|
if warmup_ok == 0:
|
|
print(" WARNING: All warmup requests failed!")
|
|
for r in warmup_results:
|
|
if r.error:
|
|
print(f" {r.error[:200]}")
|
|
print(f" Warmup done ({warmup_ok}/{num_warmups} succeeded).")
|
|
|
|
# Build request list
|
|
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
|
|
|
# Run
|
|
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
|
semaphore = asyncio.Semaphore(max_concurrency)
|
|
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
|
|
|
async def limited_request(prompt: str) -> RequestResult:
|
|
async with semaphore:
|
|
return await send_streaming_request(
|
|
session,
|
|
api_url,
|
|
create_payload_fn(prompt),
|
|
sample_rate,
|
|
sample_width,
|
|
pbar,
|
|
)
|
|
|
|
start_time = time.perf_counter()
|
|
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
|
results: list[RequestResult] = await asyncio.gather(*tasks)
|
|
wall_time = time.perf_counter() - start_time
|
|
pbar.close()
|
|
|
|
finally:
|
|
await session.close()
|
|
|
|
# Compute stats
|
|
bench = compute_stats(results, wall_time)
|
|
bench.concurrency = max_concurrency
|
|
bench.num_prompts = num_prompts
|
|
|
|
print_benchmark_results(bench)
|
|
|
|
# Print sample errors
|
|
failed = [r for r in results if not r.success]
|
|
if failed:
|
|
for r in failed[:3]:
|
|
print(f" [ERROR] {r.error[:200]}")
|
|
|
|
return bench
|
|
|
|
|
|
async def run_benchmark_sweep(
|
|
host: str,
|
|
port: int,
|
|
num_prompts: int,
|
|
concurrency_levels: list[int],
|
|
create_payload_fn: Callable[[str], dict],
|
|
sample_rate: int,
|
|
sample_width: int = 2,
|
|
num_warmups: int = 3,
|
|
request_timeout_s: float = 120.0,
|
|
config_name: str = "benchmark",
|
|
result_dir: str = "results",
|
|
) -> list[dict]:
|
|
"""Run benchmarks across multiple concurrency levels and save results."""
|
|
all_results = []
|
|
|
|
for concurrency in concurrency_levels:
|
|
result = await run_benchmark(
|
|
host=host,
|
|
port=port,
|
|
num_prompts=num_prompts,
|
|
max_concurrency=concurrency,
|
|
create_payload_fn=create_payload_fn,
|
|
sample_rate=sample_rate,
|
|
sample_width=sample_width,
|
|
num_warmups=num_warmups,
|
|
request_timeout_s=request_timeout_s,
|
|
)
|
|
result.config_name = config_name
|
|
all_results.append(asdict(result))
|
|
|
|
save_results(all_results, result_dir, config_name)
|
|
return all_results
|