210 lines
8.2 KiB
Python
210 lines
8.2 KiB
Python
"""Lightweight benchmark for TTS/ASR queueing and API throughput.
|
|
|
|
This benchmark uses the FastAPI app with a mocked upstream speech API so it
|
|
measures this project's queueing, request handling, and SSE delivery cost
|
|
without requiring a real external model endpoint.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import os
|
|
import statistics
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
|
|
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
|
if str(BACKEND_DIR) not in __import__("sys").path:
|
|
__import__("sys").path.insert(0, str(BACKEND_DIR))
|
|
|
|
import main # noqa: E402
|
|
import tts_asr # noqa: E402
|
|
from job_system import reset_job_manager # noqa: E402
|
|
|
|
|
|
def _wav_bytes(duration_ms: int = 320) -> bytes:
|
|
sample_rate = 16000
|
|
frames = max(1, int(sample_rate * duration_ms / 1000))
|
|
data = b"".join((i % 32768).to_bytes(2, "little", signed=False) for i in range(frames))
|
|
data_size = len(data)
|
|
return (
|
|
b"RIFF" + (36 + data_size).to_bytes(4, "little")
|
|
+ b"WAVE"
|
|
+ b"fmt " + (16).to_bytes(4, "little")
|
|
+ (1).to_bytes(2, "little")
|
|
+ (1).to_bytes(2, "little")
|
|
+ sample_rate.to_bytes(4, "little")
|
|
+ sample_rate.to_bytes(4, "little")
|
|
+ (2).to_bytes(2, "little")
|
|
+ (16).to_bytes(2, "little")
|
|
+ b"data" + data_size.to_bytes(4, "little")
|
|
+ data
|
|
)
|
|
|
|
|
|
def _parse_sse_done(text: str) -> dict:
|
|
for chunk in reversed([item for item in text.split("\n\n") if item.strip()]):
|
|
event = ""
|
|
data = ""
|
|
for line in chunk.splitlines():
|
|
if line.startswith("event:"):
|
|
event = line.split(":", 1)[1].strip()
|
|
elif line.startswith("data:"):
|
|
data = line.split(":", 1)[1].strip()
|
|
if event == "done" and data:
|
|
payload = json.loads(data)
|
|
result = dict(payload.get("result") or {})
|
|
for key in ("queue_ms", "run_ms", "total_ms", "queued_count", "running_count", "busy_level", "busy_ratio"):
|
|
if key in payload:
|
|
result[key] = payload[key]
|
|
return result
|
|
raise RuntimeError("done event not found")
|
|
|
|
|
|
def _percentile(values: list[float], q: float) -> float:
|
|
if not values:
|
|
return 0.0
|
|
if len(values) == 1:
|
|
return values[0]
|
|
index = (len(values) - 1) * q
|
|
lower = int(index)
|
|
upper = min(lower + 1, len(values) - 1)
|
|
if lower == upper:
|
|
return values[lower]
|
|
weight = index - lower
|
|
return values[lower] * (1 - weight) + values[upper] * weight
|
|
|
|
|
|
async def _build_mock_client(tts_delay_ms: int, asr_delay_ms: int) -> httpx.AsyncClient:
|
|
async def transport(request: httpx.Request):
|
|
if request.url.path.endswith("/audio/speech"):
|
|
await asyncio.sleep(tts_delay_ms / 1000.0)
|
|
return httpx.Response(200, content=_wav_bytes(420), headers={"x-request-id": "bench-tts"}, request=request)
|
|
await asyncio.sleep(asr_delay_ms / 1000.0)
|
|
return httpx.Response(200, json={"text": "benchmark transcript", "language": "zh"}, headers={"x-request-id": "bench-asr"}, request=request)
|
|
|
|
return httpx.AsyncClient(
|
|
base_url="https://benchmark.example/v1/",
|
|
transport=httpx.MockTransport(transport),
|
|
)
|
|
|
|
|
|
async def _run_case(case_name: str, concurrency: int, request_count: int, audio_b64: str | None = None) -> dict:
|
|
results: list[dict] = []
|
|
latencies: list[float] = []
|
|
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.ASGITransport(app=main.app),
|
|
base_url="http://testserver",
|
|
timeout=120.0,
|
|
) as client:
|
|
semaphore = asyncio.Semaphore(concurrency)
|
|
|
|
async def fire(index: int) -> None:
|
|
async with semaphore:
|
|
started = time.perf_counter()
|
|
if case_name == "tts":
|
|
response = await client.post(
|
|
"/v1/tts-asr/tts",
|
|
json={"text": f"第 {index} 条基准文本", "speaker": "Vivian", "format": "wav"},
|
|
)
|
|
else:
|
|
response = await client.post(
|
|
"/v1/tts-asr/asr",
|
|
json={"audio_base64": audio_b64, "language": "zh-CN"},
|
|
)
|
|
response.raise_for_status()
|
|
payload = _parse_sse_done(response.text)
|
|
latencies.append((time.perf_counter() - started) * 1000.0)
|
|
results.append(payload)
|
|
|
|
await asyncio.gather(*(fire(index) for index in range(request_count)))
|
|
|
|
queue_values = sorted(float(item.get("queue_ms", 0) or 0) for item in results)
|
|
run_values = sorted(float(item.get("run_ms", 0) or 0) for item in results)
|
|
total_values = sorted(float(item.get("total_ms", 0) or 0) for item in results)
|
|
latency_values = sorted(latencies)
|
|
elapsed_sum_ms = sum(latency_values)
|
|
return {
|
|
"case": case_name,
|
|
"requests": request_count,
|
|
"concurrency": concurrency,
|
|
"avg_latency_ms": round(statistics.fmean(latency_values), 2),
|
|
"p95_latency_ms": round(_percentile(latency_values, 0.95), 2),
|
|
"avg_queue_ms": round(statistics.fmean(queue_values), 2),
|
|
"p95_queue_ms": round(_percentile(queue_values, 0.95), 2),
|
|
"avg_run_ms": round(statistics.fmean(run_values), 2),
|
|
"p95_run_ms": round(_percentile(run_values, 0.95), 2),
|
|
"avg_total_ms": round(statistics.fmean(total_values), 2),
|
|
"p95_total_ms": round(_percentile(total_values, 0.95), 2),
|
|
"throughput_rps_estimate": round((request_count * 1000.0) / max(latency_values[-1], elapsed_sum_ms / max(request_count, 1)), 2),
|
|
}
|
|
|
|
|
|
async def main_async(args) -> None:
|
|
os.environ["JOB_BACKEND"] = "memory"
|
|
os.environ["JOB_TTS_CONCURRENCY"] = str(args.tts_workers)
|
|
os.environ["JOB_TTS_MAX_QUEUE"] = str(max(args.tts_requests, args.tts_workers))
|
|
os.environ["JOB_ASR_CONCURRENCY"] = str(args.asr_workers)
|
|
os.environ["JOB_ASR_MAX_QUEUE"] = str(max(args.asr_requests, args.asr_workers))
|
|
reset_job_manager()
|
|
|
|
mock_client = await _build_mock_client(args.tts_delay_ms, args.asr_delay_ms)
|
|
tts_asr._httpx_client = mock_client
|
|
try:
|
|
audio_b64 = base64.b64encode(_wav_bytes(args.audio_duration_ms)).decode("utf-8")
|
|
tts_stats = await _run_case("tts", args.tts_concurrency, args.tts_requests)
|
|
asr_stats = await _run_case("asr", args.asr_concurrency, args.asr_requests, audio_b64=audio_b64)
|
|
finally:
|
|
await mock_client.aclose()
|
|
tts_asr._httpx_client = None
|
|
reset_job_manager()
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"benchmark_date": time.strftime("%Y-%m-%d %H:%M:%S"),
|
|
"assumptions": {
|
|
"upstream_tts_delay_ms": args.tts_delay_ms,
|
|
"upstream_asr_delay_ms": args.asr_delay_ms,
|
|
"job_backend": "memory",
|
|
},
|
|
"tts": tts_stats,
|
|
"asr": asr_stats,
|
|
"recommended_defaults": {
|
|
"JOB_TTS_CONCURRENCY": args.tts_workers,
|
|
"JOB_TTS_MAX_QUEUE": max(16, args.tts_workers * 4),
|
|
"JOB_ASR_CONCURRENCY": args.asr_workers,
|
|
"JOB_ASR_MAX_QUEUE": max(8, args.asr_workers * 4),
|
|
"TTS_ASR_MAX_CONNECTIONS": max(24, (args.tts_workers + args.asr_workers) * 4),
|
|
"TTS_ASR_MAX_KEEPALIVE_CONNECTIONS": max(12, (args.tts_workers + args.asr_workers) * 2),
|
|
},
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--tts-delay-ms", type=int, default=120)
|
|
parser.add_argument("--asr-delay-ms", type=int, default=280)
|
|
parser.add_argument("--tts-workers", type=int, default=4)
|
|
parser.add_argument("--asr-workers", type=int, default=2)
|
|
parser.add_argument("--tts-concurrency", type=int, default=8)
|
|
parser.add_argument("--asr-concurrency", type=int, default=4)
|
|
parser.add_argument("--tts-requests", type=int, default=32)
|
|
parser.add_argument("--asr-requests", type=int, default=16)
|
|
parser.add_argument("--audio-duration-ms", type=int, default=320)
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main_async(parse_args()))
|