939 lines
40 KiB
Python
939 lines
40 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Strict PersonaPlex validation for the unified OpenAI Realtime duplex path."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import time
|
|
import uuid
|
|
import wave
|
|
from collections.abc import Awaitable, Callable, Sequence
|
|
from pathlib import Path
|
|
from typing import Protocol
|
|
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
import websockets
|
|
from websockets.exceptions import ConnectionClosed, WebSocketException
|
|
except ImportError as exc: # pragma: no cover - driver dependency.
|
|
raise SystemExit("Install websockets first: pip install websockets") from exc
|
|
|
|
from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats
|
|
from vllm_omni.clients.duplex import EventCollector, write_pcm16_wav
|
|
from vllm_omni.clients.duplex import (
|
|
wait_for_condition as wait_for,
|
|
)
|
|
|
|
SAMPLE_RATE_HZ = 24_000
|
|
FRAME_SAMPLES = 1_920
|
|
FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ
|
|
AUDIO_DELTA_EVENT_TYPES = frozenset({"response.audio.delta", "response.output_audio.delta"})
|
|
|
|
|
|
class _AudioAcceptanceError(AssertionError):
|
|
"""A local acceptance failure with numeric diagnostics safe to report."""
|
|
|
|
def __init__(self, check: str, message: str, stats: dict[str, object]) -> None:
|
|
super().__init__(message)
|
|
self.check = check
|
|
self.stats = dict(stats)
|
|
|
|
|
|
class _ProbeTransport(Protocol):
|
|
async def recv(self) -> str | bytes: ...
|
|
|
|
async def send(self, message: str) -> None: ...
|
|
|
|
async def close(self) -> None: ...
|
|
|
|
|
|
class RawRealtimeProbe:
|
|
"""Raw-event Realtime probe: one websocket plus an :class:`EventCollector`.
|
|
|
|
This driver speaks the wire protocol directly (it validates handshake
|
|
errors such as admission limits, which the public
|
|
:class:`vllm_omni.clients.duplex.DuplexClient` treats as failures), so it
|
|
keeps a transport-only client and reads everything from ``.events``.
|
|
"""
|
|
|
|
def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024, close_timeout_s: float = 10.0) -> None:
|
|
self.url = url
|
|
self.max_size = max_size
|
|
self.close_timeout_s = close_timeout_s
|
|
self.events = EventCollector()
|
|
self._ws: _ProbeTransport | None = None
|
|
self._reader_task: asyncio.Task[None] | None = None
|
|
|
|
async def __aenter__(self) -> RawRealtimeProbe:
|
|
self._ws = await websockets.connect(self.url, max_size=self.max_size, close_timeout=self.close_timeout_s)
|
|
self._reader_task = asyncio.create_task(self._read_events())
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, traceback) -> None:
|
|
try:
|
|
if self._ws is not None:
|
|
await self._ws.close()
|
|
finally:
|
|
if self._reader_task is not None:
|
|
self._reader_task.cancel()
|
|
try:
|
|
await self._reader_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
async def _read_events(self) -> None:
|
|
assert self._ws is not None
|
|
try:
|
|
while True:
|
|
raw = await self._ws.recv()
|
|
if not isinstance(raw, str):
|
|
continue
|
|
event = json.loads(raw)
|
|
if isinstance(event, dict):
|
|
self.events.add(event)
|
|
except ConnectionClosed:
|
|
return
|
|
|
|
async def send(self, event: dict[str, object]) -> None:
|
|
assert self._ws is not None
|
|
await self._ws.send(json.dumps(event))
|
|
|
|
|
|
def _input_identity(
|
|
path: Path,
|
|
*,
|
|
expected_sha256: str | None,
|
|
) -> dict[str, str]:
|
|
resolved = path.resolve()
|
|
digest = hashlib.sha256()
|
|
with resolved.open("rb") as input_file:
|
|
for chunk in iter(lambda: input_file.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
actual = digest.hexdigest()
|
|
if expected_sha256 is not None and actual != expected_sha256.lower():
|
|
raise ValueError(
|
|
f"input WAV SHA-256 mismatch: expected={expected_sha256.lower()}, actual={actual}, path={resolved}"
|
|
)
|
|
return {"path": str(resolved), "sha256": actual}
|
|
|
|
|
|
def _realtime_url(base_url: str, model: str, session_id: str) -> str:
|
|
parts = urlsplit(base_url)
|
|
query = dict(parse_qsl(parts.query, keep_blank_values=True))
|
|
query.update(duplex="1", model=model, autostart="0", session_id=session_id)
|
|
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
|
|
|
|
|
|
def _read_wav_as_float32(path: Path) -> np.ndarray:
|
|
with wave.open(str(path), "rb") as wav_file:
|
|
channels = wav_file.getnchannels()
|
|
sample_width = wav_file.getsampwidth()
|
|
source_rate = wav_file.getframerate()
|
|
frames = wav_file.readframes(wav_file.getnframes())
|
|
dtypes = {1: np.uint8, 2: np.dtype("<i2"), 4: np.dtype("<i4")}
|
|
if sample_width not in dtypes:
|
|
raise ValueError(f"unsupported WAV sample width: {sample_width}")
|
|
pcm = np.frombuffer(frames, dtype=dtypes[sample_width]).astype(np.float32)
|
|
if sample_width == 1:
|
|
pcm = (pcm - 128.0) / 128.0
|
|
else:
|
|
pcm /= float(1 << (sample_width * 8 - 1))
|
|
if channels > 1:
|
|
pcm = pcm.reshape(-1, channels).mean(axis=1)
|
|
if source_rate != SAMPLE_RATE_HZ and pcm.size:
|
|
output_samples = max(1, round(pcm.size * SAMPLE_RATE_HZ / source_rate))
|
|
pcm = np.interp(
|
|
np.linspace(0, pcm.size - 1, output_samples),
|
|
np.arange(pcm.size),
|
|
pcm,
|
|
).astype(np.float32)
|
|
pcm = np.ascontiguousarray(pcm, dtype="<f4")
|
|
if not pcm.size or not np.isfinite(pcm).all():
|
|
raise ValueError("input WAV must contain finite audio")
|
|
return pcm
|
|
|
|
|
|
def _events(client: RawRealtimeProbe, event_type: str) -> list[dict[str, object]]:
|
|
return [event for event in client.events.events if event.get("type") == event_type]
|
|
|
|
|
|
def _audio_events(client: RawRealtimeProbe) -> list[dict[str, object]]:
|
|
"""Return audio deltas from both current and legacy Realtime event names."""
|
|
return [event for event in client.events.events if event.get("type") in AUDIO_DELTA_EVENT_TYPES]
|
|
|
|
|
|
def _validated_audio_event(event: dict[str, object]) -> bytes:
|
|
"""Decode one wire audio delta without depending on client event aliases."""
|
|
encoded = event.get("delta")
|
|
if not isinstance(encoded, str):
|
|
raise ValueError("missing_audio_delta")
|
|
try:
|
|
chunk = base64.b64decode(encoded, validate=True)
|
|
except ValueError:
|
|
raise ValueError("invalid_audio_base64") from None
|
|
if chunk:
|
|
if len(chunk) % 2:
|
|
raise ValueError("unaligned_pcm16_packet")
|
|
if event.get("sample_rate_hz") != SAMPLE_RATE_HZ:
|
|
raise ValueError("invalid_audio_sample_rate")
|
|
if not isinstance(event.get("response_id"), str) or not event.get("response_id"):
|
|
raise ValueError("missing_audio_response_id")
|
|
return chunk
|
|
|
|
|
|
def _validated_audio_chunks(client: RawRealtimeProbe) -> list[tuple[dict[str, object], bytes]]:
|
|
return [(event, _validated_audio_event(event)) for event in _audio_events(client)]
|
|
|
|
|
|
def _audio_bytes(client: RawRealtimeProbe) -> bytes:
|
|
return b"".join(chunk for _, chunk in _validated_audio_chunks(client) if chunk)
|
|
|
|
|
|
async def _open_session(
|
|
args: argparse.Namespace,
|
|
*,
|
|
session_id: str,
|
|
persona: str,
|
|
expect_error: bool = False,
|
|
) -> tuple[RawRealtimeProbe, dict[str, object]]:
|
|
client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id))
|
|
await client.__aenter__()
|
|
await client.send(_session_update(args, session_id=session_id, persona=persona))
|
|
event_type = "error" if expect_error else "session.created"
|
|
await wait_for(
|
|
lambda: client.events.count(event_type) > 0,
|
|
timeout_s=args.timeout_s,
|
|
label=event_type,
|
|
)
|
|
return client, _events(client, event_type)[-1]
|
|
|
|
|
|
def _session_update(args: argparse.Namespace, *, session_id: str, persona: str) -> dict[str, object]:
|
|
return {
|
|
"type": "session.update",
|
|
"session": {
|
|
"session_id": session_id,
|
|
"model": args.model,
|
|
"modalities": ["audio", "text"],
|
|
"input_audio_format": "pcm_f32le",
|
|
"output_audio_format": "pcm16",
|
|
"voice": args.voice,
|
|
"instructions": persona,
|
|
"turn_detection": None,
|
|
"extra_body": {"auto_response": True},
|
|
},
|
|
}
|
|
|
|
|
|
async def _close_session(client: RawRealtimeProbe, *, timeout_s: float) -> None:
|
|
await client.send({"type": "session.close"})
|
|
await wait_for(
|
|
lambda: client.events.count("session.closed") > 0,
|
|
timeout_s=timeout_s,
|
|
label="session.closed",
|
|
)
|
|
await client.__aexit__(None, None, None)
|
|
|
|
|
|
async def _stream_frames(
|
|
client: RawRealtimeProbe,
|
|
pcm: np.ndarray,
|
|
*,
|
|
max_frames: int | None = None,
|
|
) -> int:
|
|
frame_count = math.ceil(pcm.size / FRAME_SAMPLES)
|
|
if max_frames is not None:
|
|
if max_frames < 0:
|
|
raise ValueError(f"max_frames must be non-negative, got {max_frames}")
|
|
if max_frames > 0:
|
|
frame_count = min(frame_count, max_frames)
|
|
for seq in range(frame_count):
|
|
frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
|
|
await client.send(_frame_event(frame, seq))
|
|
await asyncio.sleep(FRAME_PERIOD_S)
|
|
return frame_count
|
|
|
|
|
|
def _frame_event(frame: np.ndarray, seq: int) -> dict[str, object]:
|
|
frame = np.ascontiguousarray(np.pad(frame, (0, FRAME_SAMPLES - frame.size)), dtype="<f4")
|
|
return {
|
|
"type": "input_audio_buffer.append",
|
|
"audio": base64.b64encode(frame).decode("ascii"),
|
|
"format": "pcm_f32le",
|
|
"sample_rate_hz": SAMPLE_RATE_HZ,
|
|
"duration_ms": 80,
|
|
"audio_end_ms": (seq + 1) * 80,
|
|
"is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
|
|
}
|
|
|
|
|
|
def _audio_frame_stats(
|
|
raw: bytes,
|
|
*,
|
|
input_frames: int,
|
|
voiced_frame_rms_threshold: float,
|
|
) -> dict[str, int | float]:
|
|
pcm = np.frombuffer(raw, dtype="<i2")
|
|
if pcm.size % FRAME_SAMPLES:
|
|
raise AssertionError(
|
|
"PersonaPlex output must contain whole 80 ms codec frames: "
|
|
f"samples={pcm.size}, samples_per_frame={FRAME_SAMPLES}"
|
|
)
|
|
output_frames = pcm.size // FRAME_SAMPLES
|
|
frames = pcm.astype(np.float32).reshape(output_frames, FRAME_SAMPLES) / 32768.0
|
|
voiced_frames = int(np.count_nonzero(np.sqrt(np.mean(np.square(frames), axis=1)) >= voiced_frame_rms_threshold))
|
|
return {
|
|
"output_frames": int(output_frames),
|
|
"voiced_frames": voiced_frames,
|
|
"silent_frames": int(output_frames - voiced_frames),
|
|
"frame_deficit": int(input_frames - output_frames),
|
|
"frame_coverage_ratio": float(output_frames / input_frames),
|
|
}
|
|
|
|
|
|
def _response_ids(client: RawRealtimeProbe) -> set[str]:
|
|
return {str(event["response_id"]) for event in _audio_events(client) if isinstance(event.get("response_id"), str)}
|
|
|
|
|
|
def _session_result(
|
|
client: RawRealtimeProbe,
|
|
*,
|
|
input_frames: int,
|
|
args: argparse.Namespace,
|
|
minimum_chunks: int,
|
|
) -> tuple[bytes, set[str], dict[str, object]]:
|
|
if isinstance(client, RawRealtimeProbe):
|
|
audio_chunks = _validated_audio_chunks(client)
|
|
chunks = sum(bool(chunk) for _, chunk in audio_chunks)
|
|
raw = b"".join(chunk for _, chunk in audio_chunks if chunk)
|
|
else:
|
|
# Unit fixtures may provide an already-decoded EventCollector without
|
|
# retaining the base64 wire field. Real websocket probes always take
|
|
# the strict wire-validation path above.
|
|
chunks = sum(len(items) for items in client.events.response_audio.values())
|
|
raw = client.events.audio_bytes()
|
|
stats: dict[str, object] = {"audio_chunks": chunks, "audio_samples": len(raw) // 2}
|
|
if not raw or len(raw) % 2:
|
|
raise _AudioAcceptanceError("pcm16_output", f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}", stats)
|
|
pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
|
|
rms = float(np.sqrt(np.mean(np.square(pcm))))
|
|
stats["audio_rms"] = rms
|
|
# This low floor only catches corrupt, non-finite, or effectively zero PCM.
|
|
# Audible speech is checked per frame below with the higher 1e-3 default.
|
|
if not np.isfinite(pcm).all() or rms <= 1e-5:
|
|
raise _AudioAcceptanceError(
|
|
"audio_rms", f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}", stats
|
|
)
|
|
rates = {rate for event in _audio_events(client) if isinstance(rate := event.get("sample_rate_hz"), int)}
|
|
if rates != {SAMPLE_RATE_HZ}:
|
|
raise _AudioAcceptanceError("sample_rate", f"unexpected output sample rates: {sorted(rates)}", stats)
|
|
if input_frames < 1:
|
|
raise _AudioAcceptanceError("input_frames", f"audio received without input frames: input={input_frames}", stats)
|
|
try:
|
|
frame_stats = _audio_frame_stats(
|
|
raw,
|
|
input_frames=input_frames,
|
|
voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
|
|
)
|
|
except AssertionError as exc:
|
|
# _audio_frame_stats only formats local sample counts, not wire data.
|
|
raise _AudioAcceptanceError("whole_codec_frames", str(exc), stats) from exc
|
|
stats.update(frame_stats)
|
|
if chunks < minimum_chunks:
|
|
raise _AudioAcceptanceError(
|
|
"audio_chunks", f"insufficient audio chunks: chunks={chunks}, minimum={minimum_chunks}", stats
|
|
)
|
|
if not 0 <= int(frame_stats["frame_deficit"]) <= args.max_frame_deficit:
|
|
raise _AudioAcceptanceError(
|
|
"frame_deficit",
|
|
f"invalid synchronous frame accounting: input={input_frames}, maximum_deficit={args.max_frame_deficit}, "
|
|
f"stats={stats}",
|
|
stats,
|
|
)
|
|
min_voiced_frames = int(args.min_voiced_frames)
|
|
if min_voiced_frames < 1:
|
|
raise ValueError(f"min_voiced_frames must be positive, got {min_voiced_frames}")
|
|
if int(frame_stats["voiced_frames"]) < min_voiced_frames:
|
|
raise _AudioAcceptanceError(
|
|
"voiced_frames",
|
|
f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}",
|
|
stats,
|
|
)
|
|
runtime_metadata = [
|
|
metadata["vllm_omni"]
|
|
for event in _audio_events(client)
|
|
if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict)
|
|
]
|
|
if not any(
|
|
item.get("runtime_impl") == "scheduler_data_plane"
|
|
and item.get("uses_model_runner_scheduler") is True
|
|
and item.get("runner_kv_backed") is True
|
|
for item in runtime_metadata
|
|
):
|
|
raise _AudioAcceptanceError(
|
|
"scheduler_data_plane", "audio events did not prove the scheduler data plane", stats
|
|
)
|
|
response_ids = _response_ids(client)
|
|
if len(response_ids) != 1:
|
|
raise _AudioAcceptanceError(
|
|
"response_identity", f"continuous stream requires one response: count={len(response_ids)}", stats
|
|
)
|
|
return (
|
|
raw,
|
|
response_ids,
|
|
{
|
|
**stats,
|
|
"response_ids": sorted(response_ids),
|
|
},
|
|
)
|
|
|
|
|
|
def _capabilities(created: dict[str, object]) -> dict[str, object]:
|
|
session = created.get("session")
|
|
capabilities = session.get("capabilities") if isinstance(session, dict) else None
|
|
if not isinstance(capabilities, dict):
|
|
raise AssertionError(f"session.created omitted capabilities: {created}")
|
|
return capabilities
|
|
|
|
|
|
def _save(output_dir: Path, name: str, client: RawRealtimeProbe, raw: bytes) -> None:
|
|
(output_dir / f"{name}-events.jsonl").write_text(
|
|
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
|
|
encoding="utf-8",
|
|
)
|
|
write_pcm16_wav(output_dir / f"{name}-output.wav", raw, sample_rate_hz=SAMPLE_RATE_HZ)
|
|
|
|
|
|
async def run(args: argparse.Namespace) -> dict[str, object]:
|
|
if getattr(args, "sessions", None) is not None:
|
|
return await run_load(args)
|
|
input_path = Path(args.input_wav)
|
|
input_identity = _input_identity(
|
|
input_path,
|
|
expected_sha256=args.expected_input_sha256,
|
|
)
|
|
pcm = _read_wav_as_float32(input_path)
|
|
if args.tail_s > 0:
|
|
pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
ids = {name: f"personaplex-{name}-{uuid.uuid4().hex}" for name in ("primary", "secondary", "replacement")}
|
|
primary, created = await _open_session(args, session_id=ids["primary"], persona=args.persona)
|
|
secondary, secondary_created = await _open_session(
|
|
args,
|
|
session_id=ids["secondary"],
|
|
persona=args.secondary_persona,
|
|
)
|
|
capabilities = _capabilities(created)
|
|
if _capabilities(secondary_created) != capabilities:
|
|
raise AssertionError("concurrent sessions returned different capabilities")
|
|
expected_capabilities = {
|
|
"implementation_level": "model_native_duplex",
|
|
"chunk_period_ms": 80,
|
|
"supports_multi_session": True,
|
|
"supports_multi_session_same_replica": True,
|
|
"supports_barge_in": False,
|
|
}
|
|
if any(capabilities.get(key) != value for key, value in expected_capabilities.items()):
|
|
raise AssertionError(f"unexpected PersonaPlex capabilities: {capabilities}")
|
|
|
|
overflow, error = await _open_session(
|
|
args,
|
|
session_id=f"personaplex-overflow-{uuid.uuid4().hex}",
|
|
persona=args.persona,
|
|
expect_error=True,
|
|
)
|
|
error_body = error.get("error")
|
|
overflow_code = error_body.get("code") if isinstance(error_body, dict) else error.get("code")
|
|
await overflow.__aexit__(None, None, None)
|
|
if overflow_code != "resource_exhausted":
|
|
raise AssertionError(f"overflow session was not rejected: {error}")
|
|
|
|
primary_frames, secondary_frames = await asyncio.gather(
|
|
_stream_frames(primary, pcm),
|
|
_stream_frames(secondary, pcm),
|
|
)
|
|
await asyncio.gather(
|
|
wait_for(
|
|
lambda: len(_audio_bytes(primary)) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
|
|
timeout_s=args.timeout_s,
|
|
label="primary frame coverage",
|
|
),
|
|
wait_for(
|
|
lambda: (len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit),
|
|
timeout_s=args.timeout_s,
|
|
label="secondary frame coverage",
|
|
),
|
|
)
|
|
await asyncio.sleep(args.drain_s)
|
|
primary_audio, primary_response_ids, primary_stats = _session_result(
|
|
primary,
|
|
input_frames=primary_frames,
|
|
args=args,
|
|
minimum_chunks=args.minimum_audio_chunks,
|
|
)
|
|
secondary_audio, secondary_response_ids, _ = _session_result(
|
|
secondary,
|
|
input_frames=secondary_frames,
|
|
args=args,
|
|
minimum_chunks=args.minimum_audio_chunks,
|
|
)
|
|
if primary_response_ids & secondary_response_ids:
|
|
raise AssertionError("concurrent sessions shared a response id")
|
|
_save(output_dir, "primary", primary, primary_audio)
|
|
await _close_session(primary, timeout_s=args.timeout_s)
|
|
|
|
replacement, replacement_created = await _open_session(
|
|
args,
|
|
session_id=ids["replacement"],
|
|
persona=args.replacement_persona,
|
|
)
|
|
if _capabilities(replacement_created) != capabilities:
|
|
raise AssertionError("replacement session returned different capabilities")
|
|
continuation_frames, replacement_frames = await asyncio.gather(
|
|
_stream_frames(secondary, pcm, max_frames=args.replacement_frames),
|
|
_stream_frames(replacement, pcm, max_frames=args.replacement_frames),
|
|
)
|
|
secondary_total_frames = secondary_frames + continuation_frames
|
|
await asyncio.gather(
|
|
wait_for(
|
|
lambda: (
|
|
len(_audio_bytes(secondary)) // (2 * FRAME_SAMPLES) >= secondary_total_frames - args.max_frame_deficit
|
|
),
|
|
timeout_s=args.timeout_s,
|
|
label="survivor frame coverage after slot recycle",
|
|
),
|
|
wait_for(
|
|
lambda: (
|
|
len(_audio_bytes(replacement)) // (2 * FRAME_SAMPLES) >= replacement_frames - args.max_frame_deficit
|
|
),
|
|
timeout_s=args.timeout_s,
|
|
label="replacement frame coverage",
|
|
),
|
|
)
|
|
await asyncio.sleep(args.drain_s)
|
|
secondary_audio, secondary_after_ids, secondary_stats = _session_result(
|
|
secondary,
|
|
input_frames=secondary_total_frames,
|
|
args=args,
|
|
minimum_chunks=args.minimum_audio_chunks + 1,
|
|
)
|
|
replacement_audio, replacement_response_ids, replacement_stats = _session_result(
|
|
replacement,
|
|
input_frames=replacement_frames,
|
|
args=args,
|
|
minimum_chunks=1,
|
|
)
|
|
if secondary_after_ids != secondary_response_ids:
|
|
raise AssertionError("survivor response changed while another session slot was recycled")
|
|
if replacement_response_ids & (primary_response_ids | secondary_response_ids):
|
|
raise AssertionError("replacement session reused another session's response id")
|
|
_save(output_dir, "secondary", secondary, secondary_audio)
|
|
_save(output_dir, "replacement", replacement, replacement_audio)
|
|
await asyncio.gather(
|
|
_close_session(secondary, timeout_s=args.timeout_s),
|
|
_close_session(replacement, timeout_s=args.timeout_s),
|
|
)
|
|
|
|
errors = primary.events.errors() + secondary.events.errors() + replacement.events.errors()
|
|
result = {
|
|
"ok": not errors,
|
|
"model": args.model,
|
|
"input": {**input_identity, "tail_s": args.tail_s},
|
|
"capabilities": capabilities,
|
|
"overflow_error_code": overflow_code,
|
|
"primary": {"session_id": ids["primary"], "input_frames": primary_frames, **primary_stats},
|
|
"secondary": {
|
|
"session_id": ids["secondary"],
|
|
"input_frames": secondary_total_frames,
|
|
"continuation_frames_after_recycle": continuation_frames,
|
|
**secondary_stats,
|
|
},
|
|
"replacement": {
|
|
"session_id": ids["replacement"],
|
|
"input_frames": replacement_frames,
|
|
**replacement_stats,
|
|
},
|
|
"errors": errors,
|
|
"output_dir": str(output_dir),
|
|
}
|
|
(output_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
if errors:
|
|
raise AssertionError(f"Realtime sessions emitted errors: {errors}")
|
|
return result
|
|
|
|
|
|
def _distribution(values: list[float]) -> dict[str, int | float | None]:
|
|
if not values:
|
|
return {"count": 0, "median": None, "p99": None, "max": None, "argmax_index": None}
|
|
return {
|
|
"count": len(values),
|
|
"median": float(np.median(values)),
|
|
"p99": float(np.quantile(values, 0.99)),
|
|
"max": max(values),
|
|
"argmax_index": int(np.argmax(values)),
|
|
}
|
|
|
|
|
|
def _load_metrics(client: RawRealtimeProbe, sends: list[tuple[float, float, float]]) -> dict[str, object]:
|
|
# Use the collector's local clock, never a timestamp supplied by the server.
|
|
packets: list[dict[str, object]] = []
|
|
response_ids: set[str] = set()
|
|
packet_times: list[float] = []
|
|
packet_bytes: list[int] = []
|
|
invalid_packets: list[dict[str, object]] = []
|
|
samples = 0
|
|
for index, (raw, received) in enumerate(zip(client.events.events, client.events.event_received_at_s, strict=True)):
|
|
if raw.get("type") not in AUDIO_DELTA_EVENT_TYPES:
|
|
continue
|
|
try:
|
|
chunk = _validated_audio_event(raw)
|
|
except ValueError as exc:
|
|
# Only fixed diagnostic codes from the validator are exported,
|
|
# never audio payloads or details from the remote endpoint.
|
|
invalid_packets.append({"event_index": index, "received_at_s": received, "code": str(exc)})
|
|
continue
|
|
if chunk:
|
|
samples += len(chunk) // 2
|
|
packet_times.append(received)
|
|
packet_bytes.append(len(chunk))
|
|
response_id = raw.get("response_id")
|
|
packets.append({"received_at_s": received, "samples": len(chunk) // 2, "response_id": response_id})
|
|
if isinstance(response_id, str) and response_id:
|
|
response_ids.add(response_id)
|
|
intervals = [(b - a) * 1000 for a, b in zip(packet_times, packet_times[1:])]
|
|
first_send = sends[0][1] if sends else None
|
|
audio_duration = samples / SAMPLE_RATE_HZ
|
|
first_audio_ms = (packet_times[0] - first_send) * 1000 if packet_times and first_send is not None else None
|
|
client_rtf = (packet_times[-1] - first_send) / audio_duration if audio_duration and first_send is not None else None
|
|
continuity = compute_continuity_stats(packet_times, packet_bytes, SAMPLE_RATE_HZ) if packet_times else None
|
|
lateness = _distribution([max(0.0, sent - planned) * 1000 for planned, sent, _ in sends])
|
|
p99_lateness = lateness["p99"]
|
|
return {
|
|
"input_frames": len(sends),
|
|
"output_samples": samples,
|
|
# Populate frame accounting from _session_result, including on failure.
|
|
"frame_deficit": None,
|
|
"response_ids": sorted(response_ids),
|
|
"client_first_audio_after_stream_start_ms": first_audio_ms,
|
|
"client_stream_rtf": client_rtf,
|
|
"client_audio_packet_interval_ms": _distribution(intervals),
|
|
"client_send_lateness_ms": lateness,
|
|
"client_pacing_warning": isinstance(p99_lateness, float) and p99_lateness > FRAME_PERIOD_S * 1000,
|
|
"audio_underrun_s": continuity.max_underrun_s if continuity else None,
|
|
"audio_underrun_event_count": continuity.underrun_event_count if continuity else None,
|
|
"audio_continuity_ok": continuity.is_continuous if continuity else None,
|
|
"input_timeline": [
|
|
{"frame_index": i, "planned_at_s": planned, "send_started_at_s": sent, "send_completed_at_s": completed}
|
|
for i, (planned, sent, completed) in enumerate(sends)
|
|
],
|
|
"audio_packet_timeline": packets,
|
|
"invalid_audio_packets": invalid_packets,
|
|
}
|
|
|
|
|
|
def _check_load_connection(client: RawRealtimeProbe) -> None:
|
|
if client.events.errors():
|
|
raise RuntimeError("server emitted a Realtime error (see the session error count)")
|
|
if client._reader_task is not None and client._reader_task.done():
|
|
client._reader_task.result() # Propagate malformed JSON/reader failures.
|
|
raise RuntimeError("connection closed before session completion")
|
|
|
|
|
|
async def _wait_load_event(client: RawRealtimeProbe, event_type: str, timeout_s: float) -> None:
|
|
def received() -> bool:
|
|
if client.events.count(event_type):
|
|
return True
|
|
_check_load_connection(client)
|
|
return False
|
|
|
|
await wait_for(received, timeout_s=timeout_s, label=event_type)
|
|
|
|
|
|
async def _paced_load_frames(
|
|
client: RawRealtimeProbe,
|
|
pcm: np.ndarray,
|
|
*,
|
|
epoch: float,
|
|
timeout_s: float,
|
|
sends: list[tuple[float, float, float]],
|
|
clock: Callable[[], float] = time.monotonic,
|
|
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
|
) -> None:
|
|
next_send = epoch
|
|
for seq in range(math.ceil(pcm.size / FRAME_SAMPLES)):
|
|
await sleep(max(0.0, next_send - clock()))
|
|
_check_load_connection(client)
|
|
event = _frame_event(pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES], seq)
|
|
started = clock()
|
|
await asyncio.wait_for(client.send(event), timeout=timeout_s)
|
|
completed = clock()
|
|
sends.append((epoch + seq * FRAME_PERIOD_S, started, completed))
|
|
# Do not burst to catch up after a stalled send; report accumulated lag.
|
|
next_send = max(epoch + (seq + 1) * FRAME_PERIOD_S, completed + FRAME_PERIOD_S)
|
|
|
|
|
|
async def _request_load_close(client: RawRealtimeProbe, timeout_s: float) -> None:
|
|
await client.send({"type": "session.close"})
|
|
await _wait_load_event(client, "session.closed", timeout_s)
|
|
|
|
|
|
async def _run_load_session(
|
|
args: argparse.Namespace,
|
|
pcm: np.ndarray,
|
|
index: int,
|
|
ready: asyncio.Future[None],
|
|
start: asyncio.Future[float],
|
|
) -> dict[str, object]:
|
|
session_id = f"personaplex-load-{index}-{uuid.uuid4().hex}"
|
|
client = RawRealtimeProbe(_realtime_url(args.url, args.model, session_id), close_timeout_s=args.cleanup_timeout_s)
|
|
sends: list[tuple[float, float, float]] = []
|
|
error: str | None = None
|
|
cleanup_error: str | None = None
|
|
metrics: dict[str, object] | None = None
|
|
session_stats: dict[str, object] = {}
|
|
acceptance_check: str | None = None
|
|
phase = "admission"
|
|
try:
|
|
async with client:
|
|
try:
|
|
await asyncio.wait_for(
|
|
client.send(_session_update(args, session_id=session_id, persona=args.persona)), args.timeout_s
|
|
)
|
|
await _wait_load_event(client, "session.created", args.timeout_s)
|
|
capabilities = _capabilities(_events(client, "session.created")[-1])
|
|
if (
|
|
capabilities.get("chunk_period_ms") != 80
|
|
or capabilities.get("supports_multi_session_same_replica") is not True
|
|
):
|
|
raise AssertionError("server did not advertise PersonaPlex multi-session capabilities")
|
|
ready.set_result(None)
|
|
phase = "stream"
|
|
epoch = await asyncio.shield(start)
|
|
await _paced_load_frames(client, pcm, epoch=epoch, timeout_s=args.timeout_s, sends=sends)
|
|
phase = "drain"
|
|
# Fixed, reported observation window. Closing must not extend
|
|
# it and let a slow session pass by flushing after the cutoff.
|
|
await asyncio.sleep(args.drain_s)
|
|
_check_load_connection(client)
|
|
_, _, session_stats = _session_result(
|
|
client, input_frames=len(sends), args=args, minimum_chunks=args.minimum_audio_chunks
|
|
)
|
|
finally:
|
|
metrics = _load_metrics(client, sends)
|
|
if client.events.count("session.created") and not client.events.count("session.closed"):
|
|
try:
|
|
# Send and acknowledgement share one short cleanup budget,
|
|
# independent of the potentially long inference timeout.
|
|
await asyncio.wait_for(
|
|
_request_load_close(client, args.cleanup_timeout_s), args.cleanup_timeout_s
|
|
)
|
|
except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, WebSocketException) as exc:
|
|
cleanup_error = f"cleanup: {type(exc).__name__}"
|
|
except (OSError, asyncio.TimeoutError, ValueError, RuntimeError, AssertionError, WebSocketException) as exc:
|
|
if isinstance(exc, _AudioAcceptanceError):
|
|
acceptance_check = exc.check
|
|
session_stats = exc.stats
|
|
error = f"{phase}: AssertionError: {exc}"
|
|
else:
|
|
# Other assertions/errors may embed handshake credentials or wire
|
|
# metadata. Only our local numeric acceptance diagnostics are safe.
|
|
error = f"{phase}: {type(exc).__name__}"
|
|
finally:
|
|
# Admission failure must not leave the other sessions at the barrier.
|
|
if not ready.done():
|
|
ready.set_result(None)
|
|
if metrics is None:
|
|
metrics = _load_metrics(client, sends)
|
|
if cleanup_error:
|
|
error = f"{error}; {cleanup_error}" if error else cleanup_error
|
|
if client.events.errors():
|
|
error = error or "server emitted a Realtime error"
|
|
if metrics["invalid_audio_packets"]:
|
|
error = error or "invalid audio packets"
|
|
first_audio = metrics["client_first_audio_after_stream_start_ms"]
|
|
if isinstance(first_audio, float) and first_audio < 0:
|
|
error = error or "audio arrived before the first input frame"
|
|
rtf = metrics["client_stream_rtf"]
|
|
if args.max_client_rtf is not None and isinstance(rtf, float) and rtf > args.max_client_rtf:
|
|
error = error or f"client_stream_rtf exceeds {args.max_client_rtf}"
|
|
return {
|
|
"index": index,
|
|
"session_id": session_id,
|
|
"ok": error is None,
|
|
"error": error,
|
|
"acceptance_check": acceptance_check,
|
|
"server_error_count": len(client.events.errors()),
|
|
"server_error_codes": [
|
|
str(body.get("code", "unknown")) if isinstance(body := event.get("error"), dict) else "unknown"
|
|
for event in client.events.errors()
|
|
],
|
|
**metrics,
|
|
**session_stats,
|
|
}
|
|
|
|
|
|
async def run_load(args: argparse.Namespace) -> dict[str, object]:
|
|
input_path = Path(args.input_wav)
|
|
identity = _input_identity(input_path, expected_sha256=args.expected_input_sha256)
|
|
pcm = _read_wav_as_float32(input_path)
|
|
if args.tail_s:
|
|
pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
|
|
pcm = pcm[: args.load_frames * FRAME_SAMPLES]
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
# Refuse to silently overwrite evidence from a previous sweep.
|
|
result_path = output_dir / "load-result.json"
|
|
if result_path.exists():
|
|
raise FileExistsError(result_path)
|
|
loop = asyncio.get_running_loop()
|
|
start: asyncio.Future[float] = loop.create_future()
|
|
ready: list[asyncio.Future[None]] = [loop.create_future() for _ in range(args.sessions)]
|
|
tasks = [asyncio.create_task(_run_load_session(args, pcm, i, r, start)) for i, r in enumerate(ready)]
|
|
try:
|
|
await asyncio.gather(*ready)
|
|
start.set_result(time.monotonic())
|
|
rows = await asyncio.gather(*tasks)
|
|
finally:
|
|
for task in tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
seen: set[str] = set()
|
|
shared: set[str] = set()
|
|
for row in rows:
|
|
assert isinstance(row["response_ids"], list)
|
|
response_ids = set(row["response_ids"])
|
|
shared.update(seen & response_ids)
|
|
seen.update(response_ids)
|
|
for row in rows:
|
|
assert isinstance(row["response_ids"], list)
|
|
if shared.intersection(row["response_ids"]):
|
|
row.update(ok=False, error="concurrent sessions shared a response id")
|
|
result = {
|
|
"schema_version": 1,
|
|
"mode": "paced_multi_session",
|
|
"model": args.model,
|
|
"requested_sessions": args.sessions,
|
|
"passed_sessions": sum(row["ok"] is True for row in rows),
|
|
"ok": all(row["ok"] is True for row in rows),
|
|
"input": {**identity, "tail_s": args.tail_s, "load_frames": args.load_frames},
|
|
"frame_period_s": FRAME_PERIOD_S,
|
|
"drain_s": args.drain_s,
|
|
"cleanup_timeout_s": args.cleanup_timeout_s,
|
|
"client_pacing_warning_threshold_ms": FRAME_PERIOD_S * 1000,
|
|
"client_pacing_warning_sessions": [row["index"] for row in rows if row["client_pacing_warning"]],
|
|
"max_client_rtf": args.max_client_rtf,
|
|
"acceptance": {
|
|
"max_frame_deficit": args.max_frame_deficit,
|
|
"min_voiced_frames": args.min_voiced_frames,
|
|
"minimum_audio_chunks": args.minimum_audio_chunks,
|
|
"voiced_frame_rms_threshold": args.voiced_frame_rms_threshold,
|
|
},
|
|
"shared_response_ids": sorted(shared),
|
|
"sessions": rows,
|
|
"reported_server": {"revision": args.server_revision, "hardware": args.server_hardware},
|
|
"measurement_origin": "client monotonic clock; admission excluded; streaming and drain retained",
|
|
"interval_definition": "nonempty audio packet arrivals, not server ticks or per-frame inference latency",
|
|
}
|
|
with result_path.open("x", encoding="utf-8") as stream:
|
|
json.dump(result, stream, ensure_ascii=False, indent=2, allow_nan=False)
|
|
return result
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
|
|
parser.add_argument("--model", required=True)
|
|
parser.add_argument("--input-wav", required=True)
|
|
parser.add_argument(
|
|
"--expected-input-sha256",
|
|
help="Fail before opening sessions when the input WAV does not match this SHA-256.",
|
|
)
|
|
parser.add_argument("--output-dir", default="/tmp/personaplex-realtime-duplex")
|
|
parser.add_argument("--voice", default="NATF2.pt")
|
|
parser.add_argument("--persona", default="You are a concise and helpful assistant.")
|
|
parser.add_argument("--secondary-persona", default="You are a calm and precise assistant.")
|
|
parser.add_argument("--replacement-persona", default="You are a cheerful travel assistant.")
|
|
parser.add_argument("--tail-s", type=float, default=2.0)
|
|
parser.add_argument("--drain-s", type=float, default=2.0)
|
|
parser.add_argument("--timeout-s", type=float, default=180.0)
|
|
parser.add_argument(
|
|
"--cleanup-timeout-s",
|
|
type=float,
|
|
default=5.0,
|
|
help="Load mode: shared session-close/ack budget; also bounds the websocket close handshake.",
|
|
)
|
|
parser.add_argument(
|
|
"--max-frame-deficit",
|
|
type=int,
|
|
default=4,
|
|
help=(
|
|
"Maximum unflushed tail frames. PersonaPlex Code2Wav emits five-frame "
|
|
"chunks, so a session cutoff can leave at most four pending frames."
|
|
),
|
|
)
|
|
parser.add_argument("--voiced-frame-rms-threshold", type=float, default=1e-3)
|
|
parser.add_argument(
|
|
"--min-voiced-frames",
|
|
type=int,
|
|
default=5,
|
|
help=(
|
|
"Require at least one normal five-frame Code2Wav chunk (400 ms) above "
|
|
"the per-frame RMS threshold; this detects inaudible output, not speech quality."
|
|
),
|
|
)
|
|
parser.add_argument("--minimum-audio-chunks", type=int, default=2)
|
|
parser.add_argument(
|
|
"--replacement-frames",
|
|
type=int,
|
|
default=0,
|
|
help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.",
|
|
)
|
|
parser.add_argument(
|
|
"--sessions", type=int, help="Opt into paced load mode with N sessions; omit for the lifecycle test."
|
|
)
|
|
parser.add_argument(
|
|
"--load-frames", type=int, default=1000, help="Maximum input frames per load session (no repetition)."
|
|
)
|
|
parser.add_argument("--max-client-rtf", type=float, help="Optional client stream-window RTF acceptance ceiling.")
|
|
parser.add_argument("--server-revision", help="User-reported server commit/version; not verified by this client.")
|
|
parser.add_argument(
|
|
"--server-hardware", help="User-reported server hardware; not inferred from the client machine."
|
|
)
|
|
args = parser.parse_args(argv)
|
|
if args.sessions is not None:
|
|
if not 1 <= args.sessions <= 64 or not 1 <= args.load_frames <= 10000:
|
|
parser.error("load mode requires sessions in [1,64] and load-frames in [1,10000]")
|
|
for name in ("timeout_s", "cleanup_timeout_s", "drain_s", "tail_s"):
|
|
value = getattr(args, name)
|
|
if not math.isfinite(value) or value < 0 or (name in ("timeout_s", "cleanup_timeout_s") and value == 0):
|
|
parser.error(f"{name} must be finite and non-negative (timeout must be positive)")
|
|
if args.max_client_rtf is not None and (not math.isfinite(args.max_client_rtf) or args.max_client_rtf <= 0):
|
|
parser.error("max-client-rtf must be finite and positive")
|
|
if args.max_frame_deficit < 0 or args.min_voiced_frames < 1 or args.minimum_audio_chunks < 1:
|
|
parser.error("frame-deficit must be non-negative; voiced frames and audio chunks must be positive")
|
|
if not math.isfinite(args.voiced_frame_rms_threshold) or args.voiced_frame_rms_threshold < 0:
|
|
parser.error("voiced-frame-rms-threshold must be finite and non-negative")
|
|
return args
|
|
|
|
|
|
def main() -> None:
|
|
result = asyncio.run(run(parse_args()))
|
|
print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False))
|
|
if result.get("mode") == "paced_multi_session" and not result["ok"]:
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|