Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
418 lines
16 KiB
Python
418 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Measure streaming TTS with an auditable incremental-text client."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import csv
|
|
import json
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import wave
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
LOG_ROOT = REPO_ROOT / "logs" / "streaming_test"
|
|
DEFAULT_SERVER_BIN = REPO_ROOT / "build" / "debug" / "bin" / "audiocpp_server"
|
|
DEFAULT_SAMPLE_RATE = 48000
|
|
|
|
|
|
def timestamp_slug() -> str:
|
|
return datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
|
|
|
|
def load_json(path: Path) -> dict[str, Any]:
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
return json.load(handle)
|
|
|
|
|
|
def write_json(path: Path, value: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
|
|
|
|
|
|
def path_is_under(path: Path, root: Path) -> bool:
|
|
try:
|
|
path.resolve().relative_to(root.resolve())
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def split_incremental_text(text: str, max_chars: int) -> list[str]:
|
|
if max_chars <= 0:
|
|
raise RuntimeError("--max-chars must be positive")
|
|
sentences = [part.strip() for part in re.split(r"(?<=[.!?])\s+", text.strip()) if part.strip()]
|
|
chunks: list[str] = []
|
|
current = ""
|
|
for sentence in sentences:
|
|
if not current:
|
|
current = sentence
|
|
elif len(current) + 1 + len(sentence) <= max_chars:
|
|
current = current + " " + sentence
|
|
else:
|
|
chunks.append(current)
|
|
current = sentence
|
|
while len(current) > max_chars:
|
|
split_at = current.rfind(" ", 0, max_chars + 1)
|
|
if split_at <= 0:
|
|
split_at = max_chars
|
|
chunks.append(current[:split_at].strip())
|
|
current = current[split_at:].strip()
|
|
if current:
|
|
chunks.append(current)
|
|
if not chunks:
|
|
raise RuntimeError("input text produced no chunks")
|
|
return chunks
|
|
|
|
|
|
def gpu_memory_mib() -> int:
|
|
output = subprocess.check_output(
|
|
["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
|
|
text=True,
|
|
timeout=5,
|
|
)
|
|
values = [int(line.strip()) for line in output.splitlines() if line.strip()]
|
|
if not values:
|
|
raise RuntimeError("nvidia-smi returned no GPU memory values")
|
|
return max(values)
|
|
|
|
|
|
class VramSampler:
|
|
def __init__(self, path: Path, interval_s: float) -> None:
|
|
self.path = path
|
|
self.interval_s = interval_s
|
|
self.stop_event = threading.Event()
|
|
self.thread: threading.Thread | None = None
|
|
|
|
def __enter__(self) -> "VramSampler":
|
|
if self.interval_s <= 0.0:
|
|
raise RuntimeError("--vram-sample-ms must be positive")
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.thread = threading.Thread(target=self._run, daemon=True)
|
|
self.thread.start()
|
|
return self
|
|
|
|
def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
|
|
self.stop_event.set()
|
|
if self.thread is not None:
|
|
self.thread.join(timeout=5)
|
|
|
|
def _run(self) -> None:
|
|
start = time.perf_counter()
|
|
with self.path.open("w", newline="", encoding="utf-8") as handle:
|
|
writer = csv.writer(handle)
|
|
writer.writerow(["elapsed_ms", "vram_mib"])
|
|
while not self.stop_event.is_set():
|
|
writer.writerow([round((time.perf_counter() - start) * 1000.0, 3), gpu_memory_mib()])
|
|
handle.flush()
|
|
time.sleep(self.interval_s)
|
|
|
|
|
|
def read_vram_summary(path: Path) -> dict[str, Any]:
|
|
values: list[int] = []
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for row in csv.DictReader(handle):
|
|
values.append(int(row["vram_mib"]))
|
|
if not values:
|
|
raise RuntimeError(f"VRAM sampler wrote no samples: {path}")
|
|
return {
|
|
"vram_start_mib": values[0],
|
|
"vram_peak_mib": max(values),
|
|
"vram_end_mib": values[-1],
|
|
"vram_samples": len(values),
|
|
}
|
|
|
|
|
|
def wait_for_health(base_url: str, timeout_s: float) -> None:
|
|
deadline = time.monotonic() + timeout_s
|
|
last_error: Exception | None = None
|
|
while time.monotonic() < deadline:
|
|
try:
|
|
with urllib.request.urlopen(base_url + "/health", timeout=1) as response:
|
|
if response.status == 200:
|
|
return
|
|
except Exception as exc:
|
|
last_error = exc
|
|
time.sleep(0.5)
|
|
raise RuntimeError(f"server health timeout for {base_url}: {last_error}")
|
|
|
|
|
|
def stream_speech_request(base_url: str, model: str, text: str, sample_rate: int) -> dict[str, Any]:
|
|
payload = {
|
|
"model": model,
|
|
"input": text,
|
|
"response_format": "pcm",
|
|
"stream_format": "sse",
|
|
"options": {"retry_badcase": False},
|
|
}
|
|
request = urllib.request.Request(
|
|
base_url + "/v1/audio/speech",
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json", "Accept": "text/event-stream"},
|
|
method="POST",
|
|
)
|
|
start = time.perf_counter()
|
|
pcm = bytearray()
|
|
delta_events = 0
|
|
first_delta_ms: float | None = None
|
|
done_ms: float | None = None
|
|
ttft_ms: float | None = None
|
|
done_event: dict[str, Any] | None = None
|
|
errors: list[str] = []
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=1200) as response:
|
|
for raw_line in response:
|
|
line = raw_line.decode("utf-8", "replace").strip()
|
|
if not line or not line.startswith("data:"):
|
|
continue
|
|
data = line[5:].strip()
|
|
if data == "[DONE]":
|
|
break
|
|
event = json.loads(data)
|
|
event_ms = (time.perf_counter() - start) * 1000.0
|
|
event_type = event.get("type")
|
|
if event_type == "error":
|
|
errors.append(event.get("error", {}).get("message", json.dumps(event)))
|
|
elif event_type == "speech.audio.delta":
|
|
delta_events += 1
|
|
if first_delta_ms is None:
|
|
first_delta_ms = event_ms
|
|
pcm.extend(base64.b64decode(event["audio"]))
|
|
elif event_type == "speech.audio.done":
|
|
done_event = event
|
|
done_ms = event_ms
|
|
ttft_ms = event.get("timing", {}).get("ttft_ms")
|
|
except urllib.error.HTTPError as exc:
|
|
body = exc.read().decode("utf-8", "replace")
|
|
raise RuntimeError(f"HTTP {exc.code}: {body}") from exc
|
|
elapsed_ms = (time.perf_counter() - start) * 1000.0
|
|
if errors:
|
|
raise RuntimeError("; ".join(errors))
|
|
if ttft_ms is None:
|
|
raise RuntimeError("streaming speech response did not include ttft_ms")
|
|
if not pcm:
|
|
raise RuntimeError("streaming speech response produced no audio")
|
|
return {
|
|
"request": payload,
|
|
"elapsed_ms": elapsed_ms,
|
|
"ttft_ms": ttft_ms,
|
|
"client_first_delta_ms": first_delta_ms,
|
|
"client_done_ms": done_ms,
|
|
"audio_duration_s": len(pcm) / (2.0 * sample_rate),
|
|
"delta_events": delta_events,
|
|
"pcm_bytes": len(pcm),
|
|
"done_event": done_event,
|
|
"pcm": bytes(pcm),
|
|
}
|
|
|
|
|
|
def write_wav(path: Path, pcm: bytes, sample_rate: int = DEFAULT_SAMPLE_RATE) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with wave.open(str(path), "wb") as wav:
|
|
wav.setnchannels(1)
|
|
wav.setsampwidth(2)
|
|
wav.setframerate(sample_rate)
|
|
wav.writeframes(pcm)
|
|
|
|
|
|
def percentile(values: list[float], q: float) -> float:
|
|
if not values:
|
|
raise RuntimeError("percentile requires values")
|
|
ordered = sorted(values)
|
|
index = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * q)))
|
|
return ordered[index]
|
|
|
|
|
|
def run_measurement(args: argparse.Namespace) -> dict[str, Any]:
|
|
output_dir = args.output_dir
|
|
if not path_is_under(output_dir, LOG_ROOT):
|
|
raise RuntimeError(f"--output-dir must be under {LOG_ROOT}")
|
|
if args.max_chars <= 0:
|
|
raise RuntimeError("--max-chars must be positive")
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
config = load_json(args.server_config)
|
|
host = args.host or config.get("host", "127.0.0.1")
|
|
port = int(args.port or config["port"])
|
|
models = config.get("models") or []
|
|
if not models and not args.model:
|
|
raise RuntimeError("server config contains no models and --model was not provided")
|
|
model = args.model or models[0]["id"]
|
|
base_url = f"http://{host}:{port}"
|
|
write_json(output_dir / "server_config_snapshot.json", config)
|
|
|
|
input_text = args.input_text.read_text(encoding="utf-8")
|
|
chunks = split_incremental_text(input_text, args.max_chars)
|
|
(output_dir / "input_text.txt").write_text(input_text, encoding="utf-8")
|
|
write_json(
|
|
output_dir / "input_chunks.json",
|
|
{
|
|
"source": str(args.input_text),
|
|
"input_chars": len(input_text),
|
|
"max_chars": args.max_chars,
|
|
"chunks": [{"index": index, "chars": len(chunk), "text": chunk} for index, chunk in enumerate(chunks)],
|
|
},
|
|
)
|
|
|
|
requests_path = output_dir / "requests.jsonl"
|
|
if requests_path.exists():
|
|
requests_path.unlink()
|
|
|
|
server: subprocess.Popen[str] | None = None
|
|
server_log_handle = None
|
|
server_command: list[str] | None = None
|
|
try:
|
|
if args.start_server:
|
|
server_command = [
|
|
str(args.server_bin),
|
|
"--config",
|
|
str(args.server_config),
|
|
"--log-file",
|
|
str(output_dir / "framework.log"),
|
|
]
|
|
server_log_handle = (output_dir / "server_stdout.log").open("w", encoding="utf-8")
|
|
server = subprocess.Popen(
|
|
server_command,
|
|
cwd=args.repo_root,
|
|
stdout=server_log_handle,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
wait_for_health(base_url, args.health_timeout_s)
|
|
warmup = stream_speech_request(base_url, model, args.warmup_text, args.sample_rate)
|
|
write_wav(output_dir / "warmup_discard.wav", warmup["pcm"], args.sample_rate)
|
|
warmup_meta = {key: value for key, value in warmup.items() if key != "pcm"}
|
|
write_json(output_dir / "warmup_discard.json", warmup_meta)
|
|
|
|
measured: list[dict[str, Any]] = []
|
|
merged_pcm = bytearray()
|
|
sequence_start = time.perf_counter()
|
|
with VramSampler(output_dir / "measured_vram.csv", args.vram_sample_ms / 1000.0):
|
|
for index, chunk in enumerate(chunks):
|
|
result = stream_speech_request(base_url, model, chunk, args.sample_rate)
|
|
chunk_pcm = result.pop("pcm")
|
|
merged_pcm.extend(chunk_pcm)
|
|
write_wav(output_dir / f"chunk_{index:03d}.wav", chunk_pcm, args.sample_rate)
|
|
record = {
|
|
"index": index,
|
|
"input_chars": len(chunk),
|
|
"input_text": chunk,
|
|
**result,
|
|
}
|
|
measured.append(record)
|
|
with requests_path.open("a", encoding="utf-8") as handle:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
sequence_elapsed_ms = (time.perf_counter() - sequence_start) * 1000.0
|
|
write_wav(output_dir / "merged.wav", bytes(merged_pcm), args.sample_rate)
|
|
|
|
ttfts = [float(item["ttft_ms"]) for item in measured]
|
|
summary = {
|
|
"measurement_kind": "incremental_client_text_chunks_to_streaming_sse_tts",
|
|
"repo": str(args.repo_root),
|
|
"server_config": str(args.server_config),
|
|
"server_command": server_command,
|
|
"model": model,
|
|
"memory_saver": args.memory_saver_label,
|
|
"host": host,
|
|
"port": port,
|
|
"input_text": str(output_dir / "input_text.txt"),
|
|
"input_text_source": str(args.input_text),
|
|
"input_chars": len(input_text),
|
|
"chunk_max_chars": args.max_chars,
|
|
"chunk_count": len(chunks),
|
|
"warmup": warmup_meta,
|
|
"measured_sequence_elapsed_ms": sequence_elapsed_ms,
|
|
"sample_rate": args.sample_rate,
|
|
"total_audio_duration_s": len(merged_pcm) / (2.0 * args.sample_rate),
|
|
"ttft_first_chunk_ms": ttfts[0],
|
|
"ttft_min_ms": min(ttfts),
|
|
"ttft_p50_ms": percentile(ttfts, 0.50),
|
|
"ttft_p95_ms": percentile(ttfts, 0.95),
|
|
"ttft_max_ms": max(ttfts),
|
|
"ttft_mean_ms": sum(ttfts) / len(ttfts),
|
|
"client_first_delta_first_chunk_ms": measured[0]["client_first_delta_ms"],
|
|
"chunks": measured,
|
|
"artifacts": {
|
|
"summary": str(output_dir / "summary.json"),
|
|
"server_config_snapshot": str(output_dir / "server_config_snapshot.json"),
|
|
"framework_log": str(output_dir / "framework.log") if args.start_server else None,
|
|
"server_stdout": str(output_dir / "server_stdout.log") if args.start_server else None,
|
|
"input_text": str(output_dir / "input_text.txt"),
|
|
"chunks": str(output_dir / "input_chunks.json"),
|
|
"requests": str(requests_path),
|
|
"warmup_wav": str(output_dir / "warmup_discard.wav"),
|
|
"warmup_json": str(output_dir / "warmup_discard.json"),
|
|
"merged_wav": str(output_dir / "merged.wav"),
|
|
"vram_csv": str(output_dir / "measured_vram.csv"),
|
|
},
|
|
}
|
|
summary.update(read_vram_summary(output_dir / "measured_vram.csv"))
|
|
write_json(output_dir / "summary.json", summary)
|
|
return summary
|
|
finally:
|
|
if server is not None:
|
|
server.terminate()
|
|
try:
|
|
server.wait(timeout=10)
|
|
except subprocess.TimeoutExpired:
|
|
server.kill()
|
|
server.wait(timeout=10)
|
|
if server_log_handle is not None:
|
|
server_log_handle.close()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
default_output = LOG_ROOT / f"tts_streaming_{timestamp_slug()}"
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--repo-root", type=Path, default=REPO_ROOT)
|
|
parser.add_argument("--server-bin", type=Path, default=DEFAULT_SERVER_BIN)
|
|
parser.add_argument("--server-config", type=Path, required=True)
|
|
parser.add_argument("--start-server", action=argparse.BooleanOptionalAction, default=True)
|
|
parser.add_argument("--host", default="")
|
|
parser.add_argument("--port", type=int, default=0)
|
|
parser.add_argument("--model", default="")
|
|
parser.add_argument("--input-text", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, default=default_output)
|
|
parser.add_argument("--memory-saver-label", required=True)
|
|
parser.add_argument("--max-chars", type=int, default=480)
|
|
parser.add_argument("--warmup-text", default="Short warmup request for streaming timing.")
|
|
parser.add_argument("--sample-rate", type=int, default=DEFAULT_SAMPLE_RATE)
|
|
parser.add_argument("--health-timeout-s", type=float, default=180.0)
|
|
parser.add_argument("--vram-sample-ms", type=float, default=200.0)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
try:
|
|
summary = run_measurement(parse_args())
|
|
except Exception as exc:
|
|
print(f"ERROR: {exc}", file=sys.stderr)
|
|
return 1
|
|
keys = (
|
|
"memory_saver",
|
|
"chunk_count",
|
|
"measured_sequence_elapsed_ms",
|
|
"total_audio_duration_s",
|
|
"ttft_first_chunk_ms",
|
|
"ttft_p50_ms",
|
|
"ttft_p95_ms",
|
|
"ttft_max_ms",
|
|
"client_first_delta_first_chunk_ms",
|
|
"vram_peak_mib",
|
|
)
|
|
print(json.dumps({key: summary[key] for key in keys}, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|