diff --git a/.gitignore b/.gitignore index 5f579fff..2c2fcabb 100644 --- a/.gitignore +++ b/.gitignore @@ -274,6 +274,11 @@ traces/ benchmark_output/ benchmark_outputs/ microbench_output* +capacity_search_output/ + +# local/runtime files +.env +configs/ # uv lock file -uv.lock \ No newline at end of file +uv.lock diff --git a/README.md b/README.md index 0beb9e8a..ff52e7fc 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ microbenchmarks. One tool, any workload. **From isolated requests to complex agentic sessions, Veeksha captures the full complexity of modern LLM workloads.** -πŸ‘‰ **[Why Veeksha?](https://project-vajra.github.io/veeksha/getting_started/why_veeksha.html)** β€” Learn what sets Veeksha apart +πŸ‘‰ **[Why Veeksha?](https://project-vajra.github.io/veeksha/getting_started/why_veeksha.html)** β€” Learn what sets Veeksha apart πŸ“š **[Documentation](https://project-vajra.github.io/veeksha)** β€” Full guides and API reference ## Quick start diff --git a/configs/stt_vajra.yaml b/configs/stt_vajra.yaml new file mode 100644 index 00000000..ad0ce07a --- /dev/null +++ b/configs/stt_vajra.yaml @@ -0,0 +1,58 @@ +seed: 42 +output_dir: "benchmark_output/stt_vajra" + +client: + type: stt + provider: vajra_openai_realtime + api_base: "http://localhost:8003" + # Drives Vajra through the OpenAI-compatible realtime transcription WebSocket. + model: mistralai/Voxtral-Mini-4B-Realtime-2602 + sample_rate: 16000 + ws_chunk_size: 2048 + ws_realtime_pacing: true + request_timeout: 7200 + +session_generator: + type: trace + trace_file: traces/asr/aa_public/manifest.jsonl + wrap_mode: true + flavor: + type: audio + audio_dir: "" + +traffic_scheduler: + type: concurrent + target_concurrent_sessions: 8 + rampup_seconds: 0 + +evaluators: + - type: performance + target_channels: ["audio"] + stream_metrics: true + stream_metrics_interval: 10.0 + slos: + # Capacity gate: p99 first-partial latency climbs once the engine can't + # keep up at 1x pacing. (RTF is >=1 by construction when paced.) + - name: "P99 time-to-first-chunk under 500ms" + metric: ttfc + percentile: 0.99 + value: 500 + type: constant + +runtime: + max_sessions: -1 + benchmark_timeout: 300 + post_timeout_grace_seconds: 120 + num_client_threads: 2 + num_dispatcher_threads: 2 + num_completion_threads: 2 + +trace_recorder: + enabled: false + +wandb: + enabled: false + project: "stt-benchmark" + run_name: "stt_vajra_openai_realtime_voxtral" + tags: ["stt", "vajra", "voxtral", "openai_realtime"] + log_artifacts: false diff --git a/configs/stt_vllm_realtime.yaml b/configs/stt_vllm_realtime.yaml new file mode 100644 index 00000000..30f9b06f --- /dev/null +++ b/configs/stt_vllm_realtime.yaml @@ -0,0 +1,57 @@ +seed: 42 +output_dir: "benchmark_output/stt_vllm_realtime" + +client: + type: stt + provider: vllm_realtime + api_base: "http://localhost:8025" + model: mistralai/Voxtral-Mini-4B-Realtime-2602 + sample_rate: 16000 + ws_chunk_size: 2048 + ws_realtime_pacing: true + request_timeout: 7200 + +session_generator: + type: trace + trace_file: traces/asr/aa_public/manifest.jsonl + wrap_mode: true + flavor: + type: audio + audio_dir: "" + +traffic_scheduler: + type: concurrent + target_concurrent_sessions: 8 + rampup_seconds: 0 + +evaluators: + - type: performance + target_channels: ["audio"] + stream_metrics: true + stream_metrics_interval: 10.0 + slos: + # Capacity gate: p99 first-partial latency climbs once the engine can't + # keep up at 1x pacing. (RTF is >=1 by construction when paced.) + - name: "P99 time-to-first-chunk under 500ms" + metric: ttfc + percentile: 0.99 + value: 500 + type: constant + +runtime: + max_sessions: -1 + benchmark_timeout: 300 + post_timeout_grace_seconds: 120 + num_client_threads: 2 + num_dispatcher_threads: 2 + num_completion_threads: 2 + +trace_recorder: + enabled: false + +wandb: + enabled: false + project: "stt-benchmark" + run_name: "stt_vllm_voxtral" + tags: ["stt", "vllm_realtime", "voxtral"] + log_artifacts: false diff --git a/docs/ASR_BENCHMARKING.md b/docs/ASR_BENCHMARKING.md new file mode 100644 index 00000000..ab14283d --- /dev/null +++ b/docs/ASR_BENCHMARKING.md @@ -0,0 +1,215 @@ +# ASR Benchmarking + +## Trace Generation + +Generate the public ASR trace with: + +```bash +.venv/bin/python scripts/prepare_audio_traces.py --clips-per-dataset 128 +``` + +This writes `traces/asr/aa_public/manifest.jsonl` plus WAV files under +`traces/asr/aa_public/audio/`. + +`--max-duration` controls final clip length and defaults to 30 seconds. For +timestamped clips, longer source files are split on word boundaries. With +`--without-word-timestamping`, longer untimestamped clips are skipped. + +```bash +.venv/bin/python scripts/prepare_audio_traces.py \ + --clips-per-dataset 128 \ + --max-duration 30 +``` + +Use `--target-duration` to build clips around a requested duration instead of +splitting by `--max-duration`. Source clips are repeated before NeMo alignment +until they cover the target duration, then the aligned result is truncated to +the requested length. `--target-duration` and `--max-duration` are mutually +exclusive. + +```bash +.venv/bin/python scripts/prepare_audio_traces.py \ + --clips-per-dataset 128 \ + --target-duration 600 +``` + +Word-level reference timings for the interactivity metric are generated by +default with NeMo forced alignment in Docker: + +```bash +.venv/bin/python scripts/prepare_audio_traces.py --clips-per-dataset 128 +``` + +To skip word timestamping: + +```bash +.venv/bin/python scripts/prepare_audio_traces.py \ + --clips-per-dataset 128 \ + --without-word-timestamping +``` + +AMI can be prepared through the same script. The script downloads and caches the +AMI manual word annotations and Mix-Headset WAVs under `benchmark_output/`: + +```bash +.venv/bin/python scripts/prepare_audio_traces.py \ + --datasets ami_word_timed +``` + +This writes `traces/asr/ami_word_timed/manifest.jsonl` plus WAV files under +`traces/asr/ami_word_timed/audio/`. + +The trace uses the public Artificial Analysis cleaned datasets, VoxPopuli and +Earnings22, as a recognizable external point of reference. Earnings22 examples +are kept as full, potentially long requests rather than being chunked. The trace +does not attempt to reproduce the full Artificial Analysis benchmark exactly: +the proprietary AA-AgentTalk dataset is not available, their custom normalizer +is not open sourced, and Veeksha measures behavior across engines so any point +of difference is fine as long as it's consistent for all engines. + +## Request Metrics + +Request-level metrics are written to `request_level_metrics.jsonl`. + +Common audio metrics: + +- `ttfc`: time from the first audio byte sent to the first transcript delta + whose own payload carries text after control-token cleaning (or to the final + transcript when no such delta arrives), in milliseconds. Empty + progress/keepalive deltas do not count. +- `end_to_end_latency`: time from client request start to request completion, in + milliseconds. +- `generated_audio_duration`: audio duration represented by the request, in + milliseconds. For STT this is the input clip duration. +- `rtf`: real-time factor, computed as + `end_to_end_latency / generated_audio_duration`. +- `chunk_count`: number of transcript deltas observed, or `1` when only a final + transcript is returned. +- `input_tokens`: whitespace token count of the final transcript. + +ASR-specific metrics: + +- `time_to_first_visible_text`: time from the first audio byte sent to the + first moment the assembled (concatenated then cleaned) transcript is + non-empty β€” when a user watching the live transcript would first see text β€” + in milliseconds. Usually equal to `ttfc`; the two differ only when cleaning a + single delta disagrees with cleaning the assembled transcript (e.g. a control + token split across deltas). +- `time_to_first_partial`: time from the client sending EOF/commit for the audio + stream to the first non-empty partial transcript received after that point, in + milliseconds. This is only present when the provider emits such a partial. +- `time_to_final_transcript`: time from the client sending EOF/commit for the + audio stream to the final transcript, in milliseconds. +- `partial_transcript`: first non-empty post-EOF partial transcript used for + partial WER. +- `final_transcript`: final transcript returned by the provider. +- `expected_transcript`: reference transcript from the trace row. +- `interactivity`: mean latency between when matched reference words finish in + the audio and when they first appear in the streamed client transcript, in + milliseconds. This requires `reference_word_timestamps` in the manifest and is + most meaningful with realtime audio pacing enabled. +- `interactivity_word_count`: number of matched words used for the per-request + `interactivity` value. +- `partial_wer`: WER for `partial_transcript` against `expected_transcript`. + Present when a partial transcript is available. +- `final_wer`: WER for `final_transcript` against `expected_transcript`. + Present for every completed STT request with a reference transcript. + +## Replay Viewer + +For runs with `reference_word_timestamps` and `transcript_snapshots`, export a +browser replay bundle: + +```bash +.venv/bin/python tools/asr_replay_viewer/export_asr_replay.py \ + benchmark_output/ +``` + +Then serve the repo root and open the viewer: + +```bash +.venv/bin/python -m http.server 8765 --bind 127.0.0.1 +``` + +Open: + +```text +http://127.0.0.1:8765/tools/asr_replay_viewer/?data=benchmark_output//metrics/asr_replay.json +``` + +The viewer shows the ground-truth word timeline, the streamed transcript +timeline, and an interactive latency graph. The graph loads ECharts from a +pinned CDN URL at runtime, so the viewer needs network access when opened. + +## Partials and Finals + +Streaming providers may emit many transcript deltas before the audio stream is +committed. Veeksha concatenates those deltas to track +`time_to_first_visible_text` and the eventual transcript, but `partial_transcript` is specifically the first non-empty +transcript observed after EOF/commit. `final_transcript` is the provider's final +message when available, otherwise the concatenated deltas are used as a fallback. + +## Aggregate Metrics + +Aggregate metrics are emitted in the performance summary. + +General audio aggregates are percentile summaries over request-level values: + +- `ttfc (Mean/P50/P90/P99)` +- `end_to_end_latency (Mean/P50/P90/P99)` +- `generated_audio_duration (Mean/P50/P90/P99)` +- `rtf (Mean/P50/P90/P99)` +- `chunk_count (Mean/P50/P90/P99)` +- `time_to_first_visible_text (Mean/P50/P90/P99)`, for STT runs when available +- `time_to_first_partial (Mean/P50/P90/P99)`, for STT runs when available +- `time_to_final_transcript (Mean/P50/P90/P99)`, for STT runs when available +- `interactivity (Mean/P50/P90/P99)`, for STT runs with word-timestamped + references + +ASR WER aggregates: + +- `asr_final_sample_count`: number of scored ASR request samples. +- `asr_final_sample_mean_wer`: unweighted mean WER across scored samples. +- `asr_final_corpus_wer`: corpus-level WER from summed edit counts and summed + reference words. +- `asr_final_duration_weighted_wer`: WER weighted by audio duration. +- `asr_partial_*`: same aggregate modes for partial transcripts, over samples + where partial WER is available. +- `asr_dataset__final_*` and `asr_dataset__partial_*`: + dataset-specific versions of the same metrics. + +For comparisons between serving engines, prefer the same aggregate metric across +runs, with `asr_final_duration_weighted_wer` as the closest analogue to the +duration-weighted WER used by public ASR leaderboards. + +## Manifest Fields + +Each manifest row represents one audio request. Key fields: + +- `session_id`: trace session identifier. +- `audio_file`: WAV path relative to the manifest directory. +- `dataset`: source dataset key, such as `aa_voxpopuli` or `aa_earnings22`. +- `expected_transcript`: reference transcript used for WER. +- `duration_s`: clip duration in seconds. +- `sample_id`: unique row identifier. +- `reference_word_timestamps`: optional list of reference word timings relative + to the request audio start, e.g. + `{"word": "hello", "start_ms": 120.0, "end_ms": 340.0}`. + +When the audio trace flavor sets `target_duration_s`, each row is trimmed at +generation time: only the first `target_duration_s` seconds of audio are +streamed (the STT client slices the decoded PCM using the +`input_audio_start_ms` / `input_audio_end_ms` request metadata written by the +generator), and `expected_transcript` / `reference_word_timestamps` are trimmed +to the words that end within that prefix. Rows must provide +`reference_word_timestamps`, and every clip must be at least +`target_duration_s` long. + +## References + +- Artificial Analysis Speech to Text Methodology: + https://artificialanalysis.ai/speech-to-text/methodology +- VoxPopuli-Cleaned-AA: + https://huggingface.co/datasets/ArtificialAnalysis/VoxPopuli-Cleaned-AA +- Earnings22-Cleaned-AA: + https://huggingface.co/datasets/ArtificialAnalysis/Earnings22-Cleaned-AA diff --git a/docs/getting_started/installation.rst b/docs/getting_started/installation.rst index 2721deef..e8739b60 100644 --- a/docs/getting_started/installation.rst +++ b/docs/getting_started/installation.rst @@ -1,9 +1,10 @@ Installation ============ -You can run Veeksha without explicit installation using ``uvx``.: +You can run Veeksha without explicit installation using ``uvx``: .. code-block:: bash + uvx -p 3.14t veeksha benchmark ... Veeksha requires free-threaded Python ``>=3.14``. diff --git a/docs/index.rst b/docs/index.rst index ddd76e84..491af245 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -31,7 +31,7 @@ Key features as directed acyclic graphs with history inheritance, capturing real chat context accumulation - **Shared prefix testing**: Generate workloads with configurable prefix sharing to benchmark KV-cache efficiency - - **Trace replay**: Replay production traces (Claude Code, RAG, conversational) with + - **Trace replay**: Replay production traces (coding assistants, RAG, conversational) with preserved timing and token distributions **Flexible traffic generation** diff --git a/docs/user_guide/index.rst b/docs/user_guide/index.rst index 2e09739a..696d0c85 100644 --- a/docs/user_guide/index.rst +++ b/docs/user_guide/index.rst @@ -10,6 +10,7 @@ integrating with external tools. configuration output_files trace_flavors + tts_benchmarking capacity_search sweeps server_management diff --git a/docs/user_guide/tts_benchmarking.rst b/docs/user_guide/tts_benchmarking.rst new file mode 100644 index 00000000..24791c9b --- /dev/null +++ b/docs/user_guide/tts_benchmarking.rst @@ -0,0 +1,297 @@ +Streaming TTS benchmarks +======================== + +Veeksha measures TTS from the client: it paces text into the provider, records +when decoded PCM becomes playable, and derives every latency and continuity +metric from one monotonic request timeline. Provider clocks are not used. + +Benchmark structure +------------------- + +The TTS path has five layers: + +1. A trace flavor turns source text into one-request TTS sessions. +2. The traffic scheduler controls arrival rate or concurrent sessions. +3. A TTS client implements the provider protocol and records text and PCM events. +4. The audio performance evaluator derives latency, RTF, overlap, stalls, and + fluidity from those events. +5. The audio quality evaluator can save WAV files and compute WER and UTMOS. + +The provider adapters intentionally retain their native lifecycle: + +.. list-table:: + :header-rows: 1 + + * - ``client.type`` + - Transport + - Input behavior + * - ``realtime_tts`` + - OpenAI Realtime-compatible WebSocket + - Complete-text or duplex ``response.create`` scheduling + * - ``elevenlabs_streaming_tts`` + - ElevenLabs ``/stream-input`` WebSocket + - Paced partial text followed by the provider finalization message + * - ``deepgram_flux_streaming_tts`` + - Deepgram Flux ``/v2/speak`` WebSocket + - Paced ``Speak`` messages; audio may arrive before ``Flush`` + * - ``deepgram_aura_streaming_tts`` + - Deepgram Aura ``/v1/speak`` WebSocket + - Paced ``Speak`` messages followed by ``Flush`` + * - ``elevenlabs_http_tts`` + - ElevenLabs complete-response HTTP + - All text and then all audio + * - ``deepgram_flux_http_tts`` + - Deepgram Flux ``POST /v2/speak`` + - All text and then all audio + +The native protocols follow the vendor specifications for `ElevenLabs +stream-input `_, +`Deepgram Flux `_, +and `Deepgram Aura `_. + +Metrics +------- + +.. list-table:: + :header-rows: 1 + + * - Metric + - Meaning + * - ``trigger_to_first_playable_audio_ms`` + - Synthesis trigger to the first complete PCM playback frame on the active + connection. This is the primary steady-state TTFA: ``response.create`` + for Realtime TTS and the first real synthesis-eligible text message for + native streaming APIs. Protocol setup messages are excluded. + * - ``first_input_to_first_playable_audio_ms`` + - First real streamed text delta to the first complete PCM playback frame. + Unlike trigger TTFA, this intentionally exposes any client or provider + lookahead before synthesis is triggered. + * - ``request_start_to_first_playable_audio_ms`` + - WebSocket-connect initiation to the first complete PCM playback frame. + Report this separately as cold-session or connection-inclusive latency. + * - ``ttfc`` + - Request start to the first wire audio chunk. A tiny partial chunk may not + yet be playable, so this is retained as a transport diagnostic. + * - ``rtf`` + - End-to-end request time divided by generated audio duration. It includes + paced upstream text input. + * - ``streaming_rtf`` + - Wall time from the first to last audio arrival divided by audio delivered + after the first chunk. It isolates output delivery after startup. + * - ``audio_before_commit_ratio`` + - Fraction of output PCM received before the final text input was sent. + * - ``duplex_overlap_observed`` + - Whether a complete playable frame arrived before final text input. + * - ``required_startup_delay_ms`` + - Smallest fixed playback delay that would eliminate all underruns in the + captured stream. + * - ``zero_delay_*`` + - Stall count and duration when playback starts immediately. + * - ``user_audio_fluidity_index`` + - Fraction of accepted fixed-frame playback deadlines, including stalls + caused anywhere in the end-to-end path. + * - ``tts_service_fluidity_index`` + - The same score only when the timeline provides enough evidence to blame + misses on TTS rather than missing upstream text. + +Fluidity is inspired by the deadline, slack, and reset semantics in `Etalon +`_. Veeksha first converts raw PCM supply into +fixed playback frames (20 ms by default). Early frames accumulate playable +buffer. A late frame consumes that buffer; if it is still late, every elapsed +20 ms playback deadline is a miss and the buffer resets. The score is accepted +deadlines divided by all deadlines. + +The primary score defaults to zero artificial startup delay. Configure +``fluidity_startup_delay_ms`` only when the real playback client deliberately +buffers by that amount, and always report the delay with the score. Veeksha also +emits policy-specific scores such as +``user_audio_fluidity_index_d100ms`` when 100 ms is included in +``startup_delay_ms_values``. + +In duplex mode, no provider-independent method can know whether a silent period +means that TTS stalled or that the upstream LLM supplied no synthesis-eligible +text. Therefore: + +- ``user_audio_fluidity_index`` is always the observed user experience. +- In ``conservative`` attribution mode, service fluidity is emitted only when + all text arrived before playback began. +- ``source_oversupplied`` may be used only by a controlled workload that + guarantees enough eligible text throughout playback. + +Batch HTTP output has all audio buffered when playback begins, so its fluidity +is trivially one. Compare batch models on first-playable latency, end-to-end +latency, RTF, cost, and quality; do not use batch fluidity to rank streaming +behavior. + +Trace sources +------------- + +For publishable TTS quality runs, use ``seed_tts_text``. Its current default is +the English split of the ``TwinkStart/Seed-TTS-Eval`` Hugging Face mirror. Each +row supplies target synthesis text; Veeksha records the dataset, subset, split, +and source row in request metadata. The original benchmark is maintained by +`BytedanceSpeech/seed-tts-eval +`_. Pin or locally archive the +exact dataset revision used in a published comparison. + +``sharegpt`` is also supported when a local ShareGPT-format JSON/JSONL file is +provided. It extracts assistant turns as TTS text. Veeksha does not ship a +ShareGPT dataset. + +There is no bundled Claude Code trace and no Claude-specific trace flavor. +``timed_synthetic_session`` can replay a privacy-safe coding-assistant trace of +token lengths, dependencies, and think times, but that is an LLM serving +workloadβ€”not a canonical TTS quality corpus. + +Run a benchmark +--------------- + +This complete example uses the Seed-TTS text trace and ElevenLabs streaming. +Set ``ELEVENLABS_API_KEY`` in the environment and replace ``voice_id``. + +.. code-block:: yaml + + seed: 42 + output_dir: benchmark_output/tts_elevenlabs_streaming + + client: + type: elevenlabs_streaming_tts + api_base: https://api.elevenlabs.io + model: eleven_flash_v2_5 + voice_id: YOUR_VOICE_ID + api_key_env: ELEVENLABS_API_KEY + sample_rate: 24000 + pacing: + tokens_per_second: 20 + tokens_per_delta: 1 + gap_distribution: fixed + + session_generator: + type: trace + wrap_mode: false + flavor: + type: seed_tts_text + min_tokens: 20 + max_tokens: 150 + + traffic_scheduler: + type: concurrent + target_concurrent_sessions: 1 + rampup_seconds: 0 + + evaluators: + - type: performance + target_channels: [audio] + audio_channel: + interactivity_enabled: true + fluidity_frame_ms: 20 + fluidity_startup_delay_ms: 0 + startup_delay_ms_values: [0, 100, 300] + fluidity_attribution_mode: conservative + persist_raw_timing: true + slos: + - type: constant + name: P90 trigger-to-first-playable audio under 1 second + metric: trigger_to_first_playable_audio_ms + percentile: 0.90 + value: 1000 + - type: constant + name: P1 user fluidity at least 0.99 + metric: user_audio_fluidity_index + percentile: 0.01 + value: 0.99 + - type: audio_quality + target_channels: [audio] + save_audio_files: true + verification: + wer: + enabled: true + threshold: 0.05 + whisper: + model: large-v3 + device: cpu + compute_type: int8 + language: en + beam_size: 5 + utmos: + enabled: false + + runtime: + max_sessions: 100 + benchmark_timeout: 1800 + +Run it with: + +.. code-block:: console + + uvx -p 3.14t veeksha benchmark --config tts_benchmark.veeksha.yml + +WER requires the optional ``audio-verification`` dependencies (including +faster-whisper). UTMOS requires its corresponding optional dependency group. +Install those groups in the Veeksha environment before enabling the quality +checks. + +Correctness metric contract +--------------------------- + +ASR and TTS WER intentionally use different published protocols and therefore +different units: + +- ASR ``final_wer`` and aggregate ``asr_*_wer`` fields are percentages in the + range 0--100 for ordinary cases. Text is normalized with the Open ASR + Leaderboard English normalizer. Prefer corpus WER for the primary comparison; + sample-mean and duration-weighted WER are reported as diagnostics. +- TTS verification ``wer`` fields are ratios where 0.05 means five percent. + They use Seed-TTS-style punctuation normalization and a configured + faster-whisper judge. Keep the judge checkpoint, language, beam size, device, + compute type, corpus revision, and text normalization identical across every + provider. This WER measures intelligibility, not naturalness, speaker + similarity, emotion, or human preference. +- UTMOS is a predicted naturalness score. Treat it as a scalable regression + signal, not a substitute for a blinded human preference study. + +With ``fail_on_threshold: true``, verification fails closed: a WER threshold +violation, missing audio file, transcription failure, unavailable UTMOS model, +or run-level verification error fails the benchmark rather than silently +removing that request from the quality sample. + +For Deepgram Flux streaming, replace only the client block: + +.. code-block:: yaml + + client: + type: deepgram_flux_streaming_tts + api_base: https://api.deepgram.com + model: flux-alexis-en + api_key_env: DEEPGRAM_API_KEY + sample_rate: 24000 + pacing: + tokens_per_second: 20 + tokens_per_delta: 1 + gap_distribution: fixed + +Use ``deepgram_aura_streaming_tts`` with an Aura model for the ``/v1/speak`` +lane. Use ``elevenlabs_http_tts`` or ``deepgram_flux_http_tts`` for the +complete-text/complete-audio controls. For Vajra, use ``realtime_tts`` and set +``input_output_mode: duplex``; the server must consume conversation items added +after an active ``response.create``. + +Do not treat ``response.create`` ordering or output overlap alone as proof of +semantic duplex synthesis. A conforming run must also pass full-reference TTS +WER (or a stronger text-coverage check) so a server that speaks only the prefix +available at trigger time cannot be reported as successful streaming. The +``duplex_start_after_tokens`` value is a configurable workload threshold, not a +special eight-token protocol rule. + +Keep the trace seed, input texts, pacing, PCM format, region, concurrency sweep, +and retry policy identical across providers. Report request failures and rate +limits rather than silently retrying them away. + +Outputs +------- + +The performance evaluator writes aggregate summaries, request-level JSONL, CDF +CSVs/plots, and optional ``audio_raw_timing.jsonl``. The quality evaluator writes +WAV files and verification summaries. Provider API keys, local ``.env`` files, +downloaded traces, and benchmark output directories are not source artifacts and +must not be committed. diff --git a/pyproject.toml b/pyproject.toml index f9ad4299..56d3ceb6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "numpy", "jinja2", "datasets", - "lm-eval", + "lm-eval>=0.4.11,<0.4.12", "evaluate", "pytablewriter", "Pillow", @@ -40,10 +40,16 @@ dependencies = [ "nvidia-ml-py", "rich", "vidhi", + "websockets>=14", + "jiwer", + "regex", + "soundfile", + "librosa", ] [project.scripts] veeksha = "veeksha.__main__:main" +"veeksha-launcher" = "veeksha.launcher:main" [project.urls] Homepage = "https://github.com/project-vajra/veeksha" @@ -65,6 +71,16 @@ test = [ "rich", "jinja2", ] +audio-verification = [ + "faster-whisper", + "jiwer", +] +utmos = [ + "torch", + "scipy", + "soundfile", + "huggingface_hub", +] [tool.setuptools] include-package-data = true @@ -94,6 +110,7 @@ known_third_party = "wandb" extend_skip = ["veeksha/_version.py"] [tool.pyright] +pythonVersion = "3.14" include = ["veeksha"] exclude = [ "**/__pycache__", diff --git a/scripts/prepare_audio_traces.py b/scripts/prepare_audio_traces.py new file mode 100755 index 00000000..9d44e188 --- /dev/null +++ b/scripts/prepare_audio_traces.py @@ -0,0 +1,1291 @@ +#!/usr/bin/env python3 +"""Build portable ASR traces from public datasets. + +Examples: + .venv/bin/python scripts/prepare_audio_traces.py --clips-per-dataset 128 + + .venv/bin/python scripts/prepare_audio_traces.py \ + --datasets aa_voxpopuli,aa_earnings22 \ + --clips-per-dataset 128 \ + --max-duration 30 + + .venv/bin/python scripts/prepare_audio_traces.py \ + --clips-per-dataset 128 \ + --without-word-timestamping + + .venv/bin/python scripts/prepare_audio_traces.py \ + --datasets ami_word_timed \ + --clips-per-dataset 128 +""" + +from __future__ import annotations + +import argparse +import io +import json +import math +import os +import re +import shutil +import subprocess +import sys +import urllib.request +import xml.etree.ElementTree as ET +import zipfile +from abc import ABC, abstractmethod +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any, Iterable, Sequence + +import librosa +import numpy as np +import soundfile as sf +from datasets import load_dataset +from huggingface_hub import hf_hub_download + +######################################################################## +# Constants +######################################################################## + +REPO_ROOT = Path(__file__).resolve().parent.parent +TRACES_ROOT = REPO_ROOT / "traces" +AA_TRACE_OUTPUT_DIR = TRACES_ROOT / "asr" / "aa_public" +AMI_TRACE_OUTPUT_DIR = TRACES_ROOT / "asr" / "ami_word_timed" +DEFAULT_DATASETS = "aa_voxpopuli,aa_earnings22" +MANIFEST_NAME = "manifest.jsonl" + +DEFAULT_SAMPLE_RATE = 16000 +DEFAULT_MIN_DURATION_S = 0.25 +DEFAULT_MAX_DURATION_S = 30.0 +DEFAULT_SHUFFLE_BUFFER = 0 +DEFAULT_SEED = 42 +TARGET_DURATION_DEVIATION_FRACTION = 0.10 +AMI_MAX_GAP_S = 1.5 +AMI_CACHE_DIR = REPO_ROOT / "benchmark_output" / "ami_cache" +AMI_BASE_URL = "https://groups.inf.ed.ac.uk/ami" +AMI_ANNOTATIONS_ARCHIVE = "ami_public_manual_1.6.2.zip" +AMI_ANNOTATIONS_DIR = "ami_public_manual_1.6.2" +AMI_MIX_HEADSET = "{meeting_id}.Mix-Headset.wav" +AMI_DATASET_KEY = "ami_word_timed" +AA_DATASET_KEYS = frozenset(("aa_voxpopuli", "aa_earnings22")) + +NEMO_MODEL = "stt_en_fastconformer_hybrid_large_pc" +NEMO_DOCKER_IMAGE = "nvcr.io/nvidia/nemo:26.02" +NEMO_DOCKER_GPUS = "all" +NEMO_CACHE_DIR = REPO_ROOT / "benchmark_output" / "nemo_docker_cache" +NEMO_EXTRA_MOUNTS: tuple[str, ...] = () +NEMO_ALIGN_SCRIPT = "/opt/NeMo/tools/nemo_forced_aligner/align.py" +NEMO_MANIFEST_NAME = "nemo_manifest.jsonl" +NEMO_OUTPUT_DIR_NAME = "nemo_output" +NEMO_SOURCE_AUDIO_DIR_NAME = "source_audio" + +NEMO_CONTAINER_SCRIPT = r""" +set -euo pipefail + +cleanup() { + paths=("$NEMO_DOCKER_CACHE_DIR" "$NEMO_ALIGNMENT_OUTPUT_DIR") + chown -R "$HOST_UID:$HOST_GID" "${paths[@]}" 2>/dev/null || true +} +trap cleanup EXIT + +python /opt/NeMo/tools/nemo_forced_aligner/align.py "$@" +""" + +DATASETS: dict[str, dict[str, str]] = { + "aa_voxpopuli": { + "repo": "ArtificialAnalysis/VoxPopuli-Cleaned-AA", + "split": "test", + }, + "aa_earnings22": { + "repo": "ArtificialAnalysis/Earnings22-Cleaned-AA", + "split": "test", + }, + "ami_word_timed": { + "repo": "AMI local word XML", + "split": "local", + }, +} + +######################################################################## +# Data Models +######################################################################## + + +@dataclass(frozen=True) +class TraceSourceOptions: + sample_rate: int = DEFAULT_SAMPLE_RATE + min_duration_s: float = DEFAULT_MIN_DURATION_S + max_duration_s: float | None = None + shuffle_buffer: int = DEFAULT_SHUFFLE_BUFFER + seed: int = DEFAULT_SEED + + +@dataclass(frozen=True) +class WordTiming: + word: str + start_ms: float + end_ms: float + + +@dataclass(frozen=True) +class TraceClip: + audio: np.ndarray + transcript: str + duration_s: float + metadata: dict[str, Any] + word_timestamps: list[WordTiming] | None = None + target_duration_s: float | None = None + + +@dataclass(frozen=True) +class NemoItem: + row_index: int + audio_path: Path + text: str + + +######################################################################## +# CLI +######################################################################## + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + dataset_keys = supported_dataset_keys() + parser.add_argument( + "--datasets", + default=DEFAULT_DATASETS, + help=f"Comma-separated subset to build. Supported: {', '.join(dataset_keys)}.", + ) + parser.add_argument( + "--clips-per-dataset", + type=int, + default=16, + help="Requested final clips per dataset. Use 0 for all clips.", + ) + duration_group = parser.add_mutually_exclusive_group() + duration_group.add_argument( + "--max-duration", + type=float, + default=None, + help=( + "Maximum final clip duration in seconds. Timestamped clips are " + "split on word boundaries; untimestamped longer clips are skipped." + ), + ) + duration_group.add_argument( + "--target-duration", + type=float, + default=None, + help=( + "Target final clip duration in seconds. Source clips are repeated " + "before NeMo alignment and truncated after alignment." + ), + ) + parser.add_argument( + "--without-word-timestamping", + action="store_true", + help="Skip NeMo word timestamping.", + ) + args = parser.parse_args() + if args.target_duration is None and args.max_duration is None: + args.max_duration = DEFAULT_MAX_DURATION_S + return args + + +def validate_args(args: argparse.Namespace) -> list[str]: + if args.clips_per_dataset < 0: + raise SystemExit("--clips-per-dataset must be >= 0") + if args.max_duration is not None and args.max_duration <= 0: + raise SystemExit("--max-duration must be positive when set") + if args.target_duration is not None and args.target_duration <= 0: + raise SystemExit("--target-duration must be positive when set") + if args.target_duration is not None and args.without_word_timestamping: + raise SystemExit("--target-duration requires word timestamping") + dataset_keys = selected_dataset_keys(args.datasets) + output_dir_for_dataset_keys(dataset_keys) + return dataset_keys + + +######################################################################## +# Dataset Sources +######################################################################## + + +class ASRTraceSource(ABC): + """Base class for datasets that produce source ASR clips.""" + + def __init__(self, key: str, options: TraceSourceOptions) -> None: + self.key = key + self.options = options + self.spec = DATASETS[key] + + @property + def repo(self) -> str: + return self.spec["repo"] + + @property + def split(self) -> str: + return self.spec["split"] + + @abstractmethod + def iter_clips(self) -> Iterable[TraceClip]: + """Yield source clips before optional timestamping and chunking.""" + + +class AATraceSource(ASRTraceSource): + """Artificial Analysis public cleaned ASR datasets.""" + + def iter_clips(self) -> Iterable[TraceClip]: + dataset = load_dataset(self.repo, split=self.split, streaming=True) + if self.options.shuffle_buffer > 0: + dataset = dataset.shuffle( + seed=self.options.seed, + buffer_size=self.options.shuffle_buffer, + ) + + for row_index, sample in enumerate(dataset): + transcript = clean_text(sample.get("transcript")) + if not transcript: + continue + + try: + audio = decode_audio(fetch_aa_audio(sample, self.repo), self.options) + except Exception as exc: + print( + f" WARNING: skipping {self.key} row {row_index}: {exc}", + file=sys.stderr, + ) + continue + + duration_s = len(audio) / float(self.options.sample_rate) + if duration_s < self.options.min_duration_s: + continue + + yield self._build_clip( + audio=audio, + transcript=transcript, + duration_s=duration_s, + row={"row_index": row_index, **sample}, + ) + + def _build_clip( + self, + *, + audio: np.ndarray, + transcript: str, + duration_s: float, + row: dict[str, Any], + ) -> TraceClip: + sample_id = source_id(row) + return TraceClip( + audio=audio, + transcript=transcript, + duration_s=duration_s, + metadata={ + "source_dataset": self.repo, + "source_split": self.split, + "source_id": sample_id, + "sample_id": sample_id, + }, + ) + + +class AMITraceSource(ASRTraceSource): + """Local AMI audio plus AMI ``*.words.xml`` word timing annotations.""" + + AUDIO_DIR = "" + WORDS_DIR = "" + CACHE_DIR = AMI_CACHE_DIR + BASE_URL = AMI_BASE_URL + AUDIO_GLOB = "{meeting_id}*.wav" + + def iter_clips(self) -> Iterable[TraceClip]: + audio_dir = self.audio_dir + words_dir = self.words_dir + word_files = sorted(words_dir.rglob("*.words.xml")) + if not word_files: + raise SystemExit(f"No AMI *.words.xml files found under {words_dir}") + + audio_cache: dict[Path, np.ndarray] = {} + for words_file in word_files: + meeting_id, speaker_id = parse_ami_ids(words_file) + try: + audio_path = self.find_audio_file(audio_dir, meeting_id, speaker_id) + words = parse_ami_words(words_file) + if not words: + continue + full_audio = audio_cache.get(audio_path) + if full_audio is None: + full_audio = decode_audio(str(audio_path), self.options) + audio_cache[audio_path] = full_audio + except Exception as exc: + print(f" WARNING: skipping {words_file}: {exc}", file=sys.stderr) + continue + + for clip_index, clip_words in enumerate( + chunk_timed_words( + words, + max_duration_s=( + self.options.max_duration_s or DEFAULT_MAX_DURATION_S + ), + max_gap_s=AMI_MAX_GAP_S, + ) + ): + clip = self._build_clip( + full_audio=full_audio, + meeting_id=meeting_id, + speaker_id=speaker_id, + clip_index=clip_index, + words=clip_words, + ) + if clip is not None: + yield clip + + def _build_clip( + self, + *, + full_audio: np.ndarray, + meeting_id: str, + speaker_id: str, + clip_index: int, + words: list[WordTiming], + ) -> TraceClip | None: + clip_start_ms = words[0].start_ms + clip_end_ms = words[-1].end_ms + duration_s = (clip_end_ms - clip_start_ms) / 1000 + if duration_s < self.options.min_duration_s: + return None + + start_sample = ms_to_sample(clip_start_ms, self.options.sample_rate) + end_sample = min( + len(full_audio), + ms_to_sample(clip_end_ms, self.options.sample_rate), + ) + if end_sample <= start_sample: + return None + + source = f"{meeting_id}:{speaker_id}:{clip_index}" + relative_words = [ + WordTiming( + word=word.word, + start_ms=round(word.start_ms - clip_start_ms, 3), + end_ms=round(word.end_ms - clip_start_ms, 3), + ) + for word in words + ] + return TraceClip( + audio=full_audio[start_sample:end_sample], + transcript=" ".join(word.word for word in words), + duration_s=duration_s, + metadata={ + "source_dataset": "AMI", + "source_split": "local", + "source_id": source, + "sample_id": source, + "meeting_id": meeting_id, + "speaker_id": speaker_id, + }, + word_timestamps=relative_words, + ) + + @property + def audio_dir(self) -> Path: + if self.AUDIO_DIR: + return Path(self.AUDIO_DIR) + return self.CACHE_DIR / "wav_db" + + @property + def words_dir(self) -> Path: + if self.WORDS_DIR: + return Path(self.WORDS_DIR) + return self.ensure_annotations() + + def find_audio_file( + self, + audio_dir: Path, + meeting_id: str, + speaker_id: str, + ) -> Path: + if not self.AUDIO_DIR: + return self.ensure_meeting_audio(meeting_id) + + glob_pattern = self.AUDIO_GLOB.format( + meeting_id=meeting_id, + speaker_id=speaker_id, + ) + matches = sorted(audio_dir.glob(glob_pattern)) + if not matches: + raise FileNotFoundError( + f"No AMI audio matched {glob_pattern!r} under {audio_dir}" + ) + return matches[0] + + def ensure_annotations(self) -> Path: + words_dir = self.find_cached_words_dir() + if words_dir.exists() and any(words_dir.glob("*.words.xml")): + return words_dir + + archive_path = self.CACHE_DIR / AMI_ANNOTATIONS_ARCHIVE + download_file( + f"{self.BASE_URL}/AMICorpusAnnotations/{AMI_ANNOTATIONS_ARCHIVE}", + archive_path, + ) + safe_extract_zip(archive_path, self.CACHE_DIR) + words_dir = self.find_cached_words_dir() + if words_dir.exists() and any(words_dir.glob("*.words.xml")): + return words_dir + raise FileNotFoundError("Downloaded AMI annotations did not contain words XML") + + def find_cached_words_dir(self) -> Path: + candidates = [ + self.CACHE_DIR / "words", + self.CACHE_DIR / AMI_ANNOTATIONS_DIR / "words", + ] + for candidate in candidates: + if candidate.exists() and any(candidate.rglob("*.words.xml")): + return candidate + return candidates[0] + + def ensure_meeting_audio(self, meeting_id: str) -> Path: + wav_name = AMI_MIX_HEADSET.format(meeting_id=meeting_id) + wav_path = self.CACHE_DIR / "wav_db" / meeting_id / "audio" / wav_name + if wav_path.exists(): + return wav_path + + download_file( + f"{self.BASE_URL}/AMICorpusMirror/amicorpus/" + f"{meeting_id}/audio/{wav_name}", + wav_path, + ) + return wav_path + + +def selected_dataset_keys(raw_datasets: str) -> list[str]: + keys = [key.strip() for key in raw_datasets.split(",") if key.strip()] + if not keys: + raise SystemExit("--datasets must include at least one dataset") + unknown = [key for key in keys if key not in DATASETS] + if unknown: + raise SystemExit( + f"Unknown dataset key(s): {unknown}. Supported: {', '.join(DATASETS)}" + ) + return keys + + +def output_dir_for_dataset_keys(dataset_keys: Sequence[str]) -> Path: + key_set = set(dataset_keys) + if key_set == {AMI_DATASET_KEY}: + return AMI_TRACE_OUTPUT_DIR + if key_set.issubset(AA_DATASET_KEYS): + return AA_TRACE_OUTPUT_DIR + raise SystemExit( + "Run AMI separately from Artificial Analysis datasets: " + f"{AMI_DATASET_KEY} writes to {AMI_TRACE_OUTPUT_DIR}, while " + f"{', '.join(sorted(AA_DATASET_KEYS))} write to {AA_TRACE_OUTPUT_DIR}." + ) + + +def supported_dataset_keys() -> list[str]: + return list(DATASETS) + + +def build_trace_source(key: str, options: TraceSourceOptions) -> ASRTraceSource: + if key in AA_DATASET_KEYS: + return AATraceSource(key, options) + if key == AMI_DATASET_KEY: + return AMITraceSource(key, options) + raise ValueError(f"Unsupported ASR trace source: {key!r}") + + +######################################################################## +# Audio And Text Helpers +######################################################################## + + +def clean_text(value: Any) -> str: + return re.sub(r"\s+", " ", str(value or "")).strip() + + +def source_id(row: dict[str, Any]) -> str: + for key in ("id", "audio_id", "file", "file_name", "path", "url"): + value = row.get(key) + if value: + return str(value) + return f"source-{row['row_index']}" + + +def fetch_aa_audio(row: dict[str, Any], repo: str) -> str | bytes: + url = row.get("url") + if not url: + raise ValueError("row has no url column") + + url = str(url) + if url.startswith(("http://", "https://")): + with urllib.request.urlopen(url, timeout=60) as response: + return response.read() + + return hf_hub_download(repo_id=repo, repo_type="dataset", filename=url) + + +def download_file(url: str, target: Path) -> None: + if target.exists(): + return + + target.parent.mkdir(parents=True, exist_ok=True) + partial = target.with_name(f"{target.name}.part") + print(f" Downloading {url}") + try: + with urllib.request.urlopen(url, timeout=120) as response: + with partial.open("wb") as out: + shutil.copyfileobj(response, out) + partial.replace(target) + finally: + if partial.exists() and not target.exists(): + partial.unlink() + + +def safe_extract_zip(archive_path: Path, output_dir: Path) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + output_root = output_dir.resolve() + with zipfile.ZipFile(archive_path) as archive: + for member in archive.infolist(): + target = (output_root / member.filename).resolve() + if not is_relative_to(target, output_root): + raise ValueError( + f"Refusing to extract unsafe zip path: {member.filename}" + ) + archive.extractall(output_root) + + +def decode_audio(source: str | bytes, options: TraceSourceOptions) -> np.ndarray: + audio_input: str | io.BytesIO + audio_input = io.BytesIO(source) if isinstance(source, bytes) else source + + try: + audio, sample_rate = sf.read(audio_input, dtype="float32") + except Exception: + audio, sample_rate = librosa.load(audio_input, sr=None, mono=False) + + if audio.ndim > 1: + axis = 0 if audio.shape[0] <= 8 and audio.shape[0] < audio.shape[-1] else 1 + audio = audio.mean(axis=axis) + return resample(np.asarray(audio, dtype=np.float32), int(sample_rate), options) + + +def resample( + audio: np.ndarray, + source_sample_rate: int, + options: TraceSourceOptions, +) -> np.ndarray: + if source_sample_rate == options.sample_rate: + return audio.astype(np.float32, copy=False) + + return librosa.resample( + audio.astype(np.float32, copy=False), + orig_sr=source_sample_rate, + target_sr=options.sample_rate, + ).astype(np.float32) + + +def ms_to_sample(ms: float, sample_rate: int) -> int: + return max(0, int(round(ms * sample_rate / 1000))) + + +######################################################################## +# AMI Parsing +######################################################################## + + +def parse_ami_ids(words_file: Path) -> tuple[str, str]: + parts = words_file.name.split(".") + meeting_id = parts[0] + speaker_id = parts[1] if len(parts) > 2 else "" + return meeting_id, speaker_id + + +def parse_ami_words(words_file: Path) -> list[WordTiming]: + root = ET.parse(words_file).getroot() + words: list[WordTiming] = [] + for elem in root.iter(): + if elem.tag.rsplit("}", 1)[-1] != "w": + continue + text = clean_text("".join(elem.itertext())) + if not text: + continue + start_s = attr_float(elem, "starttime") + end_s = attr_float(elem, "endtime") + if start_s is None or end_s is None: + raise ValueError(f"AMI word missing start/end timing in {words_file}") + words.append( + WordTiming( + word=text, + start_ms=round(start_s * 1000, 3), + end_ms=round(end_s * 1000, 3), + ) + ) + return sorted(words, key=lambda word: (word.start_ms, word.end_ms)) + + +def attr_float(elem: ET.Element, suffix: str) -> float | None: + for key, value in elem.attrib.items(): + if key.rsplit("}", 1)[-1].lower() == suffix: + return float(value) + return None + + +######################################################################## +# Word Timestamping +######################################################################## + + +class NeMoDockerWordTimestampProvider: + def __init__(self, aligner: "NeMoDockerAligner | None" = None) -> None: + self.aligner = aligner or NeMoDockerAligner() + + def annotate( + self, + *, + dataset_key: str, + clips: list[TraceClip], + alignment_output_dir: Path, + options: TraceSourceOptions, + ) -> list[TraceClip]: + if not clips or all(clip.word_timestamps is not None for clip in clips): + return clips + + dataset_alignment_dir = alignment_output_dir / dataset_key + source_audio_dir = dataset_alignment_dir / NEMO_SOURCE_AUDIO_DIR_NAME + source_audio_dir.mkdir(parents=True, exist_ok=True) + nemo_manifest = dataset_alignment_dir / NEMO_MANIFEST_NAME + nemo_output_dir = dataset_alignment_dir / NEMO_OUTPUT_DIR_NAME + + items = write_nemo_inputs( + clips=clips, + source_audio_dir=source_audio_dir, + nemo_manifest=nemo_manifest, + options=options, + ) + self.aligner.run( + manifest_path=nemo_manifest, + output_dir=nemo_output_dir, + alignment_output_dir=dataset_alignment_dir, + ) + word_timings = read_nemo_word_timings(nemo_output_dir, nemo_manifest, items) + + annotated: list[TraceClip] = [] + for index, clip in enumerate(clips): + if clip.word_timestamps is not None: + annotated.append(clip) + else: + annotated.append(replace(clip, word_timestamps=word_timings[index])) + return annotated + + +class NeMoDockerAligner: + """Runs NeMo forced alignment inside NVIDIA's Docker image.""" + + def __init__( + self, + *, + image: str = NEMO_DOCKER_IMAGE, + gpus: str = NEMO_DOCKER_GPUS, + cache_dir: Path = NEMO_CACHE_DIR, + extra_mounts: Sequence[str] = NEMO_EXTRA_MOUNTS, + repo_root: Path = REPO_ROOT, + ) -> None: + self.image = image + self.gpus = gpus + self.cache_dir = cache_dir.resolve() + self.extra_mounts = tuple(mount for mount in extra_mounts if mount) + self.repo_root = repo_root.resolve() + self.host_uid = os.getuid() + self.host_gid = os.getgid() + + def run( + self, + *, + manifest_path: Path, + output_dir: Path, + alignment_output_dir: Path, + ) -> None: + self.prepare_cache_dir() + subprocess.run( + self.command( + manifest_path=manifest_path, + output_dir=output_dir, + alignment_output_dir=alignment_output_dir, + ), + check=True, + ) + + def prepare_cache_dir(self) -> None: + for subdir in ("home", "hf", "torch", "torchinductor", "tmp"): + (self.cache_dir / subdir).mkdir(parents=True, exist_ok=True) + + def command( + self, + *, + manifest_path: Path, + output_dir: Path, + alignment_output_dir: Path, + ) -> list[str]: + command = [ + "docker", + "run", + "--rm", + "--gpus", + self.gpus, + "--ipc=host", + "--ulimit", + "memlock=-1", + "--ulimit", + "stack=67108864", + "-e", + f"HOST_UID={self.host_uid}", + "-e", + f"HOST_GID={self.host_gid}", + "-e", + f"HOME={self.cache_dir / 'home'}", + "-e", + f"HF_HOME={self.cache_dir / 'hf'}", + "-e", + f"TORCH_HOME={self.cache_dir / 'torch'}", + "-e", + f"TORCHINDUCTOR_CACHE_DIR={self.cache_dir / 'torchinductor'}", + "-e", + f"TMPDIR={self.cache_dir / 'tmp'}", + "-e", + f"NEMO_DOCKER_CACHE_DIR={self.cache_dir}", + "-e", + f"NEMO_ALIGNMENT_OUTPUT_DIR={alignment_output_dir}", + "-v", + f"{self.repo_root}:{self.repo_root}", + "-w", + str(self.repo_root), + ] + + if not is_relative_to(self.cache_dir, self.repo_root): + command.extend(["-v", f"{self.cache_dir}:{self.cache_dir}"]) + for mount in self.extra_mounts: + command.extend(["-v", mount]) + + return command + [ + self.image, + "bash", + "-lc", + NEMO_CONTAINER_SCRIPT, + "nemo_align", + f"manifest_filepath={manifest_path}", + f"output_dir={output_dir}", + 'save_output_file_formats=["ctm"]', + f"pretrained_name={NEMO_MODEL}", + ] + + +def write_nemo_inputs( + *, + clips: list[TraceClip], + source_audio_dir: Path, + nemo_manifest: Path, + options: TraceSourceOptions, +) -> list[NemoItem]: + items: list[NemoItem] = [] + with nemo_manifest.open("w", encoding="utf-8") as manifest: + for row_index, clip in enumerate(clips): + audio_path = source_audio_dir / f"source_{row_index:06d}.wav" + sf.write( + str(audio_path), + np.clip(clip.audio, -1.0, 1.0), + options.sample_rate, + ) + item = NemoItem( + row_index=row_index, + audio_path=audio_path, + text=clip.transcript, + ) + items.append(item) + manifest.write( + json.dumps( + { + "audio_filepath": str(audio_path.absolute()), + "text": clip.transcript, + }, + ensure_ascii=False, + ) + + "\n" + ) + return items + + +def read_nemo_word_timings( + nemo_output_dir: Path, + nemo_manifest: Path, + items: list[NemoItem], +) -> dict[int, list[WordTiming]]: + output_manifest = find_nemo_output_manifest(nemo_output_dir, nemo_manifest) + nemo_rows = read_jsonl(output_manifest) + if len(nemo_rows) != len(items): + raise ValueError( + f"NeMo output row count {len(nemo_rows)} does not match input " + f"{len(items)}." + ) + + word_timings: dict[int, list[WordTiming]] = {} + for item, nemo_row in zip(items, nemo_rows): + ctm_path = find_word_ctm_path(nemo_row, output_manifest.parent) + word_timings[item.row_index] = parse_word_ctm(ctm_path) + return word_timings + + +def find_nemo_output_manifest(nemo_output_dir: Path, nemo_manifest: Path) -> Path: + stem = nemo_manifest.stem + candidates = [ + nemo_output_dir / f"{stem}_with_output_file_paths.json", + nemo_output_dir / f"{stem}_with_ctm_paths.json", + ] + for candidate in candidates: + if candidate.exists(): + return candidate + matches = sorted(nemo_output_dir.glob(f"{stem}_with*.json")) + if matches: + return matches[0] + raise FileNotFoundError(f"Could not find NeMo output manifest in {nemo_output_dir}") + + +def find_word_ctm_path(row: dict[str, Any], base_dir: Path) -> Path: + for key, value in row.items(): + if "word" in key and "ctm" in key and value: + path = Path(str(value)) + if path.is_absolute() or path.exists(): + return path + return base_dir / path + raise ValueError(f"NeMo row has no word-level CTM path: {row}") + + +def parse_word_ctm(path: Path) -> list[WordTiming]: + words: list[WordTiming] = [] + with path.open("r", encoding="utf-8") as f: + for line in f: + parts = line.strip().split() + if len(parts) < 5: + continue + start_s = float(parts[2]) + duration_s = float(parts[3]) + # Round to microsecond precision so second->ms conversion doesn't + # leak float artifacts (0.1 + 0.2 -> 300.00000000000006 ms) into + # manifests. + words.append( + WordTiming( + word=parts[4], + start_ms=round(start_s * 1000, 3), + end_ms=round((start_s + duration_s) * 1000, 3), + ) + ) + return words + + +def read_jsonl(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + with path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + rows.append(json.loads(line)) + return rows + + +def is_relative_to(path: Path, base: Path) -> bool: + try: + path.relative_to(base) + except ValueError: + return False + return True + + +######################################################################## +# Chunking +######################################################################## + + +def finalize_clips( + clips: Iterable[TraceClip], + *, + max_duration_s: float | None, + target_duration_s: float | None = None, + sample_rate: int, +) -> Iterable[TraceClip]: + for clip in clips: + clip_target_duration_s = clip.target_duration_s or target_duration_s + if clip_target_duration_s is not None: + target_clip = truncate_clip_to_target_duration( + clip, + target_duration_s=clip_target_duration_s, + sample_rate=sample_rate, + ) + if target_clip is not None: + yield target_clip + continue + + yield from split_or_filter_clip( + clip, + max_duration_s=max_duration_s, + sample_rate=sample_rate, + ) + + +def repeat_clip_to_cover_target_duration( + clip: TraceClip, + *, + target_duration_s: float, + sample_rate: int, +) -> TraceClip: + if len(clip.audio) == 0: + return replace( + clip, + word_timestamps=None, + target_duration_s=target_duration_s, + ) + + source_duration_s = len(clip.audio) / sample_rate + repeat_count = max(1, math.ceil(target_duration_s / source_duration_s)) + audio = np.tile(clip.audio, repeat_count) + return replace( + clip, + audio=audio, + transcript=" ".join([clip.transcript] * repeat_count), + duration_s=len(audio) / sample_rate, + word_timestamps=None, + target_duration_s=target_duration_s, + ) + + +def truncate_clip_to_target_duration( + clip: TraceClip, + *, + target_duration_s: float, + sample_rate: int, +) -> TraceClip | None: + target_samples = max(1, int(round(target_duration_s * sample_rate))) + end_sample = min(len(clip.audio), target_samples) + if end_sample <= 0: + return None + + duration_s = end_sample / sample_rate + word_timestamps = clip.word_timestamps + transcript = clip.transcript + if word_timestamps is not None: + target_end_ms = duration_s * 1000 + word_timestamps = [ + word for word in word_timestamps if word.end_ms <= target_end_ms + ] + transcript = " ".join(word.word for word in word_timestamps) + + return replace( + clip, + audio=clip.audio[:end_sample], + transcript=transcript, + duration_s=duration_s, + word_timestamps=word_timestamps, + ) + + +def split_or_filter_clip( + clip: TraceClip, + *, + max_duration_s: float | None, + sample_rate: int, +) -> Iterable[TraceClip]: + if max_duration_s is None or clip.duration_s <= max_duration_s: + yield clip + return + + if not clip.word_timestamps: + return + + for chunk_index, words in enumerate( + chunk_timed_words( + clip.word_timestamps, + max_duration_s=max_duration_s, + max_gap_s=AMI_MAX_GAP_S, + ) + ): + chunk = build_timed_chunk( + clip=clip, + words=words, + chunk_index=chunk_index, + sample_rate=sample_rate, + ) + if chunk is not None: + yield chunk + + +def build_timed_chunk( + *, + clip: TraceClip, + words: list[WordTiming], + chunk_index: int, + sample_rate: int, +) -> TraceClip | None: + start_ms = words[0].start_ms + end_ms = words[-1].end_ms + start_sample = ms_to_sample(start_ms, sample_rate) + end_sample = min(len(clip.audio), ms_to_sample(end_ms, sample_rate)) + if end_sample <= start_sample: + return None + + relative_words = [ + WordTiming( + word=word.word, + start_ms=round(word.start_ms - start_ms, 3), + end_ms=round(word.end_ms - start_ms, 3), + ) + for word in words + ] + sample_id = str(clip.metadata.get("sample_id", "source")) + metadata = { + **clip.metadata, + "source_parent_sample_id": sample_id, + "chunk_index": chunk_index, + "sample_id": f"{sample_id}:chunk-{chunk_index:05d}", + } + return TraceClip( + audio=clip.audio[start_sample:end_sample], + transcript=" ".join(word.word for word in words), + duration_s=(end_sample - start_sample) / sample_rate, + metadata=metadata, + word_timestamps=relative_words, + ) + + +def chunk_timed_words( + words: list[WordTiming], + *, + max_duration_s: float, + max_gap_s: float, +) -> Iterable[list[WordTiming]]: + current: list[WordTiming] = [] + max_duration_ms = max_duration_s * 1000 + max_gap_ms = max_gap_s * 1000 + for word in words: + if current: + clip_start_ms = current[0].start_ms + previous_end_ms = current[-1].end_ms + if ( + word.start_ms - previous_end_ms > max_gap_ms + or word.end_ms - clip_start_ms > max_duration_ms + ): + yield current + current = [] + current.append(word) + if current: + yield current + + +######################################################################## +# Manifest Writing +######################################################################## + + +def write_trace_dataset( + *, + dataset_key: str, + clips: Iterable[TraceClip], + output_dir: Path, + audio_root: Path, + rows: list[dict[str, Any]], + start_session_id: int, + clip_limit: int | None, + options: TraceSourceOptions, +) -> tuple[int, int]: + dataset_audio_dir = audio_root / dataset_key + dataset_audio_dir.mkdir(parents=True, exist_ok=True) + + produced = 0 + session_id = start_session_id + for clip in clips: + if clip_limit is not None and produced >= clip_limit: + break + + wav_path = dataset_audio_dir / f"clip_{produced:05d}.wav" + sf.write( + str(wav_path), + np.clip(clip.audio, -1.0, 1.0), + options.sample_rate, + ) + row = { + "session_id": session_id, + "audio_file": wav_path.relative_to(output_dir).as_posix(), + "expected_transcript": clip.transcript, + "dataset": dataset_key, + "duration_s": round(clip.duration_s, 3), + "sample_rate": options.sample_rate, + **clip.metadata, + } + if clip.word_timestamps is not None: + row["reference_word_timestamps"] = [ + { + "word": word.word, + "start_ms": round(word.start_ms, 3), + "end_ms": round(word.end_ms, 3), + } + for word in clip.word_timestamps + ] + rows.append(row) + produced += 1 + session_id += 1 + print( + f" [{produced}] {wav_path.relative_to(output_dir)} " + f"({clip.duration_s:.2f}s)" + ) + + return produced, session_id + + +def write_manifest(path: Path, rows: list[dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8") as manifest: + for row in rows: + manifest.write(json.dumps(row, ensure_ascii=False) + "\n") + + +######################################################################## +# Orchestration +######################################################################## + + +def main() -> None: + args = parse_args() + dataset_keys = validate_args(args) + + output_dir = output_dir_for_dataset_keys(dataset_keys) + audio_root = output_dir / "audio" + manifest_path = output_dir / MANIFEST_NAME + output_dir.mkdir(parents=True, exist_ok=True) + audio_root.mkdir(parents=True, exist_ok=True) + + clip_limit = None if args.clips_per_dataset == 0 else args.clips_per_dataset + target_duration_s = args.target_duration + target_duration_deviation_s = ( + target_duration_s * TARGET_DURATION_DEVIATION_FRACTION + if target_duration_s is not None + else None + ) + source_options = TraceSourceOptions(max_duration_s=args.max_duration) + target_duration_rng = np.random.default_rng(source_options.seed) + timestamping_enabled = not args.without_word_timestamping + timestamp_provider = NeMoDockerWordTimestampProvider() + rows: list[dict[str, Any]] = [] + session_id = 0 + + print(f"Building ASR trace: {', '.join(dataset_keys)}") + print(f" Output: {output_dir}") + print(f" Manifest: {manifest_path}") + if target_duration_s is not None and target_duration_deviation_s is not None: + print( + " Target-duration clips will be generated with " + f"mean {target_duration_s:g}s and standard deviation " + f"{target_duration_deviation_s:g}s" + ) + + for dataset_key in dataset_keys: + source = build_trace_source(dataset_key, source_options) + print(f" Loading {dataset_key}: {source.repo} split={source.split}") + + source_clips: list[TraceClip] = [] + collect_all = clip_limit is None or ( + timestamping_enabled + and dataset_key == "aa_earnings22" + and target_duration_s is None + ) + for clip in source.iter_clips(): + source_clips.append(clip) + if ( + not collect_all + and clip_limit is not None + and len(source_clips) >= clip_limit + ): + break + + if target_duration_s is not None: + print( + f" Repeating {dataset_key} clips to cover " + "sampled target durations before alignment" + ) + target_durations_s = [ + float(max(DEFAULT_MIN_DURATION_S, duration_s)) + for duration_s in target_duration_rng.normal( + target_duration_s, + target_duration_deviation_s, + len(source_clips), + ) + ] + source_clips = [ + repeat_clip_to_cover_target_duration( + clip, + target_duration_s=sampled_target_duration_s, + sample_rate=source_options.sample_rate, + ) + for clip, sampled_target_duration_s in zip( + source_clips, + target_durations_s, + ) + ] + + needs_nemo = timestamping_enabled and ( + target_duration_s is not None + or any(clip.word_timestamps is None for clip in source_clips) + ) + if needs_nemo: + print(f" Aligning {dataset_key} with NeMo Docker") + source_clips = timestamp_provider.annotate( + dataset_key=dataset_key, + clips=source_clips, + alignment_output_dir=output_dir / "alignment", + options=source_options, + ) + elif timestamping_enabled and source_clips: + print(f" Using native word timestamps for {dataset_key}") + + produced, session_id = write_trace_dataset( + dataset_key=dataset_key, + clips=finalize_clips( + source_clips, + max_duration_s=source_options.max_duration_s, + target_duration_s=target_duration_s, + sample_rate=source_options.sample_rate, + ), + output_dir=output_dir, + audio_root=audio_root, + rows=rows, + start_session_id=session_id, + clip_limit=clip_limit, + options=source_options, + ) + + if clip_limit is not None and produced < clip_limit: + duration_filter = ( + f" matching --max-duration <= {source_options.max_duration_s:g}s" + if source_options.max_duration_s is not None + else ( + f" matching --target-duration {target_duration_s:g}s" + if target_duration_s is not None + else "" + ) + ) + raise SystemExit( + f"{dataset_key} produced {produced} eligible clip(s)" + f"{duration_filter}, fewer than --clips-per-dataset {clip_limit}." + ) + if clip_limit is None and produced == 0: + print(f" WARNING: no clips produced for {dataset_key}", file=sys.stderr) + + write_manifest(manifest_path, rows) + print(f"\nWrote {len(rows)} clips -> {manifest_path}") + + +if __name__ == "__main__": + main() + # Some free-threaded Python builds abort in pyarrow/datasets finalizers. + sys.stdout.flush() + sys.stderr.flush() + os._exit(0) diff --git a/tests/unit/client/test_native_http_tts_client.py b/tests/unit/client/test_native_http_tts_client.py new file mode 100644 index 00000000..80cb1aa4 --- /dev/null +++ b/tests/unit/client/test_native_http_tts_client.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import asyncio +import importlib.util +import sys +import types +from typing import Any + +import httpx +import pytest + +if ( + "transformers" not in sys.modules + and importlib.util.find_spec("transformers") is None +): + transformers_stub = types.ModuleType("transformers") + transformers_stub.AutoTokenizer = object # type: ignore[attr-defined] + sys.modules["transformers"] = transformers_stub + +from veeksha.client.native_http_tts import ( + DeepgramFluxHTTPClient, + DeepgramFluxHTTPProtocol, + ElevenLabsHTTPProtocol, + ElevenLabsHTTPTTSClient, +) +from veeksha.config.client import ( + DeepgramFluxHTTPClientConfig, + ElevenLabsHTTPTTSClientConfig, +) +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.request import Request +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.types import ChannelModality + + +class _FakeAsyncClient: + def __init__( + self, + response: httpx.Response, + ) -> None: + self.response = response + self.calls: list[dict[str, Any]] = [] + + async def post(self, url: str, **kwargs: Any) -> httpx.Response: + self.calls.append({"url": url, **kwargs}) + return self.response + + +def _request() -> Request: + return Request( + id=17, + channels={ + ChannelModality.TEXT: TextChannelRequestContent( + input_text="A real complete-text TTS request." + ) + }, + ) + + +@pytest.mark.unit +@pytest.mark.parametrize("provider", ["elevenlabs", "deepgram"]) +def test_native_http_clients_buffer_complete_audio_response(provider: str) -> None: + if provider == "elevenlabs": + config = ElevenLabsHTTPTTSClientConfig( + api_base="https://api.elevenlabs.io", + api_key="test-key", + model="eleven_flash_v2_5", + voice_id="test-voice", + ) + client: Any = ElevenLabsHTTPTTSClient(config) + else: + config = DeepgramFluxHTTPClientConfig( + api_base="https://api.deepgram.com", + api_key="test-key", + model="flux-alexis-en", + ) + client = DeepgramFluxHTTPClient(config) + + request = httpx.Request("POST", "https://provider.example/speak") + response = httpx.Response( + 200, + content=bytes(4_800), + headers={"Content-Type": "audio/l16"}, + request=request, + ) + fake_client = _FakeAsyncClient(response) + client._get_client = lambda: fake_client + + sent_count = 0 + + def on_sent() -> None: + nonlocal sent_count + sent_count += 1 + + result = asyncio.run( + client.send_request(_request(), session_id=1, on_request_sent=on_sent) + ) + + assert result.success + assert sent_count == 1 + assert len(fake_client.calls) == 1 + metrics = result.channels[ChannelModality.AUDIO].metrics + assert metrics[AudioMetricKey.PROVIDER.value] == provider + assert metrics[AudioMetricKey.PROVIDER_MODEL.value] == config.model + assert metrics[AudioMetricKey.RAW_PCM.value] + assert metrics[AudioMetricKey.CHUNK_COUNT.value] == 1 + assert ( + metrics[AudioMetricKey.TTFC.value] + == metrics[AudioMetricKey.END_TO_END_LATENCY.value] + ) + assert metrics[AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value] == [ + [metrics[AudioMetricKey.TTFC.value], 4_800] + ] + + +@pytest.mark.unit +def test_native_http_protocols_use_raw_pcm_and_provider_auth() -> None: + eleven_config = ElevenLabsHTTPTTSClientConfig( + api_base="https://api.elevenlabs.io", + api_key="eleven-key", + model="eleven_flash_v2_5", + voice_id="voice/id", + ) + eleven = ElevenLabsHTTPProtocol(eleven_config, "eleven-key") + eleven_request = eleven.build_request(str(eleven_config.api_base), "hello") + assert "/voice%2Fid?" in eleven_request.url + assert "output_format=pcm_24000" in eleven_request.url + assert eleven_request.headers["xi-api-key"] == "eleven-key" + assert eleven_request.payload["text"] == "hello" + + deepgram_config = DeepgramFluxHTTPClientConfig( + api_base="https://api.deepgram.com", + api_key="deepgram-key", + model="flux-alexis-en", + ) + deepgram = DeepgramFluxHTTPProtocol(deepgram_config, "deepgram-key") + deepgram_request = deepgram.build_request(str(deepgram_config.api_base), "hello") + assert "/v2/speak?" in deepgram_request.url + assert "encoding=linear16" in deepgram_request.url + assert "container=none" in deepgram_request.url + assert deepgram_request.headers["Authorization"] == "Token deepgram-key" + + +@pytest.mark.unit +def test_native_http_client_preserves_provider_error() -> None: + config = DeepgramFluxHTTPClientConfig( + api_base="https://api.deepgram.com", + api_key="test-key", + model="flux-alexis-en", + ) + client: Any = DeepgramFluxHTTPClient(config) + request = httpx.Request("POST", "https://provider.example/speak") + client._get_client = lambda: _FakeAsyncClient( + httpx.Response(429, text="rate limited", request=request) + ) + + result = asyncio.run(client.send_request(_request(), session_id=1)) + + assert not result.success + assert result.error_code == 429 + assert result.channels == {} diff --git a/tests/unit/client/test_native_streaming_tts_client.py b/tests/unit/client/test_native_streaming_tts_client.py new file mode 100644 index 00000000..a5d7425b --- /dev/null +++ b/tests/unit/client/test_native_streaming_tts_client.py @@ -0,0 +1,223 @@ +from __future__ import annotations + +import asyncio +import base64 +import importlib.util +import json +import sys +import types +from typing import Any + +import pytest + +if ( + "transformers" not in sys.modules + and importlib.util.find_spec("transformers") is None +): + transformers_stub = types.ModuleType("transformers") + transformers_stub.AutoTokenizer = object # type: ignore[attr-defined] + sys.modules["transformers"] = transformers_stub + +from veeksha.client.native_streaming_tts import ( + DeepgramAuraStreamingProtocol, + DeepgramAuraStreamingTTSClient, + DeepgramFluxStreamingProtocol, + DeepgramFluxStreamingTTSClient, + ElevenLabsStreamingProtocol, + ElevenLabsStreamingTTSClient, +) +from veeksha.config.client import ( + DeepgramAuraStreamingTTSClientConfig, + DeepgramFluxStreamingTTSClientConfig, + ElevenLabsStreamingTTSClientConfig, + TextPacingConfig, +) +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.request import Request +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.types import ChannelModality + + +class _FakeWebSocket: + def __init__(self, provider: str) -> None: + self.provider = provider + self.sent: list[str | bytes] = [] + self.events: asyncio.Queue[str | bytes] = asyncio.Queue() + self.speech_started = False + if provider == "deepgram": + self.events.put_nowait(json.dumps({"type": "Connected"})) + + async def send(self, raw: str | bytes) -> None: + self.sent.append(raw) + event = json.loads(raw) + if self.provider == "elevenlabs": + if event.get("text") and event.get("text") != " ": + audio = base64.b64encode(bytes(960)).decode("ascii") + await self.events.put(json.dumps({"audio": audio, "isFinal": False})) + elif event.get("text") == "": + await self.events.put(json.dumps({"isFinal": True})) + elif self.provider == "deepgram_aura" and event.get("type") == "Flush": + await self.events.put(bytes(960)) + await self.events.put(json.dumps({"type": "Flushed", "sequence_id": 0})) + elif event.get("type") == "Speak" and self.provider == "deepgram": + if not self.speech_started: + await self.events.put(json.dumps({"type": "SpeechStarted"})) + self.speech_started = True + await self.events.put(bytes(960)) + elif event.get("type") == "Flush": + await self.events.put(json.dumps({"type": "Flushed"})) + await self.events.put( + json.dumps( + { + "type": "SpeechMetadata", + "audio_duration_ms": 60, + "input_character_count": 13, + } + ) + ) + + async def recv(self) -> str | bytes: + return await self.events.get() + + +class _FakeConnection: + def __init__(self, websocket: _FakeWebSocket) -> None: + self.websocket = websocket + + async def __aenter__(self) -> _FakeWebSocket: + return self.websocket + + async def __aexit__(self, *args: object) -> None: + return None + + +def _request() -> Request: + return Request( + id=11, + channels={ + ChannelModality.TEXT: TextChannelRequestContent(input_text="one two three") + }, + ) + + +def _pacing() -> TextPacingConfig: + return TextPacingConfig(tokens_per_second=1000, tokens_per_delta=1) + + +@pytest.mark.unit +@pytest.mark.parametrize("provider", ["elevenlabs", "deepgram"]) +def test_native_clients_stream_text_and_audio_concurrently(provider: str) -> None: + if provider == "elevenlabs": + config = ElevenLabsStreamingTTSClientConfig( + api_base="https://api.elevenlabs.io", + api_key="test-key", + model="eleven_flash_v2_5", + voice_id="test-voice", + pacing=_pacing(), + ) + client: Any = ElevenLabsStreamingTTSClient(config) + else: + config = DeepgramFluxStreamingTTSClientConfig( + api_base="https://api.deepgram.com", + api_key="test-key", + model="flux-alexis-en", + pacing=_pacing(), + ) + client = DeepgramFluxStreamingTTSClient(config) + + websocket = _FakeWebSocket(provider) + client._connect = lambda: _FakeConnection(websocket) # type: ignore[method-assign] + result = asyncio.run(client.send_request(_request(), session_id=1)) + + assert result.success + metrics = result.channels[ChannelModality.AUDIO].metrics + assert metrics[AudioMetricKey.PROVIDER.value] == provider + assert metrics[AudioMetricKey.PROVIDER_MODEL.value] == config.model + assert len(metrics[AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value]) == 3 + assert metrics[AudioMetricKey.CHUNK_COUNT.value] == 3 + assert metrics[AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value] == pytest.approx( + metrics[AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value][0][0], abs=0.001 + ) + assert ( + metrics[AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value][0][0] + < metrics[AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value] + ) + + +@pytest.mark.unit +def test_native_protocol_urls_use_raw_pcm_and_provider_auth() -> None: + eleven_config = ElevenLabsStreamingTTSClientConfig( + api_base="https://api.elevenlabs.io", + api_key="eleven-key", + model="eleven_flash_v2_5", + voice_id="voice/id", + ) + eleven = ElevenLabsStreamingProtocol(eleven_config, "eleven-key") + assert "/voice%2Fid/stream-input" in eleven.build_ws_url( + str(eleven_config.api_base) + ) + assert "output_format=pcm_24000" in eleven.build_ws_url(str(eleven_config.api_base)) + assert eleven.headers() == {"xi-api-key": "eleven-key"} + + deepgram_config = DeepgramFluxStreamingTTSClientConfig( + api_base="https://api.deepgram.com", + api_key="deepgram-key", + model="flux-alexis-en", + ) + deepgram = DeepgramFluxStreamingProtocol(deepgram_config, "deepgram-key") + assert "/v2/speak?" in deepgram.build_ws_url(str(deepgram_config.api_base)) + assert "encoding=linear16" in deepgram.build_ws_url(str(deepgram_config.api_base)) + assert "speed=" not in deepgram.build_ws_url(str(deepgram_config.api_base)) + assert deepgram.headers() == {"Authorization": "Token deepgram-key"} + + aura_config = DeepgramAuraStreamingTTSClientConfig( + api_base="https://api.deepgram.com", + api_key="deepgram-key", + model="aura-2-thalia-en", + ) + aura = DeepgramAuraStreamingProtocol(aura_config, "deepgram-key") + assert "/v1/speak?" in aura.build_ws_url(str(aura_config.api_base)) + assert "encoding=linear16" in aura.build_ws_url(str(aura_config.api_base)) + assert "speed=1.0" in aura.build_ws_url(str(aura_config.api_base)) + assert aura.headers() == {"Authorization": "Token deepgram-key"} + + +@pytest.mark.unit +def test_deepgram_aura_adapter_handles_audio_after_flush() -> None: + config = DeepgramAuraStreamingTTSClientConfig( + api_base="https://api.deepgram.com", + api_key="test-key", + model="aura-2-thalia-en", + pacing=_pacing(), + ) + client = DeepgramAuraStreamingTTSClient(config) + websocket = _FakeWebSocket("deepgram_aura") + client._connect = lambda: _FakeConnection(websocket) # type: ignore[method-assign] + + result = asyncio.run(client.send_request(_request(), session_id=1)) + + assert result.success + metrics = result.channels[ChannelModality.AUDIO].metrics + assert metrics[AudioMetricKey.PROVIDER_PROTOCOL.value] == "v1_aura_speak" + first_audio_ms = metrics[AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value][0][0] + commit_ms = metrics[AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value] + trigger_ms = metrics[AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value] + assert trigger_ms < commit_ms + assert first_audio_ms >= commit_ms + assert [json.loads(raw)["type"] for raw in websocket.sent] == [ + "Speak", + "Speak", + "Speak", + "Flush", + ] + + +@pytest.mark.unit +def test_native_client_requires_provider_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DEEPGRAM_API_KEY", raising=False) + config = DeepgramFluxStreamingTTSClientConfig( + api_base="https://api.deepgram.com", + model="flux-alexis-en", + ) + with pytest.raises(ValueError, match="DEEPGRAM_API_KEY"): + DeepgramFluxStreamingTTSClient(config) diff --git a/tests/unit/client/test_realtime_tts_client.py b/tests/unit/client/test_realtime_tts_client.py new file mode 100644 index 00000000..ffc0d919 --- /dev/null +++ b/tests/unit/client/test_realtime_tts_client.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import asyncio +import base64 +import importlib.util +import json +import sys +import types +from typing import Any + +import pytest + +# The focused free-threaded unit-test environment omits the heavyweight +# tokenizer dependency. Realtime TTS uses the built-in word tokenizer, but the +# client package imports the generic tokenizer module during initialization. +if ( + "transformers" not in sys.modules + and importlib.util.find_spec("transformers") is None +): + transformers_stub = types.ModuleType("transformers") + transformers_stub.AutoTokenizer = object # type: ignore[attr-defined] + sys.modules["transformers"] = transformers_stub + +from veeksha.client.realtime_tts import RealtimeTTSClient +from veeksha.config.client import RealtimeTTSClientConfig, TextPacingConfig +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.request import Request +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.types import ChannelModality + + +class _FakeWebSocket: + def __init__(self) -> None: + self.sent: list[dict[str, Any]] = [] + self._events: asyncio.Queue[str] = asyncio.Queue() + + async def send(self, raw: str) -> None: + event = json.loads(raw) + self.sent.append(event) + if event["type"] == "session.update": + await self._events.put( + json.dumps( + { + "type": "session.updated", + "session": { + "audio": { + "output": { + "format": {"type": "audio/pcm", "rate": 24000} + } + } + }, + } + ) + ) + elif event["type"] == "response.create": + audio = base64.b64encode(bytes(960)).decode("ascii") + for response_event in ( + {"type": "response.created"}, + {"type": "response.output_audio.delta", "delta": audio}, + {"type": "response.output_audio.done"}, + { + "type": "response.done", + "response": {"status": "completed"}, + }, + ): + await self._events.put(json.dumps(response_event)) + + async def recv(self) -> str: + return await self._events.get() + + +class _FakeConnection: + def __init__(self, websocket: _FakeWebSocket) -> None: + self.websocket = websocket + + async def __aenter__(self) -> _FakeWebSocket: + return self.websocket + + async def __aexit__(self, *args: object) -> None: + return None + + +def _request() -> Request: + return Request( + id=7, + channels={ + ChannelModality.TEXT: TextChannelRequestContent(input_text="one two three") + }, + ) + + +def _client(mode: str) -> tuple[RealtimeTTSClient, _FakeWebSocket]: + config = RealtimeTTSClientConfig( + api_base="http://example.test", + model="test-tts", + input_output_mode=mode, + duplex_start_after_tokens=1, + pacing=TextPacingConfig( + tokens_per_second=1000.0, + tokens_per_delta=1, + ), + ) + client = RealtimeTTSClient(config) + websocket = _FakeWebSocket() + client._connect = lambda: _FakeConnection(websocket) # type: ignore[method-assign] + return client, websocket + + +@pytest.mark.unit +def test_duplex_triggers_response_before_input_is_complete() -> None: + client, websocket = _client("duplex") + + result = asyncio.run(client.send_request(_request(), session_id=1)) + + event_types = [event["type"] for event in websocket.sent] + assert event_types == [ + "session.update", + "conversation.item.create", + "response.create", + "conversation.item.create", + "conversation.item.create", + ] + metrics = result.channels[ChannelModality.AUDIO].metrics + assert ( + metrics[AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value] + < metrics[AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value] + ) + assert ( + metrics[AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value][0][0] + < metrics[AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value] + ) + + +@pytest.mark.unit +def test_complete_text_triggers_response_after_last_input_delta() -> None: + client, websocket = _client("complete_text") + + result = asyncio.run(client.send_request(_request(), session_id=1)) + + event_types = [event["type"] for event in websocket.sent] + assert event_types == [ + "session.update", + "conversation.item.create", + "conversation.item.create", + "conversation.item.create", + "response.create", + ] + metrics = result.channels[ChannelModality.AUDIO].metrics + assert ( + metrics[AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value] + >= metrics[AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value] + ) + + +@pytest.mark.unit +def test_realtime_tts_config_rejects_unknown_input_output_mode() -> None: + with pytest.raises(ValueError, match="input_output_mode"): + RealtimeTTSClientConfig( + api_base="http://example.test", + model="test-tts", + input_output_mode="not-a-mode", + ) diff --git a/tests/unit/client/test_stt_client.py b/tests/unit/client/test_stt_client.py new file mode 100644 index 00000000..4ac4b300 --- /dev/null +++ b/tests/unit/client/test_stt_client.py @@ -0,0 +1,209 @@ +import asyncio +import json + +import pytest +import websockets + +from veeksha.client.stt import ( + VajraOpenAIRealtimeSTTClient, + VllmRealtimeSTTClient, + _slice_pcm16_bytes, +) +from veeksha.config.client import STTClientConfig +from veeksha.core.request import Request + + +@pytest.mark.unit +def test_slice_pcm16_bytes_uses_millisecond_offsets() -> None: + pcm = bytes(range(20)) + + sliced = _slice_pcm16_bytes(pcm, 1000, start_ms=2.0, end_ms=6.0) + + assert sliced == pcm[4:12] + + +@pytest.mark.unit +def test_slice_pcm16_bytes_passthrough_without_offsets() -> None: + pcm = bytes(range(20)) + + assert _slice_pcm16_bytes(pcm, 1000, start_ms=None, end_ms=None) is pcm + + +@pytest.mark.unit +def test_slice_pcm16_bytes_open_ended_slices() -> None: + pcm = bytes(range(20)) + + assert _slice_pcm16_bytes(pcm, 1000, start_ms=None, end_ms=6.0) == pcm[:12] + assert _slice_pcm16_bytes(pcm, 1000, start_ms=2.0, end_ms=None) == pcm[4:] + + +@pytest.mark.unit +def test_slice_pcm16_bytes_rejects_negative_start() -> None: + pcm = bytes(range(20)) + + with pytest.raises(ValueError, match="must be non-negative"): + _slice_pcm16_bytes(pcm, 1000, start_ms=-1.0, end_ms=6.0) + + +@pytest.mark.unit +def test_slice_pcm16_bytes_rejects_end_not_after_start() -> None: + pcm = bytes(range(20)) + + with pytest.raises(ValueError, match="must be greater than"): + _slice_pcm16_bytes(pcm, 1000, start_ms=6.0, end_ms=6.0) + + +@pytest.mark.unit +def test_slice_pcm16_bytes_rejects_end_past_clip() -> None: + pcm = bytes(range(20)) + + with pytest.raises(ValueError, match="exceeds decoded clip length"): + _slice_pcm16_bytes(pcm, 1000, start_ms=0.0, end_ms=11.0) + + +@pytest.mark.unit +@pytest.mark.parametrize("field_name", ["ws_ping_interval_s", "ws_ping_timeout_s"]) +def test_stt_client_config_rejects_nonpositive_ws_ping(field_name: str) -> None: + with pytest.raises(ValueError, match=f"{field_name} must be > 0 or None"): + STTClientConfig( + provider="vajra_openai_realtime", + model="mistralai/Voxtral-Mini-4B-Realtime-2602", + api_base="http://localhost:8003", + **{field_name: 0}, + ) + + +def _vajra_realtime_client() -> VajraOpenAIRealtimeSTTClient: + config = STTClientConfig( + provider="vajra_openai_realtime", + model="mistralai/Voxtral-Mini-4B-Realtime-2602", + api_base="http://localhost:8003", + ) + return VajraOpenAIRealtimeSTTClient(config) + + +@pytest.mark.unit +def test_vajra_openai_realtime_parses_transcription_events() -> None: + client = _vajra_realtime_client() + + assert client._parse_message( + {"type": "conversation.item.input_audio_transcription.delta", "delta": "hi"} + ) == ("delta", "hi") + assert client._parse_message( + { + "type": "conversation.item.input_audio_transcription.completed", + "transcript": "hi there", + } + ) == ("done", "hi there") + assert client._parse_message({"type": "input_audio_buffer.committed"}) == ("", "") + + +@pytest.mark.unit +def test_vajra_openai_realtime_maps_failed_item_to_error() -> None: + client = _vajra_realtime_client() + + kind, text = client._parse_message( + { + "type": "conversation.item.input_audio_transcription.failed", + "error": {"type": "server_error", "message": "boom"}, + } + ) + + assert (kind, text) == ("error", "boom") + + +@pytest.mark.unit +def test_vajra_openai_realtime_maps_session_error() -> None: + client = _vajra_realtime_client() + + assert client._parse_message( + {"type": "error", "error": {"message": "bad request"}} + ) == ("error", "bad request") + + +def _vllm_realtime_client() -> VllmRealtimeSTTClient: + config = STTClientConfig( + provider="vllm_realtime", + model="mistralai/Voxtral-Mini-4B-Realtime-2602", + api_base="http://localhost:8025", + ) + return VllmRealtimeSTTClient(config) + + +@pytest.mark.unit +def test_stt_send_request_finishes_lifecycle_callbacks_on_invalid_audio() -> None: + client = _vllm_realtime_client() + events: list[str] = [] + + result = asyncio.run( + client.send_request( + Request(id=99, channels={}), + session_id=7, + on_request_sent=lambda: events.append("sent"), + on_request_dispatched=lambda: events.append("dispatched"), + ) + ) + + assert result.success is False + assert result.error_code == 400 + assert events == ["dispatched", "sent"] + + +class _FakeVllmWebSocket: + def __init__(self) -> None: + self._recv_index = 0 + self._audio_sent = asyncio.Event() + + async def recv(self) -> str: + if self._recv_index == 0: + self._recv_index += 1 + return json.dumps({"type": "session.created"}) + + await self._audio_sent.wait() + if self._recv_index == 1: + self._recv_index += 1 + return json.dumps({"type": "transcription.delta", "delta": "hello"}) + self._recv_index += 1 + return json.dumps({"type": "transcription.done", "text": "hello"}) + + async def send(self, message: str | bytes) -> None: + if isinstance(message, str): + payload = json.loads(message) + if payload.get("type") == "input_audio_buffer.append": + self._audio_sent.set() + + +class _FakeConnection: + def __init__(self, websocket: _FakeVllmWebSocket) -> None: + self._websocket = websocket + + async def __aenter__(self) -> _FakeVllmWebSocket: + return self._websocket + + async def __aexit__(self, *_args) -> None: + return None + + +@pytest.mark.unit +def test_vllm_stream_fires_callbacks_after_handshake_and_first_content( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = _vllm_realtime_client() + websocket = _FakeVllmWebSocket() + monkeypatch.setattr( + websockets, + "connect", + lambda *_args, **_kwargs: _FakeConnection(websocket), + ) + events: list[str] = [] + + result = asyncio.run( + client._stream( + b"\x00\x00", + on_request_sent=lambda: events.append("sent"), + on_request_dispatched=lambda: events.append("dispatched"), + ) + ) + + assert result.final_transcript == "hello" + assert events == ["dispatched", "sent"] diff --git a/tests/unit/client/test_stt_snapshots.py b/tests/unit/client/test_stt_snapshots.py new file mode 100644 index 00000000..64e30b0d --- /dev/null +++ b/tests/unit/client/test_stt_snapshots.py @@ -0,0 +1,19 @@ +import pytest + +from veeksha.client.stt import TranscriptSnapshotRecorder + + +@pytest.mark.unit +def test_transcript_snapshot_recorder_dedupes_and_uses_audio_start() -> None: + recorder = TranscriptSnapshotRecorder() + + recorder.mark_audio_started(101.0) + recorder.add(100.5, "") + recorder.add(101.25, "hello") + recorder.add(101.30, "hello") + recorder.add(101.50, "hello world") + + assert recorder.snapshots == [ + {"elapsed_ms": 250.0, "transcript": "hello"}, + {"elapsed_ms": 500.0, "transcript": "hello world"}, + ] diff --git a/tests/unit/config/test_native_tts_client_config.py b/tests/unit/config/test_native_tts_client_config.py new file mode 100644 index 00000000..054d9316 --- /dev/null +++ b/tests/unit/config/test_native_tts_client_config.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import pytest +from vidhi import create_class_from_dict + +from veeksha.config.benchmark import BenchmarkConfig + + +@pytest.mark.parametrize( + ("client_type", "model", "extra", "config_class"), + [ + ( + "elevenlabs_streaming_tts", + "eleven_flash_v2_5", + {"voice_id": "voice"}, + "ElevenLabsStreamingTTSClientConfig", + ), + ( + "deepgram_flux_streaming_tts", + "flux-alexis-en", + {}, + "DeepgramFluxStreamingTTSClientConfig", + ), + ( + "deepgram_aura_streaming_tts", + "aura-2-thalia-en", + {}, + "DeepgramAuraStreamingTTSClientConfig", + ), + ( + "elevenlabs_http_tts", + "eleven_flash_v2_5", + {"voice_id": "voice"}, + "ElevenLabsHTTPTTSClientConfig", + ), + ( + "deepgram_flux_http_tts", + "flux-alexis-en", + {}, + "DeepgramFluxHTTPClientConfig", + ), + ], +) +def test_native_tts_client_type_deserializes_from_public_name( + client_type: str, + model: str, + extra: dict[str, str], + config_class: str, +) -> None: + config = create_class_from_dict( + BenchmarkConfig, + { + "client": { + "type": client_type, + "api_base": "https://provider.example", + "model": model, + **extra, + } + }, + ) + + assert type(config.client).__name__ == config_class + assert str(config.client.get_type()) == client_type diff --git a/tests/unit/config/test_slo_config.py b/tests/unit/config/test_slo_config.py index f8161c6a..361f6287 100644 --- a/tests/unit/config/test_slo_config.py +++ b/tests/unit/config/test_slo_config.py @@ -1,6 +1,8 @@ import unittest + from veeksha.config.slo import ConstantSloConfig + class TestSloConfig(unittest.TestCase): def test_supported_metrics(self): # Should not raise @@ -8,14 +10,18 @@ def test_supported_metrics(self): ConstantSloConfig(metric="tbc", value=1.0) ConstantSloConfig(metric="tpot", value=1.0) ConstantSloConfig(metric="e2e", value=1.0) + ConstantSloConfig(metric="user_audio_fluidity_index", value=0.99) def test_unsupported_metric(self): with self.assertRaisesRegex(ValueError, "metric 'invalid' is not supported"): ConstantSloConfig(metric="invalid", value=1.0) def test_invalid_value(self): - with self.assertRaisesRegex(ValueError, "value must be specified and must be > 0"): + with self.assertRaisesRegex( + ValueError, "value must be specified and must be > 0" + ): ConstantSloConfig(metric="ttfc", value=-1.0) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/evaluator/accuracy/test_audio_quality.py b/tests/unit/evaluator/accuracy/test_audio_quality.py new file mode 100644 index 00000000..b09c270e --- /dev/null +++ b/tests/unit/evaluator/accuracy/test_audio_quality.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +import io +import wave + +import pytest + +from veeksha.evaluator.accuracy.audio import _make_wav_header + + +@pytest.mark.unit +def test_pcm_wav_header_preserves_wire_sample_rate_and_duration() -> None: + sample_rate = 24000 + pcm = b"\x00\x00" * 480 # 20 ms of mono PCM16 at 24 kHz. + wav = io.BytesIO(_make_wav_header(len(pcm), sample_rate) + pcm) + + with wave.open(wav, "rb") as reader: + assert reader.getnchannels() == 1 + assert reader.getsampwidth() == 2 + assert reader.getframerate() == sample_rate + assert reader.getnframes() == 480 + assert reader.readframes(reader.getnframes()) == pcm diff --git a/tests/unit/evaluator/performance/test_asr_correctness.py b/tests/unit/evaluator/performance/test_asr_correctness.py new file mode 100644 index 00000000..2249bf38 --- /dev/null +++ b/tests/unit/evaluator/performance/test_asr_correctness.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import pytest + +from veeksha.evaluator.performance.asr import ( + ASRMetricAccumulator, + WERStats, + compute_wer_stats, +) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "hypothesis", + [ + "alpha delta charlie", # substitution + "alpha charlie", # deletion + "alpha bravo extra charlie", # insertion + ], +) +def test_asr_wer_counts_each_edit_class_against_a_manual_oracle( + hypothesis: str, +) -> None: + stats = compute_wer_stats("alpha bravo charlie", hypothesis) + + assert stats.errors == 1 + assert stats.reference_words == 3 + assert stats.wer == pytest.approx(100 / 3) + + +@pytest.mark.unit +def test_asr_wer_applies_the_leaderboard_normalizer_before_alignment() -> None: + stats = compute_wer_stats( + "Mr. Smith won't pay twenty-one dollars.", + "mister smith will not pay 21 dollars", + ) + + assert stats == WERStats(errors=0, reference_words=6, wer=0.0) + + +@pytest.mark.unit +def test_asr_aggregates_distinguish_sample_corpus_and_duration_weighted_wer() -> None: + accumulator = ASRMetricAccumulator() + accumulator.add_clip_sample( + dataset="short", + duration_s=1.0, + final_stats=WERStats(errors=1, reference_words=2, wer=50.0), + partial_stats=None, + ) + accumulator.add_clip_sample( + dataset="long", + duration_s=3.0, + final_stats=WERStats(errors=1, reference_words=8, wer=12.5), + partial_stats=None, + ) + + summary = accumulator.get_summary() + + assert summary["asr_final_sample_mean_wer"] == pytest.approx(31.25) + assert summary["asr_final_corpus_wer"] == pytest.approx(20.0) + assert summary["asr_final_duration_weighted_wer"] == pytest.approx(21.875) + assert summary["asr_dataset_short_final_corpus_wer"] == pytest.approx(50.0) + assert summary["asr_dataset_long_final_corpus_wer"] == pytest.approx(12.5) + + +@pytest.mark.unit +def test_asr_empty_reference_has_explicit_zero_or_full_error_semantics() -> None: + assert compute_wer_stats("", "") == WERStats(errors=0, reference_words=0, wer=0.0) + assert compute_wer_stats("", "unexpected") == WERStats( + errors=1, reference_words=0, wer=100.0 + ) + + accumulator = ASRMetricAccumulator() + accumulator.add_clip_sample( + dataset="empty", + duration_s=1.0, + final_stats=compute_wer_stats("", "unexpected"), + partial_stats=None, + ) + assert accumulator.get_summary()["asr_final_corpus_wer"] == 100.0 diff --git a/tests/unit/evaluator/performance/test_asr_interactivity.py b/tests/unit/evaluator/performance/test_asr_interactivity.py new file mode 100644 index 00000000..42646cc6 --- /dev/null +++ b/tests/unit/evaluator/performance/test_asr_interactivity.py @@ -0,0 +1,178 @@ +import random +from typing import Any, Dict + +import pytest + +from veeksha.evaluator.performance.asr import ( + ASRMetricAccumulator, + compute_interactivity_stats, + score_asr_request, +) + + +def _streaming_fixture() -> Dict[str, Any]: + """Deterministic synthetic streaming-STT request. + + Exercises growing partial transcripts, unstable tail rewrites, duplicated + snapshots, filler/number/contraction normalization, and imperfect + hypotheses. The expected metric values in the regression test below were + produced by the original (pre-optimization) scoring implementation. + """ + rng = random.Random(42) + vocab = ( + "the quick brown fox jumps over lazy dog and a half twenty one " + "dollars mr smith won't say uh hundred and three point five st " + "colour center o'clock nineteen sixty" + ).split() + ref_words = [vocab[rng.randrange(len(vocab))] for _ in range(120)] + reference_word_timestamps = [ + {"word": word, "start_ms": 250.0 * i, "end_ms": 250.0 * i + 200.0} + for i, word in enumerate(ref_words) + ] + hyp_words = [ + word if rng.random() > 0.1 else vocab[rng.randrange(len(vocab))] + for word in ref_words + ] + snapshots = [] + shown_count = 0 + elapsed = 0.0 + while shown_count < len(hyp_words): + shown_count = min(len(hyp_words), shown_count + rng.randrange(1, 4)) + elapsed += rng.uniform(300.0, 700.0) + shown = list(hyp_words[:shown_count]) + if rng.random() < 0.3 and shown_count > 3: + shown[-1] = vocab[rng.randrange(len(vocab))] # unstable tail + snapshots.append( + {"elapsed_ms": round(elapsed, 3), "transcript": " ".join(shown)} + ) + if rng.random() < 0.15: # repeated snapshot at a later timestamp + elapsed += 120.0 + snapshots.append( + {"elapsed_ms": round(elapsed, 3), "transcript": " ".join(shown)} + ) + return { + "final_transcript": " ".join(hyp_words), + "expected_transcript": " ".join(ref_words), + "partial_transcript": " ".join(hyp_words[:80]), + "dataset": "fixture", + "reference_word_timestamps": reference_word_timestamps, + "transcript_snapshots": snapshots, + } + + +@pytest.mark.unit +def test_compute_interactivity_stats_from_snapshots() -> None: + stats = compute_interactivity_stats( + { + "reference_word_timestamps": [ + {"word": "hello", "start_ms": 0, "end_ms": 100}, + {"word": "world", "start_ms": 200, "end_ms": 300}, + ], + "transcript_snapshots": [ + {"elapsed_ms": 150, "transcript": "hello"}, + {"elapsed_ms": 320, "transcript": "hello world"}, + ], + } + ) + + assert stats is not None + assert stats.word_count == 2 + assert stats.latencies_ms == [50, 20] + assert stats.mean_latency_ms == 35 + + +@pytest.mark.unit +def test_compute_interactivity_expands_and_skips_normalized_reference_words() -> None: + stats = compute_interactivity_stats( + { + "reference_word_timestamps": [ + {"word": "don't", "start_ms": 0, "end_ms": 100}, + {"word": "uh", "start_ms": 120, "end_ms": 180}, + {"word": "stop", "start_ms": 200, "end_ms": 300}, + ], + "transcript_snapshots": [ + {"elapsed_ms": 150, "transcript": "do not"}, + {"elapsed_ms": 340, "transcript": "do not stop"}, + ], + } + ) + + assert stats is not None + assert stats.word_count == 3 + assert stats.latencies_ms == [50, 50, 40] + assert stats.mean_latency_ms == pytest.approx(46.6666667) + + +@pytest.mark.unit +def test_compute_interactivity_stats_empty_snapshots_returns_none() -> None: + stats = compute_interactivity_stats( + { + "reference_word_timestamps": [ + {"word": "hello", "start_ms": 0, "end_ms": 100}, + ], + "transcript_snapshots": [], + } + ) + + assert stats is None + + +@pytest.mark.unit +def test_score_asr_request_writes_interactivity_row_fields() -> None: + metrics = score_asr_request( + request_id=0, + channel_metrics={ + "final_transcript": "hello world", + "expected_transcript": "hello world", + "dataset": "toy", + "reference_word_timestamps": [ + {"word": "hello", "start_ms": 0, "end_ms": 100}, + {"word": "world", "start_ms": 200, "end_ms": 300}, + ], + "transcript_snapshots": [ + {"elapsed_ms": 150, "transcript": "hello"}, + {"elapsed_ms": 320, "transcript": "hello world"}, + ], + }, + duration_s=0.3, + accumulator=ASRMetricAccumulator(), + ) + + row = metrics.to_request_row() + assert row["interactivity"] == 35 + assert row["interactivity_word_count"] == 2 + + +@pytest.mark.unit +def test_streaming_fixture_matches_pre_optimization_values() -> None: + """Pin exact metric values produced by the original scoring code. + + The interactivity/WER implementations were optimized for speed; this + fixture asserts bit-identical numerical results against values captured + from the pre-optimization implementation on the same input. + """ + channel_metrics = _streaming_fixture() + assert len(channel_metrics["transcript_snapshots"]) == 70 + + stats = compute_interactivity_stats(channel_metrics) + assert stats is not None + assert stats.word_count == 96 + assert stats.mean_latency_ms == 2057.4468854166666 + assert sum(stats.latencies_ms) == 197514.90099999998 + + accumulator = ASRMetricAccumulator() + metrics = score_asr_request( + request_id=0, + channel_metrics=channel_metrics, + duration_s=30.0, + accumulator=accumulator, + ) + assert metrics.final_wer == 13.513513513513514 + assert metrics.partial_wer == 42.34234234234234 + + summary = accumulator.get_summary() + assert summary["asr_final_corpus_wer"] == 13.513513513513514 + assert summary["asr_final_sample_mean_wer"] == 13.513513513513514 + assert summary["asr_final_duration_weighted_wer"] == 13.513513513513514 + assert summary["asr_partial_corpus_wer"] == 42.34234234234234 + assert summary["asr_final_sample_count"] == 1.0 diff --git a/tests/unit/evaluator/performance/test_asr_normalizer.py b/tests/unit/evaluator/performance/test_asr_normalizer.py new file mode 100644 index 00000000..cf1e9b43 --- /dev/null +++ b/tests/unit/evaluator/performance/test_asr_normalizer.py @@ -0,0 +1,78 @@ +"""Regression tests pinning EnglishTextNormalizer outputs. + +The expected strings were produced by the original (pre-optimization) +open-asr-leaderboard normalizer implementation; the optimized version must +reproduce them exactly so WER and interactivity metrics stay numerically +identical across runs. +""" + +import pytest + +from veeksha.evaluator.performance.asr_normalizer import EnglishTextNormalizer +from veeksha.evaluator.performance.asr_normalizer.normalizer import ( + remove_symbols, + remove_symbols_and_diacritics, +) + +_PINNED_CASES = [ + ( + "Mr. O'Brien's cafΓ© costs $20 million and Β’7, twenty-one and a half mhm", + "mister 0 brien is cafe costs $20000000 and Β’721.5", + ), + ( + "one oh one double seven point five percent 1,234.56 3rd 1960s", + "10177.5% 1234.56 3rd 1960s", + ), + ( + "won't can't y'all's it's don't we're I'ma [noise] (laughs) ", + "will not can not you all is it is do not we are i am going to", + ), + ( + "minus five degrees plus positive Β£3 euros Ε“Γ¦ΓŸΓ° Β²", + "-5 degrees +positive €3 oeaessd 2", + ), + ( + "one hundred and twenty three thousand four hundred fifty six", + "123456", + ), + ( + "two and a half percent of $0.5 st dr colour flavour theatre", + "2.5% of Β’5 saint doctor color flavor theater", + ), + ( + "The DOCTOR said: it's ten thirty--five, o'clock!", + "the doctor said it is 10350 clock", + ), +] + + +@pytest.mark.unit +class TestEnglishTextNormalizerRegression: + @pytest.mark.parametrize("raw,expected", _PINNED_CASES) + def test_pinned_outputs(self, raw: str, expected: str) -> None: + assert EnglishTextNormalizer()(raw) == expected + + def test_empty_and_whitespace_inputs(self) -> None: + normalizer = EnglishTextNormalizer() + assert normalizer("") == "" + assert normalizer(" ") == "" + + def test_normalizer_is_deterministic_across_calls(self) -> None: + normalizer = EnglishTextNormalizer() + raw = _PINNED_CASES[0][0] + assert normalizer(raw) == normalizer(raw) + + +@pytest.mark.unit +class TestSymbolRemoval: + def test_remove_symbols_and_diacritics_keep_set(self) -> None: + assert ( + remove_symbols_and_diacritics("cafΓ© $5.0, 100%!", keep=".%$’€£") + == "cafe $5.0 100% " + ) + + def test_remove_symbols_and_diacritics_default_keep(self) -> None: + assert remove_symbols_and_diacritics("Ε“Γ†ΓŸ naΓ―ve!") == "oeAEss naive " + + def test_remove_symbols_keeps_diacritics(self) -> None: + assert remove_symbols("naΓ―ve, test!") == "naΓ―ve test " diff --git a/tests/unit/evaluator/performance/test_audio_evaluator_tasks.py b/tests/unit/evaluator/performance/test_audio_evaluator_tasks.py new file mode 100644 index 00000000..da61b041 --- /dev/null +++ b/tests/unit/evaluator/performance/test_audio_evaluator_tasks.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import pytest + +from veeksha.config.evaluator import ( + AudioChannelPerformanceConfig, + PerformanceEvaluatorConfig, +) +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.response import ChannelResponse, RequestResult +from veeksha.evaluator.performance.audio import AudioPerformanceEvaluator +from veeksha.types import AudioTask, ChannelModality + + +def _evaluator(*, interactivity_enabled: bool) -> AudioPerformanceEvaluator: + audio_config = AudioChannelPerformanceConfig( + interactivity_enabled=interactivity_enabled, + startup_delay_ms_values=[0.0], + startup_buffer_ms_values=[0.0], + ) + config = PerformanceEvaluatorConfig( + target_channels=[ChannelModality.AUDIO], + slos=[], + audio_channel=audio_config, + ) + return AudioPerformanceEvaluator( + config, + channel_config=audio_config, + benchmark_start_time=100.0, + ) + + +def _record( + evaluator: AudioPerformanceEvaluator, + *, + request_id: int, + content: bytes | str, + metrics: dict, +) -> None: + evaluator.register_request( + request_id=request_id, + session_id=request_id, + dispatched_at=100.0, + content=None, + ) + evaluator.record_request_completed( + request_id=request_id, + session_id=request_id, + completed_at=101.0, + response=RequestResult( + request_id=request_id, + session_id=request_id, + channels={ + ChannelModality.AUDIO: ChannelResponse( + modality=ChannelModality.AUDIO, + content=content, + metrics=metrics, + ) + }, + scheduler_dispatched_at=100.0, + client_completed_at=101.0, + ), + ) + + +def test_tts_task_preserves_playback_interactivity_metrics() -> None: + evaluator = _evaluator(interactivity_enabled=True) + + _record( + evaluator, + request_id=1, + content=b"\0" * 4800, + metrics={ + "audio_task": AudioTask.TTS, + AudioMetricKey.TTFC.value: 100.0, + AudioMetricKey.END_TO_END_LATENCY.value: 250.0, + AudioMetricKey.CHUNK_COUNT.value: 2, + AudioMetricKey.RAW_PCM.value: True, + AudioMetricKey.SAMPLE_RATE.value: 24000, + AudioMetricKey.INPUT_CHARS.value: 5, + AudioMetricKey.INPUT_TOKENS.value: 1, + AudioMetricKey.INPUT_TEXT.value: "hello", + AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value: [[20.0, 5]], + AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value: [ + [100.0, 2400], + [180.0, 2400], + ], + AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value: 70.0, + AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value: 150.0, + AudioMetricKey.AUDIO_DONE_OFFSET_MS.value: 230.0, + AudioMetricKey.RESPONSE_DONE_OFFSET_MS.value: 240.0, + }, + ) + + row = evaluator._export_request_rows()[0] + assert row["audio_task"] == "tts" + assert row[AudioMetricKey.GENERATED_AUDIO_DURATION.value] == 100.0 + assert row[AudioMetricKey.ZERO_DELAY_STALL_COUNT.value] == 1 + assert row[AudioMetricKey.REQUEST_START_TO_FIRST_PLAYABLE_AUDIO_MS.value] == 100.0 + assert row[AudioMetricKey.FIRST_INPUT_TO_FIRST_PLAYABLE_AUDIO_MS.value] == 80.0 + assert row[AudioMetricKey.TRIGGER_TO_FIRST_PLAYABLE_AUDIO_MS.value] == 30.0 + assert row[AudioMetricKey.USER_AUDIO_FLUIDITY_INDEX.value] == pytest.approx( + 4 / 7, abs=1e-5 + ) + assert row[AudioMetricKey.TTS_SERVICE_FLUIDITY_ELIGIBLE.value] == 0 + assert AudioMetricKey.TTS_SERVICE_FLUIDITY_INDEX.value not in row + assert row[AudioMetricKey.UNATTRIBUTED_MISSED_DEADLINES.value] == 3 + assert row[AudioMetricKey.DUPLEX_OVERLAP_OBSERVED.value] == 1 + assert row[AudioMetricKey.AUDIO_PLAYABLE_FRAME_COUNT.value] == 5 + assert row[AudioMetricKey.AUDIO_FLUIDITY_FRAME_MS.value] == 20.0 + assert row[AudioMetricKey.AUDIO_FLUIDITY_STARTUP_DELAY_MS.value] == 0.0 + assert row[AudioMetricKey.INPUT_TEXT.value] == "hello" + + +def test_stt_task_uses_input_pcm_duration_and_exports_asr_metrics() -> None: + evaluator = _evaluator(interactivity_enabled=True) + + _record( + evaluator, + request_id=2, + content="hello world", + metrics={ + "audio_task": "stt", + AudioMetricKey.TTFC.value: 120.0, + AudioMetricKey.END_TO_END_LATENCY.value: 1000.0, + AudioMetricKey.CHUNK_COUNT.value: 4, + AudioMetricKey.PCM_BYTE_COUNT.value: 32000, + AudioMetricKey.RAW_PCM.value: True, + AudioMetricKey.SAMPLE_RATE.value: 16000, + "final_transcript": "hello world", + "partial_transcript": "hello", + "expected_transcript": "hello world", + "dataset": "fixture", + "time_to_first_visible_text": 120.0, + "time_to_first_partial": 150.0, + "time_to_final_transcript": 900.0, + "transcript_snapshots": [], + }, + ) + + row = evaluator._export_request_rows()[0] + summary = evaluator.get_summary() + assert row["audio_task"] == "stt" + assert row[AudioMetricKey.PCM_BYTE_COUNT.value] == 32000 + assert row[AudioMetricKey.GENERATED_AUDIO_DURATION.value] == 1000.0 + assert row["final_wer"] == 0.0 + assert summary["asr_final_sample_count"] == 1.0 + assert AudioMetricKey.ZERO_DELAY_STALL_COUNT.value not in row + + +def test_audio_response_without_task_is_rejected() -> None: + evaluator = _evaluator(interactivity_enabled=False) + + with pytest.raises(ValueError, match="unknown audio_task"): + _record( + evaluator, + request_id=3, + content=b"\0" * 320, + metrics={ + AudioMetricKey.RAW_PCM.value: True, + AudioMetricKey.SAMPLE_RATE.value: 16000, + }, + ) diff --git a/tests/unit/evaluator/performance/test_audio_interactivity.py b/tests/unit/evaluator/performance/test_audio_interactivity.py new file mode 100644 index 00000000..c78408d1 --- /dev/null +++ b/tests/unit/evaluator/performance/test_audio_interactivity.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import pytest + +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.evaluator.performance.audio_interactivity import ( + RequestTiming, + compute_audio_fluidity, + compute_interactivity_metrics, + parse_request_timing, +) + + +def _chunk(receipt_ms: float, duration_ms: float) -> tuple[float, float, float]: + # 16-bit mono PCM at 24 kHz: 48 decoded bytes per millisecond. + return receipt_ms, duration_ms, duration_ms * 48 + + +def test_parse_timing_coalesces_transport_frames_and_uses_wire_sample_rate() -> None: + timing = parse_request_timing( + { + AudioMetricKey.SAMPLE_RATE.value: 16000, + AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value: [[5.0, 4]], + AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value: [ + [30.0, 320], + [30.0, 320], + [50.0, 640], + ], + AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value: 10.0, + AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value: 25.0, + } + ) + + assert timing is not None + assert timing.sample_rate == 16000 + assert timing.response_trigger_ms == 10.0 + assert timing.audio_chunks == [ + (30.0, 20.0, 640.0), + (50.0, 20.0, 640.0), + ] + + +def test_audio_fluidity_carries_early_frame_slack_forward() -> None: + result = compute_audio_fluidity( + [_chunk(100.0, 60.0), _chunk(150.0, 20.0)], + frame_duration_ms=20.0, + startup_delay_ms=0.0, + ) + + assert result is not None + assert result.playable_frame_count == 4 + assert result.missed_deadlines == 0 + assert result.total_deadlines == 4 + assert result.fluidity_index == 1.0 + + +def test_audio_fluidity_counts_misses_and_resets_after_a_stall() -> None: + chunks = [ + _chunk(100.0, 20.0), + _chunk(145.0, 20.0), + _chunk(165.0, 20.0), + ] + + zero_delay = compute_audio_fluidity( + chunks, + frame_duration_ms=20.0, + startup_delay_ms=0.0, + ) + buffered = compute_audio_fluidity( + chunks, + frame_duration_ms=20.0, + startup_delay_ms=25.0, + ) + + assert zero_delay is not None + assert zero_delay.playable_frame_count == 3 + assert zero_delay.missed_deadlines == 2 + assert zero_delay.total_deadlines == 4 + assert zero_delay.fluidity_index == pytest.approx(0.5) + assert buffered is not None + assert buffered.fluidity_index == 1.0 + + +def test_audio_fluidity_is_invariant_to_transport_fragmentation() -> None: + coarse_chunks = [_chunk(100.0, 40.0), _chunk(140.0, 40.0)] + fragmented_chunks = [ + _chunk(90.0, 10.0), + _chunk(100.0, 30.0), + _chunk(130.0, 10.0), + _chunk(140.0, 30.0), + ] + + coarse = compute_audio_fluidity( + coarse_chunks, + frame_duration_ms=20.0, + startup_delay_ms=0.0, + ) + fragmented = compute_audio_fluidity( + fragmented_chunks, + frame_duration_ms=20.0, + startup_delay_ms=0.0, + ) + + assert coarse is not None + assert fragmented is not None + assert coarse.playable_frame_count == fragmented.playable_frame_count == 4 + assert coarse.total_deadlines == fragmented.total_deadlines + assert coarse.missed_deadlines == fragmented.missed_deadlines + assert coarse.fluidity_index == fragmented.fluidity_index + + +def test_audio_fluidity_ignores_an_incomplete_tail_frame() -> None: + assert ( + compute_audio_fluidity( + [_chunk(100.0, 19.0)], + frame_duration_ms=20.0, + startup_delay_ms=0.0, + ) + is None + ) + + +def test_playback_policies_match_a_manually_solved_packet_timeline() -> None: + timing = RequestTiming( + text_deltas=[(10.0, 5)], + audio_chunks=[ + _chunk(100.0, 100.0), + _chunk(250.0, 20.0), + _chunk(400.0, 20.0), + ], + response_trigger_ms=20.0, + commit_ms=20.0, + audio_done_ms=420.0, + response_done_ms=425.0, + sample_rate=24000, + ) + + metrics = compute_interactivity_metrics( + timing, + startup_delay_ms_values=[0.0, 180.0], + startup_buffer_ms_values=[100.0, 140.0], + min_reportable_stall_ms=0.0, + fluidity_frame_ms=20.0, + fluidity_startup_delay_ms=0.0, + fluidity_attribution_mode="conservative", + ) + + assert metrics.required_startup_delay_ms == 180.0 + assert metrics.fixed_delay_playback[0.0].stall_count == 2 + assert metrics.fixed_delay_playback[0.0].total_stall_ms == 180.0 + assert metrics.fixed_delay_playback[180.0].stall_free + assert not metrics.buffer_target_playback[100.0].stall_free + assert metrics.buffer_target_playback[140.0].stall_free + + +def test_initial_audio_reservoir_can_make_required_extra_delay_zero() -> None: + timing = RequestTiming( + text_deltas=[(10.0, 5)], + audio_chunks=[_chunk(100.0, 1000.0), _chunk(500.0, 20.0)], + response_trigger_ms=20.0, + commit_ms=20.0, + audio_done_ms=520.0, + response_done_ms=525.0, + sample_rate=24000, + ) + + metrics = compute_interactivity_metrics( + timing, + startup_delay_ms_values=[0.0], + startup_buffer_ms_values=[0.0], + min_reportable_stall_ms=0.0, + fluidity_frame_ms=20.0, + fluidity_startup_delay_ms=0.0, + fluidity_attribution_mode="conservative", + ) + + assert metrics.required_startup_delay_ms == 0.0 + assert metrics.fixed_delay_playback[0.0].stall_free + + +def test_trigger_ttfa_is_not_inferred_when_protocol_did_not_record_a_trigger() -> None: + timing = RequestTiming( + text_deltas=[(20.0, 5)], + audio_chunks=[_chunk(100.0, 20.0)], + response_trigger_ms=None, + commit_ms=20.0, + audio_done_ms=120.0, + response_done_ms=125.0, + sample_rate=24000, + ) + + metrics = compute_interactivity_metrics( + timing, + startup_delay_ms_values=[0.0], + startup_buffer_ms_values=[0.0], + min_reportable_stall_ms=0.0, + fluidity_frame_ms=20.0, + fluidity_startup_delay_ms=0.0, + fluidity_attribution_mode="conservative", + ) + + assert metrics.first_input_to_first_playable_audio_ms == 80.0 + assert metrics.trigger_to_first_playable_audio_ms is None + + +def test_interactivity_distinguishes_first_byte_from_first_playable_frame() -> None: + timing = RequestTiming( + text_deltas=[(20.0, 5)], + audio_chunks=[_chunk(100.0, 10.0), _chunk(110.0, 10.0)], + response_trigger_ms=80.0, + commit_ms=80.0, + audio_done_ms=120.0, + response_done_ms=125.0, + sample_rate=24000, + ) + + metrics = compute_interactivity_metrics( + timing, + startup_delay_ms_values=[0.0, 100.0], + startup_buffer_ms_values=[0.0], + min_reportable_stall_ms=10.0, + fluidity_frame_ms=20.0, + fluidity_startup_delay_ms=100.0, + fluidity_attribution_mode="conservative", + ) + + assert metrics.request_start_to_first_audio_ms == 100.0 + assert metrics.request_start_to_first_playable_audio_ms == 110.0 + assert metrics.first_input_to_first_playable_audio_ms == 90.0 + assert metrics.trigger_to_first_playable_audio_ms == 30.0 + assert metrics.user_audio_fluidity is not None + assert metrics.user_audio_fluidity.startup_delay_ms == 100.0 + assert metrics.user_audio_fluidity.fluidity_index == 1.0 + + +def test_duplex_user_fluidity_is_not_attributed_to_service_without_source_proof() -> ( + None +): + timing = RequestTiming( + text_deltas=[(0.0, 5), (200.0, 5)], + audio_chunks=[ + _chunk(50.0, 20.0), + _chunk(115.0, 20.0), + ], + response_trigger_ms=10.0, + commit_ms=200.0, + audio_done_ms=140.0, + response_done_ms=210.0, + sample_rate=24000, + ) + + metrics = compute_interactivity_metrics( + timing, + startup_delay_ms_values=[0.0], + startup_buffer_ms_values=[0.0], + min_reportable_stall_ms=10.0, + fluidity_frame_ms=20.0, + fluidity_startup_delay_ms=0.0, + fluidity_attribution_mode="conservative", + ) + + assert metrics.duplex_overlap_observed + assert metrics.user_audio_fluidity is not None + assert metrics.user_audio_fluidity.missed_deadlines > 0 + assert metrics.tts_service_fluidity is None + assert not metrics.tts_service_fluidity_eligible + assert metrics.unattributed_missed_deadlines > 0 + + +def test_oversupplied_duplex_attributes_fluidity_to_service() -> None: + timing = RequestTiming( + text_deltas=[(0.0, 5), (200.0, 5)], + audio_chunks=[ + _chunk(50.0, 20.0), + _chunk(115.0, 20.0), + ], + response_trigger_ms=10.0, + commit_ms=200.0, + audio_done_ms=140.0, + response_done_ms=210.0, + sample_rate=24000, + ) + + metrics = compute_interactivity_metrics( + timing, + startup_delay_ms_values=[0.0], + startup_buffer_ms_values=[0.0], + min_reportable_stall_ms=10.0, + fluidity_frame_ms=20.0, + fluidity_startup_delay_ms=0.0, + fluidity_attribution_mode="source_oversupplied", + ) + + assert metrics.tts_service_fluidity is metrics.user_audio_fluidity + assert metrics.tts_service_fluidity_eligible + assert metrics.unattributed_missed_deadlines == 0 + + +@pytest.mark.parametrize( + ("frame_ms", "startup_ms", "message"), + [ + (0.0, 0.0, "frame_duration_ms must be > 0"), + (20.0, -1.0, "startup_delay_ms must be >= 0"), + ], +) +def test_audio_fluidity_rejects_invalid_policy( + frame_ms: float, startup_ms: float, message: str +) -> None: + with pytest.raises(ValueError, match=message): + compute_audio_fluidity( + [_chunk(100.0, 20.0)], + frame_duration_ms=frame_ms, + startup_delay_ms=startup_ms, + ) diff --git a/tests/unit/generator/session/trace/test_audio_trace_generator.py b/tests/unit/generator/session/trace/test_audio_trace_generator.py new file mode 100644 index 00000000..b5974bd6 --- /dev/null +++ b/tests/unit/generator/session/trace/test_audio_trace_generator.py @@ -0,0 +1,190 @@ +import json +from unittest.mock import MagicMock + +import pytest + +from veeksha.config.generator.session import ( + AudioTraceFlavorConfig, + TraceSessionGeneratorConfig, +) +from veeksha.core.seeding import SeedManager +from veeksha.generator.session.trace.audio import AudioTraceFlavorGenerator + + +@pytest.mark.unit +def test_audio_trace_metadata_preserves_reference_word_timestamp_list(tmp_path) -> None: + audio_path = tmp_path / "clip.wav" + audio_path.write_bytes(b"") + reference_word_timestamps = [ + {"word": "hello", "start_ms": 0.0, "end_ms": 100.0}, + {"word": "world", "start_ms": 120.0, "end_ms": 300.0}, + ] + manifest_path = tmp_path / "manifest.jsonl" + manifest_path.write_text( + json.dumps( + { + "session_id": 0, + "audio_file": audio_path.name, + "expected_transcript": "hello world", + "reference_word_timestamps": reference_word_timestamps, + } + ) + + "\n", + encoding="utf-8", + ) + + flavor_config = AudioTraceFlavorConfig() + config = TraceSessionGeneratorConfig( + trace_file=str(manifest_path), + flavor=flavor_config, + wrap_mode=False, + ) + tokenizer_provider = MagicMock() + tokenizer_provider.for_modality.return_value = MagicMock() + + generator = AudioTraceFlavorGenerator( + config, + flavor_config, + SeedManager(seed=42), + tokenizer_provider, + ) + session = generator.generate_session() + + assert session.requests[0].metadata["reference_word_timestamps"] == ( + reference_word_timestamps + ) + + +@pytest.mark.unit +def test_audio_trace_target_duration_trims_reference_and_adds_slice_metadata( + tmp_path, +) -> None: + audio_path = tmp_path / "clip.wav" + audio_path.write_bytes(b"") + reference_word_timestamps = [ + {"word": "alpha", "start_ms": 0.0, "end_ms": 100.0}, + {"word": "beta", "start_ms": 150.0, "end_ms": 900.0}, + {"word": "gamma", "start_ms": 1200.0, "end_ms": 1400.0}, + ] + manifest_path = tmp_path / "manifest.jsonl" + manifest_path.write_text( + json.dumps( + { + "session_id": 0, + "audio_file": audio_path.name, + "expected_transcript": "alpha beta gamma", + "reference_word_timestamps": reference_word_timestamps, + } + ) + + "\n", + encoding="utf-8", + ) + + flavor_config = AudioTraceFlavorConfig(target_duration_s=1.0) + config = TraceSessionGeneratorConfig( + trace_file=str(manifest_path), + flavor=flavor_config, + wrap_mode=False, + ) + tokenizer_provider = MagicMock() + tokenizer_provider.for_modality.return_value = MagicMock() + + generator = AudioTraceFlavorGenerator( + config, + flavor_config, + SeedManager(seed=42), + tokenizer_provider, + ) + session = generator.generate_session() + metadata = session.requests[0].metadata + + assert metadata["expected_transcript"] == "alpha beta" + assert metadata["reference_word_timestamps"] == reference_word_timestamps[:2] + assert metadata["input_audio_start_ms"] == 0.0 + assert metadata["input_audio_end_ms"] == 1000.0 + assert metadata["duration_s"] == 1.0 + + +@pytest.mark.unit +def test_audio_trace_target_duration_requires_word_timestamps(tmp_path) -> None: + audio_path = tmp_path / "clip.wav" + audio_path.write_bytes(b"") + manifest_path = tmp_path / "manifest.jsonl" + manifest_path.write_text( + json.dumps( + { + "session_id": 0, + "audio_file": audio_path.name, + "expected_transcript": "alpha beta", + } + ) + + "\n", + encoding="utf-8", + ) + + flavor_config = AudioTraceFlavorConfig(target_duration_s=1.0) + config = TraceSessionGeneratorConfig( + trace_file=str(manifest_path), + flavor=flavor_config, + wrap_mode=False, + ) + tokenizer_provider = MagicMock() + tokenizer_provider.for_modality.return_value = MagicMock() + + generator = AudioTraceFlavorGenerator( + config, + flavor_config, + SeedManager(seed=42), + tokenizer_provider, + ) + + with pytest.raises(ValueError, match="reference_word_timestamps"): + generator.generate_session() + + +@pytest.mark.unit +def test_audio_trace_target_duration_rejects_empty_trimmed_transcript( + tmp_path, +) -> None: + audio_path = tmp_path / "clip.wav" + audio_path.write_bytes(b"") + manifest_path = tmp_path / "manifest.jsonl" + manifest_path.write_text( + json.dumps( + { + "session_id": 0, + "audio_file": audio_path.name, + "expected_transcript": "alpha", + "reference_word_timestamps": [ + {"word": "alpha", "start_ms": 1100.0, "end_ms": 1400.0}, + ], + } + ) + + "\n", + encoding="utf-8", + ) + + flavor_config = AudioTraceFlavorConfig(target_duration_s=1.0) + config = TraceSessionGeneratorConfig( + trace_file=str(manifest_path), + flavor=flavor_config, + wrap_mode=False, + ) + tokenizer_provider = MagicMock() + tokenizer_provider.for_modality.return_value = MagicMock() + + generator = AudioTraceFlavorGenerator( + config, + flavor_config, + SeedManager(seed=42), + tokenizer_provider, + ) + + with pytest.raises(ValueError, match="expected_transcript would"): + generator.generate_session() + + +@pytest.mark.unit +def test_audio_trace_flavor_config_rejects_nonpositive_target_duration() -> None: + with pytest.raises(ValueError, match="target_duration_s must be positive"): + AudioTraceFlavorConfig(target_duration_s=0.0) diff --git a/tests/unit/orchestration/test_benchmark_orchestrator.py b/tests/unit/orchestration/test_benchmark_orchestrator.py index 5115b43f..9f3e5c2c 100644 --- a/tests/unit/orchestration/test_benchmark_orchestrator.py +++ b/tests/unit/orchestration/test_benchmark_orchestrator.py @@ -1,79 +1,82 @@ +"""Tests for the current managed-engine orchestrator API.""" -import pytest # type: ignore[import] +import os from unittest.mock import MagicMock, patch +import pytest + +from veeksha.config.endpoint import EndpointConfig +from veeksha.config.server import SglangServerConfig, VllmServerConfig +from veeksha.orchestration.benchmark_orchestrator import ( + create_server_manager, + managed_server, +) +from veeksha.orchestration.managed_engines import ( + SglangOmniDockerRunner, + VllmOmniDockerRunner, +) + pytestmark = pytest.mark.unit -from veeksha.orchestration.benchmark_orchestrator import create_server_manager, managed_server -from veeksha.config.server import VllmServerConfig, SglangServerConfig, BaseServerConfig -from veeksha.orchestration.vllm_server import VLLMServerManager -from veeksha.orchestration.sglang_server import SGLangServerManager + + +def _vllm_config() -> VllmServerConfig: + return VllmServerConfig( + hf_model="meta/demo-model", + deploy_config="/tmp/vllm-deploy.yaml", + ) + + +def _sglang_config() -> SglangServerConfig: + return SglangServerConfig( + model_path="meta/demo-model", + deploy_config="/tmp/sglang-deploy.yaml", + bootstrap="", + ) + class TestBenchmarkOrchestrator: - def test_create_server_manager_vllm(self): - config = VllmServerConfig() - manager = create_server_manager(config, output_dir="/tmp") - assert isinstance(manager, VLLMServerManager) + manager = create_server_manager(_vllm_config(), output_dir="/tmp") + assert isinstance(manager, VllmOmniDockerRunner) def test_create_server_manager_sglang(self): - config = SglangServerConfig() - manager = create_server_manager(config, output_dir="/tmp") - assert isinstance(manager, SGLangServerManager) + manager = create_server_manager(_sglang_config(), output_dir="/tmp") + assert isinstance(manager, SglangOmniDockerRunner) @patch("veeksha.orchestration.benchmark_orchestrator.create_server_manager") def test_managed_server_context(self, mock_create): - """Test the managed_server context manager.""" - config = VllmServerConfig( - host="localhost", - port=8000, + endpoint = EndpointConfig( + engine_type="vllm", + model="meta/demo-model", + api_base="http://localhost:8000/v1", api_key="test-key", + health_url="http://localhost:8000/v1/models", + port=8000, ) - mock_manager = MagicMock() - mock_manager.launch.return_value = True - mock_manager.wait_for_ready.return_value = True + mock_manager.get_endpoint.return_value = endpoint mock_create.return_value = mock_manager - - with managed_server(config, output_dir="/tmp") as info: - assert info["api_base"] == "http://localhost:8000/v1" - assert info["api_key"] == "test-key" + + with managed_server(_vllm_config(), output_dir="/tmp") as info: + assert info["endpoint"] == endpoint + assert info["api_base"] == endpoint.api_base + assert info["api_key"] == endpoint.api_key assert info["server_manager"] == mock_manager - - mock_manager.launch.assert_called_once() - mock_manager.wait_for_ready.assert_called_once() - - # Check env vars - import os + mock_manager.start.assert_called_once() + mock_manager.get_endpoint.assert_called_once() assert os.environ["OPENAI_API_KEY"] == "test-key" assert os.environ["OPENAI_API_BASE"] == "http://localhost:8000/v1" - - mock_manager.shutdown.assert_called_once() - @patch("veeksha.orchestration.benchmark_orchestrator.create_server_manager") - def test_managed_server_launch_failure(self, mock_create): - """Test managed_server when launch fails.""" - config = VllmServerConfig() - mock_manager = MagicMock() - mock_manager.launch.return_value = False - mock_create.return_value = mock_manager - - with pytest.raises(RuntimeError, match="Failed to launch server"): - with managed_server(config, output_dir="/tmp"): - pass - - mock_manager.shutdown.assert_called_once() + mock_manager.stop.assert_called_once() @patch("veeksha.orchestration.benchmark_orchestrator.create_server_manager") - def test_managed_server_ready_failure(self, mock_create): - """Test managed_server when wait_for_ready fails.""" - config = VllmServerConfig() + def test_managed_server_start_failure_still_stops(self, mock_create): mock_manager = MagicMock() - mock_manager.launch.return_value = True - mock_manager.wait_for_ready.return_value = False + mock_manager.start.side_effect = RuntimeError("failed to launch server") mock_create.return_value = mock_manager - - with pytest.raises(RuntimeError, match="Server failed to become ready"): - with managed_server(config, output_dir="/tmp"): + + with pytest.raises(RuntimeError, match="failed to launch server"): + with managed_server(_vllm_config(), output_dir="/tmp"): pass - - mock_manager.shutdown.assert_called_once() + + mock_manager.stop.assert_called_once() diff --git a/tests/unit/orchestration/test_engine_server_commands.py b/tests/unit/orchestration/test_engine_server_commands.py index 940088e2..e52f0b8d 100644 --- a/tests/unit/orchestration/test_engine_server_commands.py +++ b/tests/unit/orchestration/test_engine_server_commands.py @@ -1,85 +1,88 @@ -"""Unit tests for engine-specific server managers with advanced configs.""" +"""Unit tests for current engine-runner command construction.""" -import json from textwrap import dedent -import pytest # type: ignore[import] +import pytest import yaml from vidhi import create_class_from_dict from veeksha.config.server import BaseServerConfig, VllmServerConfig -from veeksha.orchestration.vllm_server import VLLMServerManager +from veeksha.orchestration.managed_engines import VllmOmniDockerRunner pytestmark = pytest.mark.unit -def test_vllm_launch_command_with_advanced_configuration(): +def test_vllm_server_command_with_advanced_configuration(): config = VllmServerConfig( - model="meta/test-model", - host="0.0.0.0", - port=9001, - api_key="secret-key", - tensor_parallel_size=2, - dtype="bfloat16", - max_model_len=8192, - additional_args={ - "max-num-batched-tokens": 4096, - "trust-remote-code": True, - "disable-log-requests": False, - "some_list": ["alpha", "beta"], - "rope_scaling": {"type": "linear", "factor": 2.0}, - }, + hf_model="meta/test-model", + model="served-model", + deploy_config="/tmp/deploy.yaml", + container_deploy_config="/etc/vllm/custom.yaml", + container_port=9001, + bootstrap="", + engine_args=[ + "--max-num-batched-tokens", + "4096", + "--trust-remote-code", + ], ) - command = VLLMServerManager(config, output_dir="/tmp")._build_launch_command() - - # Base flags - assert command[:3] == ["python", "-m", "vllm.entrypoints.openai.api_server"] - assert command[command.index("--model") + 1] == "meta/test-model" - assert command[command.index("--host") + 1] == "0.0.0.0" - assert command[command.index("--port") + 1] == "9001" - assert command[command.index("--api-key") + 1] == "secret-key" - assert command[command.index("--tensor-parallel-size") + 1] == "2" - assert command[command.index("--dtype") + 1] == "bfloat16" - assert command[command.index("--max-model-len") + 1] == "8192" - - # Additional args - assert command[command.index("--max-num-batched-tokens") + 1] == "4096" - assert "--trust-remote-code" in command - assert "--disable-log-requests" not in command - - list_start = command.index("--some_list") - assert command[list_start : list_start + 3] == ["--some_list", "alpha", "beta"] - - rope_index = command.index("--rope-scaling") - assert command[rope_index + 1] == json.dumps({"type": "linear", "factor": 2.0}) - -def test_server_config_additional_args_loaded_from_yaml(tmp_path): + command = VllmOmniDockerRunner(config, output_dir="/tmp")._build_server_cmd() + + assert command == [ + "vllm", + "serve", + "meta/test-model", + "--omni", + "--port", + "9001", + "--deploy-config", + "/etc/vllm/custom.yaml", + "--max-num-batched-tokens", + "4096", + "--trust-remote-code", + ] + + +def test_server_config_loaded_from_yaml_builds_docker_command(tmp_path): + deploy_config = tmp_path / "deploy.yaml" + deploy_config.write_text("model: fixture\n") config_file = tmp_path / "server_config.yaml" config_file.write_text( dedent( - """ + f""" type: vllm - model: meta/demo-model - port: 8100 - additional_args: - trust-remote-code: true - max-num-batched-tokens: 512 - rope_scaling: - type: linear - factor: 1.25 + hf_model: meta/demo-model + model: served-model + image: vllm-omni:test + deploy_config: {deploy_config} + docker_gpus: device=0,1 + bootstrap: "" + env: + TOKEN: fixture + engine_args: + - --max-num-batched-tokens + - "512" + - --trust-remote-code """ ).strip() ) - data = yaml.safe_load(config_file.read_text()) - server_config = create_class_from_dict(BaseServerConfig, data) - - command = VLLMServerManager(server_config, output_dir="/tmp")._build_launch_command() - - assert "--trust-remote-code" in command - idx = command.index("--max-num-batched-tokens") - assert command[idx + 1] == "512" - - rope_idx = command.index("--rope-scaling") - assert json.loads(command[rope_idx + 1]) == {"type": "linear", "factor": 1.25} + server_config = create_class_from_dict( + BaseServerConfig, + yaml.safe_load(config_file.read_text()), + ) + command = VllmOmniDockerRunner( + server_config, output_dir=tmp_path + )._build_docker_run_cmd() + + assert command[:3] == ["docker", "run", "-d"] + assert command[command.index("--gpus") + 1] == "device=0,1" + assert f"{deploy_config}:/etc/vllm-omni/{deploy_config.name}:ro" in command + assert "TOKEN=fixture" in command + assert "vllm-omni:test" in command + assert command[-3:] == [ + "--max-num-batched-tokens", + "512", + "--trust-remote-code", + ] diff --git a/tests/unit/orchestration/test_server_manager.py b/tests/unit/orchestration/test_server_manager.py index fb01be5c..33065d1a 100644 --- a/tests/unit/orchestration/test_server_manager.py +++ b/tests/unit/orchestration/test_server_manager.py @@ -1,39 +1,51 @@ - -import os -import time import tempfile from pathlib import Path +from unittest.mock import ANY, MagicMock, patch import pytest # type: ignore[import] -from unittest.mock import MagicMock, patch, ANY -from veeksha.orchestration.server_manager import BaseServerManager + from veeksha.config.server import VllmServerConfig +from veeksha.orchestration.server_manager import BaseServerManager pytestmark = pytest.mark.unit -class TestServerManager(BaseServerManager): + +class FakeServerManager(BaseServerManager): """Concrete implementation for testing.""" + def _build_launch_command(self) -> list[str]: return ["python", "-m", "mock_server"] + +def _server_config(**kwargs) -> VllmServerConfig: + """Build a valid current vLLM config for BaseServerManager tests.""" + return VllmServerConfig( + hf_model="meta/test-model", + deploy_config="/tmp/test-deploy.yaml", + **kwargs, + ) + + @pytest.fixture def server_config(): - return VllmServerConfig( + return _server_config( host="localhost", port=8000, gpu_ids=[0], startup_timeout=1, - health_check_interval=0.1 + health_check_interval=0.1, ) + @pytest.fixture def manager(server_config): - mgr = TestServerManager(server_config) + mgr = FakeServerManager(server_config) mgr._is_port_in_use = MagicMock(return_value=False) return mgr + class TestBaseServerManager: - + @patch("subprocess.Popen") def test_launch_success(self, mock_popen, manager): """Test successful server launch.""" @@ -41,19 +53,19 @@ def test_launch_success(self, mock_popen, manager): mock_process.poll.return_value = None mock_process.pid = 12345 mock_popen.return_value = mock_process - + success, error = manager.launch() assert success assert error is None - + assert manager.is_running assert manager.process == mock_process mock_popen.assert_called_once() - + # Check environment variables call_args = mock_popen.call_args - env = call_args[1]['env'] - assert env['CUDA_VISIBLE_DEVICES'] == "0" + env = call_args[1]["env"] + assert env["CUDA_VISIBLE_DEVICES"] == "0" @patch("subprocess.Popen") def test_launch_already_running(self, mock_popen, manager): @@ -61,7 +73,7 @@ def test_launch_already_running(self, mock_popen, manager): manager._is_running = True manager.process = MagicMock() manager.process.poll.return_value = None - + success, error = manager.launch() assert success assert error is None @@ -71,7 +83,7 @@ def test_launch_already_running(self, mock_popen, manager): def test_launch_failure(self, mock_popen, manager): """Test launch failure.""" mock_popen.side_effect = Exception("Launch failed") - + success, error = manager.launch() assert not success assert error == "Launch failed" @@ -94,9 +106,9 @@ def test_health_check_success(self, mock_get, manager): mock_response = MagicMock() mock_response.status_code = 200 mock_get.return_value = mock_response - + assert manager.health_check() - mock_get.assert_called_with("http://localhost:8000/health", timeout=5) + mock_get.assert_called_with("http://localhost:8000/v1/models", timeout=5) @patch("requests.get") def test_health_check_failure_status(self, mock_get, manager): @@ -104,15 +116,18 @@ def test_health_check_failure_status(self, mock_get, manager): mock_response = MagicMock() mock_response.status_code = 503 mock_get.return_value = mock_response - + assert not manager.health_check() @patch("requests.get") def test_health_check_failure_exception(self, mock_get, manager): """Test health check failure due to exception.""" import requests - mock_get.side_effect = requests.exceptions.RequestException("Connection refused") - + + mock_get.side_effect = requests.exceptions.RequestException( + "Connection refused" + ) + assert not manager.health_check() @patch("time.sleep") @@ -123,15 +138,15 @@ def test_wait_for_ready_success(self, mock_get, mock_sleep, manager): manager._is_running = True manager.process = MagicMock() manager.process.poll.return_value = None - + mock_response_fail = MagicMock() mock_response_fail.status_code = 503 - + mock_response_success = MagicMock() mock_response_success.status_code = 200 - + mock_get.side_effect = [mock_response_fail, mock_response_success] - + assert manager.wait_for_ready(timeout=5) assert mock_get.call_count == 2 @@ -142,26 +157,26 @@ def test_wait_for_ready_timeout(self, mock_get, mock_sleep, manager): manager._is_running = True manager.process = MagicMock() manager.process.poll.return_value = None - + mock_response = MagicMock() mock_response.status_code = 503 mock_get.return_value = mock_response - + # Mock time.time to simulate timeout time_calls = {"value": 0} def time_mock(): time_calls["value"] += 1 return time_calls["value"] * 0.2 # 0.2, 0.4, 0.6, ... - + with patch("time.time", side_effect=time_mock): assert not manager.wait_for_ready(timeout=1) @patch("time.sleep") def test_wait_for_ready_process_died(self, mock_sleep, manager): """Test wait_for_ready when process dies.""" - manager._is_running = False # Process not running - + manager._is_running = False # Process not running + assert not manager.wait_for_ready(timeout=5) def test_shutdown(self, manager): @@ -170,9 +185,9 @@ def test_shutdown(self, manager): mock_process.poll.return_value = None manager.process = mock_process manager._is_running = True - + assert manager.shutdown() - + mock_process.terminate.assert_called_once() mock_process.wait.assert_called() assert not manager.is_running @@ -183,29 +198,29 @@ def test_shutdown_force(self, manager): mock_process.poll.return_value = None manager.process = mock_process manager._is_running = True - + assert manager.shutdown(force=True) - + mock_process.kill.assert_called_once() assert not manager.is_running def test_get_additional_args_dict(self): """Test parsing additional args.""" # Test string (JSON) - config = VllmServerConfig(additional_args='{"key": "value"}') - manager = TestServerManager(config) + config = _server_config(additional_args='{"key": "value"}') + manager = FakeServerManager(config) args = manager.get_additional_args_dict() assert args == {"key": "value"} # Test None - config = VllmServerConfig(additional_args=None) - manager = TestServerManager(config) + config = _server_config(additional_args=None) + manager = FakeServerManager(config) args = manager.get_additional_args_dict() assert args == {} # Test dict - config = VllmServerConfig(additional_args={"key": "value"}) - manager = TestServerManager(config) + config = _server_config(additional_args={"key": "value"}) + manager = FakeServerManager(config) args = manager.get_additional_args_dict() assert args == {"key": "value"} # Ensure it's a shallow copy @@ -213,8 +228,8 @@ def test_get_additional_args_dict(self): assert config.additional_args == {"key": "value"} # Test invalid JSON string - config = VllmServerConfig(additional_args='{"invalid": json}') - manager = TestServerManager(config) + config = _server_config(additional_args='{"invalid": json}') + manager = FakeServerManager(config) with pytest.raises(ValueError, match="Invalid JSON in additional_args"): manager.get_additional_args_dict() @@ -223,43 +238,46 @@ def test_get_additional_args_dict(self): def test_auto_allocation(self): """Test auto-allocation of GPUs during launch.""" # Config without explicit GPU IDs - config = VllmServerConfig(gpu_ids=None, tensor_parallel_size=2) - manager = TestServerManager(config) - + config = _server_config(gpu_ids=None, tensor_parallel_size=2) + manager = FakeServerManager(config) + # Mock resource manager manager.resource_manager = MagicMock() - manager.resource_manager.wait_for_resources.return_value = [("node1", 0), ("node1", 1)] - + manager.resource_manager.wait_for_resources.return_value = [ + ("node1", 0), + ("node1", 1), + ] + with patch("subprocess.Popen") as mock_popen: mock_process = MagicMock() mock_process.poll.return_value = None mock_popen.return_value = mock_process - + success, error = manager.launch() assert success assert error is None - + # Check that GPUs were allocated manager.resource_manager.wait_for_resources.assert_called_with( num_gpus=2, timeout=300, job_id=ANY, contiguous=True ) - + # Check that config was updated assert manager.config.gpu_ids == [0, 1] - + # Check env var call_args = mock_popen.call_args - env = call_args[1]['env'] - assert env['CUDA_VISIBLE_DEVICES'] == "0,1" + env = call_args[1]["env"] + assert env["CUDA_VISIBLE_DEVICES"] == "0,1" def test_auto_allocation_non_contiguous(self): """Ensure we can request non-contiguous GPUs when flag is disabled.""" - config = VllmServerConfig( + config = _server_config( gpu_ids=None, tensor_parallel_size=2, require_contiguous_gpus=False, ) - manager = TestServerManager(config) + manager = FakeServerManager(config) manager.resource_manager = MagicMock() manager.resource_manager.wait_for_resources.return_value = [ @@ -283,9 +301,8 @@ def test_auto_allocation_non_contiguous(self): contiguous=False, ) - def test_get_server_logs_reads_last_n_lines(self, tmp_path, manager): + def test_get_server_logs_reads_last_n_lines(self, manager): """Test that get_server_logs returns the last N lines from the log file.""" - import tempfile # Create a temporary file and write some lines tmp = tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".log") try: @@ -315,14 +332,14 @@ def test_logs_written_to_metrics_output_dir( self, mock_popen, tmp_path, monkeypatch ): """Server logs should live inside the benchmark output directory.""" - config = VllmServerConfig( + config = _server_config( host="localhost", port=8123, gpu_ids=[0], startup_timeout=1, health_check_interval=0.1, ) - manager = TestServerManager(config) + manager = FakeServerManager(config) monkeypatch.setenv("VEEKSHA_OUTPUT_DIR", str(tmp_path)) mock_process = MagicMock() @@ -341,4 +358,3 @@ def test_logs_written_to_metrics_output_dir( assert log_path.exists() monkeypatch.delenv("VEEKSHA_OUTPUT_DIR", raising=False) - diff --git a/tests/unit/scripts/test_align_audio_trace.py b/tests/unit/scripts/test_align_audio_trace.py new file mode 100644 index 00000000..4714f11e --- /dev/null +++ b/tests/unit/scripts/test_align_audio_trace.py @@ -0,0 +1,93 @@ +import json + +import pytest + +from scripts.prepare_audio_traces import ( + NEMO_MODEL, + NeMoDockerAligner, + NemoItem, + read_nemo_word_timings, +) + + +@pytest.mark.unit +def test_attach_word_timings_maps_ctm_to_manifest_row(tmp_path) -> None: + nemo_output_dir = tmp_path / "nemo_output" + nemo_output_dir.mkdir() + nemo_manifest = tmp_path / "nemo_manifest.jsonl" + nemo_manifest.write_text("", encoding="utf-8") + ctm_path = nemo_output_dir / "clip.ctm" + ctm_path.write_text( + "utt 1 0.100 0.200 hello\n" "utt 1 1.200 0.300 world\n", + encoding="utf-8", + ) + (nemo_output_dir / "nemo_manifest_with_output_file_paths.json").write_text( + json.dumps({"word_ctm": str(ctm_path)}) + "\n", + encoding="utf-8", + ) + + word_timings = read_nemo_word_timings( + nemo_output_dir, + nemo_manifest, + [NemoItem(row_index=7, audio_path=tmp_path / "clip.wav", text="hello world")], + ) + + assert [(word.word, word.start_ms, word.end_ms) for word in word_timings[7]] == [ + ("hello", 100.0, 300.0), + ("world", 1200.0, 1500.0), + ] + + +@pytest.mark.unit +def test_nemo_forced_aligner_command_owns_docker_lifecycle(tmp_path) -> None: + command = NeMoDockerAligner( + image="nemo:test", + gpus="device=0", + cache_dir=tmp_path / "cache", + extra_mounts=("/outside:/outside:ro",), + repo_root=tmp_path, + ).command( + manifest_path=tmp_path / "manifest.jsonl", + output_dir=tmp_path / "alignment" / "nemo_output", + alignment_output_dir=tmp_path / "alignment", + ) + + assert command[:2] == ["docker", "run"] + assert "--gpus" in command + assert "device=0" in command + assert "nemo:test" in command + assert "/outside:/outside:ro" in command + assert f"pretrained_name={NEMO_MODEL}" in command + + +@pytest.mark.unit +def test_parse_nemo_output_path_that_is_relative_to_working_directory( + tmp_path, monkeypatch +) -> None: + monkeypatch.chdir(tmp_path) + nemo_output_dir = tmp_path / "alignment" / "nemo_output" + nemo_output_dir.mkdir(parents=True) + nemo_manifest = tmp_path / "alignment" / "nemo_manifest.jsonl" + nemo_manifest.write_text("", encoding="utf-8") + ctm_path = nemo_output_dir / "ctm" / "words" / "clip.ctm" + ctm_path.parent.mkdir(parents=True) + ctm_path.write_text("utt 1 0.100 0.200 hello NA lex NA\n", encoding="utf-8") + (nemo_output_dir / "nemo_manifest_with_output_file_paths.json").write_text( + json.dumps( + { + "words_level_ctm_filepath": ctm_path.relative_to(tmp_path).as_posix(), + } + ) + + "\n", + encoding="utf-8", + ) + + word_timings = read_nemo_word_timings( + nemo_output_dir, + nemo_manifest, + [NemoItem(row_index=0, audio_path=tmp_path / "clip.wav", text="hello")], + ) + + assert [(word.word, word.start_ms, word.end_ms) for word in word_timings[0]] == [ + ("hello", 100.0, 300.0), + ] diff --git a/tests/unit/scripts/test_asr_trace_sources.py b/tests/unit/scripts/test_asr_trace_sources.py new file mode 100644 index 00000000..bc363882 --- /dev/null +++ b/tests/unit/scripts/test_asr_trace_sources.py @@ -0,0 +1,158 @@ +import numpy as np +import pytest +import soundfile as sf + +from scripts import prepare_audio_traces +from scripts.prepare_audio_traces import TraceSourceOptions, build_trace_source + + +@pytest.mark.unit +def test_aa_source_filters_complete_clips_by_max_duration(monkeypatch) -> None: + samples = [ + {"transcript": "short clip", "url": "short"}, + {"transcript": "edge clip", "url": "edge"}, + {"transcript": "long clip", "url": "long"}, + ] + durations = {"short": 10.0, "edge": 30.0, "long": 30.1} + + monkeypatch.setattr( + prepare_audio_traces, + "load_dataset", + lambda *args, **kwargs: samples, + ) + monkeypatch.setattr( + prepare_audio_traces, + "fetch_aa_audio", + lambda row, repo: row["url"], + ) + + def decode_audio(source, options): + return np.zeros( + int(round(durations[source] * options.sample_rate)), + dtype=np.float32, + ) + + monkeypatch.setattr(prepare_audio_traces, "decode_audio", decode_audio) + + source = build_trace_source( + "aa_voxpopuli", + TraceSourceOptions(), + ) + clips = list(source.iter_clips()) + + final_clips = list( + prepare_audio_traces.finalize_clips( + clips, + max_duration_s=30.0, + sample_rate=TraceSourceOptions().sample_rate, + ) + ) + + assert [clip.transcript for clip in final_clips] == ["short clip", "edge clip"] + assert [clip.metadata["sample_id"] for clip in final_clips] == ["short", "edge"] + + +@pytest.mark.unit +def test_target_duration_repeats_before_and_truncates_after_alignment() -> None: + sample_rate = 2 + clip = prepare_audio_traces.TraceClip( + audio=np.arange(4, dtype=np.float32), + transcript="hello world", + duration_s=2.0, + metadata={"sample_id": "source"}, + ) + + repeated = prepare_audio_traces.repeat_clip_to_cover_target_duration( + clip, + target_duration_s=5.0, + sample_rate=sample_rate, + ) + + assert repeated.duration_s == 6.0 + assert repeated.transcript == "hello world hello world hello world" + assert repeated.word_timestamps is None + assert repeated.target_duration_s == 5.0 + assert repeated.audio.tolist() == [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] + + aligned = prepare_audio_traces.TraceClip( + audio=repeated.audio, + transcript=repeated.transcript, + duration_s=repeated.duration_s, + metadata=repeated.metadata, + target_duration_s=repeated.target_duration_s, + word_timestamps=[ + prepare_audio_traces.WordTiming("hello", 0.0, 500.0), + prepare_audio_traces.WordTiming("world", 800.0, 1500.0), + prepare_audio_traces.WordTiming("hello", 2000.0, 2500.0), + prepare_audio_traces.WordTiming("world", 2800.0, 3500.0), + prepare_audio_traces.WordTiming("hello", 4000.0, 4500.0), + prepare_audio_traces.WordTiming("world", 4800.0, 5500.0), + ], + ) + + final_clips = list( + prepare_audio_traces.finalize_clips( + [aligned], + max_duration_s=None, + target_duration_s=99.0, + sample_rate=sample_rate, + ) + ) + + assert len(final_clips) == 1 + final_clip = final_clips[0] + assert final_clip.duration_s == 5.0 + assert final_clip.audio.tolist() == [0, 1, 2, 3, 0, 1, 2, 3, 0, 1] + assert final_clip.transcript == "hello world hello world hello" + assert [word.word for word in final_clip.word_timestamps or []] == [ + "hello", + "world", + "hello", + "world", + "hello", + ] + + +@pytest.mark.unit +def test_ami_word_timed_source_yields_relative_word_timestamps( + tmp_path, + monkeypatch, +) -> None: + audio_dir = tmp_path / "audio" + words_dir = tmp_path / "words" + audio_dir.mkdir() + words_dir.mkdir() + + sf.write(audio_dir / "ES2001a.wav", np.zeros(16000, dtype=np.float32), 16000) + (words_dir / "ES2001a.A.words.xml").write_text( + """ + + hello + world + + """, + encoding="utf-8", + ) + + monkeypatch.setattr( + prepare_audio_traces.AMITraceSource, "AUDIO_DIR", str(audio_dir) + ) + monkeypatch.setattr( + prepare_audio_traces.AMITraceSource, "WORDS_DIR", str(words_dir) + ) + + source = build_trace_source( + "ami_word_timed", + TraceSourceOptions(min_duration_s=0.0, max_duration_s=2.0), + ) + clips = list(source.iter_clips()) + + assert len(clips) == 1 + clip = clips[0] + assert clip.transcript == "hello world" + assert clip.metadata["meeting_id"] == "ES2001a" + assert clip.metadata["speaker_id"] == "A" + assert clip.word_timestamps == [ + prepare_audio_traces.WordTiming(word="hello", start_ms=0.0, end_ms=200.0), + prepare_audio_traces.WordTiming(word="world", start_ms=250.0, end_ms=400.0), + ] diff --git a/tests/unit/slo/test_fluidity_slo.py b/tests/unit/slo/test_fluidity_slo.py new file mode 100644 index 00000000..0c85cb3d --- /dev/null +++ b/tests/unit/slo/test_fluidity_slo.py @@ -0,0 +1,32 @@ +from veeksha.config.slo import ConstantSloConfig +from veeksha.slo.slo import ConstantSlo + + +def test_fluidity_slo_uses_lower_tail_and_higher_is_better() -> None: + slo = ConstantSlo( + ConstantSloConfig( + metric="user_audio_fluidity_index", + percentile=0.01, + value=0.90, + ) + ) + + met, observed = slo.evaluate({"user_audio_fluidity_index": [0.95, 0.97, 0.99, 1.0]}) + + assert met + assert observed > 0.90 + + +def test_fluidity_slo_fails_when_lower_tail_is_below_threshold() -> None: + slo = ConstantSlo( + ConstantSloConfig( + metric="user_audio_fluidity_index", + percentile=0.01, + value=0.99, + ) + ) + + met, observed = slo.evaluate({"user_audio_fluidity_index": [0.80, 0.99, 1.0, 1.0]}) + + assert not met + assert observed < 0.99 diff --git a/tests/unit/slo/test_metrics.py b/tests/unit/slo/test_metrics.py index aee49632..67918199 100644 --- a/tests/unit/slo/test_metrics.py +++ b/tests/unit/slo/test_metrics.py @@ -1,33 +1,49 @@ import unittest -from veeksha.slo.metrics import extract_metric_values + +from veeksha.slo.metrics import extract_metric_values, lower_is_better + class TestMetrics(unittest.TestCase): + def test_duplex_success_metrics_are_higher_is_better(self): + for metric in ( + "audio_before_commit_ratio", + "duplex_overlap_observed", + "user_audio_fluidity_index", + "tts_service_fluidity_index", + "tts_service_fluidity_eligible", + ): + self.assertFalse(lower_is_better(metric)) + self.assertTrue(lower_is_better("ttfc")) + def test_extract_metric_values(self): metrics = { "ttfc": [0.1, 0.2, 0.3], "tbc": [[0.01, 0.02], [0.03, 0.04]], "end_to_end_latency": [1.0, 2.0, 3.0], "tpot": [0.001, 0.002, 0.003], - "other": [10, 20] + "other": [10, 20], } # Test ttfc self.assertEqual(extract_metric_values("ttfc", metrics), [0.1, 0.2, 0.3]) - + # Test tbc flattening - self.assertEqual(extract_metric_values("tbc", metrics), [0.01, 0.02, 0.03, 0.04]) - + self.assertEqual( + extract_metric_values("tbc", metrics), [0.01, 0.02, 0.03, 0.04] + ) + # Test e2e mapping self.assertEqual(extract_metric_values("e2e", metrics), [1.0, 2.0, 3.0]) - + # Test tpot mapping self.assertEqual(extract_metric_values("tpot", metrics), [0.001, 0.002, 0.003]) - + # Test unknown self.assertEqual(extract_metric_values("unknown", metrics), []) - + # Test empty self.assertEqual(extract_metric_values("ttfc", {}), []) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/verification/test_audio_verification.py b/tests/unit/verification/test_audio_verification.py new file mode 100644 index 00000000..a85bd5b3 --- /dev/null +++ b/tests/unit/verification/test_audio_verification.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from veeksha.config.verification import ( + AudioVerificationConfig, + UTMOSVerifierConfig, + WERVerifierConfig, + WhisperTranscriptionConfig, +) +from veeksha.verification import audio as audio_verification +from veeksha.verification.audio import ( + AudioVerificationError, + LocalWhisperTranscriber, + compute_wer, + normalize_text, + verify_audio_outputs, +) + + +def _wer_config( + *, threshold: float = 0.3, fail_on_threshold: bool = False +) -> AudioVerificationConfig: + return AudioVerificationConfig( + fail_on_threshold=fail_on_threshold, + wer=WERVerifierConfig(enabled=True, threshold=threshold), + ) + + +def _write_quality_fixture(output_dir: Path, references: list[str]) -> None: + metrics_dir = output_dir / "metrics" + audio_dir = output_dir / "audio_files" + metrics_dir.mkdir(parents=True) + audio_dir.mkdir() + rows = [ + {"request_id": request_id, "input_text": reference} + for request_id, reference in enumerate(references) + ] + (metrics_dir / "request_level_metrics.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" + ) + for request_id in range(len(references)): + (audio_dir / f"request_{request_id}.wav").write_bytes(b"fixture") + + +@pytest.mark.unit +def test_seed_tts_style_normalization_and_wer_match_manual_edit_counts() -> None: + assert normalize_text("Hello, ROCK'N'ROLL -- world!") == ( + "hello rock'n'roll world" + ) + assert compute_wer("one two three", "one too three") == pytest.approx(1 / 3) + assert compute_wer("one two three", "one three") == pytest.approx(1 / 3) + assert compute_wer("one two three", "one two extra three") == pytest.approx(1 / 3) + assert compute_wer("Hello, world!", "hello world") == 0.0 + assert compute_wer("", "") == 0.0 + assert compute_wer("", "unexpected") == 1.0 + + +@pytest.mark.unit +def test_tts_verification_summary_matches_manual_percentiles(tmp_path: Path) -> None: + references = ["one two three four"] * 4 + hypotheses = { + "request_0.wav": "one two three four", + "request_1.wav": "one two three", + "request_2.wav": "one two", + "request_3.wav": "", + } + _write_quality_fixture(tmp_path, references) + + summary = verify_audio_outputs( + tmp_path, + _wer_config(), + transcribe_audio=lambda path: hypotheses[path.name], + ) + + assert summary.transcribed_requests == 4 + assert summary.passed_requests == 2 + assert summary.failed_requests == 2 + assert summary.error_requests == 0 + assert summary.wer_avg == pytest.approx(0.4375) + assert summary.wer_p50 == pytest.approx(0.375) + assert summary.wer_p90 == pytest.approx(0.85) + assert summary.wer_p99 == pytest.approx(0.985) + assert summary.wer_max == 1.0 + + +@pytest.mark.unit +def test_strict_tts_verification_fails_closed_on_missing_audio(tmp_path: Path) -> None: + metrics_dir = tmp_path / "metrics" + metrics_dir.mkdir(parents=True) + (metrics_dir / "request_level_metrics.jsonl").write_text( + json.dumps({"request_id": 7, "input_text": "hello"}) + "\n", + encoding="utf-8", + ) + + with pytest.raises(AudioVerificationError, match="could not be verified"): + verify_audio_outputs( + tmp_path, + _wer_config(fail_on_threshold=True), + transcribe_audio=lambda _path: "hello", + ) + + +@pytest.mark.unit +def test_strict_tts_verification_fails_on_wer_threshold(tmp_path: Path) -> None: + _write_quality_fixture(tmp_path, ["one two three"]) + + with pytest.raises(AudioVerificationError, match="exceeded WER threshold"): + verify_audio_outputs( + tmp_path, + _wer_config(threshold=0.2, fail_on_threshold=True), + transcribe_audio=lambda _path: "one", + ) + + +@pytest.mark.unit +def test_utmos_scores_and_failures_are_counted_without_sample_bias( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _write_quality_fixture(tmp_path, ["first", "second", "third"]) + scores = { + "request_0.wav": 3.0, + "request_1.wav": 4.0, + "request_2.wav": None, + } + monkeypatch.setattr( + audio_verification, + "_utmos_predict_audio_path", + lambda path, _config: scores[path.name], + ) + config = AudioVerificationConfig( + utmos=UTMOSVerifierConfig(enabled=True, device="cpu") + ) + + summary = verify_audio_outputs(tmp_path, config) + + assert summary.utmos_evaluated == 2 + assert summary.utmos_mean == 3.5 + assert summary.utmos_median == 3.5 + assert summary.utmos_failed == 1 + assert summary.error_requests == 1 + + +@pytest.mark.unit +def test_whisper_judge_uses_pinned_language_task_and_beam() -> None: + calls = [] + + class FakeModel: + def transcribe(self, path: str, **kwargs): + calls.append((path, kwargs)) + return ( + iter([SimpleNamespace(text=" hello "), SimpleNamespace(text="world")]), + None, + ) + + transcriber = object.__new__(LocalWhisperTranscriber) + transcriber.config = AudioVerificationConfig( + wer=WERVerifierConfig( + enabled=True, + whisper=WhisperTranscriptionConfig(language="en", beam_size=3), + ) + ) + transcriber.model = FakeModel() + + assert transcriber.transcribe(Path("sample.wav")) == "hello world" + assert calls == [ + ( + "sample.wav", + {"language": "en", "task": "transcribe", "beam_size": 3}, + ) + ] + + +@pytest.mark.unit +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"language": ""}, "language is required"), + ({"beam_size": 0}, "beam_size must be >= 1"), + ], +) +def test_whisper_judge_rejects_unusable_pinned_settings( + kwargs: dict, match: str +) -> None: + with pytest.raises(ValueError, match=match): + WhisperTranscriptionConfig(**kwargs) diff --git a/tools/asr_replay_viewer/app.js b/tools/asr_replay_viewer/app.js new file mode 100644 index 00000000..ae68aba5 --- /dev/null +++ b/tools/asr_replay_viewer/app.js @@ -0,0 +1,752 @@ +const state = { + data: null, + selected: null, + sortedRequests: [], + referenceTokens: [], + receivedTokens: [], + referenceWordEls: [], + receivedWordEls: [], + referenceActiveCount: 0, + receivedActiveCount: 0, + replayTimeMs: 0, + replayMaxMs: 0, + replayPlaying: false, + replayStartedAt: 0, + replayBaseMs: 0, + audioSeekToken: 0, + raf: null, + latencyChart: null, + chartSeries: null, + chartLibraryFailed: false, +}; + +const ECHARTS_URL = "https://cdn.jsdelivr.net/npm/echarts@5.6.0/dist/echarts.min.js"; + +const els = { + dataUrl: document.getElementById("data-url"), + loadUrl: document.getElementById("load-url"), + fileInput: document.getElementById("file-input"), + runSubtitle: document.getElementById("run-subtitle"), + requestTitle: document.getElementById("request-title"), + requestMeta: document.getElementById("request-meta"), + latencyValue: document.getElementById("latency-value"), + latencyDetail: document.getElementById("latency-detail"), + latencyPanel: document.getElementById("latency-panel"), + metricStrip: document.getElementById("metric-strip"), + audio: document.getElementById("audio-player"), + playToggle: document.getElementById("play-toggle"), + replayButton: document.getElementById("replay-button"), + timeSlider: document.getElementById("time-slider"), + timeLabel: document.getElementById("time-label"), + durationLabel: document.getElementById("duration-label"), + latencyChart: document.getElementById("latency-chart"), + latencyChartEmpty: document.getElementById("latency-chart-empty"), + chartLiveMean: document.getElementById("chart-live-mean"), + chartLatest: document.getElementById("chart-latest"), + chartCount: document.getElementById("chart-count"), + chartPlayhead: document.getElementById("chart-playhead"), + chartPlayheadLabel: document.getElementById("chart-playhead-label"), + referenceStream: document.getElementById("reference-stream"), + transcriptStream: document.getElementById("transcript-stream"), + replayWarning: document.getElementById("replay-warning"), + sortSelect: document.getElementById("sort-select"), + requestTable: document.getElementById("request-table"), +}; + +let chartLibraryLoad = null; + +els.loadUrl.addEventListener("click", () => loadFromUrl(els.dataUrl.value.trim())); +els.fileInput.addEventListener("change", loadFromFile); +els.playToggle.addEventListener("click", toggleReplay); +els.replayButton.addEventListener("click", replayFromStart); +els.sortSelect.addEventListener("change", () => { + sortRequests(); + renderRequestTable(); +}); +els.timeSlider.addEventListener("input", () => { + pauseReplay(); + setReplayTime(Number(els.timeSlider.value)); + seekAudioToReplayTime(); +}); +window.addEventListener("resize", () => { + state.latencyChart?.resize(); + renderLatencyChart(); + updateLatencyChartProgress(); +}); + +ensureLatencyChartLibrary().catch(() => {}); + +const queryData = new URLSearchParams(window.location.search).get("data"); +if (queryData) { + els.dataUrl.value = queryData; + loadFromUrl(queryData); +} + +async function loadFromUrl(path) { + if (!path) { + return; + } + const url = new URL(path, `${window.location.origin}/`); + const response = await fetch(url); + if (!response.ok) { + throw new Error(`Failed to load ${url}: ${response.status}`); + } + loadData(await response.json()); +} + +function loadFromFile(event) { + const file = event.target.files?.[0]; + if (!file) { + return; + } + const reader = new FileReader(); + reader.onload = () => loadData(JSON.parse(String(reader.result || "{}"))); + reader.readAsText(file); +} + +function loadData(data) { + state.data = data; + els.runSubtitle.textContent = `${data.run_dir || "run"} Β· ${data.requests?.length || 0} ASR requests`; + sortRequests(); + selectRequest(state.sortedRequests[0] || null); +} + +function sortRequests() { + const requests = [...(state.data?.requests || [])]; + const sortKey = els.sortSelect.value; + const metric = (request, name) => Number(request.metrics?.[name] ?? -Infinity); + requests.sort((a, b) => { + if (sortKey === "request_id_asc") { + return Number(a.request_id ?? 0) - Number(b.request_id ?? 0); + } + if (sortKey === "final_wer_desc") { + return metric(b, "final_wer") - metric(a, "final_wer"); + } + if (sortKey === "rtf_desc") { + return metric(b, "rtf") - metric(a, "rtf"); + } + return metric(b, "interactivity") - metric(a, "interactivity"); + }); + state.sortedRequests = requests; +} + +function renderRequestTable() { + els.requestTable.innerHTML = state.sortedRequests + .map((request) => { + const selected = state.selected?.request_id === request.request_id ? " selected" : ""; + return ` + + `; + }) + .join(""); + + for (const row of els.requestTable.querySelectorAll(".request-row")) { + row.addEventListener("click", () => { + const requestId = row.getAttribute("data-request-id"); + selectRequest(state.sortedRequests.find((item) => String(item.request_id) === requestId)); + }); + } +} + +function selectRequest(request) { + pauseReplay(); + state.selected = request || null; + state.referenceTokens = buildReferenceTokens(request || {}); + state.receivedTokens = buildReceivedTokens(request || {}); + state.referenceWordEls = []; + state.receivedWordEls = []; + state.referenceActiveCount = 0; + state.receivedActiveCount = 0; + state.chartSeries = null; + renderRequestTable(); + if (!request) { + return; + } + + els.requestTitle.textContent = "Live call"; + els.requestMeta.textContent = `Request #${request.request_id} Β· ${request.dataset || "unknown"} Β· ${request.sample_id || request.source_id || ""}`; + els.audio.src = request.audio_url ? new URL(request.audio_url, `${window.location.origin}/`).href : ""; + + configureTimeline(request); + renderMetricStrip(request); + renderLatencyPanel(request); + renderStreams(); + renderLatencyChart(true); + setReplayTime(0); + seekAudioToReplayTime(); +} + +function configureTimeline(request) { + const durationMs = Number(request.duration_ms || 0); + const lastReferenceMs = Math.max(0, ...state.referenceTokens.map((word) => word.end_ms)); + const lastReceivedMs = Math.max(0, ...state.receivedTokens.map((word) => word.time_ms)); + state.replayMaxMs = Math.max(durationMs, lastReferenceMs, lastReceivedMs) + 300; + els.timeSlider.max = String(Math.max(100, Math.ceil(state.replayMaxMs))); + els.durationLabel.textContent = formatClock(state.replayMaxMs); +} + +function renderMetricStrip(request) { + const metrics = request.metrics || {}; + const items = [ + ["TTFC", metrics.ttfc, "ms"], + ["First partial", metrics.time_to_first_partial, "ms"], + ["Final after EOF", metrics.time_to_final_transcript, "ms"], + ["Final WER", metrics.final_wer, "%"], + ["RTF", metrics.rtf], + ]; + els.metricStrip.innerHTML = items + .map( + ([label, value, unit]) => ` +
+ ${escapeHtml(label)} + ${formatMetric(value, unit)} +
+ `, + ) + .join(""); +} + +function renderLatencyPanel(request) { + const metrics = request.metrics || {}; + const interactivity = metrics.interactivity; + const wordCount = metrics.interactivity_word_count ?? state.receivedTokens.length; + els.latencyValue.textContent = formatMetric(interactivity, "ms"); + els.latencyDetail.textContent = `${wordCount} matched words Β· ${formatMetric(metrics.ttfc, "ms")} TTFC`; + els.latencyPanel.className = `latency-panel ${latencyClass(interactivity)}`; +} + +function setReplayTime(ms) { + state.replayTimeMs = Math.max(0, Math.min(ms, state.replayMaxMs)); + els.timeSlider.value = String(Math.round(state.replayTimeMs)); + els.timeLabel.textContent = formatClock(state.replayTimeMs); + updateStreamProgress(); + updateLatencyChartProgress(); +} + +function toggleReplay() { + if (state.replayPlaying) { + pauseReplay(); + } else { + startReplay(); + } +} + +function replayFromStart() { + pauseReplay(); + setReplayTime(0); + startReplay(); +} + +async function startReplay() { + if (!state.selected) { + return; + } + if (state.replayTimeMs >= state.replayMaxMs) { + setReplayTime(0); + } + await seekAudioToReplayTime(); + state.replayPlaying = true; + state.replayBaseMs = state.replayTimeMs; + state.replayStartedAt = performance.now(); + updatePlaybackControls(); + if (els.audio.src) { + els.audio.play().catch(() => {}); + } + tickReplay(); +} + +function pauseReplay() { + state.replayPlaying = false; + if (state.raf !== null) { + cancelAnimationFrame(state.raf); + state.raf = null; + } + els.audio.pause(); + updatePlaybackControls(); +} + +function tickReplay() { + if (!state.replayPlaying) { + return; + } + const nextMs = state.replayBaseMs + performance.now() - state.replayStartedAt; + setReplayTime(nextMs); + if (nextMs >= state.replayMaxMs) { + pauseReplay(); + return; + } + state.raf = requestAnimationFrame(tickReplay); +} + +async function seekAudioToReplayTime() { + if (!els.audio.src) { + return; + } + const token = ++state.audioSeekToken; + const seconds = Math.max(0, state.replayTimeMs / 1000); + if (els.audio.readyState < 1) { + els.audio.load(); + await new Promise((resolve) => { + els.audio.addEventListener("loadedmetadata", resolve, { once: true }); + }); + if (token !== state.audioSeekToken) { + return; + } + } + await seekAudioToSeconds(seconds, token); +} + +async function seekAudioToSeconds(seconds, token) { + if (!els.audio.src || token !== state.audioSeekToken) { + return; + } + const target = Number.isFinite(els.audio.duration) + ? Math.min(seconds, els.audio.duration) + : seconds; + if (Math.abs(els.audio.currentTime - target) < 0.05) { + return; + } + const seeked = new Promise((resolve) => { + const done = () => { + clearTimeout(timeout); + els.audio.removeEventListener("seeked", done); + resolve(); + }; + const timeout = setTimeout(done, 400); + els.audio.addEventListener("seeked", done); + }); + if (Number.isFinite(els.audio.duration)) { + els.audio.currentTime = target; + } else { + els.audio.currentTime = target; + } + await seeked; +} + +function updatePlaybackControls() { + els.playToggle.classList.toggle("is-playing", state.replayPlaying); + els.playToggle.setAttribute("aria-label", state.replayPlaying ? "Pause" : "Play"); + els.playToggle.title = state.replayPlaying ? "Pause" : "Play"; +} + +function renderStreams() { + els.referenceStream.innerHTML = renderTimedWords(state.referenceTokens, "reference"); + els.transcriptStream.innerHTML = renderTimedWords(state.receivedTokens, "received"); + state.referenceWordEls = [...els.referenceStream.querySelectorAll(".stream-word")]; + state.receivedWordEls = [...els.transcriptStream.querySelectorAll(".stream-word")]; + state.referenceActiveCount = 0; + state.receivedActiveCount = 0; + updateStreamProgress(); + els.replayWarning.textContent = state.selected?.has_replay + ? "" + : "This row has scalar metrics but no precomputed replay words."; +} + +function renderTimedWords(tokens, kind) { + return tokens + .map((token) => { + const active = Number(token.time_ms || 0) <= state.replayTimeMs; + const detail = kind === "reference" + ? formatClock(Number(token.time_ms || 0)) + : `${formatClock(Number(token.time_ms || 0))}${latencySuffix(token.latency_ms)}`; + return ` + + ${escapeHtml(token.word)} + + `; + }) + .join(" "); +} + +function updateStreamProgress() { + state.referenceActiveCount = updateActiveWords( + state.referenceWordEls, + state.referenceActiveCount, + upperBoundByTime(state.referenceTokens, state.replayTimeMs), + ); + state.receivedActiveCount = updateActiveWords( + state.receivedWordEls, + state.receivedActiveCount, + upperBoundByTime(state.receivedTokens, state.replayTimeMs), + ); +} + +function updateActiveWords(elements, previousCount, nextCount) { + if (nextCount > previousCount) { + for (let index = previousCount; index < nextCount; index += 1) { + elements[index]?.classList.add("active"); + } + } else if (nextCount < previousCount) { + for (let index = nextCount; index < previousCount; index += 1) { + elements[index]?.classList.remove("active"); + } + } + return nextCount; +} + +function buildReferenceTokens(request) { + return (request.reference_words || []).map((word) => ({ + word: String(word.word || ""), + time_ms: Number(word.start_ms ?? word.end_ms ?? 0), + end_ms: Number(word.end_ms ?? word.start_ms ?? 0), + })); +} + +function buildReceivedTokens(request) { + return (request.received_words || []) + .map((word, index) => ({ + index, + word: String(word.word || ""), + time_ms: Number(word.time_ms ?? 0), + reference_end_ms: Number(word.reference_end_ms ?? 0), + latency_ms: Number(word.latency_ms ?? 0), + })) + .filter((word) => Number.isFinite(word.time_ms) && Number.isFinite(word.latency_ms)) + .sort((a, b) => a.time_ms - b.time_ms); +} + +function initLatencyChart() { + if (state.latencyChart) { + return; + } + if (!window.echarts || !els.latencyChart) { + return; + } + state.chartLibraryFailed = false; + state.latencyChart = window.echarts.init(els.latencyChart, null, { + renderer: "canvas", + }); + state.latencyChart.on("click", (params) => { + const value = Array.isArray(params.value) ? params.value : null; + if (!value) { + return; + } + pauseReplay(); + setReplayTime(Number(value[0])); + seekAudioToReplayTime(); + }); +} + +function ensureLatencyChartLibrary() { + if (window.echarts) { + initLatencyChart(); + return Promise.resolve(window.echarts); + } + if (chartLibraryLoad) { + return chartLibraryLoad; + } + + chartLibraryLoad = new Promise((resolve, reject) => { + const script = document.createElement("script"); + script.src = ECHARTS_URL; + script.async = true; + script.crossOrigin = "anonymous"; + script.onload = () => { + state.chartLibraryFailed = false; + initLatencyChart(); + renderLatencyChart(); + resolve(window.echarts); + }; + script.onerror = () => { + state.chartLibraryFailed = true; + els.latencyChartEmpty.textContent = "Could not load chart library from CDN"; + els.latencyChartEmpty.classList.remove("hidden"); + reject(new Error(`Failed to load ${ECHARTS_URL}`)); + }; + document.head.appendChild(script); + }); + + return chartLibraryLoad; +} + +function renderLatencyChart() { + const series = getLatencySeries(); + els.latencyChartEmpty.textContent = state.chartLibraryFailed + ? "Could not load chart library from CDN" + : "Latency points appear as transcript words arrive"; + els.latencyChartEmpty.classList.toggle( + "hidden", + !state.chartLibraryFailed && series.points.length > 0, + ); + + if (!state.latencyChart) { + ensureLatencyChartLibrary().catch(() => {}); + } + if (!state.latencyChart) { + return; + } + + state.latencyChart.setOption( + { + backgroundColor: "#faf9f6", + animation: false, + animationEasingUpdate: "cubicOut", + grid: { + left: 68, + right: 28, + top: 74, + bottom: 40, + }, + tooltip: { + trigger: "item", + confine: true, + borderWidth: 1, + borderColor: "rgba(23, 23, 23, 0.14)", + backgroundColor: "#ffffff", + padding: 0, + extraCssText: "box-shadow:0 12px 34px rgba(23,23,23,0.16);border-radius:8px;", + formatter: formatChartTooltip, + }, + xAxis: { + type: "value", + min: 0, + max: series.xMax, + axisLabel: { + color: "#74716c", + formatter: (value) => formatClock(Number(value)), + }, + axisLine: { lineStyle: { color: "rgba(23, 23, 23, 0.44)" } }, + axisTick: { show: false }, + splitLine: { lineStyle: { color: "rgba(23, 23, 23, 0.08)" } }, + }, + yAxis: { + type: "value", + min: series.yMin, + max: series.yMax, + axisLabel: { + color: "#74716c", + formatter: (value) => `${Math.round(Number(value))} ms`, + }, + axisLine: { lineStyle: { color: "rgba(23, 23, 23, 0.44)" } }, + axisTick: { show: false }, + splitLine: { + lineStyle: { + color: "rgba(23, 23, 23, 0.08)", + type: "dashed", + }, + }, + }, + series: [ + { + name: "Word latency", + type: "scatter", + data: series.pointData, + symbolSize: (value) => Math.max(9, Math.min(18, 9 + Math.abs(Number(value[1])) / 180)), + itemStyle: { + color: "#171717", + borderColor: "#faf9f6", + borderWidth: 2, + shadowBlur: 8, + shadowColor: "rgba(23, 23, 23, 0.16)", + }, + emphasis: { + scale: 1.45, + itemStyle: { + borderColor: "#171717", + borderWidth: 2, + shadowBlur: 16, + shadowColor: "rgba(23, 23, 23, 0.28)", + }, + }, + }, + { + name: "Running mean", + type: "line", + data: series.meanData, + showSymbol: false, + smooth: 0.3, + lineStyle: { + color: "#195e53", + width: 3, + shadowBlur: 8, + shadowColor: "rgba(25, 94, 83, 0.22)", + }, + }, + ], + }, + { + notMerge: true, + lazyUpdate: true, + }, + ); + updateLatencyChartProgress(); +} + +function getLatencySeries() { + if (state.chartSeries) { + return state.chartSeries; + } + const points = runningMeanPoints(state.receivedTokens.map((token, index) => ({ + index, + word: token.word, + time_ms: token.time_ms, + latency_ms: token.latency_ms, + reference_end_ms: token.reference_end_ms, + }))); + const latencies = points.map((point) => point.latency_ms); + const maxLatency = Math.max(500, ...latencies); + state.chartSeries = { + points, + pointData: points.map((point) => [ + point.time_ms, + point.latency_ms, + point.word, + point.reference_end_ms, + point.index, + point.mean_ms, + ]), + meanData: points.map((point) => [ + point.time_ms, + point.mean_ms, + point.word, + point.latency_ms, + point.index, + ]), + xMax: Math.max(state.replayMaxMs, ...points.map((point) => point.time_ms), 1000), + yMin: 0, + yMax: Math.ceil((maxLatency + 150) / 100) * 100, + }; + return state.chartSeries; +} + +function updateLatencyChartProgress() { + const series = getLatencySeries(); + const activeCount = upperBoundByTime(series.points, state.replayTimeMs); + const activePoints = series.points.slice(0, activeCount); + updateChartStats(activePoints); + updateChartPlayhead(series); +} + +function runningMeanPoints(points) { + let total = 0; + return points.map((point, index) => { + total += point.latency_ms; + return { + ...point, + mean_ms: total / (index + 1), + }; + }); +} + +function upperBoundByTime(points, timeMs) { + let lo = 0; + let hi = points.length; + while (lo < hi) { + const mid = Math.floor((lo + hi) / 2); + if (points[mid].time_ms <= timeMs) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; +} + +function updateChartPlayhead(series) { + if (!els.chartPlayhead) { + return; + } + const chartRect = els.latencyChart.getBoundingClientRect(); + const ratio = Math.max(0, Math.min(1, state.replayTimeMs / Math.max(series.xMax, 1))); + const x = 68 + ratio * Math.max(0, chartRect.width - 96); + els.chartPlayhead.style.transform = `translateX(${Math.round(x)}px)`; + els.chartPlayheadLabel.textContent = formatClock(state.replayTimeMs); +} + +function formatChartTooltip(params) { + const value = Array.isArray(params.value) ? params.value : []; + const seriesName = params.seriesName || ""; + const receivedMs = Number(value[0]); + const yValue = Number(value[1]); + const word = value[2] || ""; + const latencyMs = seriesName === "Word latency" ? yValue : Number(value[3]); + const meanMs = seriesName === "Word latency" ? Number(value[5]) : yValue; + const referenceEndMs = seriesName === "Word latency" ? Number(value[3]) : receivedMs - latencyMs; + const rows = [ + ["spoken end", formatClock(referenceEndMs)], + ["received", formatClock(receivedMs)], + ["latency", formatMetric(latencyMs, "ms")], + ["running mean", formatMetric(meanMs, "ms")], + ]; + + return ` +
+ ${escapeHtml(word || seriesName)} + ${rows + .map(([label, detail]) => ` +
+ ${escapeHtml(label)} + ${escapeHtml(detail)} +
+ `) + .join("")} +
+ `; +} + +function updateChartStats(activePoints) { + if (!activePoints.length) { + els.chartLiveMean.textContent = "mean n/a"; + els.chartLatest.textContent = "latest n/a"; + els.chartCount.textContent = "0 words"; + return; + } + const latest = activePoints.at(-1); + const mean = activePoints.reduce((sum, point) => sum + point.latency_ms, 0) / activePoints.length; + els.chartLiveMean.textContent = `mean ${Math.round(mean)} ms`; + els.chartLatest.textContent = `latest ${Math.round(latest.latency_ms)} ms`; + els.chartCount.textContent = `${activePoints.length} words`; +} + +function latencyClass(value) { + const latency = Number(value); + if (!Number.isFinite(latency)) { + return ""; + } + if (latency < 200) { + return "good"; + } + if (latency <= 500) { + return "warn"; + } + return "bad"; +} + +function latencySuffix(value) { + if (value === null || value === undefined || !Number.isFinite(Number(value))) { + return ""; + } + return ` Β· ${Math.round(Number(value))} ms latency`; +} + +function formatMetric(value, unit = "") { + if (value === null || value === undefined || !Number.isFinite(Number(value))) { + return "n/a"; + } + const number = Number(value); + const formatted = Math.abs(number) >= 100 ? number.toFixed(0) : number.toFixed(2); + return unit ? `${formatted} ${unit}` : formatted; +} + +function formatClock(value) { + if (!Number.isFinite(value)) { + return "not seen"; + } + if (value >= 1000) { + return `${(value / 1000).toFixed(2)} s`; + } + return `${Math.round(value)} ms`; +} + +function escapeHtml(value) { + return String(value ?? "") + .replaceAll("&", "&") + .replaceAll("<", "<") + .replaceAll(">", ">") + .replaceAll('"', """) + .replaceAll("'", "'"); +} diff --git a/tools/asr_replay_viewer/export_asr_replay.py b/tools/asr_replay_viewer/export_asr_replay.py new file mode 100644 index 00000000..490ed228 --- /dev/null +++ b/tools/asr_replay_viewer/export_asr_replay.py @@ -0,0 +1,302 @@ +#!/usr/bin/env python3 +"""Export ASR benchmark rows into a replay-viewer JSON bundle.""" + +from __future__ import annotations + +import argparse +import difflib +import json +import sys +from pathlib import Path +from typing import Any + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from veeksha.evaluator.performance.asr_interactivity import ( # noqa: E402 + _expand_reference_words, + _normalize_words, + _parse_reference_words, + _parse_transcript_snapshots, +) + +DEFAULT_TRACE_MANIFESTS = ( + REPO_ROOT / "traces" / "asr" / "aa_public" / "manifest.jsonl", + REPO_ROOT / "traces" / "asr" / "ami_word_timed" / "manifest.jsonl", +) +METRICS_DIR_NAME = "metrics" +REQUEST_METRICS_NAME = "request_level_metrics.jsonl" +SUMMARY_NAME = "summary_stats.json" +REPLAY_NAME = "asr_replay.json" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "run_dir", + type=Path, + help="Benchmark run directory, metrics directory, or request_level_metrics.jsonl.", + ) + parser.add_argument( + "--trace-manifest", + action="append", + type=Path, + default=[], + help=( + "Trace manifest to use for audio/reference backfill. May be passed " + "multiple times. Defaults to known traces/asr manifests when present." + ), + ) + parser.add_argument( + "--output", + type=Path, + default=None, + help="Output JSON path. Defaults to /metrics/asr_replay.json.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + metrics_path = resolve_metrics_path(args.run_dir) + metrics_dir = metrics_path.parent + run_dir = ( + metrics_dir.parent if metrics_dir.name == METRICS_DIR_NAME else metrics_dir + ) + output_path = args.output or metrics_dir / REPLAY_NAME + + trace_rows = load_trace_rows(args.trace_manifest or existing_default_manifests()) + request_rows = read_jsonl(metrics_path) + summary_path = metrics_dir / SUMMARY_NAME + summary = read_json(summary_path) if summary_path.exists() else {} + + requests = [ + build_replay_request(row, trace_rows) + for row in sorted(request_rows, key=lambda item: item.get("request_id", 0)) + if str(row.get("audio_task") or "").lower() == "stt" + ] + + payload = { + "schema_version": 2, + "run_dir": repo_relative_or_absolute(run_dir), + "metrics_file": repo_relative_or_absolute(metrics_path), + "summary": summary, + "requests": requests, + } + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(payload, separators=(",", ":")), encoding="utf-8") + print(f"Wrote {len(requests)} ASR replay rows -> {output_path}") + + +def resolve_metrics_path(path: Path) -> Path: + candidates = [] + if path.is_file(): + candidates.append(path) + candidates.extend( + [ + path / REQUEST_METRICS_NAME, + path / METRICS_DIR_NAME / REQUEST_METRICS_NAME, + ] + ) + for candidate in candidates: + if candidate.exists(): + return candidate.resolve() + raise FileNotFoundError(f"Could not find {REQUEST_METRICS_NAME} under {path}") + + +def existing_default_manifests() -> list[Path]: + return [path for path in DEFAULT_TRACE_MANIFESTS if path.exists()] + + +def load_trace_rows(paths: list[Path]) -> dict[tuple[str, str], dict[str, Any]]: + rows: dict[tuple[str, str], dict[str, Any]] = {} + for path in paths: + manifest_dir = path.resolve().parent + for row in read_jsonl(path): + dataset = str(row.get("dataset") or "") + for key_field in ("sample_id", "source_id"): + value = row.get(key_field) + if value is not None: + rows[(dataset, str(value))] = attach_trace_audio_url( + row, + manifest_dir, + ) + return rows + + +def attach_trace_audio_url(row: dict[str, Any], manifest_dir: Path) -> dict[str, Any]: + row = dict(row) + audio_file = row.get("audio_file") + if audio_file: + audio_path = Path(str(audio_file)) + if not audio_path.is_absolute(): + audio_path = manifest_dir / audio_path + row["audio_url"] = repo_relative_or_absolute(audio_path) + return row + + +def build_replay_request( + row: dict[str, Any], + trace_rows: dict[tuple[str, str], dict[str, Any]], +) -> dict[str, Any]: + dataset = str(row.get("dataset") or "") + sample_id = optional_str(row.get("sample_id")) + source_id = optional_str(row.get("source_id")) + trace_row = None + for key in (sample_id, source_id): + if key is not None: + trace_row = trace_rows.get((dataset, key)) + if trace_row is not None: + break + + audio_file = optional_str(row.get("audio_file")) + audio_url = audio_url_from_row(audio_file) + if audio_url is None and trace_row is not None: + audio_url = optional_str(trace_row.get("audio_url")) + audio_file = audio_file or optional_str(trace_row.get("audio_file")) + + raw_reference_words = list_or_empty(row.get("reference_word_timestamps")) + if not raw_reference_words and trace_row is not None: + raw_reference_words = list_or_empty(trace_row.get("reference_word_timestamps")) + + reference_words, received_words = build_word_timings( + raw_reference_words, + list_or_empty(row.get("transcript_snapshots")), + ) + latencies = [word["latency_ms"] for word in received_words] + interactivity = ( + sum(latencies) / len(latencies) if latencies else row.get("interactivity") + ) + + request = { + "request_id": row.get("request_id"), + "session_id": row.get("session_id"), + "dataset": dataset, + "source_id": source_id, + "sample_id": sample_id, + "audio_file": audio_file, + "audio_url": audio_url, + "duration_ms": row.get("generated_audio_duration"), + "reference_words": reference_words, + "received_words": received_words, + "metrics": { + "ttfc": row.get("ttfc"), + "time_to_first_partial": row.get("time_to_first_partial"), + "time_to_final_transcript": row.get("time_to_final_transcript"), + "interactivity": interactivity, + "interactivity_word_count": len(received_words), + "partial_wer": row.get("partial_wer"), + "final_wer": row.get("final_wer"), + "rtf": row.get("rtf"), + "chunk_count": row.get("chunk_count"), + }, + } + request["has_replay"] = bool(reference_words and received_words) + return request + + +def build_word_timings( + raw_reference_words: list[dict[str, Any]], + raw_snapshots: list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + if not raw_reference_words: + return [], [] + + reference_words = _expand_reference_words( + _parse_reference_words(raw_reference_words) + ) + reference_payload = [ + { + "word": word.word, + "start_ms": round(word.start_ms, 3), + "end_ms": round(word.end_ms, 3), + } + for word in reference_words + ] + if not raw_snapshots: + return reference_payload, [] + + first_seen_ms: list[float | None] = [None] * len(reference_words) + reference_tokens = [word.word for word in reference_words] + for snapshot in sorted( + _parse_transcript_snapshots(raw_snapshots), + key=lambda item: item.elapsed_ms, + ): + matcher = difflib.SequenceMatcher( + a=reference_tokens, + b=_normalize_words(snapshot.transcript), + autojunk=False, + ) + for tag, ref_start, ref_end, _hyp_start, _hyp_end in matcher.get_opcodes(): + if tag != "equal": + continue + for ref_index in range(ref_start, ref_end): + if ( + first_seen_ms[ref_index] is None + and snapshot.elapsed_ms >= reference_words[ref_index].start_ms + ): + first_seen_ms[ref_index] = snapshot.elapsed_ms + + received_words = [] + for word, seen_ms in zip(reference_words, first_seen_ms): + if seen_ms is None: + continue + received_words.append( + { + "word": word.word, + "time_ms": round(seen_ms, 3), + "reference_start_ms": round(word.start_ms, 3), + "reference_end_ms": round(word.end_ms, 3), + "latency_ms": round(max(0.0, seen_ms - word.end_ms), 3), + } + ) + return reference_payload, received_words + + +def audio_url_from_row(audio_file: str | None) -> str | None: + if audio_file is None: + return None + return repo_relative_or_absolute(Path(audio_file)) + + +def read_jsonl(path: Path) -> list[dict[str, Any]]: + with path.open("r", encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def read_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + value = json.load(handle) + return value if isinstance(value, dict) else {} + + +def list_or_empty(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, list): + return [] + return [dict(item) for item in value if isinstance(item, dict)] + + +def optional_str(value: Any) -> str | None: + if value is None: + return None + return str(value) + + +def repo_relative_or_absolute(path: Path) -> str: + resolved = path.resolve() + if is_relative_to(resolved, REPO_ROOT): + return resolved.relative_to(REPO_ROOT).as_posix() + return str(path) + + +def is_relative_to(path: Path, base: Path) -> bool: + try: + path.relative_to(base) + return True + except ValueError: + return False + + +if __name__ == "__main__": + main() diff --git a/tools/asr_replay_viewer/index.html b/tools/asr_replay_viewer/index.html new file mode 100644 index 00000000..53cde187 --- /dev/null +++ b/tools/asr_replay_viewer/index.html @@ -0,0 +1,143 @@ + + + + + + ASR Interactivity Replay + + + +
+
+
+

Veeksha STT benchmark

+

Realtime transcript timing

+
+
+ + + +
+
+ +
+
+

No run loaded

+

Live call

+

+ Load a replay export to compare when words are spoken with when + Veeksha receives them from the server. +

+
+ +
+ Mean word latency + n/a +

No request selected

+
+
+ +
+
+
+
+ + Live call +
+ 0 ms +
+ +
+
+
+ Latency over time + word points and running mean +
+
+ mean n/a + latest n/a + 0 words +
+
+ +
+ 0 ms +
+
+ Latency points appear as transcript words arrive +
+
+ +
+ + + + 0 ms +
+ + + +
+
+
+ Ground truth + dataset word timestamps +
+
+
+ +
+
+ Veeksha transcript + first received by client +
+
+
+
+ +

+
+ + +
+
+ + + + diff --git a/tools/asr_replay_viewer/styles.css b/tools/asr_replay_viewer/styles.css new file mode 100644 index 00000000..9e32c2bc --- /dev/null +++ b/tools/asr_replay_viewer/styles.css @@ -0,0 +1,601 @@ +:root { + color-scheme: light; + --page: #f6f2ea; + --panel: #fffcf6; + --panel-soft: #faf9f6; + --ink: #171717; + --muted: #746f67; + --faint: #bdb7ad; + --line: #ddd6ca; + --line-dark: #c9c1b5; + --green: #1d6b42; + --yellow: #7b5f16; + --red: #8f2b22; + font-family: + Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", + sans-serif; +} + +* { + box-sizing: border-box; +} + +body { + margin: 0; + min-height: 100vh; + background: var(--page); + color: var(--ink); +} + +button, +input, +select { + font: inherit; +} + +button, +.file-button { + border: 1px solid var(--line-dark); + background: var(--ink); + color: #fff; + border-radius: 999px; + min-height: 36px; + padding: 0 16px; + cursor: pointer; +} + +button:hover, +.file-button:hover { + background: #000; +} + +.shell { + width: min(1480px, 100%); + margin: 0 auto; + padding: 24px; +} + +.topbar { + display: flex; + align-items: flex-start; + justify-content: space-between; + gap: 20px; + margin-bottom: 22px; +} + +.eyebrow { + margin: 0 0 8px; + color: var(--muted); + font-size: 13px; + letter-spacing: 0; +} + +h1, +h2, +h3, +p { + margin: 0; +} + +h1 { + font-size: 30px; + line-height: 1.1; + font-weight: 680; +} + +h2 { + font-size: clamp(34px, 4.2vw, 58px); + line-height: 0.98; + font-weight: 680; + letter-spacing: 0; +} + +h3 { + font-size: 13px; + color: var(--muted); + font-weight: 650; + text-transform: uppercase; + letter-spacing: 0; +} + +.loader { + display: grid; + grid-template-columns: minmax(280px, 520px) auto auto; + gap: 8px; +} + +.loader input[type="text"], +select { + min-height: 36px; + border: 1px solid var(--line-dark); + border-radius: 999px; + padding: 0 14px; + background: var(--panel); + color: var(--ink); +} + +.file-button { + display: inline-flex; + align-items: center; + justify-content: center; + position: relative; + overflow: hidden; +} + +.file-button input { + position: absolute; + inset: 0; + opacity: 0; + cursor: pointer; +} + +.hero { + display: grid; + grid-template-columns: minmax(0, 1fr) 290px; + gap: 24px; + align-items: end; + margin-bottom: 14px; +} + +.hero-copy { + max-width: 920px; +} + +#request-meta { + margin-top: 10px; + max-width: 760px; + color: var(--muted); + font-size: 17px; + line-height: 1.45; +} + +.latency-panel { + min-height: 128px; + border: 1px solid var(--line); + border-radius: 8px; + background: var(--panel); + padding: 18px; +} + +.latency-panel span, +.latency-panel p { + color: var(--muted); + font-size: 13px; +} + +.latency-panel strong { + display: block; + margin: 8px 0 6px; + font-size: 40px; + line-height: 1; + font-weight: 680; + font-variant-numeric: tabular-nums; +} + +.latency-panel.good strong { + color: var(--green); +} + +.latency-panel.warn strong { + color: var(--yellow); +} + +.latency-panel.bad strong { + color: var(--red); +} + +.demo-stage { + display: grid; + grid-template-columns: minmax(0, 1fr) 360px; + gap: 20px; + align-items: start; +} + +.call-card, +.metrics-card, +.request-picker { + border: 1px solid var(--line); + border-radius: 8px; + background: var(--panel); +} + +.call-card { + min-height: 690px; + padding: 20px 22px; +} + +.card-head, +.stream-head, +.picker-head { + display: flex; + align-items: center; + justify-content: space-between; + gap: 12px; +} + +.card-label { + color: var(--muted); + font-size: 14px; +} + +.status-dot { + display: inline-block; + width: 8px; + height: 8px; + margin-right: 7px; + border-radius: 50%; + background: #ef4444; + box-shadow: 0 0 0 5px rgba(239, 68, 68, 0.12); +} + +.time-label { + color: var(--muted); + font-size: 13px; + font-variant-numeric: tabular-nums; +} + +.latency-chart-wrap { + position: relative; + height: 320px; + margin-top: 18px; + border: 1px solid rgba(23, 23, 23, 0.08); + border-radius: 8px; + overflow: hidden; + background: var(--panel-soft); +} + +.chart-toolbar { + position: absolute; + top: 14px; + left: 16px; + right: 16px; + z-index: 2; + display: flex; + align-items: flex-start; + justify-content: space-between; + gap: 16px; + pointer-events: none; +} + +.chart-title { + display: block; + font-size: 14px; + font-weight: 680; +} + +.chart-toolbar small { + display: block; + margin-top: 2px; + color: var(--muted); + font-size: 12px; +} + +.chart-stats { + display: flex; + flex-wrap: wrap; + justify-content: flex-end; + gap: 6px; +} + +.chart-stats span { + display: inline-flex; + align-items: center; + min-height: 25px; + border: 1px solid rgba(23, 23, 23, 0.08); + border-radius: 999px; + background: rgba(255, 255, 255, 0.82); + padding: 0 9px; + color: var(--muted); + font-size: 12px; + font-variant-numeric: tabular-nums; +} + +.chart-tooltip { + min-width: 178px; + padding: 10px 12px; + color: var(--ink); + font-size: 12px; + background: #fff; +} + +.chart-tooltip strong { + display: block; + margin-bottom: 7px; + font-size: 13px; +} + +.chart-tooltip div { + display: flex; + justify-content: space-between; + gap: 14px; + padding-top: 5px; + border-top: 1px solid rgba(23, 23, 23, 0.08); +} + +.chart-tooltip span { + color: var(--muted); +} + +.chart-tooltip b { + font-variant-numeric: tabular-nums; +} + +#latency-chart { + width: 100%; + height: 100%; + display: block; +} + +.chart-playhead { + position: absolute; + top: 74px; + bottom: 40px; + left: 0; + z-index: 3; + width: 0; + border-left: 2px dashed #171717; + filter: drop-shadow(0 0 6px rgba(23, 23, 23, 0.18)); + pointer-events: none; + will-change: transform; +} + +.chart-playhead span { + position: absolute; + top: -30px; + left: 8px; + border: 1px solid rgba(23, 23, 23, 0.12); + border-radius: 8px; + background: #fff; + padding: 3px 7px; + color: #171717; + font-size: 12px; + font-variant-numeric: tabular-nums; + white-space: nowrap; +} + +.latency-chart-empty { + position: absolute; + inset: 56px 0 0; + display: grid; + place-items: center; + color: var(--muted); + background: rgba(246, 242, 234, 0.38); + pointer-events: none; +} + +.hidden { + display: none; +} + +.controls { + display: grid; + grid-template-columns: auto auto minmax(160px, 1fr) 72px; + gap: 10px; + align-items: center; + margin: 14px 0 18px; +} + +.icon-button { + display: inline-grid; + place-items: center; + width: 38px; + min-width: 38px; + height: 38px; + min-height: 38px; + padding: 0; +} + +.icon-button svg { + width: 20px; + height: 20px; + fill: currentColor; +} + +.icon-button .icon-pause, +.icon-button.is-playing .icon-play { + display: none; +} + +.icon-button.is-playing .icon-pause { + display: block; +} + +.controls input[type="range"] { + width: 100%; + accent-color: var(--ink); +} + +#duration-label { + color: var(--muted); + text-align: right; + font-size: 13px; + font-variant-numeric: tabular-nums; +} + +.stream-grid { + display: grid; + grid-template-columns: repeat(2, minmax(0, 1fr)); + gap: 16px; +} + +.stream-block { + border-top: 1px solid var(--line); + padding-top: 14px; +} + +.stream-head span { + font-size: 15px; + font-weight: 680; +} + +.stream-head small { + color: var(--muted); +} + +.word-stream { + min-height: 108px; + padding: 12px 0 2px; + font-size: clamp(24px, 2.6vw, 40px); + line-height: 1.16; + letter-spacing: 0; +} + +.stream-word { + color: var(--faint); + transition: color 140ms ease; +} + +.stream-word.active { + color: var(--ink); +} + +.replay-warning { + margin-top: 18px; + color: var(--red); + font-size: 13px; +} + +.side-panel { + display: grid; + gap: 16px; +} + +.metrics-card, +.request-picker { + padding: 16px; +} + +.metric-strip { + display: grid; + gap: 10px; + margin-top: 14px; +} + +.metric { + display: flex; + align-items: center; + justify-content: space-between; + gap: 12px; + border-top: 1px solid var(--line); + padding-top: 10px; +} + +.metric span { + color: var(--muted); + font-size: 13px; +} + +.metric strong { + font-variant-numeric: tabular-nums; + font-size: 15px; +} + +.request-picker { + max-height: calc(100vh - 320px); + min-height: 420px; + display: flex; + flex-direction: column; +} + +.picker-head { + margin-bottom: 12px; +} + +.picker-head select { + max-width: 168px; +} + +.request-table { + overflow: auto; + border-top: 1px solid var(--line); +} + +.request-row { + display: grid; + grid-template-columns: 56px minmax(0, 1fr) 80px; + gap: 10px; + align-items: center; + width: 100%; + min-height: 48px; + border: 0; + border-bottom: 1px solid var(--line); + border-radius: 0; + padding: 8px 0; + background: transparent; + color: var(--ink); + text-align: left; +} + +.request-row:hover, +.request-row.selected { + background: transparent; +} + +.request-row.selected .request-text, +.request-row.selected .request-number { + color: var(--ink); + font-weight: 700; +} + +.request-number, +.request-latency { + color: var(--muted); + font-size: 12px; + font-variant-numeric: tabular-nums; +} + +.request-text { + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + color: var(--muted); + font-size: 13px; +} + +.request-latency { + text-align: right; +} + +audio { + display: none; +} + +@media (max-width: 1120px) { + .topbar, + .hero, + .demo-stage { + grid-template-columns: 1fr; + } + + .loader { + width: 100%; + grid-template-columns: 1fr auto auto; + } + + .stream-grid { + grid-template-columns: 1fr; + } + + .request-picker { + max-height: 420px; + } +} + +@media (max-width: 680px) { + .shell { + padding: 16px; + } + + .loader, + .controls { + grid-template-columns: 1fr; + } + + #duration-label { + text-align: left; + } + + .word-stream { + font-size: 30px; + } +} diff --git a/veeksha/benchmark.py b/veeksha/benchmark.py index f5de901a..d289dc9d 100644 --- a/veeksha/benchmark.py +++ b/veeksha/benchmark.py @@ -1,4 +1,6 @@ import os +import sys +import sysconfig import threading import time from dataclasses import replace @@ -13,19 +15,15 @@ ) from veeksha.client.registry import ClientRegistry from veeksha.config.benchmark import BenchmarkConfig +from veeksha.config.endpoint import EndpointConfig from veeksha.core.seeding import SeedManager from veeksha.core.thread_pool import ThreadPoolManager -from veeksha.core.tokenizer import ( - TokenizerProvider, - build_hf_tokenizer_handle_from_model, -) from veeksha.core.trace_recorder import TraceRecorder from veeksha.generator.session.registry import SessionGeneratorRegistry from veeksha.health import HealthChecker from veeksha.logger import init_logger -from veeksha.orchestration import managed_server +from veeksha.orchestration.benchmark_orchestrator import managed_server from veeksha.traffic.registry import TrafficSchedulerRegistry -from veeksha.types import ChannelModality from veeksha.wandb_integration import ( maybe_finish_wandb_run, maybe_init_wandb_run, @@ -39,6 +37,26 @@ logger = init_logger(__name__) +def _warn_if_gil_enabled(stage: str) -> None: + """Warn when the GIL is active on a free-threaded build. + + A C extension that does not declare free-threading support re-enables the + GIL at import time unless the process runs with ``-Xgil=0`` / + ``PYTHON_GIL=0``. This can happen mid-run (e.g. the first + ``librosa.load`` lazily imports ``msgpack``), silently serializing every + client worker thread and invalidating high-concurrency measurements. + """ + if not sysconfig.get_config_var("Py_GIL_DISABLED"): + return + if sys._is_gil_enabled(): + logger.warning( + "The GIL is enabled at %s on a free-threaded Python build; " + "client worker threads serialize on it. Launch with -Xgil=0 or " + "PYTHON_GIL=0 to keep it disabled.", + stage, + ) + + def _maybe_pregenerate_sessions(benchmark_config, session_generator) -> Optional[list]: """Pre-generate sessions when enabled in runtime config.""" if not ( @@ -77,10 +95,23 @@ def _run_main_loop( ) -> None: """Run the main benchmark loop with all workers.""" logger.info("Starting main loop") + _warn_if_gil_enabled("benchmark start") if benchmark_start_time is None: benchmark_start_time = time.monotonic() - client_queues = [Queue() for _ in range(runtime_config.num_client_threads)] + num_client_threads = runtime_config.num_client_threads + if num_client_threads is None: + # Provision client workers for the offered load (the sweep planner + # already does this; direct configs get the same protection): an + # under-provisioned pool serializes per-session sends and shows up + # as phantom server-side latency at high concurrency. + target_sessions = getattr( + traffic_scheduler, "target_concurrent_sessions", None + ) or getattr(traffic_scheduler, "_target_concurrent", None) + num_client_threads = ( + max(3, -(-int(target_sessions) // 8)) if target_sessions else 3 + ) + client_queues = [Queue() for _ in range(num_client_threads)] output_queue = Queue() stop_event = threading.Event() generator_lock = threading.Lock() @@ -165,6 +196,10 @@ def _run_main_loop( logger.info("Interrupted, stopping") pending_in_flight = set() + # The GIL can flip on mid-run via lazy extension imports; re-check so a + # serialized run is at least loudly reported. + _warn_if_gil_enabled("benchmark end") + stop_event.set() pool_manager.join_pool("prefetch", timeout=1.0) pool_manager.join_pool("dispatch", timeout=1.0) @@ -197,11 +232,8 @@ def _run_benchmark( seed_manager = SeedManager(benchmark_config.seed) # get session generator - model_name = benchmark_config.client.model - tokenizer_provider = TokenizerProvider( - {ChannelModality.TEXT: build_hf_tokenizer_handle_from_model(model_name)}, - model_name=model_name, - ) + tokenizer_provider = benchmark_config.client.build_tokenizer_provider() + append_min_tokens_instruction = False if ( hasattr(benchmark_config.client, "use_min_tokens_prompt_fallback") @@ -302,12 +334,23 @@ def _run_benchmark( logger.info("Finalizing evaluator...") # finalize and save results + finalize_started_at = time.monotonic() result = evaluator.finalize() + logger.info( + "Benchmark phase 'evaluator_finalize' took %.2fs", + time.monotonic() - finalize_started_at, + ) + save_started_at = time.monotonic() evaluator.save(f"{benchmark_config.output_dir}/metrics") + logger.info( + "Benchmark phase 'evaluator_save' took %.2fs", + time.monotonic() - save_started_at, + ) # health checks logger.info("Running health checks...") + health_started_at = time.monotonic() health_checker = HealthChecker( trace_file=f"{benchmark_config.output_dir}/traces/dispatch_trace.jsonl", metrics_file=f"{benchmark_config.output_dir}/metrics/request_level_metrics.jsonl", @@ -316,10 +359,25 @@ def _run_benchmark( health_checker.run_and_save( f"{benchmark_config.output_dir}/health_check_results.txt" ) + logger.info( + "Benchmark phase 'health_checks' took %.2fs", + time.monotonic() - health_started_at, + ) return result +def _with_endpoint( + benchmark_config: BenchmarkConfig, endpoint: EndpointConfig +) -> BenchmarkConfig: + return replace( + benchmark_config, + client=endpoint.apply_to_client_config(benchmark_config.client), + endpoint=endpoint, + server=None, + ) + + def manage_benchmark_run( benchmark_config: BenchmarkConfig, ): @@ -340,35 +398,34 @@ def manage_benchmark_run( if benchmark_config.server is not None: logger.info(f"Launching {benchmark_config.server.engine} server...") - - with managed_server( - benchmark_config.server, output_dir=benchmark_config.output_dir - ) as server_info: - logger.info(f"Server ready at {server_info['api_base']}") - - # server dictates client - updated_client_config = replace( - benchmark_config.client, - api_base=server_info["api_base"], - api_key=server_info["api_key"], - model=benchmark_config.server.model, - ) - updated_benchmark_config = replace( - benchmark_config, - client=updated_client_config, - server=None, - ) - - maybe_init_wandb_run(updated_benchmark_config, run_kind="benchmark") - try: - result = _run_benchmark(updated_benchmark_config) - maybe_log_benchmark_scalars(updated_benchmark_config.output_dir) - maybe_log_benchmark_artifacts(updated_benchmark_config) - return result - finally: + updated_benchmark_config = None + result = None + try: + with managed_server( + benchmark_config.server, output_dir=benchmark_config.output_dir + ) as server_info: + endpoint = server_info["endpoint"] + logger.info(f"Server ready at {endpoint.api_base}") + + updated_benchmark_config = _with_endpoint(benchmark_config, endpoint) + + maybe_init_wandb_run(updated_benchmark_config, run_kind="benchmark") + try: + result = _run_benchmark(updated_benchmark_config) + finally: + logger.info("Server shutting down...") + + maybe_log_benchmark_scalars(updated_benchmark_config.output_dir) + maybe_log_benchmark_artifacts(updated_benchmark_config) + return result + finally: + if updated_benchmark_config is not None: maybe_finish_wandb_run(updated_benchmark_config.output_dir) - logger.info("Server shutting down...") else: + if benchmark_config.endpoint is not None: + benchmark_config = _with_endpoint( + benchmark_config, benchmark_config.endpoint + ) maybe_init_wandb_run(benchmark_config, run_kind="benchmark") try: result = _run_benchmark(benchmark_config) diff --git a/veeksha/benchmark_utils.py b/veeksha/benchmark_utils.py index 1919baa5..f31147ae 100644 --- a/veeksha/benchmark_utils.py +++ b/veeksha/benchmark_utils.py @@ -1,17 +1,19 @@ """Utilities used by the benchmark runner.""" import hashlib +import json import os import shutil import time from datetime import datetime -from typing import Any, Dict, Set, Tuple +from pathlib import Path +from typing import Any, Callable, Dict, Set, Tuple import yaml from tqdm import tqdm -from vidhi import dataclass_to_dict from veeksha.config.benchmark import BenchmarkConfig +from veeksha.config.utils import to_serializable_config_dict from veeksha.core.seeding import SeedManager from veeksha.evaluator.base import BaseEvaluator from veeksha.evaluator.composite import CompositeEvaluator @@ -40,7 +42,7 @@ def _persist_config_yaml(benchmark_config: BenchmarkConfig) -> str: Path to the persisted YAML file. """ os.makedirs(benchmark_config.output_dir, exist_ok=True) - config_dict = dataclass_to_dict(benchmark_config) + config_dict = to_serializable_config_dict(benchmark_config) config_path = os.path.join(benchmark_config.output_dir, "config.yml") with open(config_path, "w", encoding="utf-8") as config_file: yaml.safe_dump( @@ -157,6 +159,8 @@ def build_evaluator( "output_dir": f"{benchmark_config.output_dir}/metrics", "benchmark_start_time": benchmark_start_time, } + if cfg.get_type() == EvaluationType.PERFORMANCE: + kwargs["client_type"] = benchmark_config.client.get_type() if cfg.get_type() == EvaluationType.ACCURACY_LMEVAL: kwargs["session_generator"] = session_generator evaluator_instances.append(EvaluatorRegistry.get(cfg.get_type(), **kwargs)) @@ -217,6 +221,35 @@ def _update_pbar( state["last_completed"] = total_done +def _progress_writer(max_sessions: int) -> Callable[[int], None]: + """Return a callable that publishes benchmark progress as JSON. + + When ``VEEKSHA_PROGRESS_FILE`` is set (the launcher points the benchmark at + a file there), the returned callable writes ``{"completed", "total"}`` to it + atomically so consumers get structured progress without scraping the console. + Returns a no-op when the variable is unset (standalone benchmark runs). + """ + target = os.environ.get("VEEKSHA_PROGRESS_FILE") + if not target: + return lambda completed: None + + path = Path(target) + tmp_path = path.with_name(path.name + ".tmp") + total = max_sessions if max_sessions > 0 else None + + def write(completed: int) -> None: + try: + tmp_path.write_text( + json.dumps({"completed": completed, "total": total}), + encoding="utf-8", + ) + os.replace(tmp_path, path) + except OSError as exc: + logger.debug("Failed to write progress file %s: %s", path, exc) + + return write + + def _monitor_for_completion( traffic_scheduler, evaluator, @@ -238,6 +271,10 @@ def _monitor_for_completion( timeout_start: float = 0.0 in_flight_remaining: Set[str] = set() + write_progress = _progress_writer(max_sessions) + write_progress(0) + last_written = 0 + try: while True: time.sleep(0.1) @@ -247,6 +284,9 @@ def _monitor_for_completion( elapsed = time.monotonic() - benchmark_start _update_pbar(pbar, time_based_progress, elapsed, total_done, pbar_state) + if total_done != last_written: + write_progress(total_done) + last_written = total_done if ( not timeout_triggered diff --git a/veeksha/capacity_search.py b/veeksha/capacity_search.py index cf427d70..2830c46c 100644 --- a/veeksha/capacity_search.py +++ b/veeksha/capacity_search.py @@ -14,7 +14,6 @@ from typing import Any, Callable, Dict, Optional, Tuple, cast import yaml -from vidhi import dataclass_to_dict from veeksha.benchmark import manage_benchmark_run from veeksha.config.benchmark import BenchmarkConfig @@ -24,6 +23,7 @@ FixedIntervalGeneratorConfig, ) from veeksha.config.traffic import ConcurrentTrafficConfig, RateTrafficConfig +from veeksha.config.utils import to_serializable_config_dict from veeksha.logger import init_logger from veeksha.wandb_integration import ( dedup_tags, @@ -37,7 +37,7 @@ def _persist_capacity_search_config_yaml(config: CapacitySearchConfig) -> str: """Write the resolved capacity search configuration to config.yml.""" os.makedirs(config.output_dir, exist_ok=True) - config_dict = dataclass_to_dict(config) + config_dict = to_serializable_config_dict(config) config_path = os.path.join(config.output_dir, "config.yml") with open(config_path, "w", encoding="utf-8") as config_file: yaml.safe_dump( diff --git a/veeksha/client/__init__.py b/veeksha/client/__init__.py index 0b40cb69..5267b062 100644 --- a/veeksha/client/__init__.py +++ b/veeksha/client/__init__.py @@ -3,11 +3,15 @@ from veeksha.client.base import BaseLLMClient from veeksha.client.openai_chat import OpenAIChatCompletionsClient from veeksha.client.openai_router import OpenAIRouterClient +from veeksha.client.realtime_tts import RealtimeTTSClient from veeksha.client.registry import ClientRegistry +from veeksha.client.tts import TTSClient __all__ = [ "BaseLLMClient", "OpenAIChatCompletionsClient", "OpenAIRouterClient", + "RealtimeTTSClient", + "TTSClient", "ClientRegistry", ] diff --git a/veeksha/client/native_http_tts.py b/veeksha/client/native_http_tts.py new file mode 100644 index 00000000..858cac08 --- /dev/null +++ b/veeksha/client/native_http_tts.py @@ -0,0 +1,250 @@ +"""Native complete-text HTTP clients for cloud TTS providers.""" + +from __future__ import annotations + +import os +import threading +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Optional +from urllib.parse import quote, urlencode, urljoin + +import httpx + +from veeksha.client.base import BaseLLMClient +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.request import Request +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.core.response import ChannelResponse, RequestResult +from veeksha.logger import init_logger +from veeksha.types import AudioTask, ChannelModality + +if TYPE_CHECKING: + from veeksha.config.client import ( + DeepgramFluxHTTPClientConfig, + ElevenLabsHTTPTTSClientConfig, + ) + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class NativeHTTPRequest: + url: str + headers: dict[str, str] + payload: dict[str, Any] + + +class ElevenLabsHTTPProtocol: + provider = "elevenlabs" + protocol_name = "v1_text_to_speech" + + def __init__(self, config: "ElevenLabsHTTPTTSClientConfig", api_key: str) -> None: + self.config = config + self.api_key = api_key + + def build_request(self, api_base: str, text: str) -> NativeHTTPRequest: + normalized = api_base.rstrip("/") + "/" + path = f"v1/text-to-speech/{quote(self.config.voice_id, safe='')}" + query = urlencode({"output_format": f"pcm_{self.config.sample_rate}"}) + return NativeHTTPRequest( + url=f"{urljoin(normalized, path)}?{query}", + headers={ + "xi-api-key": self.api_key, + "Content-Type": "application/json", + }, + payload={ + "text": text, + "model_id": self.config.model, + "voice_settings": { + "stability": self.config.stability, + "similarity_boost": self.config.similarity_boost, + "speed": self.config.speed, + }, + "apply_text_normalization": self.config.apply_text_normalization, + }, + ) + + +class DeepgramFluxHTTPProtocol: + provider = "deepgram" + protocol_name = "v2_flux_speak_http" + + def __init__(self, config: "DeepgramFluxHTTPClientConfig", api_key: str) -> None: + self.config = config + self.api_key = api_key + + def build_request(self, api_base: str, text: str) -> NativeHTTPRequest: + normalized = api_base.rstrip("/") + "/" + query = urlencode( + { + "model": self.config.model, + "encoding": "linear16", + "container": "none", + "sample_rate": self.config.sample_rate, + "mip_opt_out": str(self.config.mip_opt_out).lower(), + } + ) + return NativeHTTPRequest( + url=f"{urljoin(normalized, 'v2/speak')}?{query}", + headers={ + "Authorization": f"Token {self.api_key}", + "Content-Type": "application/json", + }, + payload={"text": text}, + ) + + +class NativeNonStreamingTTSClient(BaseLLMClient): + """Measure a complete-input/complete-output HTTP TTS operation.""" + + def __init__( + self, config: Any, protocol_factory: Callable[[Any, str], Any] + ) -> None: + configured_key = config.api_key + super().__init__(config) + api_key = configured_key or os.environ.get(config.api_key_env) + if not api_key: + raise ValueError( + f"API key is required: set client.api_key or {config.api_key_env}" + ) + self.api_key = api_key + self._http_config = config + self._protocol = protocol_factory(config, api_key) + self._client_storage = threading.local() + + def _get_client(self) -> httpx.AsyncClient: + if not hasattr(self._client_storage, "client"): + self._client_storage.client = httpx.AsyncClient( + timeout=self._http_config.request_timeout + ) + return self._client_storage.client + + async def send_request( + self, + request: Request, + session_id: int, + session_total_requests: int = 1, + on_request_sent: Optional[Callable[[], None]] = None, + on_request_dispatched: Optional[Callable[[], None]] = None, + ) -> RequestResult: + text_content = request.channels.get(ChannelModality.TEXT) + if not isinstance(text_content, TextChannelRequestContent): + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + success=False, + error_code=400, + error_msg="No TEXT channel in request for TTS", + client_completed_at=time.monotonic(), + ) + + input_text = text_content.input_text + native_request = self._protocol.build_request(str(self.api_base), input_text) + t_start = time.monotonic() + error_code: Optional[int] = None + error_msg: Optional[str] = None + audio_data = b"" + + try: + if on_request_dispatched is not None: + on_request_dispatched() + response = await self._get_client().post( + native_request.url, + headers=native_request.headers, + json=native_request.payload, + timeout=self._http_config.request_timeout, + ) + response.raise_for_status() + audio_data = response.content + content_type = response.headers.get("Content-Type", "") + if content_type and not content_type.lower().startswith( + ("audio/", "application/octet-stream", "binary/octet-stream") + ): + raise ValueError(f"Expected audio response, got {content_type!r}") + if not audio_data: + raise ValueError("TTS provider returned an empty audio response") + if on_request_sent is not None: + on_request_sent() + except httpx.HTTPStatusError as exc: + error_code = exc.response.status_code + error_msg = str(exc) + except httpx.ConnectError as exc: + error_code = 503 + error_msg = str(exc) + except httpx.TimeoutException: + error_code = 408 + error_msg = "TTS request timed out" + except Exception as exc: # noqa: BLE001 - converted to request result. + error_code = 502 + error_msg = str(exc) + + completed_at = time.monotonic() + total_latency_ms = (completed_at - t_start) * 1000 + success = error_code is None and error_msg is None + if not success: + logger.warning( + "%s non-streaming TTS error: (%s) %s", + self._protocol.provider, + error_code, + error_msg, + ) + if on_request_sent is not None: + on_request_sent() + + channels: dict = {} + if success: + terminal_ms = round(total_latency_ms, 3) + channels[ChannelModality.AUDIO] = ChannelResponse( + modality=ChannelModality.AUDIO, + content=audio_data, + metrics={ + "audio_task": AudioTask.TTS, + AudioMetricKey.PROVIDER.value: self._protocol.provider, + AudioMetricKey.PROVIDER_MODEL.value: self._http_config.model, + AudioMetricKey.PROVIDER_PROTOCOL.value: ( + self._protocol.protocol_name + ), + AudioMetricKey.TTFC.value: terminal_ms, + AudioMetricKey.END_TO_END_LATENCY.value: terminal_ms, + AudioMetricKey.CHUNK_COUNT.value: 1, + AudioMetricKey.RAW_PCM.value: True, + AudioMetricKey.SAMPLE_RATE.value: self._http_config.sample_rate, + AudioMetricKey.INPUT_CHARS.value: len(input_text), + AudioMetricKey.INPUT_TOKENS.value: ( + text_content.target_prompt_tokens or 0 + ), + AudioMetricKey.INPUT_TEXT.value: input_text, + AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value: [ + [0.0, len(input_text)] + ], + AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value: [ + [terminal_ms, len(audio_data)] + ], + AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value: 0.0, + AudioMetricKey.AUDIO_DONE_OFFSET_MS.value: terminal_ms, + AudioMetricKey.RESPONSE_DONE_OFFSET_MS.value: terminal_ms, + }, + ) + + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + channels=channels, + success=success, + error_code=error_code, + error_msg=error_msg, + client_completed_at=completed_at, + ) + + +class ElevenLabsHTTPTTSClient(NativeNonStreamingTTSClient): + def __init__(self, config: "ElevenLabsHTTPTTSClientConfig", **kwargs) -> None: + super().__init__(config, ElevenLabsHTTPProtocol) + + +class DeepgramFluxHTTPClient(NativeNonStreamingTTSClient): + def __init__(self, config: "DeepgramFluxHTTPClientConfig", **kwargs) -> None: + super().__init__(config, DeepgramFluxHTTPProtocol) diff --git a/veeksha/client/native_streaming_tts.py b/veeksha/client/native_streaming_tts.py new file mode 100644 index 00000000..ca5af9c5 --- /dev/null +++ b/veeksha/client/native_streaming_tts.py @@ -0,0 +1,540 @@ +"""Native cloud-provider WebSocket clients for streaming text-to-speech. + +Both clients emit Veeksha's transport-independent raw PCM timing contract, so +the same first-playable, stall, RTF, duplex-overlap, and fluidity evaluators can +be applied to Vajra, ElevenLabs, and Deepgram without provider-side clocks. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import os +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Optional +from urllib.parse import quote, urlencode, urljoin + +from websockets.asyncio.client import connect +from websockets.exceptions import InvalidHandshake, InvalidStatus + +from veeksha.client.base import BaseLLMClient +from veeksha.client.utils import TextDeltaPacer, segment_text +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.request import Request +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.core.response import ChannelResponse, RequestResult +from veeksha.logger import init_logger +from veeksha.types import AudioTask, ChannelModality + +if TYPE_CHECKING: + from veeksha.config.client import ( + DeepgramAuraStreamingTTSClientConfig, + DeepgramFluxStreamingTTSClientConfig, + ElevenLabsStreamingTTSClientConfig, + ) + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class NativeProtocolEvent: + """One normalized provider event.""" + + audio: bytes = b"" + ready: bool = False + response_started: bool = False + audio_done: bool = False + terminal: bool = False + error: Optional[str] = None + + +class NativeStreamingTTSError(Exception): + """Fatal provider event received after a successful WebSocket handshake.""" + + +class ElevenLabsStreamingProtocol: + provider = "elevenlabs" + protocol_name = "v1_stream_input" + has_ready_event = False + + def __init__( + self, config: "ElevenLabsStreamingTTSClientConfig", api_key: str + ) -> None: + self.config = config + self.api_key = api_key + + def build_ws_url(self, api_base: str) -> str: + normalized = api_base.rstrip("/") + "/" + path = f"v1/text-to-speech/{quote(self.config.voice_id, safe='')}/stream-input" + query = urlencode( + { + "model_id": self.config.model, + "output_format": f"pcm_{self.config.sample_rate}", + "auto_mode": str(self.config.auto_mode).lower(), + "apply_text_normalization": self.config.apply_text_normalization, + } + ) + url = urljoin(normalized, path) + if url.startswith("https://"): + url = "wss://" + url[len("https://") :] + elif url.startswith("http://"): + url = "ws://" + url[len("http://") :] + return f"{url}?{query}" + + def headers(self) -> dict[str, str]: + return {"xi-api-key": self.api_key} + + def initial_messages(self) -> list[str]: + payload: dict[str, Any] = { + "text": " ", + "voice_settings": { + "stability": self.config.stability, + "similarity_boost": self.config.similarity_boost, + "speed": self.config.speed, + }, + } + if not self.config.auto_mode: + payload["generation_config"] = { + "chunk_length_schedule": self.config.chunk_length_schedule + } + return [json.dumps(payload)] + + def text_message(self, text: str) -> str: + return json.dumps({"text": text}) + + def finish_message(self) -> str: + # An empty text message flushes buffered text and closes the generation. + return json.dumps({"text": ""}) + + def parse(self, raw: str | bytes) -> NativeProtocolEvent: + if isinstance(raw, bytes): + return NativeProtocolEvent(error="Unexpected binary ElevenLabs frame") + try: + event = json.loads(raw) + except json.JSONDecodeError, TypeError, ValueError: + return NativeProtocolEvent() + if not isinstance(event, dict): + return NativeProtocolEvent() + error = event.get("error") + if error: + if isinstance(error, dict): + message = error.get("message") or json.dumps(error) + else: + message = str(error) + return NativeProtocolEvent(error=message) + audio = b"" + encoded_audio = event.get("audio") + if isinstance(encoded_audio, str) and encoded_audio: + try: + audio = base64.b64decode(encoded_audio, validate=True) + except (ValueError, TypeError) as exc: + return NativeProtocolEvent(error=f"Invalid ElevenLabs audio: {exc}") + terminal = bool(event.get("isFinal")) + return NativeProtocolEvent( + audio=audio, + response_started=bool(audio), + audio_done=terminal, + terminal=terminal, + ) + + +class DeepgramFluxStreamingProtocol: + provider = "deepgram" + protocol_name = "v2_flux_speak" + has_ready_event = True + + def __init__( + self, config: "DeepgramFluxStreamingTTSClientConfig", api_key: str + ) -> None: + self.config = config + self.api_key = api_key + + def build_ws_url(self, api_base: str) -> str: + normalized = api_base.rstrip("/") + "/" + query = urlencode( + { + "model": self.config.model, + "encoding": "linear16", + "sample_rate": self.config.sample_rate, + "mip_opt_out": str(self.config.mip_opt_out).lower(), + } + ) + url = urljoin(normalized, "v2/speak") + if url.startswith("https://"): + url = "wss://" + url[len("https://") :] + elif url.startswith("http://"): + url = "ws://" + url[len("http://") :] + return f"{url}?{query}" + + def headers(self) -> dict[str, str]: + return {"Authorization": f"Token {self.api_key}"} + + def initial_messages(self) -> list[str]: + return [] + + def text_message(self, text: str) -> str: + return json.dumps({"type": "Speak", "text": text}) + + def finish_message(self) -> str: + return json.dumps({"type": "Flush"}) + + def parse(self, raw: str | bytes) -> NativeProtocolEvent: + if isinstance(raw, bytes): + return NativeProtocolEvent(audio=raw, response_started=bool(raw)) + try: + event = json.loads(raw) + except json.JSONDecodeError, TypeError, ValueError: + return NativeProtocolEvent() + if not isinstance(event, dict): + return NativeProtocolEvent() + event_type = event.get("type") + if event_type == "Error": + return NativeProtocolEvent( + error=str(event.get("description") or event.get("code") or event) + ) + return NativeProtocolEvent( + ready=event_type == "Connected", + response_started=event_type == "SpeechStarted", + audio_done=event_type == "SpeechMetadata", + terminal=event_type == "SpeechMetadata", + ) + + +class DeepgramAuraStreamingProtocol: + """Deepgram Aura continuous-text WebSocket protocol. + + The shared client sends text incrementally and records whether Aura emits + audio before the final ``Flush`` instead of assuming a particular buffering + policy. + """ + + provider = "deepgram" + protocol_name = "v1_aura_speak" + has_ready_event = False + + def __init__( + self, config: "DeepgramAuraStreamingTTSClientConfig", api_key: str + ) -> None: + self.config = config + self.api_key = api_key + + def build_ws_url(self, api_base: str) -> str: + normalized = api_base.rstrip("/") + "/" + query = urlencode( + { + "model": self.config.model, + "encoding": "linear16", + "sample_rate": self.config.sample_rate, + "mip_opt_out": str(self.config.mip_opt_out).lower(), + "speed": self.config.speed, + } + ) + url = urljoin(normalized, "v1/speak") + if url.startswith("https://"): + url = "wss://" + url[len("https://") :] + elif url.startswith("http://"): + url = "ws://" + url[len("http://") :] + return f"{url}?{query}" + + def headers(self) -> dict[str, str]: + return {"Authorization": f"Token {self.api_key}"} + + def initial_messages(self) -> list[str]: + return [] + + def text_message(self, text: str) -> str: + return json.dumps({"type": "Speak", "text": text}) + + def finish_message(self) -> str: + return json.dumps({"type": "Flush"}) + + def parse(self, raw: str | bytes) -> NativeProtocolEvent: + if isinstance(raw, bytes): + return NativeProtocolEvent(audio=raw, response_started=bool(raw)) + try: + event = json.loads(raw) + except json.JSONDecodeError, TypeError, ValueError: + return NativeProtocolEvent() + if not isinstance(event, dict): + return NativeProtocolEvent() + event_type = event.get("type") + if event_type == "Error": + return NativeProtocolEvent( + error=str(event.get("description") or event.get("code") or event) + ) + return NativeProtocolEvent( + audio_done=event_type == "Flushed", + terminal=event_type == "Flushed", + ) + + +def _round_ms(value: Optional[float]) -> Optional[float]: + return round(value, 3) if value is not None else None + + +def _flatten_exception(exc: BaseException) -> BaseException: + if not isinstance(exc, BaseExceptionGroup): + return exc + leaves: list[BaseException] = [] + + def collect(node: BaseException) -> None: + if isinstance(node, BaseExceptionGroup): + for child in node.exceptions: + collect(child) + else: + leaves.append(node) + + collect(exc) + return leaves[0] if leaves else exc + + +def _map_error(exc: BaseException) -> tuple[int, str]: + if isinstance(exc, NativeStreamingTTSError): + return 500, str(exc) + if isinstance(exc, InvalidStatus): + body = exc.response.body + detail = "" + if isinstance(body, bytes) and body: + detail = body.decode("utf-8", errors="replace")[:500] + elif body: + detail = str(body)[:500] + message = str(exc) if not detail else f"{exc}: {detail}" + return exc.response.status_code, message + if isinstance(exc, TimeoutError): + return 408, "Streaming TTS request timed out" + if isinstance(exc, (InvalidHandshake, OSError)): + return 503, str(exc) + return 520, str(exc) + + +class NativeStreamingTTSClient(BaseLLMClient): + """Shared paced-input/raw-PCM measurement loop for native cloud APIs.""" + + def __init__( + self, config: Any, protocol_factory: Callable[[Any, str], Any] + ) -> None: + configured_key = config.api_key + super().__init__(config) + api_key = configured_key or os.environ.get(config.api_key_env) + if not api_key: + raise ValueError( + f"API key is required: set client.api_key or {config.api_key_env}" + ) + self.api_key = api_key + self._streaming_config = config + self._protocol = protocol_factory(config, api_key) + + def _connect(self): + open_timeout = min(self._streaming_config.request_timeout, 30) + return connect( + self._protocol.build_ws_url(str(self.api_base)), + max_size=None, + compression=None, + open_timeout=open_timeout, + additional_headers=self._protocol.headers(), + ) + + async def send_request( + self, + request: Request, + session_id: int, + session_total_requests: int = 1, + on_request_sent: Optional[Callable[[], None]] = None, + on_request_dispatched: Optional[Callable[[], None]] = None, + ) -> RequestResult: + text_content = request.channels.get(ChannelModality.TEXT) + if not isinstance(text_content, TextChannelRequestContent): + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + success=False, + error_code=400, + error_msg="No TEXT channel in request for streaming TTS", + client_completed_at=time.monotonic(), + ) + + input_text = text_content.input_text + pacing = self._streaming_config.pacing + segments = segment_text(input_text, pacing.tokens_per_delta) + pacer = TextDeltaPacer(pacing, seed=pacing.seed + request.id) + input_tokens = text_content.target_prompt_tokens or sum( + segment.n_tokens for segment in segments + ) + + audio_chunks: list[bytes] = [] + audio_chunk_ts: list[list[float]] = [] + text_delta_ts: list[list[float]] = [] + ttfc: Optional[float] = None + ws_connect_latency: Optional[float] = None + session_ready_offset: Optional[float] = None + response_trigger_offset: Optional[float] = None + response_created_offset: Optional[float] = None + input_complete_offset: Optional[float] = None + audio_done_offset: Optional[float] = None + response_done_offset: Optional[float] = None + + sent_fired = False + + def fire_sent_once() -> None: + nonlocal sent_fired + if not sent_fired and on_request_sent is not None: + on_request_sent() + sent_fired = True + + t_start = time.monotonic() + + async def send_loop(ws) -> None: + nonlocal input_complete_offset, response_trigger_offset + if pacer.initial_delay_s > 0: + await asyncio.sleep(pacer.initial_delay_s) + deadline = time.monotonic() + for segment in segments: + deadline += pacer.next_gap() + sleep_s = deadline - time.monotonic() + if sleep_s > 0: + await asyncio.sleep(sleep_s) + offset_ms = (time.monotonic() - t_start) * 1000 + if response_trigger_offset is None: + # Native streaming APIs begin accepting synthesis work with + # the first real text message. Protocol setup payloads such + # as ElevenLabs' single-space initializer are deliberately + # excluded from this semantic trigger. + response_trigger_offset = offset_ms + await ws.send(self._protocol.text_message(segment.text)) + text_delta_ts.append([offset_ms, segment.n_chars]) + input_complete_offset = (time.monotonic() - t_start) * 1000 + await ws.send(self._protocol.finish_message()) + + async def recv_loop(ws) -> None: + nonlocal ttfc, session_ready_offset, response_created_offset + nonlocal audio_done_offset, response_done_offset + while True: + raw = await ws.recv() + wire_offset_ms = (time.monotonic() - t_start) * 1000 + event = self._protocol.parse(raw) + if event.error: + raise NativeStreamingTTSError(event.error) + if event.ready and session_ready_offset is None: + session_ready_offset = wire_offset_ms + if event.response_started and response_created_offset is None: + response_created_offset = wire_offset_ms + if event.audio: + playable_offset_ms = (time.monotonic() - t_start) * 1000 + audio_chunks.append(event.audio) + audio_chunk_ts.append([playable_offset_ms, len(event.audio)]) + if ttfc is None: + ttfc = wire_offset_ms + if response_created_offset is None: + response_created_offset = wire_offset_ms + fire_sent_once() + if event.audio_done and audio_done_offset is None: + audio_done_offset = wire_offset_ms + if event.terminal: + response_done_offset = wire_offset_ms + return + + error_code: Optional[int] = None + error_msg: Optional[str] = None + try: + async with asyncio.timeout(self._streaming_config.request_timeout): + async with self._connect() as ws: + ws_connect_latency = (time.monotonic() - t_start) * 1000 + for message in self._protocol.initial_messages(): + await ws.send(message) + if not self._protocol.has_ready_event: + session_ready_offset = (time.monotonic() - t_start) * 1000 + if on_request_dispatched is not None: + on_request_dispatched() + async with asyncio.TaskGroup() as task_group: + task_group.create_task(send_loop(ws)) + task_group.create_task(recv_loop(ws)) + except TimeoutError: + error_code = 408 + error_msg = "Streaming TTS request timed out" + except Exception as exc: # noqa: BLE001 - converted to request result. + error_code, error_msg = _map_error(_flatten_exception(exc)) + logger.warning( + "%s streaming TTS error: (%s) %s", + self._protocol.provider, + error_code, + error_msg, + ) + + completed_at = time.monotonic() + total_latency_ms = (completed_at - t_start) * 1000 + success = error_code is None and error_msg is None + fire_sent_once() + + metrics = { + "audio_task": AudioTask.TTS, + AudioMetricKey.PROVIDER.value: self._protocol.provider, + AudioMetricKey.PROVIDER_MODEL.value: self._streaming_config.model, + AudioMetricKey.PROVIDER_PROTOCOL.value: self._protocol.protocol_name, + AudioMetricKey.TTFC.value: round(ttfc or 0.0, 3), + AudioMetricKey.END_TO_END_LATENCY.value: round(total_latency_ms, 3), + AudioMetricKey.CHUNK_COUNT.value: len(audio_chunks), + AudioMetricKey.RAW_PCM.value: True, + AudioMetricKey.SAMPLE_RATE.value: self._streaming_config.sample_rate, + AudioMetricKey.INPUT_CHARS.value: len(input_text), + AudioMetricKey.INPUT_TOKENS.value: input_tokens, + AudioMetricKey.INPUT_TEXT.value: input_text, + AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value: text_delta_ts, + AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value: audio_chunk_ts, + AudioMetricKey.WS_CONNECT_LATENCY_MS.value: _round_ms(ws_connect_latency), + AudioMetricKey.SESSION_READY_OFFSET_MS.value: _round_ms( + session_ready_offset + ), + AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value: _round_ms( + response_trigger_offset + ), + AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value: _round_ms( + input_complete_offset + ), + AudioMetricKey.RESPONSE_CREATED_OFFSET_MS.value: _round_ms( + response_created_offset + ), + AudioMetricKey.AUDIO_DONE_OFFSET_MS.value: _round_ms(audio_done_offset), + AudioMetricKey.RESPONSE_DONE_OFFSET_MS.value: _round_ms( + response_done_offset + ), + } + + channels: dict = {} + if success or audio_chunks or text_delta_ts: + channels[ChannelModality.AUDIO] = ChannelResponse( + modality=ChannelModality.AUDIO, + content=b"".join(audio_chunks), + metrics=metrics, + ) + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + channels=channels, + success=success, + error_code=error_code, + error_msg=error_msg, + client_completed_at=completed_at, + ) + + +class ElevenLabsStreamingTTSClient(NativeStreamingTTSClient): + def __init__(self, config: "ElevenLabsStreamingTTSClientConfig", **kwargs) -> None: + super().__init__(config, ElevenLabsStreamingProtocol) + + +class DeepgramFluxStreamingTTSClient(NativeStreamingTTSClient): + def __init__( + self, config: "DeepgramFluxStreamingTTSClientConfig", **kwargs + ) -> None: + super().__init__(config, DeepgramFluxStreamingProtocol) + + +class DeepgramAuraStreamingTTSClient(NativeStreamingTTSClient): + def __init__( + self, config: "DeepgramAuraStreamingTTSClientConfig", **kwargs + ) -> None: + super().__init__(config, DeepgramAuraStreamingProtocol) diff --git a/veeksha/client/openai_chat.py b/veeksha/client/openai_chat.py index b802243f..d124e483 100644 --- a/veeksha/client/openai_chat.py +++ b/veeksha/client/openai_chat.py @@ -1,11 +1,13 @@ from __future__ import annotations +import base64 import time from typing import TYPE_CHECKING, Any, Callable, List, Optional import httpx # type: ignore from veeksha.client.openai_base import OpenAIBaseClient +from veeksha.core.audio_contract import DEFAULT_AUDIO_SAMPLE_RATE, AudioMetricKey from veeksha.core.request import Request from veeksha.core.request_content import ( AudioChannelRequestContent, @@ -16,7 +18,7 @@ from veeksha.core.response import ChannelResponse, RequestResult from veeksha.core.tokenizer import TokenizerProvider from veeksha.logger import init_logger -from veeksha.types import ChannelModality +from veeksha.types import AudioTask, ChannelModality if TYPE_CHECKING: from veeksha.config.client import OpenAIChatCompletionsClientConfig @@ -277,7 +279,7 @@ def _build_channel_responses( channels[ChannelModality.AUDIO] = ChannelResponse( modality=ChannelModality.AUDIO, content=audio_data, - metrics={}, + metrics={"audio_task": AudioTask.LLM_AUDIO}, ) if video_data is not None: @@ -338,11 +340,16 @@ async def send_request( # multimodal response data image_data: Optional[Any] = None - audio_data: Optional[Any] = None video_data: Optional[Any] = None + # audio output tracking + audio_chunks: List[bytes] = [] + audio_ttfc: Optional[float] = None + audio_chunk_count = 0 + delta_prompt_len = 0 messages = [] + t_start = time.monotonic() try: messages, delta_prompt_len = self._build_message_content(request) @@ -382,7 +389,8 @@ async def send_request( } client = self._get_client() - most_recent_token_time = time.monotonic() + t_start = time.monotonic() + most_recent_token_time = t_start async with client.stream( "POST", self.chat_address, @@ -417,7 +425,15 @@ async def send_request( ) continue - if delta.get("content"): + modality_type = data.get("modality") + + if modality_type == "audio" and delta.get("content"): + chunk_bytes = base64.b64decode(delta["content"]) + if audio_ttfc is None: + audio_ttfc = (receive_time - t_start) * 1000 + audio_chunks.append(chunk_bytes) + audio_chunk_count += 1 + elif delta.get("content"): ( generated_text, chunks_received, @@ -440,9 +456,6 @@ async def send_request( # TODO: image deltas image_data = self._process_image_response(delta, image_data) - # TODO: audio deltas - audio_data = self._process_audio_response(delta, audio_data) - # TODO: video deltas video_data = self._process_video_response(delta, video_data) @@ -488,10 +501,27 @@ async def send_request( total_prompt_len=num_total_prompt_tokens, tokens_received=num_completion_tokens, image_data=image_data, - audio_data=audio_data, + audio_data=None, video_data=video_data, ) + # Build audio channel from streamed audio chunks + if success and audio_chunks: + audio_bytes = b"".join(audio_chunks) + total_latency_ms = (completed_at - t_start) * 1000 + channels[ChannelModality.AUDIO] = ChannelResponse( + modality=ChannelModality.AUDIO, + content=audio_bytes, + metrics={ + "audio_task": AudioTask.LLM_AUDIO, + AudioMetricKey.TTFC.value: round(audio_ttfc or 0.0, 3), + AudioMetricKey.END_TO_END_LATENCY.value: round(total_latency_ms, 3), + AudioMetricKey.CHUNK_COUNT.value: audio_chunk_count, + AudioMetricKey.RAW_PCM.value: True, + AudioMetricKey.SAMPLE_RATE.value: DEFAULT_AUDIO_SAMPLE_RATE, + }, + ) + return RequestResult( request_id=request.id, session_id=session_id, diff --git a/veeksha/client/realtime_tts.py b/veeksha/client/realtime_tts.py new file mode 100644 index 00000000..edc8a23e --- /dev/null +++ b/veeksha/client/realtime_tts.py @@ -0,0 +1,538 @@ +"""WebSocket client for text-to-speech over the OpenAI Realtime protocol. + +It sends paced text conversation items at an emulated upstream-LLM decode rate, +receives audio chunks, and records the per-event millisecond timestamps consumed +by the audio interactivity evaluator. + +Structurally mirrors the streaming HTTP :class:`~veeksha.client.tts.TTSClient` +(same metric contract), but the transport is a +WebSocket and the request is a paced sequence of ``conversation.item.create`` +events. In complete-text mode, an audio-only ``response.create`` follows the +last item. In duplex mode, it is sent after a configurable amount of initial +text and later items continue arriving while the response is active. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import time +from collections.abc import Callable +from enum import StrEnum +from typing import TYPE_CHECKING, Optional +from urllib.parse import quote, urljoin + +from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosedError, InvalidHandshake, InvalidStatus + +from veeksha.client.base import BaseLLMClient +from veeksha.client.utils import TextDeltaPacer, segment_text +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.request import Request +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.core.response import ChannelResponse, RequestResult +from veeksha.logger import init_logger +from veeksha.types import AudioTask, ChannelModality + +if TYPE_CHECKING: + from veeksha.config.client import RealtimeTTSClientConfig + +logger = init_logger(__name__) + + +class RealtimeEventKind(StrEnum): + SESSION_UPDATED = "session_updated" + RESPONSE_CREATED = "response_created" + AUDIO_DELTA = "audio_delta" + AUDIO_DONE = "audio_done" + RESPONSE_DONE = "response_done" + ERROR = "error" + OTHER = "other" + + +class RealtimeTTSProtocol: + """OpenAI Realtime wire contract used by the benchmark client.""" + + _EVENT_KINDS = { + "session.updated": RealtimeEventKind.SESSION_UPDATED, + "response.created": RealtimeEventKind.RESPONSE_CREATED, + "response.output_audio.delta": RealtimeEventKind.AUDIO_DELTA, + "response.output_audio.done": RealtimeEventKind.AUDIO_DONE, + "response.done": RealtimeEventKind.RESPONSE_DONE, + "error": RealtimeEventKind.ERROR, + } + + def __init__( + self, config: "RealtimeTTSClientConfig", api_key: Optional[str] + ) -> None: + self.config = config + self._api_key = api_key + + @property + def raw_pcm(self) -> bool: + return True + + def build_ws_url(self, api_base: str) -> str: + ws_base = api_base + if ws_base.startswith("https://"): + ws_base = "wss://" + ws_base[len("https://") :] + elif ws_base.startswith("http://"): + ws_base = "ws://" + ws_base[len("http://") :] + normalized_base = ws_base.rstrip("/") + path = "realtime" if normalized_base.endswith("/v1") else "v1/realtime" + url = urljoin(f"{normalized_base}/", path) + return f"{url}?model={quote(self.config.model, safe='')}" + + def headers(self) -> dict[str, str]: + if not self._api_key: + return {} + return {"Authorization": f"Bearer {self._api_key}"} + + def session_update_json(self) -> str: + output: dict = { + "format": {"type": "audio/pcm", "rate": self.config.sample_rate} + } + if self.config.voice_id: + output["voice"] = self.config.voice_id + session = { + "type": "realtime", + "output_modalities": ["audio"], + "audio": {"output": output}, + } + return json.dumps({"type": "session.update", "session": session}) + + def conversation_item_create_json(self, text: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": text}], + }, + } + ) + + def response_create_json(self) -> str: + return json.dumps( + { + "type": "response.create", + "response": {"output_modalities": ["audio"]}, + } + ) + + def classify(self, event_type: str) -> RealtimeEventKind: + return self._EVENT_KINDS.get(event_type, RealtimeEventKind.OTHER) + + def extract_audio(self, event: dict) -> bytes: + encoded = event.get("delta") + return base64.b64decode(encoded) if encoded else b"" + + +def _build_realtime_protocol( + config: "RealtimeTTSClientConfig", api_key: Optional[str] = None +) -> RealtimeTTSProtocol: + return RealtimeTTSProtocol( + config, api_key=config.api_key if api_key is None else api_key + ) + + +# --------------------------------------------------------------------------- +# Server-side error signalling +# --------------------------------------------------------------------------- + + +class RealtimeServerError(Exception): + """Raised when the server sends an ``error`` event mid-stream. + + Carries the decoded error event payload so the client can surface the + server-provided message. + """ + + def __init__(self, event: dict) -> None: + self.event = event + super().__init__(str(event)) + + +def _validate_completed_response(event: dict) -> None: + """Raise when a terminal Realtime response did not complete successfully.""" + response = event.get("response") + if not isinstance(response, dict) or response.get("status") != "completed": + raise RealtimeServerError(event) + + +# --------------------------------------------------------------------------- +# Client +# --------------------------------------------------------------------------- + + +def _round_ms(value: Optional[float]) -> Optional[float]: + return round(value, 3) if value is not None else None + + +def _extract_output_sample_rate(session: dict) -> Optional[int]: + """Read ``session.audio.output.format.rate`` from a GA session object.""" + audio = session.get("audio") + if isinstance(audio, dict): + output = audio.get("output") + if isinstance(output, dict): + fmt = output.get("format") + if isinstance(fmt, dict) and isinstance(fmt.get("rate"), int): + return fmt["rate"] + return None + + +class RealtimeTTSClient(BaseLLMClient): + """Async WebSocket client for OpenAI Realtime text-to-speech.""" + + def __init__(self, config: "RealtimeTTSClientConfig", **kwargs) -> None: + # **kwargs swallows tokenizer_provider (matching TTSClient); realtime + # TTS token counts come from whitespace segmentation instead. + super().__init__(config) + self._realtime_config = config + self._protocol = _build_realtime_protocol(config, api_key=self.api_key) + + def _connect(self): + """Return the websocket connect context manager. + + A seam so tests can override the transport. ``max_size=None`` lifts the + inbound-frame cap (audio deltas can be large); ``compression=None`` + keeps binary PCM uncompressed. The asyncio transport sets ``TCP_NODELAY`` + by default, so no explicit socket option is required. + """ + open_timeout = min(self._realtime_config.request_timeout, 30) + return connect( + self._protocol.build_ws_url(str(self.api_base)), + max_size=None, + compression=None, + open_timeout=open_timeout, + additional_headers=self._protocol.headers(), + ) + + async def send_request( + self, + request: Request, + session_id: int, + session_total_requests: int = 1, + on_request_sent: Optional[Callable[[], None]] = None, + on_request_dispatched: Optional[Callable[[], None]] = None, + ) -> RequestResult: + """Stream text deltas over a websocket and collect audio metrics.""" + text_content = request.channels.get(ChannelModality.TEXT) + if not isinstance(text_content, TextChannelRequestContent): + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + success=False, + error_code=400, + error_msg="No TEXT channel in request for realtime TTS", + client_completed_at=time.monotonic(), + ) + + input_text = text_content.input_text + pacing = self._realtime_config.pacing + segments = segment_text(input_text, pacing.tokens_per_delta) + pacer = TextDeltaPacer(pacing, seed=pacing.seed + request.id) + input_tokens = text_content.target_prompt_tokens or sum( + seg.n_tokens for seg in segments + ) + + logger.debug( + "[RealtimeTTS] request_id=%d session_id=%d chars=%d deltas=%d", + request.id, + session_id, + len(input_text), + len(segments), + ) + + # Collected state (shared with the nested send/recv loops via closure). + audio_chunks: list[bytes] = [] + audio_chunk_ts: list[list[float]] = [] # [[offset_ms, n_bytes], ...] + text_delta_ts: list[list[float]] = [] # [[offset_ms, n_chars], ...] + ttfc: Optional[float] = None + ws_connect_latency: Optional[float] = None + session_ready_offset: Optional[float] = None + response_created_offset: Optional[float] = None + response_trigger_offset: Optional[float] = None + input_complete_offset: Optional[float] = None + audio_done_offset: Optional[float] = None + response_done_offset: Optional[float] = None + sample_rate = self._realtime_config.sample_rate + + sent_fired = False + + def fire_sent_once() -> None: + nonlocal sent_fired + if not sent_fired and on_request_sent is not None: + on_request_sent() + sent_fired = True + + error_code: Optional[int] = None + error_msg: Optional[str] = None + + t_start = time.monotonic() + + async def send_loop(ws) -> None: + nonlocal input_complete_offset, response_trigger_offset + if pacer.initial_delay_s > 0: + await asyncio.sleep(pacer.initial_delay_s) + # Pace by absolute deadlines so ws.send backpressure never + # accumulates drift into subsequent gaps. + deadline = time.monotonic() + response_triggered = False + sent_tokens = 0 + for seg in segments: + deadline += pacer.next_gap() + sleep_s = deadline - time.monotonic() + if sleep_s > 0: + await asyncio.sleep(sleep_s) + await ws.send(self._protocol.conversation_item_create_json(seg.text)) + text_delta_ts.append([(time.monotonic() - t_start) * 1000, seg.n_chars]) + sent_tokens += seg.n_tokens + if ( + self._realtime_config.input_output_mode == "duplex" + and not response_triggered + and sent_tokens >= self._realtime_config.duplex_start_after_tokens + ): + response_trigger_offset = (time.monotonic() - t_start) * 1000 + await ws.send(self._protocol.response_create_json()) + response_triggered = True + input_complete_offset = (time.monotonic() - t_start) * 1000 + if not response_triggered: + response_trigger_offset = (time.monotonic() - t_start) * 1000 + await ws.send(self._protocol.response_create_json()) + + async def recv_loop(ws) -> None: + nonlocal ttfc, session_ready_offset, response_created_offset + nonlocal audio_done_offset, response_done_offset, sample_rate + while True: + raw = await ws.recv() + # Stamp receipt BEFORE any json/base64 decode work. + offset_ms = (time.monotonic() - t_start) * 1000 + try: + event = json.loads(raw) + except json.JSONDecodeError, TypeError, ValueError: + continue + if not isinstance(event, dict): + continue + kind = self._protocol.classify(event.get("type", "")) + + if kind is RealtimeEventKind.AUDIO_DELTA: + chunk = self._protocol.extract_audio(event) + if chunk: + # TTFC retains the historical first-wire-audio meaning, + # while the supply timeline represents decoded playable + # PCM availability at the client. + playable_offset_ms = (time.monotonic() - t_start) * 1000 + audio_chunks.append(chunk) + audio_chunk_ts.append([playable_offset_ms, len(chunk)]) + if ttfc is None: + ttfc = offset_ms + fire_sent_once() + elif kind is RealtimeEventKind.SESSION_UPDATED: + if session_ready_offset is None: + session_ready_offset = offset_ms + session = event.get("session") + if isinstance(session, dict): + server_sr = _extract_output_sample_rate(session) + if ( + isinstance(server_sr, int) + and server_sr > 0 + and server_sr != sample_rate + ): + logger.warning( + "Realtime server echoed sample_rate=%d (configured " + "%d); using server value.", + server_sr, + self._realtime_config.sample_rate, + ) + sample_rate = server_sr + elif kind is RealtimeEventKind.RESPONSE_CREATED: + if response_created_offset is None: + response_created_offset = offset_ms + elif kind is RealtimeEventKind.AUDIO_DONE: + if audio_done_offset is None: + audio_done_offset = offset_ms + elif kind is RealtimeEventKind.RESPONSE_DONE: + response_done_offset = offset_ms + _validate_completed_response(event) + return # Terminal: stop on response.done, not audio.done. + elif kind is RealtimeEventKind.ERROR: + raise RealtimeServerError(event) + # OTHER: keep receiving until response.done. + + try: + async with asyncio.timeout(self._realtime_config.request_timeout): + async with self._connect() as ws: + ws_connect_latency = (time.monotonic() - t_start) * 1000 + await ws.send(self._protocol.session_update_json()) + # Analog of the HTTP-200 ack: the scheduler's dispatch pacing + # advances on this callback. + if on_request_dispatched is not None: + on_request_dispatched() + + async with asyncio.TaskGroup() as task_group: + task_group.create_task(send_loop(ws)) + task_group.create_task(recv_loop(ws)) + except TimeoutError: + error_code = 408 + error_msg = "Realtime TTS request timed out" + logger.warning("Realtime TTS timeout: (%s) %s", error_code, error_msg) + except Exception as exc: # noqa: BLE001 - mapped to error codes below. + error_code, error_msg = _map_error(_flatten_exception(exc)) + logger.warning("Realtime TTS error: (%s) %s", error_code, error_msg) + + completed_at = time.monotonic() + total_latency_ms = (completed_at - t_start) * 1000 + success = error_code is None and error_msg is None + + # on_request_sent fallback: fire exactly once even on error paths. + # With ordering="prefill" the client_runner's DispatchTracker only + # advances via on_request_sent (or when send_request RAISES); a request + # that fails *before* first audio would otherwise deadlock the tracker, + # since we return an error result rather than raising. Firing here keeps + # the sequential launch moving. + fire_sent_once() + + audio_data = b"".join(audio_chunks) if audio_chunks else b"" + + metrics = { + "audio_task": AudioTask.TTS, + AudioMetricKey.TTFC.value: round(ttfc or 0.0, 3), + AudioMetricKey.END_TO_END_LATENCY.value: round(total_latency_ms, 3), + AudioMetricKey.CHUNK_COUNT.value: len(audio_chunks), + AudioMetricKey.RAW_PCM.value: self._protocol.raw_pcm, + AudioMetricKey.SAMPLE_RATE.value: sample_rate, + AudioMetricKey.INPUT_CHARS.value: len(input_text), + AudioMetricKey.INPUT_TOKENS.value: input_tokens, + AudioMetricKey.INPUT_TEXT.value: input_text, + AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value: text_delta_ts, + AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value: audio_chunk_ts, + AudioMetricKey.WS_CONNECT_LATENCY_MS.value: _round_ms(ws_connect_latency), + AudioMetricKey.SESSION_READY_OFFSET_MS.value: _round_ms( + session_ready_offset + ), + AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS.value: _round_ms( + response_trigger_offset + ), + AudioMetricKey.RESPONSE_CREATED_OFFSET_MS.value: _round_ms( + response_created_offset + ), + AudioMetricKey.INPUT_COMMIT_OFFSET_MS.value: _round_ms( + input_complete_offset + ), + AudioMetricKey.AUDIO_DONE_OFFSET_MS.value: _round_ms(audio_done_offset), + AudioMetricKey.RESPONSE_DONE_OFFSET_MS.value: _round_ms( + response_done_offset + ), + } + + channels: dict = {} + has_partial = bool(audio_chunks or text_delta_ts) + if success or has_partial: + channels[ChannelModality.AUDIO] = ChannelResponse( + modality=ChannelModality.AUDIO, + content=audio_data, + metrics=metrics, + ) + + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + channels=channels, + success=success, + error_code=error_code, + error_msg=error_msg, + client_completed_at=completed_at, + ) + + +# --------------------------------------------------------------------------- +# Error flattening / mapping +# --------------------------------------------------------------------------- + + +def _flatten_exception(exc: BaseException) -> BaseException: + """Reduce a (possibly nested) ExceptionGroup to its most specific leaf. + + ``asyncio.TaskGroup`` raises an ``ExceptionGroup`` bundling the failing + task exceptions; connect-phase failures propagate bare. Collect the leaves + and pick by priority so the mapping below sees the exception that actually + determines the error code. + """ + if not isinstance(exc, BaseExceptionGroup): + return exc + + leaves: list[BaseException] = [] + + def _collect(node: BaseException) -> None: + if isinstance(node, BaseExceptionGroup): + for sub in node.exceptions: + _collect(sub) + else: + leaves.append(node) + + _collect(exc) + if not leaves: + return exc + + priority = ( + RealtimeServerError, + InvalidStatus, + InvalidHandshake, + ConnectionClosedError, + TimeoutError, + OSError, + ) + for exc_type in priority: + for leaf in leaves: + if isinstance(leaf, exc_type): + return leaf + return leaves[0] + + +def _server_error_message(event: dict) -> str: + """Extract the useful message from an error or failed response event.""" + error = event.get("error") + if isinstance(error, dict): + message = error.get("message") + if isinstance(message, str) and message: + return message + + response = event.get("response") + if isinstance(response, dict): + status_details = response.get("status_details") + if isinstance(status_details, dict): + error = status_details.get("error") + if isinstance(error, dict): + message = error.get("message") + if isinstance(message, str) and message: + return message + reason = status_details.get("reason") + if isinstance(reason, str) and reason: + return reason + + status = response.get("status") + if isinstance(status, str) and status: + return f"Realtime response ended with status {status}" + + return json.dumps(event) + + +def _map_error(exc: BaseException) -> tuple[int, str]: + """Map a flattened exception to an (error_code, error_msg) pair.""" + if isinstance(exc, RealtimeServerError): + return 500, _server_error_message(exc.event) + if isinstance(exc, InvalidStatus): + # Handshake rejected: surface the server's HTTP status code. + return exc.response.status_code, str(exc) + if isinstance(exc, TimeoutError): + return 408, "Realtime TTS request timed out" + if isinstance(exc, (InvalidHandshake, OSError)): + # Connect-phase failure (DNS/refused/handshake): unreachable server. + return 503, str(exc) + # ConnectionClosedError and anything else: unclassified transport failure. + return 520, str(exc) diff --git a/veeksha/client/registry.py b/veeksha/client/registry.py index 1794640d..f83587e0 100644 --- a/veeksha/client/registry.py +++ b/veeksha/client/registry.py @@ -32,3 +32,67 @@ def get_key_from_str(cls, key_str: str) -> ClientType: "OpenAICompletionsClient", ), ) + +ClientRegistry.register( + ClientType.TTS, + _LazyLoader( + "veeksha.client.tts", + "TTSClient", + ), +) + +ClientRegistry.register( + ClientType.REALTIME_TTS, + _LazyLoader( + "veeksha.client.realtime_tts", + "RealtimeTTSClient", + ), +) + +ClientRegistry.register( + ClientType.ELEVENLABS_STREAMING_TTS, + _LazyLoader( + "veeksha.client.native_streaming_tts", + "ElevenLabsStreamingTTSClient", + ), +) + +ClientRegistry.register( + ClientType.DEEPGRAM_FLUX_STREAMING_TTS, + _LazyLoader( + "veeksha.client.native_streaming_tts", + "DeepgramFluxStreamingTTSClient", + ), +) + +ClientRegistry.register( + ClientType.DEEPGRAM_AURA_STREAMING_TTS, + _LazyLoader( + "veeksha.client.native_streaming_tts", + "DeepgramAuraStreamingTTSClient", + ), +) + +ClientRegistry.register( + ClientType.ELEVENLABS_HTTP_TTS, + _LazyLoader( + "veeksha.client.native_http_tts", + "ElevenLabsHTTPTTSClient", + ), +) + +ClientRegistry.register( + ClientType.DEEPGRAM_FLUX_HTTP_TTS, + _LazyLoader( + "veeksha.client.native_http_tts", + "DeepgramFluxHTTPClient", + ), +) + +ClientRegistry.register( + ClientType.STT, + _LazyLoader( + "veeksha.client.stt", + "STTClient", + ), +) diff --git a/veeksha/client/stt.py b/veeksha/client/stt.py new file mode 100644 index 00000000..af30c971 --- /dev/null +++ b/veeksha/client/stt.py @@ -0,0 +1,792 @@ +"""STT clients for realtime streaming speech-to-text (vajra_openai_realtime, vllm_realtime). + +Both providers stream PCM16 over a WebSocket and report transcription metrics +(ttfc, end-to-end latency, RTF; see ``STTStreamResult`` for the timing metric +definitions). They share one lifecycle in ``_STTClientBase``: +audio is paced at 1x playback when ``ws_realtime_pacing`` is on, and the send +and receive loops run concurrently so ``ttfc`` reflects when the first partial +actually arrives rather than when the upload finishes. Each provider only +supplies four small protocol hooks (open session, encode a chunk, EOF sentinel, +parse a message). + +``STTClient(config)`` is a factory returning the client for ``config.provider`` +(see ``_PROVIDERS``), so the registry needs only one entry. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import os +import re +import threading +import time +from abc import abstractmethod +from collections import OrderedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Optional + +import numpy as np + +from veeksha.client.base import BaseLLMClient +from veeksha.core.request import Request +from veeksha.core.request_content import AudioChannelRequestContent +from veeksha.core.response import ChannelResponse, RequestResult +from veeksha.logger import init_logger +from veeksha.types import AudioTask, ChannelModality + +if TYPE_CHECKING: + from veeksha.config.client import STTClientConfig + +logger = init_logger(__name__) + +BYTES_PER_SAMPLE = 2 +TranscriptSnapshotRow = dict[str, float | str] + +# Voxtral streaming tokens injected by the model that should be stripped. +_STREAMING_TOKEN_RE = re.compile(r"\[STREAMING_(?:PAD|WORD)\]") + + +def _pcm_duration_ms(pcm_bytes: int, sample_rate: int) -> float: + """Compute PCM16 mono audio duration in ms.""" + return (pcm_bytes / BYTES_PER_SAMPLE / sample_rate) * 1000 + + +def _clean_transcript(text: str) -> str: + """Strip Voxtral streaming control tokens and collapse whitespace.""" + text = _STREAMING_TOKEN_RE.sub("", text) + return re.sub(r"\s+", " ", text).strip() + + +def _audio_to_pcm16_bytes(audio_path: str, target_sr: int) -> bytes: + """Load an audio file and convert to raw PCM16 bytes at target sample rate.""" + import librosa + + audio, _ = librosa.load(audio_path, sr=target_sr, mono=True) + pcm16 = (audio * 32767).astype(np.int16) + return pcm16.tobytes() + + +def _metadata_ms(metadata: dict, key: str) -> Optional[float]: + """Read an optional millisecond offset from request metadata.""" + value = metadata.get(key) + if value is None: + return None + return float(value) + + +def _slice_pcm16_bytes( + pcm_bytes: bytes | memoryview, + sample_rate: int, + *, + start_ms: Optional[float], + end_ms: Optional[float], +) -> bytes | memoryview: + """Slice raw PCM16 mono bytes by millisecond offsets. + + Accepts a ``memoryview`` for zero-copy slicing of cached clip PCM (the + slice then aliases the cached buffer instead of duplicating ~minutes of + audio per session). + """ + if start_ms is None and end_ms is None: + return pcm_bytes + + total_samples = len(pcm_bytes) // BYTES_PER_SAMPLE + start_sample = 0 if start_ms is None else int(round(start_ms * sample_rate / 1000)) + end_sample = ( + total_samples if end_ms is None else int(round(end_ms * sample_rate / 1000)) + ) + + if start_sample < 0: + raise ValueError(f"input_audio_start_ms must be non-negative; got {start_ms}") + if end_sample <= start_sample: + raise ValueError( + "input_audio_end_ms must be greater than input_audio_start_ms; " + f"got start_ms={start_ms}, end_ms={end_ms}" + ) + if end_sample > total_samples: + raise ValueError( + "Requested audio slice exceeds decoded clip length: " + f"end_ms={end_ms}, clip_ms={_pcm_duration_ms(len(pcm_bytes), sample_rate):.3f}" + ) + + start_byte = start_sample * BYTES_PER_SAMPLE + end_byte = end_sample * BYTES_PER_SAMPLE + return pcm_bytes[start_byte:end_byte] + + +# Decoded clips (and their pre-encoded wire messages) are cached per client; +# traces replay a finite clip set, so this is bounded in practice. The FIFO +# cap guards pathological manifests with thousands of distinct clips. +_CLIP_CACHE_MAX_CLIPS = 64 + + +@dataclass(frozen=True) +class _ClipAssets: + """Deterministic per-clip artifacts shared by every session replaying it. + + ``wire_messages[i]`` is the provider append-message for the full chunk + ``pcm[i * chunk_size : (i + 1) * chunk_size]``; the trailing partial chunk + (from the clip end or a per-session slice) is encoded on the fly. + """ + + pcm: bytes + wire_messages: list[str | bytes] + + +@dataclass +class STTStreamResult: + """Provider-neutral transcript and timing output from one streaming request. + + All timings are in milliseconds, measured with ``time.monotonic`` deltas: + + - ``ttfc`` (time to first content): from the first audio byte on the wire + to the first delta whose own payload still contains text after + control-token cleaning. Empty progress/keepalive deltas and pure + control-token pads do not count. Falls back to the completion message + when the stream produces a final transcript without any content-bearing + delta. + - ``time_to_first_visible_text``: from the first audio byte on the wire to + the first moment the assembled (concatenated then cleaned) transcript is + non-empty, i.e. when a user watching the live transcript would first see + text. Matches ``ttfc`` for well-formed streams; the two differ when + cleaning a delta in isolation disagrees with cleaning the assembled + transcript (e.g. a control token split across deltas cleans non-empty on + its own but vanishes once joined). Same completion fallback as ``ttfc``. + - ``time_to_first_partial``: from end-of-audio (EOF sentinel sent) to the + first delta after EOF with a non-empty assembled transcript; ``None`` + when no such delta arrives before completion. + - ``time_to_final_transcript``: from end-of-audio to the completion + message. + """ + + ttfc: Optional[float] + time_to_first_visible_text: Optional[float] + time_to_first_partial: Optional[float] + time_to_final_transcript: Optional[float] + partial_transcript: Optional[str] + final_transcript: str + transcript_snapshots: list[TranscriptSnapshotRow] + chunk_count: int + pcm_byte_count: int + + +class TranscriptSnapshotRecorder: + """Records evolving transcripts relative to first sent audio byte.""" + + def __init__(self) -> None: + self._audio_started_at: Optional[float] = None + self._snapshots: list[TranscriptSnapshotRow] = [] + + @property + def snapshots(self) -> list[TranscriptSnapshotRow]: + return self._snapshots + + def mark_audio_started(self, now: float) -> None: + if self._audio_started_at is None: + self._audio_started_at = now + + def add(self, now: float, transcript: str) -> None: + if not transcript: + return + if self._snapshots and self._snapshots[-1]["transcript"] == transcript: + return + self._snapshots.append( + { + "elapsed_ms": round(self._elapsed_ms(now), 3), + "transcript": transcript, + } + ) + + def _elapsed_ms(self, now: float) -> float: + assert ( + self._audio_started_at is not None + ), "snapshot recorded before any audio was sent" + return (now - self._audio_started_at) * 1000 + + +class _STTClientBase(BaseLLMClient): + """Shared streaming lifecycle for realtime STT providers. + + Subclasses set ``ws_path`` and implement the four protocol hooks; the + request lifecycle, 1x pacing, concurrent stream, error mapping, and metrics + assembly live here. + """ + + #: WebSocket path appended to ``api_base`` (set by subclasses). + ws_path: str = "" + + def __init__(self, config: "STTClientConfig", **kwargs) -> None: + super().__init__(config) + self._provider = config.provider + self._model = config.model + self._sample_rate = config.sample_rate + self._ws_chunk_size = config.ws_chunk_size + self._pacing = config.ws_realtime_pacing + self._request_timeout = config.request_timeout + self._ws_ping_interval_s = config.ws_ping_interval_s + self._ws_ping_timeout_s = config.ws_ping_timeout_s + self._ws_compression = "deflate" if config.ws_permessage_deflate else None + self._ws_url = self._http_to_ws(self.ws_path) + + # Per-clip decode + wire-message cache, shared across the worker + # threads that all hold this one client instance. + self._clip_cache: OrderedDict[str, _ClipAssets] = OrderedDict() + self._clip_cache_lock = threading.Lock() + self._clip_locks: dict[str, threading.Lock] = {} + + # ------------------------------------------------------------------ + # Provider protocol hooks + # ------------------------------------------------------------------ + + @abstractmethod + async def _open_session(self, ws) -> None: + """Complete the provider handshake before audio is sent.""" + + @abstractmethod + def _encode_chunk(self, chunk: bytes | memoryview) -> str | bytes: + """Frame a PCM16 chunk into the message the provider expects.""" + raise NotImplementedError + + @abstractmethod + def _eof(self) -> str | bytes: + """Return the end-of-audio sentinel message.""" + raise NotImplementedError + + @abstractmethod + def _parse_message(self, msg: dict) -> tuple[str, str]: + """Map a server message to ``(kind, text)``. + + ``kind`` is ``"delta"``, ``"done"``, ``"error"``, or ``""`` (ignore). + """ + + # ------------------------------------------------------------------ + # Shared helpers + # ------------------------------------------------------------------ + + def _http_to_ws(self, path: str) -> str: + """Convert the http(s) api_base to a ws(s) URL with ``path`` appended.""" + base = self.api_base or "" + if base.startswith("https://"): + base = "wss://" + base[len("https://") :] + elif base.startswith("http://"): + base = "ws://" + base[len("http://") :] + return f"{base}{path}" + + async def _maybe_pace_until(self, target_at: float) -> None: + """Sleep until an absolute audio schedule time when pacing is on.""" + if not self._pacing: + return + delay_s = target_at - time.monotonic() + if delay_s > 0: + await asyncio.sleep(delay_s) + + def _clip_lock(self, audio_path: str) -> threading.Lock: + """Return the per-clip lock guarding decode/encode for one path.""" + with self._clip_cache_lock: + lock = self._clip_locks.get(audio_path) + if lock is None: + lock = self._clip_locks[audio_path] = threading.Lock() + return lock + + def _clip_assets(self, audio_path: str) -> _ClipAssets: + """Decode a clip and pre-encode its append messages, cached per clip. + + Blocking (audio decode + base64/JSON encode); callers on an event + loop must run this in an executor. The per-clip lock makes the + thundering herd at benchmark start decode each clip exactly once. + """ + with self._clip_lock(audio_path): + with self._clip_cache_lock: + assets = self._clip_cache.get(audio_path) + if assets is not None: + return assets + + pcm = _audio_to_pcm16_bytes(audio_path, self._sample_rate) + view = memoryview(pcm) + full_end = len(pcm) - len(pcm) % self._ws_chunk_size + wire_messages = [ + self._encode_chunk(view[offset : offset + self._ws_chunk_size]) + for offset in range(0, full_end, self._ws_chunk_size) + ] + assets = _ClipAssets(pcm=pcm, wire_messages=wire_messages) + + with self._clip_cache_lock: + self._clip_cache[audio_path] = assets + while len(self._clip_cache) > _CLIP_CACHE_MAX_CLIPS: + self._clip_cache.popitem(last=False) + return assets + + async def _stream( + self, + pcm_bytes: bytes | memoryview, + wire_messages: Optional[list[str | bytes]] = None, + on_request_sent: Optional[Callable[[], None]] = None, + on_request_dispatched: Optional[Callable[[], None]] = None, + ) -> STTStreamResult: + """Stream pre-decoded PCM16 to the provider and collect the transcript. + + The sender and receiver run concurrently so partial transcripts are + timed when they arrive, not after the (paced) upload completes. ``ttfc`` + is anchored to ``audio_started_at`` (the first audio byte on the wire), + so clip decode, connect, and handshake latency are excluded. + + ``wire_messages`` are the cached pre-encoded append messages for the + clip this PCM is a prefix of (chunk ``i`` at byte ``i * chunk_size``); + the trailing partial chunk is encoded on the fly. + """ + import websockets + + ttfc: Optional[float] = None + audio_started_at: Optional[float] = None + audio_end_at: Optional[float] = None + time_to_first_visible_text: Optional[float] = None + time_to_first_partial: Optional[float] = None + time_to_final_transcript: Optional[float] = None + partial_transcript: Optional[str] = None + final_transcript = "" + chunk_count = 0 + transcript_chunks: list[str] = [] + snapshots = TranscriptSnapshotRecorder() + + async with websockets.connect( + self._ws_url, + ping_interval=self._ws_ping_interval_s, + ping_timeout=self._ws_ping_timeout_s, + compression=self._ws_compression, + ) as ws: + await self._open_session(ws) + if on_request_dispatched is not None: + on_request_dispatched() + + async def _send() -> None: + nonlocal audio_end_at, audio_started_at + for byte_offset in range(0, len(pcm_bytes), self._ws_chunk_size): + if audio_started_at is None: + audio_started_at = time.monotonic() + snapshots.mark_audio_started(audio_started_at) + else: + await self._maybe_pace_until( + audio_started_at + + byte_offset / BYTES_PER_SAMPLE / self._sample_rate + ) + chunk_end = byte_offset + self._ws_chunk_size + if wire_messages is not None and chunk_end <= len(pcm_bytes): + message = wire_messages[byte_offset // self._ws_chunk_size] + else: + message = self._encode_chunk(pcm_bytes[byte_offset:chunk_end]) + await ws.send(message) + if audio_started_at is not None: + await self._maybe_pace_until( + audio_started_at + + len(pcm_bytes) / BYTES_PER_SAMPLE / self._sample_rate + ) + await ws.send(self._eof()) + audio_end_at = time.monotonic() + + send_task = asyncio.ensure_future(_send()) + try: + while True: + kind, text = self._parse_message(json.loads(await ws.recv())) + now = time.monotonic() + if kind == "delta": + # TTFC counts only deltas whose own payload carries + # transcript text after cleaning; empty progress / + # keepalive deltas (e.g. Vajra's priming delta) and + # pure control-token pads are skipped. See + # STTStreamResult for how this differs from + # time_to_first_visible_text below. + if ttfc is None and _clean_transcript(text): + assert ( + audio_started_at is not None + ), "delta arrived before any audio was sent" + ttfc = (now - audio_started_at) * 1000 + if on_request_sent is not None: + on_request_sent() + transcript_chunks.append(text) + chunk_count += 1 + current_transcript = _clean_transcript( + "".join(transcript_chunks) + ) + snapshots.add(now, current_transcript) + if time_to_first_visible_text is None and current_transcript: + assert ( + audio_started_at is not None + ), "transcript arrived before any audio was sent" + time_to_first_visible_text = (now - audio_started_at) * 1000 + if ( + audio_end_at is not None + and time_to_first_partial is None + and current_transcript + ): + time_to_first_partial = (now - audio_end_at) * 1000 + partial_transcript = current_transcript + elif kind == "done": + final_transcript = _clean_transcript( + text or "".join(transcript_chunks) + ) + snapshots.add(now, final_transcript) + if final_transcript: + if ttfc is None: + assert ( + audio_started_at is not None + ), "completion arrived before any audio was sent" + ttfc = (now - audio_started_at) * 1000 + if on_request_sent is not None: + on_request_sent() + if time_to_first_visible_text is None: + assert ( + audio_started_at is not None + ), "transcript arrived before any audio was sent" + time_to_first_visible_text = ( + now - audio_started_at + ) * 1000 + if chunk_count == 0: + chunk_count = 1 + if audio_end_at is not None: + time_to_final_transcript = (now - audio_end_at) * 1000 + break + elif kind == "error": + raise RuntimeError( + f"{self._provider} streaming error: {text or 'unknown'}" + ) + finally: + # Normal path: the sender is already done (servers only emit + # "done" after EOF). Error path: cancel it. A CancelledError is + # expected; any other error is a real send failure and must + # propagate. + if not send_task.done(): + send_task.cancel() + try: + await send_task + except asyncio.CancelledError: + pass + + if not final_transcript: + final_transcript = _clean_transcript("".join(transcript_chunks)) + + return STTStreamResult( + ttfc=ttfc, + time_to_first_visible_text=time_to_first_visible_text, + time_to_first_partial=time_to_first_partial, + time_to_final_transcript=time_to_final_transcript, + partial_transcript=partial_transcript, + final_transcript=final_transcript, + transcript_snapshots=snapshots.snapshots, + chunk_count=chunk_count, + pcm_byte_count=len(pcm_bytes), + ) + + # ------------------------------------------------------------------ + # Request lifecycle (shared) + # ------------------------------------------------------------------ + + async def send_request( + self, + request: Request, + session_id: int, + session_total_requests: int = 1, + on_request_sent: Optional[Callable[[], None]] = None, + on_request_dispatched: Optional[Callable[[], None]] = None, + ) -> RequestResult: + """Stream an audio file to the STT API and collect transcription metrics.""" + + sent_fired = False + dispatched_fired = False + + def fire_sent_once() -> None: + nonlocal sent_fired + if sent_fired: + return + sent_fired = True + if on_request_sent is not None: + on_request_sent() + + def fire_dispatched_once() -> None: + nonlocal dispatched_fired + if dispatched_fired: + return + dispatched_fired = True + if on_request_dispatched is not None: + on_request_dispatched() + + def finish_callbacks() -> None: + # Error results are returned rather than raised, so both dispatch + # orderings need a fallback to avoid blocking every later ticket. + fire_dispatched_once() + fire_sent_once() + + audio_content = request.channels.get(ChannelModality.AUDIO) + if not isinstance(audio_content, AudioChannelRequestContent): + finish_callbacks() + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + success=False, + error_code=400, + error_msg="No AUDIO channel in request for STT", + client_completed_at=time.monotonic(), + ) + + audio_path = audio_content.input_audio + + try: + os.path.getsize(audio_path) + except OSError as e: + finish_callbacks() + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + success=False, + error_code=404, + error_msg=f"Failed to read audio file {audio_path}: {e}", + client_completed_at=time.monotonic(), + ) + + logger.debug( + "[STT %s] request_id=%d session_id=%d file=%s", + self._provider, + request.id, + session_id, + audio_path, + ) + + error_msg: Optional[str] = None + error_code: Optional[int] = None + stream_result: Optional[STTStreamResult] = None + + # Decode and slice before the latency clock so file I/O and DSP don't + # inflate TTFC. Decode and append-message encoding are cached per + # clip and run in an executor so they never stall the shared event + # loop (a synchronous decode here freezes pacing for every session + # on this worker's loop). + try: + loop = asyncio.get_running_loop() + clip = await loop.run_in_executor(None, self._clip_assets, audio_path) + start_ms = _metadata_ms(request.metadata, "input_audio_start_ms") + pcm_bytes = _slice_pcm16_bytes( + memoryview(clip.pcm), + self._sample_rate, + start_ms=start_ms, + end_ms=_metadata_ms(request.metadata, "input_audio_end_ms"), + ) + # Cached messages are aligned to the clip start; a non-zero slice + # start shifts the chunk grid, so encode per chunk in that case. + wire_messages = clip.wire_messages if not start_ms else None + except Exception as e: + finish_callbacks() + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + success=False, + error_code=422, + error_msg=f"Failed to decode audio file {audio_path}: {e}", + client_completed_at=time.monotonic(), + ) + + t_start = time.monotonic() + + try: + async with asyncio.timeout(self._request_timeout): + stream_result = await self._stream( + pcm_bytes, + wire_messages, + on_request_sent=fire_sent_once, + on_request_dispatched=fire_dispatched_once, + ) + except TimeoutError: + error_code = 408 + error_msg = f"STT request timed out after {self._request_timeout}s" + except (OSError, ConnectionError) as e: + # Includes ConnectionRefusedError when the server isn't up. + error_code = 503 + error_msg = str(e) + except Exception as e: + error_code = 520 + error_msg = str(e) + logger.error( + "[STT %s] request_id=%d error: %s", + self._provider, + request.id, + e, + exc_info=True, + ) + + finish_callbacks() + completed_at = time.monotonic() + total_latency_ms = (completed_at - t_start) * 1000 + success = error_msg is None and error_code is None + + channels = {} + if success and stream_result is not None: + input_audio_duration_ms = _pcm_duration_ms( + stream_result.pcm_byte_count, self._sample_rate + ) + # Report the input clip's byte count; the evaluator derives + # duration/RTF from pcm_byte_count + sample_rate. + metrics_dict: dict = { + "audio_task": AudioTask.STT, + "ttfc": round(stream_result.ttfc or 0.0, 3), + "end_to_end_latency": round(total_latency_ms, 3), + "time_to_first_visible_text": ( + round(stream_result.time_to_first_visible_text, 3) + if stream_result.time_to_first_visible_text is not None + else None + ), + "time_to_first_partial": ( + round(stream_result.time_to_first_partial, 3) + if stream_result.time_to_first_partial is not None + else None + ), + "time_to_final_transcript": ( + round(stream_result.time_to_final_transcript, 3) + if stream_result.time_to_final_transcript is not None + else None + ), + "chunk_count": stream_result.chunk_count, + "pcm_byte_count": stream_result.pcm_byte_count, + "raw_pcm": True, + "input_tokens": len(stream_result.final_transcript.split()), + "sample_rate": self._sample_rate, + "input_audio_duration_ms": round(input_audio_duration_ms, 3), + "partial_transcript": stream_result.partial_transcript, + "final_transcript": stream_result.final_transcript, + "transcript_snapshots": stream_result.transcript_snapshots, + } + + # Ground truth and dataset metadata are guaranteed by the trace generator. + for key, value in request.metadata.items(): + metrics_dict.setdefault(key, value) + + channels[ChannelModality.AUDIO] = ChannelResponse( + modality=ChannelModality.AUDIO, + content=stream_result.final_transcript, + metrics=metrics_dict, + ) + + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + channels=channels, + success=success, + error_code=error_code, + error_msg=error_msg, + client_completed_at=completed_at, + ) + + +class VllmRealtimeSTTClient(_STTClientBase): + """vllm_realtime: WebSocket /v1/realtime β€” base64 PCM16 chunks, deltas. + + Server -> Client: ``{"type": "session.created"}`` then + ``transcription.delta`` / ``transcription.done`` / ``error``. + """ + + ws_path = "/v1/realtime" + + async def _open_session(self, ws) -> None: + msg = json.loads(await ws.recv()) + if msg.get("type") != "session.created": + raise RuntimeError(f"Expected session.created, got: {msg}") + await ws.send(json.dumps({"type": "session.update", "model": self._model})) + await ws.send(json.dumps({"type": "input_audio_buffer.commit"})) + + def _encode_chunk(self, chunk: bytes | memoryview) -> str: + return json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(chunk).decode("utf-8"), + } + ) + + def _eof(self) -> str: + return json.dumps({"type": "input_audio_buffer.commit", "final": True}) + + def _parse_message(self, msg: dict) -> tuple[str, str]: + msg_type = msg.get("type") + if msg_type == "transcription.delta": + return "delta", msg.get("delta", "") + if msg_type == "transcription.done": + return "done", msg.get("text", "") + if msg_type == "error": + return "error", str(msg.get("error", "")) + return "", "" + + +class VajraOpenAIRealtimeSTTClient(_STTClientBase): + """vajra_openai_realtime: WebSocket /openai/v1/realtime β€” OpenAI transcription. + + Drives Vajra's OpenAI-compatible realtime transcription endpoint. Server -> + Client: ``transcription_session.created`` then + ``conversation.item.input_audio_transcription.delta`` / + ``conversation.item.input_audio_transcription.completed`` / ``error``. v1 is + manual-commit, one transcript per connection (PCM16 mono 16 kHz). + """ + + ws_path = "/openai/v1/realtime?intent=transcription" + + async def _open_session(self, ws) -> None: + msg = json.loads(await ws.recv()) + if msg.get("type") != "transcription_session.created": + raise RuntimeError(f"Expected transcription_session.created, got: {msg}") + # Configure the session; unlike vllm_realtime, do NOT commit here β€” a + # commit before any audio would finalize an empty transcript. + await ws.send( + json.dumps( + { + "type": "transcription_session.update", + "input_audio_format": "pcm16", + "input_audio_transcription": {"model": self._model}, + } + ) + ) + + def _encode_chunk(self, chunk: bytes | memoryview) -> str: + return json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(chunk).decode("utf-8"), + } + ) + + def _eof(self) -> str: + return json.dumps({"type": "input_audio_buffer.commit"}) + + def _parse_message(self, msg: dict) -> tuple[str, str]: + msg_type = msg.get("type") + if msg_type == "conversation.item.input_audio_transcription.delta": + return "delta", msg.get("delta", "") + if msg_type == "conversation.item.input_audio_transcription.completed": + return "done", msg.get("transcript", "") + # ".failed" is the terminal event for a committed item whose + # transcription failed; treat it like a session-level "error" so the + # request fails fast instead of waiting for the timeout. + if msg_type in ( + "error", + "conversation.item.input_audio_transcription.failed", + ): + error = msg.get("error") + if isinstance(error, dict): + return "error", str(error.get("message", "")) + return "error", str(error or "") + return "", "" + + +_PROVIDERS: dict[str, type[_STTClientBase]] = { + "vllm_realtime": VllmRealtimeSTTClient, + "vajra_openai_realtime": VajraOpenAIRealtimeSTTClient, +} + + +def STTClient(config: "STTClientConfig", **kwargs) -> _STTClientBase: + """Factory: return the STT client for ``config.provider``.""" + try: + cls = _PROVIDERS[config.provider] + except KeyError: + raise ValueError(f"Unsupported STT provider: {config.provider}") + return cls(config, **kwargs) diff --git a/veeksha/client/tts.py b/veeksha/client/tts.py new file mode 100644 index 00000000..ed8ea17e --- /dev/null +++ b/veeksha/client/tts.py @@ -0,0 +1,225 @@ +"""TTS client for streaming HTTP-based text-to-speech APIs.""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING +from urllib.parse import urljoin + +import httpx + +from veeksha.client.base import BaseLLMClient +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.core.request import Request +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.core.response import ChannelResponse, RequestResult +from veeksha.logger import init_logger +from veeksha.types import AudioTask, ChannelModality + +if TYPE_CHECKING: + from veeksha.config.client import TTSClientConfig + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class OpenAISpeechRequest: + """Canonical OpenAI Audio Speech HTTP request.""" + + url: str + headers: dict[str, str] + payload: dict[str, str | bool] + + +class TTSProtocolError(Exception): + """Raised when a server response violates the selected Speech stream format.""" + + +def _build_audio_speech_url(api_base: str) -> str: + """Build ``/v1/audio/speech`` from an API base with or without ``/v1``.""" + normalized_base = api_base.rstrip("/") + if not normalized_base.endswith("/v1"): + normalized_base = f"{normalized_base}/v1" + return urljoin(f"{normalized_base}/", "audio/speech") + + +class TTSClient(BaseLLMClient): + """Async client for streaming HTTP-based TTS APIs.""" + + def __init__(self, config: TTSClientConfig, **kwargs) -> None: + super().__init__(config) + self._chunk_size = config.chunk_size + self._speech_url = _build_audio_speech_url(str(self.api_base)) + self._client_storage = threading.local() + + def _build_request(self, text: str) -> OpenAISpeechRequest: + """Build the one wire contract accepted by the HTTP TTS client.""" + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + return OpenAISpeechRequest( + url=self._speech_url, + headers=headers, + payload={ + "model": self.config.model, + "input": text, + "voice": self.config.voice_id, + "response_format": "pcm" if self.config.raw_pcm else "wav", + "stream": True, + "stream_format": "audio", + }, + ) + + def _validate_audio_response(self, response: httpx.Response) -> None: + content_type = response.headers.get("Content-Type", "") + media_type = content_type.partition(";")[0].strip().lower() + allowed = { + "application/octet-stream", + "application/x-wav", + "binary/octet-stream", + } + if ( + media_type + and not media_type.startswith("audio/") + and media_type not in allowed + ): + raise TTSProtocolError( + "OpenAI Audio Speech stream_format=audio requires an audio byte " + f"response, got Content-Type {content_type!r}" + ) + + def _get_client(self) -> httpx.AsyncClient: + """Return a thread-local httpx client bound to the caller's event loop.""" + if not hasattr(self._client_storage, "client"): + self._client_storage.client = httpx.AsyncClient( + timeout=self.config.request_timeout + ) + return self._client_storage.client + + async def send_request( + self, + request: Request, + session_id: int, + session_total_requests: int = 1, + on_request_sent: Callable[[], None] | None = None, + on_request_dispatched: Callable[[], None] | None = None, + ) -> RequestResult: + """Send a streaming TTS request and collect audio metrics.""" + text_content = request.channels.get(ChannelModality.TEXT) + if not isinstance(text_content, TextChannelRequestContent): + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + success=False, + error_code=400, + error_msg="No TEXT channel in request for TTS", + client_completed_at=time.monotonic(), + ) + input_text = text_content.input_text + speech_request = self._build_request(input_text) + + logger.debug( + "[TTS] request_id=%d session_id=%d chars=%d text=%.80r", + request.id, + session_id, + len(input_text), + input_text, + ) + + error_msg: str | None = None + error_code: int | None = None + ttfc: float | None = None + chunk_count = 0 + audio_chunks: list[bytes] = [] + + t_start = time.monotonic() + + try: + async with self._get_client().stream( + "POST", + speech_request.url, + headers=speech_request.headers, + json=speech_request.payload, + timeout=self.config.request_timeout, + ) as response: + response.raise_for_status() + self._validate_audio_response(response) + if on_request_dispatched is not None: + on_request_dispatched() + + sent_notified = False + async for chunk in response.aiter_bytes(chunk_size=self._chunk_size): + if not chunk: + continue + receive_time = time.monotonic() + if ttfc is None: + ttfc = (receive_time - t_start) * 1000 + if not sent_notified and on_request_sent is not None: + on_request_sent() + sent_notified = True + + audio_chunks.append(chunk) + chunk_count += 1 + + if not sent_notified and on_request_sent is not None: + on_request_sent() + + except TTSProtocolError as e: + error_code = 502 + error_msg = str(e) + logger.warning("TTS protocol error: (%s) %s", error_code, error_msg) + except httpx.HTTPStatusError as e: + error_code = e.response.status_code if e.response else 500 + error_msg = str(e) + logger.warning("HTTP Error: status=%s msg=%s", error_code, error_msg) + except httpx.ConnectError as e: + error_code = 503 + error_msg = str(e) + logger.warning("Connection Error: (%s) %s", error_code, error_msg) + except httpx.TimeoutException: + error_code = 408 + error_msg = "TTS request timed out" + logger.warning("Timeout Error: (%s) %s", error_code, error_msg) + except Exception as e: + error_code = 520 + error_msg = str(e) + logger.exception("Unexpected error: (%s) %s", error_code, error_msg) + + completed_at = time.monotonic() + total_latency_ms = (completed_at - t_start) * 1000 + success = error_msg is None and error_code is None + audio_data = b"".join(audio_chunks) if audio_chunks else b"" + + channels = {} + if success: + channels[ChannelModality.AUDIO] = ChannelResponse( + modality=ChannelModality.AUDIO, + content=audio_data, + metrics={ + "audio_task": AudioTask.TTS, + AudioMetricKey.TTFC.value: round(ttfc or 0.0, 3), + AudioMetricKey.END_TO_END_LATENCY.value: round(total_latency_ms, 3), + AudioMetricKey.CHUNK_COUNT.value: chunk_count, + AudioMetricKey.RAW_PCM.value: self.config.raw_pcm, + AudioMetricKey.SAMPLE_RATE.value: self.config.sample_rate, + AudioMetricKey.INPUT_CHARS.value: len(input_text), + AudioMetricKey.INPUT_TOKENS.value: text_content.target_prompt_tokens + or 0, + AudioMetricKey.INPUT_TEXT.value: input_text, + }, + ) + + return RequestResult( + request_id=request.id, + session_id=session_id, + session_total_requests=session_total_requests, + channels=channels, + success=success, + error_code=error_code, + error_msg=error_msg, + client_completed_at=completed_at, + ) diff --git a/veeksha/client/utils.py b/veeksha/client/utils.py new file mode 100644 index 00000000..ebf1dc6a --- /dev/null +++ b/veeksha/client/utils.py @@ -0,0 +1,112 @@ +"""Shared client-side helpers for paced Realtime text-to-speech. + +Text segmentation (emulating an upstream LLM's per-token output) and delta +pacing (emulating its decode cadence) are transport-agnostic: they turn a prompt +string into paced ``conversation.item.create`` payloads. They live here rather +than in the WebSocket client so the client stays focused on the +transport and metric contract. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +from veeksha.config.generator.interval import ( + FixedIntervalGeneratorConfig, + PoissonIntervalGeneratorConfig, +) +from veeksha.generator.interval.fixed import FixedIntervalGenerator +from veeksha.generator.interval.poisson import PoissonIntervalGenerator + +if TYPE_CHECKING: + from veeksha.config.client import TextPacingConfig + + +# --------------------------------------------------------------------------- +# Text segmentation (LLM-token emulation) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class TextSegment: + """Text for one paced Realtime conversation item. + + ``text`` is the exact substring streamed for this delta (whitespace + preserved); ``n_tokens`` is the number of whitespace tokens it groups. + """ + + text: str + n_chars: int + n_tokens: int + + +def segment_text(text: str, tokens_per_delta: int) -> list[TextSegment]: + """Split ``text`` into paced deltas of ``tokens_per_delta`` whitespace tokens. + + Each token is a ``\\S+`` run plus its trailing whitespace (via + ``re.findall(r"\\S+\\s*")``); any leading whitespace is prepended to the + first piece. Consecutive pieces are grouped ``tokens_per_delta`` at a time. + + HARD INVARIANT: ``"".join(s.text for s in segments) == text`` β€” the WER + verification path reconstructs the exact input by concatenating deltas, so + no character may be dropped or added. + """ + if not text: + return [] + + pieces = re.findall(r"\S+\s*", text) + if not pieces: + # Whitespace-only text: emit it as a single (zero-token) segment so the + # round-trip invariant still holds. + return [TextSegment(text=text, n_chars=len(text), n_tokens=0)] + + # Prepend any leading whitespace (which ``\S+\s*`` does not capture) to the + # first piece so the concatenation is byte-for-byte identical to ``text``. + leading = text[: len(text) - len(text.lstrip())] + if leading: + pieces[0] = leading + pieces[0] + + segments: list[TextSegment] = [] + for start in range(0, len(pieces), tokens_per_delta): + group = pieces[start : start + tokens_per_delta] + chunk = "".join(group) + segments.append( + TextSegment(text=chunk, n_chars=len(chunk), n_tokens=len(group)) + ) + return segments + + +# --------------------------------------------------------------------------- +# Delta pacing (upstream decode-rate emulation) +# --------------------------------------------------------------------------- + + +class TextDeltaPacer: + """Emits inter-delta gaps emulating an upstream LLM's decode cadence. + + ``mean_gap_s = tokens_per_delta / tokens_per_second``. A ``poisson`` + distribution jitters the gaps around that mean (seeded, so runs are + reproducible); ``fixed`` returns the mean gap every time. + """ + + def __init__(self, pacing: "TextPacingConfig", seed: int) -> None: + self.initial_delay_s = pacing.initial_delay_s + mean_gap_s = pacing.tokens_per_delta / pacing.tokens_per_second + if pacing.gap_distribution == "poisson": + self._generator = PoissonIntervalGenerator( + PoissonIntervalGeneratorConfig(arrival_rate=1.0 / mean_gap_s), + np.random.RandomState(seed), + ) + else: + # FixedIntervalGenerator's signature declares ``rng: None``. + self._generator = FixedIntervalGenerator( + FixedIntervalGeneratorConfig(interval=mean_gap_s), None + ) + + def next_gap(self) -> float: + """Return the next inter-delta gap in seconds.""" + return self._generator.get_next_interval() diff --git a/veeksha/config/benchmark.py b/veeksha/config/benchmark.py index 26799bdf..f95d0c48 100644 --- a/veeksha/config/benchmark.py +++ b/veeksha/config/benchmark.py @@ -7,6 +7,7 @@ BaseClientConfig, OpenAIChatCompletionsClientConfig, ) +from veeksha.config.endpoint import EndpointConfig from veeksha.config.evaluator import ( BaseEvaluatorConfig, PerformanceEvaluatorConfig, @@ -60,7 +61,11 @@ class BenchmarkConfig(VeekshaCommand, name="benchmark", default=True): ) server: Optional[BaseServerConfig] = field( None, - help="Server configuration for managed servers. If set, client.model, client.api_key and client.api_base will be overwritten.", + help="Server configuration for managed servers. If set, it produces endpoint.", + ) + endpoint: Optional[EndpointConfig] = field( + None, + help="Client-facing endpoint contract used to overwrite client.model, client.api_key and client.api_base.", ) wandb: WandbConfig = field( default_factory=WandbConfig, @@ -80,3 +85,5 @@ def create_from_cli_args(cls): def __post_init__(self): if not self.evaluators: raise ValueError("BenchmarkConfig.evaluators must be non-empty.") + if self.server is not None and self.endpoint is not None: + raise ValueError("BenchmarkConfig cannot set both server and endpoint") diff --git a/veeksha/config/client.py b/veeksha/config/client.py index 1fad731f..fd0b2548 100644 --- a/veeksha/config/client.py +++ b/veeksha/config/client.py @@ -3,6 +3,7 @@ from vidhi import BasePolyConfig, field, frozen_dataclass +from veeksha.core.audio_contract import DEFAULT_AUDIO_SAMPLE_RATE from veeksha.logger import init_logger from veeksha.types import ClientType @@ -41,6 +42,23 @@ def __post_init__(self): self.additional_sampling_params ) + def build_tokenizer_provider(self): + """Build a TokenizerProvider for this client config. + + Default implementation uses a HuggingFace tokenizer based on self.model. + Subclasses can override for non-HF models. + """ + from veeksha.core.tokenizer import ( + TokenizerProvider, + build_hf_tokenizer_handle_from_model, + ) + from veeksha.types import ChannelModality + + return TokenizerProvider( + {ChannelModality.TEXT: build_hf_tokenizer_handle_from_model(self.model)}, + model_name=self.model, + ) + @frozen_dataclass class OpenAIChatCompletionsClientConfig(BaseClientConfig): @@ -123,3 +141,507 @@ class OpenAIRouterClientConfig(OpenAIChatCompletionsClientConfig): @classmethod def get_type(cls) -> ClientType: return ClientType.OPENAI_ROUTER + + +@frozen_dataclass +class TTSClientConfig(BaseClientConfig): + """TTS client configuration for the OpenAI Audio Speech API. + + Every server used with ``client.type: tts`` must implement the canonical + ``POST /v1/audio/speech`` request and raw-audio streaming response. + """ + + voice_id: str = field( + "", + help="Required OpenAI Audio Speech voice identifier.", + ) + sample_rate: int = field(DEFAULT_AUDIO_SAMPLE_RATE, help="Audio sample rate in Hz.") + chunk_size: int = field( + 1024, help="Chunk size in bytes for streaming audio response." + ) + raw_pcm: bool = field( + False, + help="Request response_format=pcm (True) or response_format=wav (False).", + ) + model: str = field("", help="The TTS model ID.") + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.TTS + + def __post_init__(self): + super().__post_init__() + + # Skip validation when instantiated with defaults by the flat_dataclass + # framework for non-selected polymorphic children. + if not self.model and self.api_base is None and not self.voice_id: + return + if not self.model: + raise ValueError("TTSClientConfig.model is required.") + if not self.voice_id: + raise ValueError("TTSClientConfig.voice_id is required.") + if self.api_base is None: + raise ValueError("TTSClientConfig.api_base is required.") + if self.sample_rate <= 0: + raise ValueError("TTSClientConfig.sample_rate must be > 0") + if self.chunk_size <= 0: + raise ValueError("TTSClientConfig.chunk_size must be > 0") + + def build_tokenizer_provider(self): + """TTS models use a simple word-split tokenizer.""" + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) + + +@frozen_dataclass +class TextPacingConfig: + """LLM decode-rate emulation for paced Realtime text conversation items.""" + + tokens_per_second: float = field( + 20.0, help="Emulated upstream LLM decode rate (whitespace tokens/sec)." + ) + tokens_per_delta: int = field(1, help="Whitespace tokens per input append event.") + gap_distribution: str = field( + "fixed", help="Inter-delta gap distribution: fixed | poisson." + ) + initial_delay_s: float = field( + 0.0, help="Delay before the first delta (upstream TTFT emulation)." + ) + seed: int = field( + 42, + help="Base seed for per-request gap jitter (per-request seed = seed + request_id).", + ) + + def __post_init__(self): + if self.tokens_per_second <= 0: + raise ValueError("TextPacingConfig.tokens_per_second must be > 0") + if self.tokens_per_delta < 1: + raise ValueError("TextPacingConfig.tokens_per_delta must be >= 1") + if self.gap_distribution not in ("fixed", "poisson"): + raise ValueError( + "TextPacingConfig.gap_distribution must be one of " + "('fixed', 'poisson'), " + f"got '{self.gap_distribution}'" + ) + if self.initial_delay_s < 0: + raise ValueError("TextPacingConfig.initial_delay_s must be >= 0") + + +@frozen_dataclass +class RealtimeTTSClientConfig(BaseClientConfig): + """WebSocket TTS client implementing the OpenAI Realtime event contract. + + `client.type: realtime_tts` opens a websocket to a realtime TTS server, + sends text conversation items at an emulated LLM decode rate, and measures + streaming interactivity metrics from the audio chunks it receives back. + """ + + voice_id: str = field( + "", + help="Optional Realtime output voice identifier.", + ) + sample_rate: int = field(DEFAULT_AUDIO_SAMPLE_RATE, help="Audio sample rate in Hz.") + raw_pcm: bool = field(True, help="The Realtime protocol streams raw PCM16.") + model: str = field("", help="The realtime TTS model ID.") + pacing: TextPacingConfig = field( + default_factory=TextPacingConfig, + help="Text pacing (LLM decode-rate emulation) configuration.", + ) + input_output_mode: str = field( + "complete_text", + help=( + "Realtime input/output scheduling: 'complete_text' sends " + "response.create after all text deltas; 'duplex' sends it once " + "duplex_start_after_tokens have arrived and continues appending text " + "while audio is generated. Duplex mode requires a server that consumes " + "conversation items added to an active response." + ), + ) + duplex_start_after_tokens: int = field( + 1, + help=( + "Minimum paced input tokens to send before response.create in duplex " + "mode. The trigger fires after the delta that crosses this threshold." + ), + ) + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.REALTIME_TTS + + def __post_init__(self): + super().__post_init__() + + # Skip validation when instantiated with defaults by the flat_dataclass + # framework for non-selected polymorphic children. + if not self.model and self.api_base is None: + return + if not self.model: + raise ValueError("RealtimeTTSClientConfig.model is required.") + if self.api_base is None: + raise ValueError("RealtimeTTSClientConfig.api_base is required.") + if self.sample_rate <= 0: + raise ValueError("RealtimeTTSClientConfig.sample_rate must be > 0") + if self.input_output_mode not in ("complete_text", "duplex"): + raise ValueError( + "RealtimeTTSClientConfig.input_output_mode must be one of " + "('complete_text', 'duplex')" + ) + if self.duplex_start_after_tokens < 1: + raise ValueError( + "RealtimeTTSClientConfig.duplex_start_after_tokens must be >= 1" + ) + + def build_tokenizer_provider(self): + """Realtime TTS models use a simple word-split tokenizer.""" + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) + + +@frozen_dataclass +class ElevenLabsStreamingTTSClientConfig(BaseClientConfig): + """Native ElevenLabs ``stream-input`` WebSocket TTS client.""" + + voice_id: str = field("", help="ElevenLabs voice identifier.") + sample_rate: int = field(DEFAULT_AUDIO_SAMPLE_RATE, help="PCM sample rate in Hz.") + model: str = field("", help="ElevenLabs model ID, e.g. eleven_flash_v2_5.") + api_key_env: str = field( + "ELEVENLABS_API_KEY", + help="Environment variable used when client.api_key is omitted.", + ) + pacing: TextPacingConfig = field( + default_factory=TextPacingConfig, + help="Upstream LLM text pacing configuration.", + ) + chunk_length_schedule: list[int] = field( + default_factory=lambda: [120, 160, 250, 290], + help="Provider character thresholds that trigger successive audio chunks.", + ) + stability: float = field(0.5, help="ElevenLabs voice stability.") + similarity_boost: float = field(0.8, help="ElevenLabs similarity boost.") + speed: float = field(1.0, help="ElevenLabs speaking-rate multiplier.") + auto_mode: bool = field( + False, + help="Use ElevenLabs auto mode instead of the configured chunk schedule.", + ) + apply_text_normalization: str = field( + "off", + help="ElevenLabs text normalization mode: auto | on | off.", + ) + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.ELEVENLABS_STREAMING_TTS + + def __post_init__(self): + super().__post_init__() + if not self.model and self.api_base is None and not self.voice_id: + return + if self.api_base is None: + raise ValueError("ElevenLabsStreamingTTSClientConfig.api_base is required") + if not self.model: + raise ValueError("ElevenLabsStreamingTTSClientConfig.model is required") + if not self.voice_id: + raise ValueError("ElevenLabsStreamingTTSClientConfig.voice_id is required") + if self.sample_rate <= 0: + raise ValueError( + "ElevenLabsStreamingTTSClientConfig.sample_rate must be > 0" + ) + if not self.api_key_env: + raise ValueError( + "ElevenLabsStreamingTTSClientConfig.api_key_env is required" + ) + if not self.chunk_length_schedule or any( + value <= 0 for value in self.chunk_length_schedule + ): + raise ValueError("chunk_length_schedule must contain positive values") + if self.chunk_length_schedule != sorted(self.chunk_length_schedule): + raise ValueError("chunk_length_schedule must be non-decreasing") + if not 0.7 <= self.speed <= 1.2: + raise ValueError("ElevenLabs speed must be between 0.7 and 1.2") + if self.apply_text_normalization not in ("auto", "on", "off"): + raise ValueError( + "apply_text_normalization must be one of ('auto', 'on', 'off')" + ) + + def build_tokenizer_provider(self): + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) + + +@frozen_dataclass +class DeepgramFluxStreamingTTSClientConfig(BaseClientConfig): + """Native Deepgram Flux ``/v2/speak`` streaming TTS client.""" + + sample_rate: int = field(DEFAULT_AUDIO_SAMPLE_RATE, help="PCM sample rate in Hz.") + model: str = field("", help="Deepgram Flux model ID, e.g. flux-alexis-en.") + api_key_env: str = field( + "DEEPGRAM_API_KEY", + help="Environment variable used when client.api_key is omitted.", + ) + pacing: TextPacingConfig = field( + default_factory=TextPacingConfig, + help="Upstream LLM text pacing configuration.", + ) + mip_opt_out: bool = field( + False, + help="Opt out of Deepgram's Model Improvement Program.", + ) + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.DEEPGRAM_FLUX_STREAMING_TTS + + def __post_init__(self): + super().__post_init__() + if not self.model and self.api_base is None: + return + if self.api_base is None: + raise ValueError( + "DeepgramFluxStreamingTTSClientConfig.api_base is required" + ) + if not self.model: + raise ValueError("DeepgramFluxStreamingTTSClientConfig.model is required") + if self.sample_rate <= 0: + raise ValueError( + "DeepgramFluxStreamingTTSClientConfig.sample_rate must be > 0" + ) + if not self.api_key_env: + raise ValueError( + "DeepgramFluxStreamingTTSClientConfig.api_key_env is required" + ) + + def build_tokenizer_provider(self): + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) + + +@frozen_dataclass +class DeepgramAuraStreamingTTSClientConfig(BaseClientConfig): + """Native Deepgram Aura ``/v1/speak`` WebSocket TTS client. + + Aura accepts incremental ``Speak`` messages and an explicit ``Flush`` to + drain the queued turn. This remains a distinct lane from Flux's + streaming-first, turn-based ``/v2/speak`` lifecycle. + """ + + sample_rate: int = field(DEFAULT_AUDIO_SAMPLE_RATE, help="PCM sample rate in Hz.") + model: str = field("", help="Deepgram Aura model ID, e.g. aura-2-thalia-en.") + api_key_env: str = field( + "DEEPGRAM_API_KEY", + help="Environment variable used when client.api_key is omitted.", + ) + pacing: TextPacingConfig = field( + default_factory=TextPacingConfig, + help="Upstream LLM text pacing configuration.", + ) + mip_opt_out: bool = field( + False, + help="Opt out of Deepgram's Model Improvement Program.", + ) + speed: float = field(1.0, help="Deepgram Aura speaking-rate multiplier.") + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.DEEPGRAM_AURA_STREAMING_TTS + + def __post_init__(self): + super().__post_init__() + if not self.model and self.api_base is None: + return + if self.api_base is None: + raise ValueError( + "DeepgramAuraStreamingTTSClientConfig.api_base is required" + ) + if not self.model: + raise ValueError("DeepgramAuraStreamingTTSClientConfig.model is required") + if self.sample_rate <= 0: + raise ValueError( + "DeepgramAuraStreamingTTSClientConfig.sample_rate must be > 0" + ) + if not self.api_key_env: + raise ValueError( + "DeepgramAuraStreamingTTSClientConfig.api_key_env is required" + ) + if not 0.7 <= self.speed <= 1.5: + raise ValueError("Deepgram Aura speed must be between 0.7 and 1.5") + + def build_tokenizer_provider(self): + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) + + +@frozen_dataclass +class ElevenLabsHTTPTTSClientConfig(BaseClientConfig): + """ElevenLabs complete-text, non-streaming HTTP TTS client.""" + + voice_id: str = field("", help="ElevenLabs voice identifier.") + sample_rate: int = field(DEFAULT_AUDIO_SAMPLE_RATE, help="PCM sample rate in Hz.") + model: str = field("", help="ElevenLabs model ID.") + api_key_env: str = field( + "ELEVENLABS_API_KEY", + help="Environment variable used when client.api_key is omitted.", + ) + stability: float = field(0.5, help="ElevenLabs voice stability.") + similarity_boost: float = field(0.8, help="ElevenLabs similarity boost.") + speed: float = field(1.0, help="ElevenLabs speaking-rate multiplier.") + apply_text_normalization: str = field( + "off", help="ElevenLabs text normalization mode: auto | on | off." + ) + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.ELEVENLABS_HTTP_TTS + + def __post_init__(self): + super().__post_init__() + if not self.model and self.api_base is None and not self.voice_id: + return + if self.api_base is None or not self.model or not self.voice_id: + raise ValueError( + "ElevenLabsHTTPTTSClientConfig requires api_base, model, and voice_id" + ) + if self.sample_rate <= 0: + raise ValueError("ElevenLabsHTTPTTSClientConfig.sample_rate must be > 0") + if not self.api_key_env: + raise ValueError("ElevenLabsHTTPTTSClientConfig.api_key_env is required") + if not 0.7 <= self.speed <= 1.2: + raise ValueError("ElevenLabs speed must be between 0.7 and 1.2") + if self.apply_text_normalization not in ("auto", "on", "off"): + raise ValueError( + "apply_text_normalization must be one of ('auto', 'on', 'off')" + ) + + def build_tokenizer_provider(self): + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) + + +@frozen_dataclass +class DeepgramFluxHTTPClientConfig(BaseClientConfig): + """Deepgram Flux complete-text, non-streaming HTTP TTS client.""" + + sample_rate: int = field(DEFAULT_AUDIO_SAMPLE_RATE, help="PCM sample rate in Hz.") + model: str = field("", help="Deepgram Flux model ID.") + api_key_env: str = field( + "DEEPGRAM_API_KEY", + help="Environment variable used when client.api_key is omitted.", + ) + mip_opt_out: bool = field( + False, help="Opt out of Deepgram's Model Improvement Program." + ) + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.DEEPGRAM_FLUX_HTTP_TTS + + def __post_init__(self): + super().__post_init__() + if not self.model and self.api_base is None: + return + if self.api_base is None or not self.model: + raise ValueError("DeepgramFluxHTTPClientConfig requires api_base and model") + if self.sample_rate <= 0: + raise ValueError("DeepgramFluxHTTPClientConfig.sample_rate must be > 0") + if not self.api_key_env: + raise ValueError("DeepgramFluxHTTPClientConfig.api_key_env is required") + + def build_tokenizer_provider(self): + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) + + +@frozen_dataclass +class STTClientConfig(BaseClientConfig): + """STT client configuration for realtime streaming speech-to-text APIs.""" + + provider: str = field( + "", + help=( + "STT provider name. Supported: 'vajra_openai_realtime', " "'vllm_realtime'." + ), + ) + sample_rate: int = field(16000, help="Expected audio sample rate in Hz.") + ws_chunk_size: int = field( + 4096, + help=( + "Bytes of raw PCM audio per WebSocket message. Client CPU scales " + "with concurrency * sample_rate * 2 / ws_chunk_size, so prefer " + "larger chunks at high concurrency." + ), + ) + ws_permessage_deflate: bool = field( + False, + help=( + "Negotiate WebSocket permessage-deflate compression. Disabled by " + "default because base64 PCM is high entropy and compression adds " + "substantial client and server CPU." + ), + ) + ws_realtime_pacing: bool = field( + False, + help=( + "Sleep between WebSocket audio chunks to simulate realtime input. " + "Enable for live-audio SLO measurements; disable for engine-bound " + "throughput measurements." + ), + ) + ws_ping_interval_s: Optional[int] = field( + 20, help="WebSocket ping interval in seconds; None disables pings." + ) + ws_ping_timeout_s: Optional[int] = field( + None, + help=( + "WebSocket ping timeout in seconds. None disables keepalive timeout " + "while preserving request_timeout." + ), + ) + model: str = field("", help="The STT model ID.") + + _SUPPORTED_PROVIDERS = ("vllm_realtime", "vajra_openai_realtime") + + @classmethod + def get_type(cls) -> ClientType: + return ClientType.STT + + def __post_init__(self): + super().__post_init__() + + # Skip validation for the default polymorphic child instantiated by Vidhi. + if not self.provider and not self.model and self.api_base is None: + return + if not self.provider: + raise ValueError( + "STTClientConfig.provider is required. " + f"Supported: {', '.join(self._SUPPORTED_PROVIDERS)}" + ) + if self.provider not in self._SUPPORTED_PROVIDERS: + raise ValueError( + f"Unsupported STT provider: {self.provider}. " + f"Supported: {', '.join(self._SUPPORTED_PROVIDERS)}" + ) + if not self.model: + raise ValueError("STTClientConfig.model is required.") + if self.api_base is None: + raise ValueError("STTClientConfig.api_base is required.") + if self.sample_rate <= 0: + raise ValueError("STTClientConfig.sample_rate must be > 0") + if self.ws_chunk_size <= 0: + raise ValueError("STTClientConfig.ws_chunk_size must be > 0") + if self.ws_ping_interval_s is not None and self.ws_ping_interval_s <= 0: + raise ValueError("STTClientConfig.ws_ping_interval_s must be > 0 or None") + if self.ws_ping_timeout_s is not None and self.ws_ping_timeout_s <= 0: + raise ValueError("STTClientConfig.ws_ping_timeout_s must be > 0 or None") + + def build_tokenizer_provider(self): + """STT models use a simple word-split tokenizer.""" + from veeksha.core.tokenizer import build_word_split_tokenizer_provider + + return build_word_split_tokenizer_provider(self.model) diff --git a/veeksha/config/endpoint.py b/veeksha/config/endpoint.py new file mode 100644 index 00000000..fbf12f63 --- /dev/null +++ b/veeksha/config/endpoint.py @@ -0,0 +1,75 @@ +"""Inference endpoint contract shared by launchers and benchmarks.""" + +from __future__ import annotations + +from dataclasses import asdict, fields, replace +from typing import Any, Optional + +from vidhi import field, frozen_dataclass + +from veeksha.types import ServerType + + +@frozen_dataclass +class EndpointConfig: + """Client-facing contract for an inference endpoint.""" + + engine_type: str = field("", help="Endpoint engine type: vajra, vllm, or sglang.") + model: str = field("", help="Model name exposed by the endpoint.") + api_base: str = field("", help="OpenAI-compatible API base URL.") + api_key: Optional[str] = field(None, help="API key for the endpoint.") + health_url: Optional[str] = field(None, help="Health endpoint URL.") + host: str = field("localhost", help="Host used for local port ownership checks.") + port: int = field(0, help="Host port owned by the endpoint.") + + def __post_init__(self) -> None: + self.engine_type = _normalize_engine_type(self.engine_type) + if not self.model: + raise ValueError("endpoint.model must be non-empty") + if not self.api_base: + raise ValueError("endpoint.api_base must be non-empty") + self.api_base = self.api_base.rstrip("/") + if self.health_url is not None: + self.health_url = self.health_url.rstrip("/") + if not self.host: + raise ValueError("endpoint.host must be non-empty") + if self.port < 0: + raise ValueError("endpoint.port must be >= 0") + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + def client_overrides_for(self, client_config: Any) -> dict[str, Any]: + client_fields = {config_field.name for config_field in fields(client_config)} + overrides: dict[str, Any] = { + "api_base": self.api_base, + "api_key": self.api_key, + "model": self.model, + } + return { + key: value + for key, value in overrides.items() + if key in client_fields and value is not None + } + + def apply_to_client_config(self, client_config: Any) -> Any: + return replace(client_config, **self.client_overrides_for(client_config)) + + def apply_to_client_mapping(self, client_mapping: dict[str, Any]) -> None: + client_mapping["api_base"] = self.api_base + client_mapping["model"] = self.model + if self.api_key is not None: + client_mapping["api_key"] = self.api_key + + +def _normalize_engine_type(engine_type: str | ServerType) -> str: + if isinstance(engine_type, ServerType): + return str(engine_type) + if not isinstance(engine_type, str) or not engine_type: + raise ValueError("endpoint.engine_type must be one of: vajra, vllm, sglang") + try: + return str(ServerType.from_str(engine_type)) + except KeyError as exc: + raise ValueError( + "endpoint.engine_type must be one of: vajra, vllm, sglang" + ) from exc diff --git a/veeksha/config/evaluator.py b/veeksha/config/evaluator.py index f2c664bf..6ede6b2b 100644 --- a/veeksha/config/evaluator.py +++ b/veeksha/config/evaluator.py @@ -7,13 +7,15 @@ - BaseEvaluatorConfig (abstract base) - PerformanceEvaluatorConfig (latency, throughput) - LMEvalAccuracyEvaluatorConfig (task-specific correctness - lm-eval) + - AudioQualityEvaluatorConfig (generated audio quality checks) """ -from typing import Optional, Union +from typing import List, Optional, Union, cast from vidhi import BasePolyConfig, field, frozen_dataclass from veeksha.config.slo import BaseSloConfig, ConstantSloConfig +from veeksha.config.verification import AudioVerificationConfig from veeksha.types import ChannelModality, EvaluationType @@ -67,9 +69,6 @@ def __post_init__(self): ) -# ---- Channel-specific performance configs ---- - - @frozen_dataclass class BaseChannelPerformanceConfig(BasePolyConfig): """Base config for channel-specific performance""" @@ -105,12 +104,83 @@ def get_type(cls) -> ChannelModality: @frozen_dataclass class AudioChannelPerformanceConfig(BaseChannelPerformanceConfig): - """Audio channel performance configuration""" + """Audio channel performance configuration for TTS benchmarking.""" + + interactivity_enabled: bool = field( + True, + help="Compute streaming interactivity metrics when timestamp lists are present.", + ) + startup_delay_ms_values: List[float] = field( + default_factory=lambda: [0.0, 100.0, 300.0], + help="Fixed delays after first audio to simulate, in milliseconds.", + ) + startup_buffer_ms_values: List[float] = field( + default_factory=lambda: [0.0, 100.0, 300.0], + help="Playable-audio targets to simulate before playback, in milliseconds.", + ) + min_reportable_stall_ms: float = field( + 10.0, + help="Playback gaps at or below this duration are treated as transport noise.", + ) + fluidity_frame_ms: float = field( + 20.0, + help=( + "Duration of each complete playable PCM frame used by the " + "Etalon-inspired audio fluidity metric." + ), + ) + fluidity_startup_delay_ms: float = field( + 0.0, + help=( + "Playback slack after the first complete frame for the primary " + "user_audio_fluidity_index metric. Set this explicitly when the " + "playback client intentionally buffers before starting." + ), + ) + fluidity_attribution_mode: str = field( + "conservative", + help=( + "How to attribute fluidity misses to the TTS service. " + "'conservative' emits tts_service_fluidity_index only when all " + "text was sent before the first playable frame. " + "'source_oversupplied' treats every miss as service-caused and must " + "only be used with a trace that guarantees synthesis-eligible text " + "throughout playback. User fluidity is always reported." + ), + ) + persist_raw_timing: bool = field( + False, + help="Write metrics/audio_raw_timing.jsonl with raw per-event timestamps.", + ) @classmethod def get_type(cls) -> ChannelModality: return ChannelModality.AUDIO + def __post_init__(self): + for field_name in ("startup_delay_ms_values", "startup_buffer_ms_values"): + values = getattr(self, field_name) + if any(value < 0 for value in values): + raise ValueError(f"{field_name} must contain only values >= 0") + normalized = list(dict.fromkeys(float(value) for value in values)) + if 0.0 not in normalized: + normalized.insert(0, 0.0) + object.__setattr__(self, field_name, normalized) + if self.min_reportable_stall_ms < 0: + raise ValueError("min_reportable_stall_ms must be >= 0") + if self.fluidity_frame_ms <= 0: + raise ValueError("fluidity_frame_ms must be > 0") + if self.fluidity_startup_delay_ms < 0: + raise ValueError("fluidity_startup_delay_ms must be >= 0") + if self.fluidity_attribution_mode not in ( + "conservative", + "source_oversupplied", + ): + raise ValueError( + "fluidity_attribution_mode must be one of " + "('conservative', 'source_oversupplied')" + ) + @frozen_dataclass class VideoChannelPerformanceConfig(BaseChannelPerformanceConfig): @@ -121,9 +191,6 @@ def get_type(cls) -> ChannelModality: return ChannelModality.VIDEO -# ---- Base evaluator config ---- - - def _default_slos() -> list[BaseSloConfig]: return [ ConstantSloConfig( @@ -141,6 +208,16 @@ def _default_slos() -> list[BaseSloConfig]: ] +def _normalize_channel_modality(channel: object) -> ChannelModality: + if isinstance(channel, ChannelModality): + return channel + if isinstance(channel, str): + return cast(ChannelModality, ChannelModality.from_str(channel)) + if isinstance(channel, int) and not isinstance(channel, bool): + return ChannelModality(channel) + raise ValueError(f"Invalid target channel modality: {channel!r}") + + @frozen_dataclass class BaseEvaluatorConfig(BasePolyConfig): """Base configuration for all evaluators (performance, accuracy)""" @@ -149,12 +226,10 @@ class BaseEvaluatorConfig(BasePolyConfig): default_factory=lambda: ["text"], help="List of ChannelModality values to evaluate.", ) - slos: list[BaseSloConfig] = field( default_factory=_default_slos, help="List of SLO definitions to evaluate against request-level metrics.", ) - stream_metrics: bool = field(True, help="Enable real-time metric streaming") stream_metrics_interval: float = field( 5.0, help="Interval for streaming metrics in seconds" @@ -162,16 +237,11 @@ class BaseEvaluatorConfig(BasePolyConfig): def __post_init__(self): if self.target_channels: - converted = [] - for ch in self.target_channels: - if isinstance(ch, str): - converted.append(ChannelModality.from_str(ch)) - else: - converted.append(ch) - object.__setattr__(self, "target_channels", converted) - - -# ---- Performance evaluator config ---- + object.__setattr__( + self, + "target_channels", + [_normalize_channel_modality(ch) for ch in self.target_channels], + ) @frozen_dataclass @@ -197,10 +267,18 @@ class PerformanceEvaluatorConfig(BaseEvaluatorConfig): def get_type(cls) -> EvaluationType: return EvaluationType.PERFORMANCE + def __post_init__(self): + super().__post_init__() + if ChannelModality.AUDIO in self.target_channels and self.audio_channel is None: + object.__setattr__(self, "audio_channel", AudioChannelPerformanceConfig()) + if ChannelModality.VIDEO in self.target_channels and self.video_channel is None: + object.__setattr__(self, "video_channel", VideoChannelPerformanceConfig()) + def get_channel_config( - self, channel: ChannelModality + self, channel: ChannelModality | str | int ) -> Optional[BaseChannelPerformanceConfig]: """Get the performance config for a specific channel.""" + channel = _normalize_channel_modality(channel) if channel == ChannelModality.TEXT: return self.text_channel elif channel == ChannelModality.IMAGE: @@ -212,18 +290,18 @@ def get_channel_config( return None -# ---- Accuracy evaluator config(s) ---- - - @frozen_dataclass class LMEvalAccuracyEvaluatorConfig(BaseEvaluatorConfig): - """Configuration for lm-eval accuracy evaluation (task-specific correctness). - - IMPORTANT: For lm-eval accuracy evaluation, the content generation must use - `LMEvalSessionGenerator`. The generator owns the lm-eval Task/Instance objects, - and the evaluator binds responses to instances for evaluation. - """ + """Configuration for lm-eval accuracy evaluation (task-specific correctness).""" + slos: list[BaseSloConfig] = field( + default_factory=list, + help="Accuracy evaluators do not use performance SLOs.", + ) + stream_metrics: bool = field( + False, + help="Accuracy evaluation does not stream incremental metrics.", + ) bootstrap_iters: int = field( 100000, help="Bootstrap iterations for confidence intervals" ) @@ -231,3 +309,33 @@ class LMEvalAccuracyEvaluatorConfig(BaseEvaluatorConfig): @classmethod def get_type(cls) -> EvaluationType: return EvaluationType.ACCURACY_LMEVAL + + +@frozen_dataclass +class AudioQualityEvaluatorConfig(BaseEvaluatorConfig): + """Configuration for generated audio quality evaluation.""" + + target_channels: list = field( + default_factory=lambda: ["audio"], + help="List of modalities whose output quality should be evaluated.", + ) + slos: list[BaseSloConfig] = field( + default_factory=list, + help="Accuracy evaluators do not use performance SLOs.", + ) + stream_metrics: bool = field( + False, + help="Accuracy evaluation does not stream incremental metrics.", + ) + verification: AudioVerificationConfig = field( + default_factory=AudioVerificationConfig, + help="Generated audio verification configuration.", + ) + save_audio_files: bool = field( + True, + help="Whether to persist audio artifacts for quality evaluation.", + ) + + @classmethod + def get_type(cls) -> EvaluationType: + return EvaluationType.AUDIO_QUALITY diff --git a/veeksha/config/generator/session.py b/veeksha/config/generator/session.py index e610af1f..cff1018c 100644 --- a/veeksha/config/generator/session.py +++ b/veeksha/config/generator/session.py @@ -153,24 +153,15 @@ class UntimedContentMultiTurnTraceFlavorConfig(BaseTraceFlavorConfig): """ conversation_column: str = field( - default="conversations", - metadata={"help": "Column containing the list of conversation messages"}, - ) - role_key: str = field( - default="from", - metadata={"help": "Key for the role field in each message dict"}, + "conversations", help="Column containing the list of conversation messages" ) + role_key: str = field("from", help="Key for the role field in each message dict") content_key: str = field( - default="value", - metadata={"help": "Key for the content field in each message dict"}, - ) - user_role_value: str = field( - default="human", - metadata={"help": "Role value indicating user messages"}, + "value", help="Key for the content field in each message dict" ) + user_role_value: str = field("human", help="Role value indicating user messages") assistant_role_value: str = field( - default="gpt", - metadata={"help": "Role value indicating assistant messages"}, + "gpt", help="Role value indicating assistant messages" ) @classmethod @@ -178,6 +169,225 @@ def get_type(cls): return TraceFlavorType.UNTIMED_CONTENT_MULTI_TURN +@frozen_dataclass +class ShareGPTTraceFlavorConfig(BaseTraceFlavorConfig): + """ShareGPT conversation trace flavor configuration. + + Reads ShareGPT-format conversations and uses assistant turn text + as TTS input. Each assistant turn becomes a single-request session. + """ + + assistant_role: str = field( + "gpt", + help="Role name for assistant turns in the ShareGPT data " + "(common values: 'gpt', 'assistant').", + ) + min_tokens: int = field( + 20, + help="Minimum input token count. Turns shorter than this are skipped. " + "Ignored when min_chars/max_chars are set.", + ) + max_tokens: int = field( + 100, + help="Maximum input token count. Text is truncated to sampled length. " + "Ignored when min_chars/max_chars are set.", + ) + min_chars: int = field( + -1, + help="If >= 0, enables char-based input length control (mutually exclusive " + "with min_tokens/max_tokens). Turns with fewer than this many chars are skipped.", + ) + max_chars: int = field( + -1, + help="If >= 0, enables char-based input length control (mutually exclusive " + "with min_tokens/max_tokens). Text is truncated to a char length sampled " + "uniformly in [min_chars, max_chars].", + ) + min_alpha_ratio: float = field( + 0.5, + help="Minimum ratio of alphabetic characters to total non-space characters. " + "Filters out junk entries like number sequences or code snippets. " + "Set to 0.0 to disable.", + ) + + @property + def use_chars(self) -> bool: + return self.min_chars >= 0 and self.max_chars >= 0 + + def __post_init__(self): + one_set = (self.min_chars >= 0) != (self.max_chars >= 0) + if one_set: + raise ValueError( + "min_chars and max_chars must both be set (>= 0) or both left " + f"unset (-1); got min_chars={self.min_chars}, max_chars={self.max_chars}" + ) + if self.use_chars: + if self.min_chars > self.max_chars: + raise ValueError( + f"min_chars ({self.min_chars}) must be <= max_chars ({self.max_chars})" + ) + elif self.min_tokens > self.max_tokens: + raise ValueError( + f"min_tokens ({self.min_tokens}) must be <= max_tokens ({self.max_tokens})" + ) + + @classmethod + def get_type(cls): + return TraceFlavorType.SHAREGPT + + +@frozen_dataclass +class AudioTraceFlavorConfig(BaseTraceFlavorConfig): + """Audio trace flavor configuration for STT benchmarking.""" + + audio_dir: str = field( + "", + help="Optional base directory prepended to relative audio_file paths.", + ) + target_duration_s: float | None = field( + None, + help=( + "Optional per-session audio duration to stream. The trace must " + "provide reference_word_timestamps so expected transcripts can be " + "trimmed to the streamed prefix." + ), + ) + target_duration_spread_s: float | None = field( + None, + help=( + "Hard half-width of a clipped-Gaussian duration spread around " + "target_duration_s. Requires target_duration_s." + ), + ) + target_duration_sigma_s: float | None = field( + None, + help=( + "Standard deviation of the clipped-Gaussian duration draw. Defaults " + "to half the spread and must not exceed the spread." + ), + ) + + @classmethod + def get_type(cls): + return TraceFlavorType.AUDIO + + def __post_init__(self): + if self.target_duration_s is not None and self.target_duration_s <= 0: + raise ValueError( + "target_duration_s must be positive when set; " + f"got {self.target_duration_s}" + ) + if self.target_duration_spread_s is not None: + if self.target_duration_s is None: + raise ValueError("target_duration_spread_s requires target_duration_s") + if not 0 < self.target_duration_spread_s < self.target_duration_s: + raise ValueError( + "target_duration_spread_s must be in (0, " + f"target_duration_s); got {self.target_duration_spread_s} " + f"with target_duration_s={self.target_duration_s}" + ) + if self.target_duration_sigma_s is not None: + if self.target_duration_spread_s is None: + raise ValueError( + "target_duration_sigma_s requires target_duration_spread_s" + ) + if self.target_duration_sigma_s <= 0: + raise ValueError( + "target_duration_sigma_s must be positive; got " + f"{self.target_duration_sigma_s}" + ) + if self.target_duration_sigma_s > self.target_duration_spread_s: + raise ValueError( + "target_duration_sigma_s must be <= " + "target_duration_spread_s; got sigma=" + f"{self.target_duration_sigma_s} spread=" + f"{self.target_duration_spread_s}" + ) + + +@frozen_dataclass +class SeedTTSTextTraceFlavorConfig(BaseTraceFlavorConfig): + """Seed TTS eval text dataset configuration. + + Loads one text-only TTS request per dataset row. The default source is the + English split of the Hugging Face Seed-TTS eval mirror. + """ + + dataset_name: str = field( + "TwinkStart/Seed-TTS-Eval", + help="Hugging Face dataset name used when local_path is empty.", + ) + subset: str = field("en", help="Dataset subset/config name. Defaults to English.") + split: str = field("train", help="Dataset split to load.") + text_column: str = field( + "text", help="Column containing the target TTS synthesis text." + ) + id_column: str = field( + "filename", + help="Optional source row identifier column copied into request metadata.", + ) + local_path: str = field( + "", + help="Optional local dataset path. Supports saved HF datasets and common " + "data files such as JSON/JSONL, CSV, and Parquet.", + ) + min_tokens: int = field( + 20, + help="Minimum input word count. Rows with fewer words are skipped. " + "Ignored when min_chars/max_chars are set.", + ) + max_tokens: int = field( + 150, + help="Maximum input word count. Text is truncated to a sampled word " + "count in [min_tokens, max_tokens]. Ignored when min_chars/max_chars " + "are set.", + ) + min_chars: int = field( + -1, + help="If >= 0, enables char-based input length control. Rows with " + "fewer chars are skipped.", + ) + max_chars: int = field( + -1, + help="If >= 0, enables char-based input length control. Text is " + "truncated to a sampled char count in [min_chars, max_chars].", + ) + + @property + def use_chars(self) -> bool: + return self.min_chars >= 0 and self.max_chars >= 0 + + @classmethod + def get_type(cls): + return TraceFlavorType.SEED_TTS_TEXT + + def __post_init__(self): + if not self.text_column: + raise ValueError("SeedTTSTextTraceFlavorConfig.text_column is required.") + if not self.local_path and not self.dataset_name: + raise ValueError( + "SeedTTSTextTraceFlavorConfig requires dataset_name or local_path." + ) + one_char_bound_set = (self.min_chars >= 0) != (self.max_chars >= 0) + if one_char_bound_set: + raise ValueError( + "min_chars and max_chars must both be set (>= 0) or both left " + f"unset (-1); got min_chars={self.min_chars}, max_chars={self.max_chars}" + ) + if self.use_chars: + if self.min_chars > self.max_chars: + raise ValueError( + f"min_chars ({self.min_chars}) must be <= max_chars ({self.max_chars})" + ) + else: + if self.min_tokens < 0 or self.max_tokens < 0: + raise ValueError("min_tokens and max_tokens must be non-negative.") + if self.min_tokens > self.max_tokens: + raise ValueError( + f"min_tokens ({self.min_tokens}) must be <= max_tokens ({self.max_tokens})" + ) + + # ----- Trace Session Generator Config ----- diff --git a/veeksha/config/launcher.py b/veeksha/config/launcher.py new file mode 100644 index 00000000..907a407e --- /dev/null +++ b/veeksha/config/launcher.py @@ -0,0 +1,357 @@ +"""Configuration model for orchestrated Veeksha launcher runs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, fields, is_dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Mapping, Optional + +import yaml +from vidhi import create_class_from_dict +from vidhi.utils import get_all_subclasses + +from veeksha.config.endpoint import EndpointConfig +from veeksha.config.server import BaseServerConfig, ManagedServerConfig +from veeksha.sweeps import planner as sweep_planner + + +class LauncherConfigError(ValueError): + """Raised when a launcher YAML file is invalid.""" + + +_SERVER_TYPE_ERROR = "server.type must be one of: vajra, vllm, sglang" + + +@dataclass(frozen=True) +class RetryConfig: + max_attempts_per_run: int = 2 + restart_engine_before_retry: bool = True + fail_sweep_after_exhausted_retries: bool = True + + def __post_init__(self) -> None: + if self.max_attempts_per_run <= 0: + raise LauncherConfigError("retry.max_attempts_per_run must be positive") + + +@dataclass(frozen=True) +class LauncherConfig: + sweep: sweep_planner.SweepConfig + server: Optional[ManagedServerConfig] = None + endpoint: Optional[EndpointConfig] = None + retry: RetryConfig = field(default_factory=RetryConfig) + output_dir: str = field(default_factory=lambda: _default_output_dir()) + + def __post_init__(self) -> None: + if self.server is not None and self.endpoint is not None: + raise LauncherConfigError( + "launcher config accepts either server or endpoint, not both" + ) + + @classmethod + def from_file(cls, path: str | Path) -> "LauncherConfig": + config_path = Path(path) + with config_path.open("r", encoding="utf-8") as f: + raw = yaml.safe_load(f) + if not isinstance(raw, Mapping): + raise LauncherConfigError("launcher config must be a YAML mapping") + return cls.from_mapping(raw) + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any]) -> "LauncherConfig": + _reject_unknown_keys( + raw, {"server", "endpoint", "sweep", "retry", "output_dir"}, "launcher" + ) + + sweep_raw = raw.get("sweep") + if not isinstance(sweep_raw, Mapping): + raise LauncherConfigError("launcher config requires a sweep mapping") + try: + sweep = sweep_planner.SweepConfig.from_mapping(sweep_raw) + sweep, _ = sweep_planner.resolve_sweep_config(sweep) + except sweep_planner.SweepConfigError as exc: + raise LauncherConfigError(str(exc)) from exc + + server_raw = raw.get("server") + endpoint_raw = raw.get("endpoint") + if server_raw is not None and endpoint_raw is not None: + raise LauncherConfigError( + "launcher config accepts either server or endpoint, not both" + ) + + server = None + endpoint = None + if server_raw is not None: + if not isinstance(server_raw, Mapping): + raise LauncherConfigError("server must be a mapping when provided") + server = _server_config_from_mapping(server_raw) + _validate_endpoint_matches_sweep( + server.get_endpoint(), sweep.engine, "server.type" + ) + if endpoint_raw is not None: + if not isinstance(endpoint_raw, Mapping): + raise LauncherConfigError("endpoint must be a mapping when provided") + endpoint = _endpoint_config_from_mapping(endpoint_raw) + _validate_endpoint_matches_sweep( + endpoint, sweep.engine, "endpoint.engine_type" + ) + + retry_raw = raw.get("retry", {}) + if not isinstance(retry_raw, Mapping): + raise LauncherConfigError("retry must be a mapping") + retry = _retry_config_from_mapping(retry_raw) + + output_dir = raw.get("output_dir") or _default_output_dir() + if not isinstance(output_dir, str): + raise LauncherConfigError("output_dir must be a string") + + return cls( + sweep=sweep, + server=server, + endpoint=endpoint, + retry=retry, + output_dir=output_dir, + ) + + def to_dict(self) -> dict[str, Any]: + return _to_plain_data(self) + + +def _default_output_dir() -> str: + timestamp = datetime.now(timezone.utc).strftime("%d_%m_%Y-%H_%M_%S") + return f"benchmark_output/veeksha_launcher_{timestamp}" + + +def _server_config_from_mapping(raw: Mapping[str, Any]) -> ManagedServerConfig: + engine_type = raw.get("type") + config_cls = _server_config_class_from_type(engine_type) + data = _coerce_server_dict(raw, config_cls) + data.pop("type", None) + try: + return create_class_from_dict(config_cls, data) + except (TypeError, ValueError) as exc: + raise LauncherConfigError(str(exc)) from exc + + +def _endpoint_config_from_mapping(raw: Mapping[str, Any]) -> EndpointConfig: + _reject_unknown_keys( + raw, {field.name for field in fields(EndpointConfig)}, "endpoint" + ) + try: + return create_class_from_dict(EndpointConfig, dict(raw)) + except (TypeError, ValueError) as exc: + raise LauncherConfigError(str(exc)) from exc + + +def _validate_endpoint_matches_sweep( + endpoint: EndpointConfig, sweep_engine: str, source: str +) -> None: + if endpoint.engine_type != sweep_engine: + raise LauncherConfigError( + f"{source} must match sweep.engine " + f"({source}={endpoint.engine_type}, sweep.engine={sweep_engine})" + ) + + +def _server_config_class_from_type(engine_type: object) -> type[BaseServerConfig]: + if not isinstance(engine_type, str): + raise LauncherConfigError(_SERVER_TYPE_ERROR) + normalized_engine_type = engine_type.lower() + for subclass in get_all_subclasses(BaseServerConfig): + try: + subtype = subclass.get_type() + except NotImplementedError: + continue + if getattr(subtype, "name", "").lower() == normalized_engine_type: + return subclass + if str(subtype).lower() == normalized_engine_type: + return subclass + raise LauncherConfigError(_SERVER_TYPE_ERROR) + + +def _coerce_server_dict( + raw: Mapping[str, Any], config_cls: type[BaseServerConfig] +) -> dict[str, Any]: + allowed = {field.name for field in fields(config_cls)} | {"type"} + _reject_unknown_keys(raw, allowed, "server") + data = dict(raw) + + for key in ("host", "type", "api_key", "dtype"): + _coerce_present_str(data, key) + for key in ( + "api_base", + "health_url", + "setup_dir", + "env_path", + "max_model_len", + ): + if key == "max_model_len": + _coerce_optional_int(data, key) + else: + _coerce_optional_str(data, key) + for key in ("image", "hf_model", "model_path", "deploy_config"): + _coerce_default_str(data, key) + for key in ("container_name", "docker_gpus", "docker_runtime", "ipc_mode"): + _coerce_optional_str(data, key) + for key in ( + "container_deploy_config", + "shm_size", + "model_name", + "source_dir", + "container_source_dir", + "venv_path", + "bootstrap", + "model", + ): + _coerce_optional_str(data, key) + _coerce_present_int(data, "port") + for key in ("max_restarts", "tensor_parallel_size"): + _coerce_default_int(data, key) + _coerce_optional_int(data, "container_port") + for key in ("startup_timeout", "health_check_interval"): + _coerce_default_float(data, key) + for key in ("command", "engine_args", "volumes", "pass_env", "docker_run_args"): + _coerce_string_list(data, key) + _coerce_string_mapping(data, "env") + _coerce_int_list(data, "gpu_ids") + _coerce_bool(data, "require_contiguous_gpus") + return data + + +def _retry_config_from_mapping(raw: Mapping[str, Any]) -> RetryConfig: + _reject_unknown_keys( + raw, + { + "max_attempts_per_run", + "restart_engine_before_retry", + "fail_sweep_after_exhausted_retries", + }, + "retry", + ) + data = dict(raw) + value = data.get("max_attempts_per_run") + if value is not None and (isinstance(value, bool) or not isinstance(value, int)): + raise LauncherConfigError("retry.max_attempts_per_run must be an integer") + for key in ("restart_engine_before_retry", "fail_sweep_after_exhausted_retries"): + value = data.get(key) + if value is not None and not isinstance(value, bool): + raise LauncherConfigError(f"retry.{key} must be a boolean") + return RetryConfig(**data) + + +def _reject_unknown_keys( + raw: Mapping[str, Any], allowed: set[str], section: str +) -> None: + unknown = sorted(set(raw) - allowed) + if unknown: + raise LauncherConfigError( + f"unsupported {section} config keys: {', '.join(unknown)}" + ) + + +def _coerce_present_str(data: dict[str, Any], key: str) -> None: + if key not in data: + return + if not isinstance(data[key], str): + raise LauncherConfigError(f"server.{key} must be a string") + + +def _coerce_optional_str(data: dict[str, Any], key: str) -> None: + if key not in data or data[key] is None: + return + if not isinstance(data[key], str): + raise LauncherConfigError(f"server.{key} must be a string") + + +def _coerce_default_str(data: dict[str, Any], key: str) -> None: + if data.get(key) is None: + data.pop(key, None) + return + if not isinstance(data[key], str): + raise LauncherConfigError(f"server.{key} must be a string") + + +def _coerce_present_int(data: dict[str, Any], key: str) -> None: + if key not in data: + return + if isinstance(data[key], bool) or not isinstance(data[key], int): + raise LauncherConfigError(f"server.{key} must be an integer") + + +def _coerce_optional_int(data: dict[str, Any], key: str) -> None: + if key not in data or data[key] is None: + return + if isinstance(data[key], bool) or not isinstance(data[key], int): + raise LauncherConfigError(f"server.{key} must be an integer") + + +def _coerce_default_int(data: dict[str, Any], key: str) -> None: + if data.get(key) is None: + data.pop(key, None) + return + if isinstance(data[key], bool) or not isinstance(data[key], int): + raise LauncherConfigError(f"server.{key} must be an integer") + + +def _coerce_default_float(data: dict[str, Any], key: str) -> None: + if data.get(key) is None: + data.pop(key, None) + return + if isinstance(data[key], bool) or not isinstance(data[key], (int, float)): + raise LauncherConfigError(f"server.{key} must be a number") + data[key] = float(data[key]) + + +def _coerce_string_list(data: dict[str, Any], key: str) -> None: + value = data.get(key) + if value is None: + data.pop(key, None) + return + if not isinstance(value, list) or not all(isinstance(item, str) for item in value): + raise LauncherConfigError(f"server.{key} must be a list of strings") + + +def _coerce_int_list(data: dict[str, Any], key: str) -> None: + value = data.get(key) + if value is None: + return + if not isinstance(value, list) or not all( + isinstance(item, int) and not isinstance(item, bool) for item in value + ): + raise LauncherConfigError(f"server.{key} must be a list of integers") + + +def _coerce_bool(data: dict[str, Any], key: str) -> None: + value = data.get(key) + if value is None: + return + if not isinstance(value, bool): + raise LauncherConfigError(f"server.{key} must be a boolean") + + +def _coerce_string_mapping(data: dict[str, Any], key: str) -> None: + value = data.get(key) + if value is None: + data.pop(key, None) + return + if not isinstance(value, Mapping): + raise LauncherConfigError(f"server.{key} must be a mapping") + data[key] = {str(k): str(v) for k, v in value.items()} + + +def _to_plain_data(value: Any) -> Any: + if is_dataclass(value) and not isinstance(value, type): + data = { + field.name: _to_plain_data(getattr(value, field.name)) + for field in fields(value) + } + if hasattr(value, "get_type") and callable(getattr(value, "get_type")): + data["type"] = str(value.get_type()) + return data + if isinstance(value, tuple): + return [_to_plain_data(item) for item in value] + if isinstance(value, list): + return [_to_plain_data(item) for item in value] + if isinstance(value, dict): + return {key: _to_plain_data(item) for key, item in value.items()} + return value diff --git a/veeksha/config/runtime.py b/veeksha/config/runtime.py index 691c7f5a..beb0a781 100644 --- a/veeksha/config/runtime.py +++ b/veeksha/config/runtime.py @@ -1,3 +1,5 @@ +from typing import Optional + from vidhi import field, frozen_dataclass @@ -15,10 +17,21 @@ class RuntimeConfig: 2, help="Number of threads for dispatching requests to workers." ) num_completion_threads: int = field( - 2, help="Number of threads for processing completed requests." + 8, + help=( + "Number of threads for processing completed requests. Completion " + "workers also run per-request ASR scoring concurrently, so " + "under-provisioning stretches the post-run drain." + ), ) - num_client_threads: int = field( - 3, help="Number of async worker threads for making concurrent requests." + num_client_threads: Optional[int] = field( + None, + help=( + "Number of async worker threads for making concurrent requests. " + "None provisions one thread per eight target-concurrent sessions, " + "with a minimum of three. Under-provisioning stretches realtime " + "send cadence and invalidates high-concurrency latency metrics." + ), ) pregenerate_sessions: bool = field( False, diff --git a/veeksha/config/server.py b/veeksha/config/server.py index cc0dfb81..5c2c684c 100644 --- a/veeksha/config/server.py +++ b/veeksha/config/server.py @@ -1,69 +1,112 @@ """ -Server configuration for LLM inference systems. +Server configuration for launcher-managed inference systems. """ -from typing import Any, Dict, List, Optional, Union +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional, TypeAlias, Union from vidhi import BasePolyConfig, field, frozen_dataclass +from veeksha.config.endpoint import EndpointConfig from veeksha.types import ServerType +_DEFAULT_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct" + +VLLM_OMNI_DEFAULT_IMAGE = "vllm-omni:0.21-local" +SGLANG_OMNI_DEFAULT_IMAGE = "frankleeeee/sglang-omni:dev" + +SGLANG_OMNI_DEFAULT_BOOTSTRAP = """set -e +if [ ! -x {venv}/bin/sgl-omni ]; then + uv venv {venv} -p 3.12 + . {venv}/bin/activate + cd {src} + uv pip install -e . + uv pip install transformers==4.57.3 accelerate==1.12.0 sox einops + uv pip install --no-deps qwen-tts==0.1.1 +else + . {venv}/bin/activate +fi""" + @frozen_dataclass class BaseServerConfig(BasePolyConfig): - """Base configuration for server launch and management.""" + """Base configuration for a managed inference server.""" env_path: Optional[str] = field( None, help="Path to a Python environment directory (virtualenv/conda)." ) - model: str = field( - "meta-llama/Meta-Llama-3-8B-Instruct", help="Model name or path." + _DEFAULT_MODEL, help="Model name exposed to the benchmark client." ) - host: str = field("localhost", help="Host address for the server") - port: int = field(8000, help="Port number for the server") - api_key: str = field("token-abc123", help="API key for server authentication") - gpu_ids: Optional[List[int]] = field( None, help="List of GPU IDs to use (None means auto-assign)" ) - - startup_timeout: int = field(300, help="Timeout in seconds for server startup") - + startup_timeout: float = field(300.0, help="Timeout in seconds for server startup") health_check_interval: float = field( 2.0, help="Interval in seconds between health checks" ) - require_contiguous_gpus: bool = field( True, help="Require contiguous GPU allocation (e.g., GPUs 0,1,2 instead of 0,2,5)", ) - - # engine arguments - tensor_parallel_size: int = field(1, help="Number of GPUs for tensor parallelism") - dtype: str = field( "auto", help="Data type for model weights (auto, float16, bfloat16, etc.)", ) - max_model_len: Optional[int] = field(None, help="Maximum model context length") - additional_args: Union[str, Dict[str, Any], None] = field( "{}", help="Additional engine-specific arguments as JSON string, dict, or None.", ) + api_base: Optional[str] = field(None, help="External API base URL for the server.") + health_url: Optional[str] = field(None, help="Health endpoint URL for the server.") + setup_dir: Optional[str] = field( + None, help="Source checkout used by subprocess engines." + ) + max_restarts: int = field(3, help="Maximum server restarts per sweep run.") + + def __post_init__(self) -> None: + if self.port <= 0: + raise ValueError("server.port must be a positive integer") + if self.startup_timeout <= 0: + raise ValueError("server.startup_timeout must be positive") + if self.health_check_interval <= 0: + raise ValueError("server.health_check_interval must be positive") + if self.max_restarts < 0: + raise ValueError("server.max_restarts must be >= 0") + if self.tensor_parallel_size <= 0: + raise ValueError("server.tensor_parallel_size must be positive") + if self.gpu_ids is not None and any(gpu_id < 0 for gpu_id in self.gpu_ids): + raise ValueError("server.gpu_ids must contain non-negative IDs") def get_api_base_url(self) -> str: + if self.api_base: + return self.api_base.rstrip("/") return f"http://{self.host}:{self.port}/v1" def get_health_check_url(self) -> str: + if self.health_url: + return self.health_url return f"http://{self.host}:{self.port}/health" + def get_endpoint(self) -> EndpointConfig: + return EndpointConfig( + engine_type=self.get_type(), + model=self.model, + api_base=self.get_api_base_url(), + api_key=self.api_key, + health_url=self.get_health_check_url(), + host=self.host, + port=self.port, + ) + def get_gpu_env_var(self) -> Optional[str]: """Get CUDA_VISIBLE_DEVICES value if gpu_ids is specified.""" if self.gpu_ids is not None: @@ -79,18 +122,238 @@ def get_num_gpus(self) -> int: @property def engine(self) -> str: """Get the engine name for logging/compat.""" - return self.get_type().name.lower() + engine_type = self.get_type() + return getattr(engine_type, "name", str(engine_type)).lower() + + @property + def type(self) -> str: + """String discriminator used by launcher-compatible config payloads.""" + return str(self.get_type()) + + @property + def api_base_url(self) -> str: + """Accessor for the server API base URL.""" + return self.get_api_base_url() + + @property + def health_check_url(self) -> str: + """Accessor for the server health check URL.""" + return self.get_health_check_url() + + +def _validate_vajra_command_interpreter(command: List[str]) -> None: + """Reject Vajra module launches through the current Veeksha interpreter.""" + if len(command) < 3 or "-m" not in command: + return + module_index = command.index("-m") + 1 + if module_index >= len(command) or not command[module_index].startswith("vajra"): + return + + interpreter = Path(command[0]).expanduser() + current_interpreter = Path(sys.executable).expanduser() + try: + same_interpreter = interpreter.resolve() == current_interpreter.resolve() + except OSError: + same_interpreter = str(interpreter) == str(current_interpreter) + + ambiguous_interpreter = command[0] in {"python", "python3"} + if same_interpreter or ambiguous_interpreter: + raise ValueError( + "vajra server.command must use the Vajra Python environment, not the " + "current Veeksha interpreter or an ambiguous python executable. Use an " + "absolute path like /path/to/vajra/env/bin/python for server.command[0]." + ) + + +@frozen_dataclass +class VajraServerConfig(BaseServerConfig): + """Vajra subprocess server config.""" + + command: List[str] = field(default_factory=list) + env: Optional[Dict[str, str]] = field(None) + + @classmethod + def get_type(cls) -> ServerType: + return ServerType.VAJRA + + def __post_init__(self) -> None: + super().__post_init__() + if not self.command: + raise ValueError("vajra requires server.command") + if not self.setup_dir: + raise ValueError( + "vajra requires server.setup_dir (the Vajra source checkout, " + "used to record the engine git commit)" + ) + if self.model == _DEFAULT_MODEL: + raise ValueError("vajra requires server.model for the endpoint") + _validate_vajra_command_interpreter(self.command) + + def get_api_base_url(self) -> str: + if self.api_base: + return self.api_base.rstrip("/") + return f"http://{self.host}:{self.port}" @frozen_dataclass class VllmServerConfig(BaseServerConfig): + """vLLM Omni Docker server config.""" + + image: str = field(VLLM_OMNI_DEFAULT_IMAGE) + container_name: Optional[str] = None + container_port: Optional[int] = None + engine_args: List[str] = field(default_factory=list) + volumes: List[str] = field(default_factory=list) + docker_gpus: Optional[str] = None + docker_runtime: Optional[str] = "nvidia" + ipc_mode: Optional[str] = "host" + hf_model: str = "" + deploy_config: str = "" + container_deploy_config: Optional[str] = None + env: Optional[Dict[str, str]] = field(None) + pass_env: List[str] = field(default_factory=list) + bootstrap: Optional[str] = None + @classmethod def get_type(cls) -> ServerType: return ServerType.VLLM + def __post_init__(self) -> None: + if self.model == _DEFAULT_MODEL and self.hf_model: + self.model = self.hf_model + super().__post_init__() + _validate_docker_engine_config(self) + if not self.hf_model: + raise ValueError("vllm requires server.hf_model") + if not self.deploy_config: + raise ValueError("vllm requires server.deploy_config") + + def get_api_base_url(self) -> str: + if self.api_base: + return self.api_base.rstrip("/") + return f"http://{self.host}:{self.port}/v1" + + def get_health_check_url(self) -> str: + if self.health_url: + return self.health_url + return f"{self.get_api_base_url()}/models" + + @property + def container_name_prefix(self) -> str: + return self.container_name or "veeksha-vllm-omni" + + @property + def resolved_container_port(self) -> int: + return self.container_port or self.port + + @property + def resolved_container_deploy_config(self) -> str: + if self.container_deploy_config: + return self.container_deploy_config + return f"/etc/vllm-omni/{Path(self.deploy_config).name}" + + @property + def uses_bootstrap(self) -> bool: + """Whether a bootstrap snippet runs before ``vllm serve``.""" + return self.resolved_bootstrap != "" + + @property + def resolved_bootstrap(self) -> str: + return self.bootstrap or "" + @frozen_dataclass class SglangServerConfig(BaseServerConfig): + """SGLang Omni Docker server config.""" + + image: str = field(SGLANG_OMNI_DEFAULT_IMAGE) + container_name: Optional[str] = None + container_port: Optional[int] = None + engine_args: List[str] = field(default_factory=list) + volumes: List[str] = field(default_factory=list) + docker_gpus: Optional[str] = None + docker_runtime: Optional[str] = "nvidia" + ipc_mode: Optional[str] = "host" + shm_size: Optional[str] = "32g" + docker_run_args: List[str] = field(default_factory=list) + model_path: str = "" + model_name: Optional[str] = None + deploy_config: str = "" + container_deploy_config: Optional[str] = None + source_dir: Optional[str] = None + container_source_dir: str = "/sglang-omni" + venv_path: str = "/opt/sglomni/.venv" + bootstrap: Optional[str] = None + env: Optional[Dict[str, str]] = field(None) + pass_env: List[str] = field(default_factory=list) + @classmethod def get_type(cls) -> ServerType: return ServerType.SGLANG + + def __post_init__(self) -> None: + if self.model == _DEFAULT_MODEL and (self.model_name or self.model_path): + self.model = self.model_name or self.model_path + super().__post_init__() + _validate_docker_engine_config(self) + if not self.model_path: + raise ValueError("sglang requires server.model_path") + if not self.deploy_config: + raise ValueError("sglang requires server.deploy_config") + if self.uses_bootstrap and not self.source_dir: + raise ValueError( + "sglang requires server.source_dir when bootstrap is enabled " + "(the default bootstrap installs sglang-omni from the mounted " + "source checkout); set server.bootstrap to '' to disable" + ) + + def get_api_base_url(self) -> str: + if self.api_base: + return self.api_base.rstrip("/") + return f"http://{self.host}:{self.port}/v1" + + def get_health_check_url(self) -> str: + if self.health_url: + return self.health_url + return f"http://{self.host}:{self.port}/health" + + @property + def container_name_prefix(self) -> str: + return self.container_name or "veeksha-sglang-omni" + + @property + def resolved_container_port(self) -> int: + return self.container_port or self.port + + @property + def resolved_container_deploy_config(self) -> str: + if self.container_deploy_config: + return self.container_deploy_config + return f"/etc/sglang-omni/{Path(self.deploy_config).name}" + + @property + def uses_bootstrap(self) -> bool: + """Whether a bootstrap snippet runs before ``sgl-omni serve``.""" + return self.resolved_bootstrap != "" + + @property + def resolved_bootstrap(self) -> str: + if self.bootstrap is not None: + return self.bootstrap + return SGLANG_OMNI_DEFAULT_BOOTSTRAP.format( + src=self.container_source_dir, venv=self.venv_path + ) + + +def _validate_docker_engine_config( + config: VllmServerConfig | SglangServerConfig, +) -> None: + if config.container_port is not None and config.container_port <= 0: + raise ValueError("server.container_port must be positive") + if config.docker_gpus is not None and config.gpu_ids is not None: + raise ValueError("use either server.docker_gpus or server.gpu_ids, not both") + + +ManagedServerConfig: TypeAlias = ( + VajraServerConfig | VllmServerConfig | SglangServerConfig +) diff --git a/veeksha/config/slo.py b/veeksha/config/slo.py index ca155975..cc295b2c 100644 --- a/veeksha/config/slo.py +++ b/veeksha/config/slo.py @@ -4,7 +4,29 @@ from veeksha.types import SloType -SUPPORTED_SLO_METRICS = {"ttfc", "tbc", "tpot", "e2e"} +SUPPORTED_SLO_METRICS = { + "ttfc", + "tbc", + "tpot", + "e2e", + "generated_audio_duration", + "rtf", + "first_input_to_first_audio_ms", + "first_input_to_first_playable_audio_ms", + "trigger_to_first_playable_audio_ms", + "request_start_to_first_playable_audio_ms", + "audio_before_commit_ratio", + "duplex_overlap_observed", + "post_commit_audio_delivery_ms", + "required_startup_delay_ms", + "zero_delay_stall_count", + "zero_delay_total_stall_ms", + "zero_delay_longest_stall_ms", + "streaming_rtf", + "user_audio_fluidity_index", + "tts_service_fluidity_index", + "tts_service_fluidity_eligible", +} @frozen_dataclass diff --git a/veeksha/config/utils.py b/veeksha/config/utils.py index 3e8f211f..55577af1 100644 --- a/veeksha/config/utils.py +++ b/veeksha/config/utils.py @@ -4,11 +4,14 @@ import logging import os import time +from enum import Enum from importlib.resources.abc import Traversable from typing import Any, Dict import yaml +from veeksha.types import ChannelModality + logger = logging.getLogger(__name__) @@ -67,6 +70,48 @@ def scrub(obj): return hashlib.blake2s(stable_json.encode()).hexdigest()[:8] +def serialize_config_value(value: Any) -> Any: + """Return a YAML/JSON-safe config value with enums rendered by schema names.""" + if isinstance(value, Enum): + if isinstance(value.value, str): + return value.value + return str(value) + if isinstance(value, dict): + converted: Dict[Any, Any] = {} + for key, item in value.items(): + converted_key = serialize_config_value(key) + if converted_key == "target_channels" and isinstance(item, (list, tuple)): + converted[converted_key] = [ + ( + str(ChannelModality(channel)) + if isinstance(channel, int) and not isinstance(channel, bool) + else serialize_config_value(channel) + ) + for channel in item + ] + else: + converted[converted_key] = serialize_config_value(item) + return converted + if isinstance(value, list): + return [serialize_config_value(item) for item in value] + if isinstance(value, tuple): + return [serialize_config_value(item) for item in value] + return value + + +def to_serializable_config_dict(config: Any) -> Dict[str, Any]: + """Convert a vidhi/dataclass config object into schema-friendly plain data.""" + from vidhi import dataclass_to_dict + + config_as_dict = dataclass_to_dict(config) + assert isinstance( + config_as_dict, dict + ), f"Expected dict, got {type(config_as_dict)}" + serialized = serialize_config_value(config_as_dict) + assert isinstance(serialized, dict), f"Expected dict, got {type(serialized)}" + return serialized + + def _build_unique_output_dir(root: str, model_name: str, config_hash: str) -> str: """Return a unique timestamped output directory path. @@ -85,8 +130,6 @@ def prepare_benchmark_output_dir(benchmark_config) -> None: named with model and config-hash plus a high-entropy timestamp. - Save both `config.json` and `config.yml` in the final output directory. """ - from vidhi import dataclass_to_dict - current_output_dir = benchmark_config.metrics_config.output_dir existing_config_path = os.path.join(current_output_dir, "config.json") if os.path.isfile(existing_config_path): @@ -99,10 +142,7 @@ def prepare_benchmark_output_dir(benchmark_config) -> None: base_output_dir = benchmark_config.metrics_config.output_dir model_name = benchmark_config.client_config.model.split("/")[-1] - config_as_dict = dataclass_to_dict(benchmark_config) - assert isinstance( - config_as_dict, dict - ), f"Expected dict, got {type(config_as_dict)}" + config_as_dict = to_serializable_config_dict(benchmark_config) cfg_hash = get_config_hash(config_as_dict) unique_dir = _build_unique_output_dir(base_output_dir, model_name, cfg_hash) object.__setattr__(benchmark_config.metrics_config, "output_dir", unique_dir) diff --git a/veeksha/config/verification.py b/veeksha/config/verification.py new file mode 100644 index 00000000..3b3855a5 --- /dev/null +++ b/veeksha/config/verification.py @@ -0,0 +1,134 @@ +from vidhi import BasePolyConfig, field, frozen_dataclass + +from veeksha.types import VerificationType + + +@frozen_dataclass +class BaseVerificationConfig(BasePolyConfig): + """Base class for post-run verification configuration.""" + + fail_on_threshold: bool = field( + False, + help="If True, fail the run when verification thresholds fail.", + ) + + def is_enabled(self) -> bool: + return False + + +@frozen_dataclass +class WhisperTranscriptionConfig: + """Whisper transcription configuration used by WER verification.""" + + model: str = field( + "large-v3", + help="Whisper model identifier passed to faster-whisper.", + ) + device: str = field( + "cuda", + help="Device passed to faster-whisper, e.g. 'cuda' or 'cpu'.", + ) + compute_type: str = field( + "float16", + help="Compute type passed to faster-whisper, e.g. 'float16' or 'int8'.", + ) + language: str = field( + "en", + help=( + "Language code passed to faster-whisper. Pin this for comparable WER " + "runs instead of relying on automatic language detection." + ), + ) + beam_size: int = field( + 5, + help="Beam size passed to faster-whisper for deterministic WER scoring.", + ) + + def __post_init__(self): + if not self.model: + raise ValueError("WhisperTranscriptionConfig.model is required") + if not self.device: + raise ValueError("WhisperTranscriptionConfig.device is required") + if not self.compute_type: + raise ValueError("WhisperTranscriptionConfig.compute_type is required") + if not self.language: + raise ValueError("WhisperTranscriptionConfig.language is required") + if self.beam_size < 1: + raise ValueError("WhisperTranscriptionConfig.beam_size must be >= 1") + + +@frozen_dataclass +class WERVerifierConfig: + """WER verifier configuration for generated speech.""" + + enabled: bool = field( + False, + help="Enable WER verification using inline Whisper transcription.", + ) + threshold: float = field( + 0.05, + help="Per-request WER threshold for pass/fail classification.", + ) + whisper: WhisperTranscriptionConfig = field( + default_factory=WhisperTranscriptionConfig, + help="Inline Whisper transcription configuration.", + ) + + def __post_init__(self): + if self.threshold < 0: + raise ValueError("WERVerifierConfig.threshold must be >= 0") + + +@frozen_dataclass +class UTMOSVerifierConfig: + """UTMOS predicted MOS verifier configuration.""" + + enabled: bool = field( + False, + help="Enable UTMOS predicted MOS scoring for generated audio.", + ) + hf_repo: str = field( + "balacoon/utmos", + help="Hugging Face model repo containing the UTMOS JIT file.", + ) + jit_file: str = field( + "utmos.jit", + help="TorchScript filename to load from the UTMOS HF repo.", + ) + device: str = field( + "cuda:0", + help="Device for UTMOS TorchScript inference, e.g. 'cuda:0' or 'cpu'.", + ) + + def __post_init__(self): + if not self.hf_repo: + raise ValueError("UTMOSVerifierConfig.hf_repo is required") + if not self.jit_file: + raise ValueError("UTMOSVerifierConfig.jit_file is required") + if not self.device: + raise ValueError("UTMOSVerifierConfig.device is required") + + +@frozen_dataclass +class AudioVerificationConfig(BaseVerificationConfig): + """Post-run verification for generated audio artifacts.""" + + max_requests: int = field( + 2000, + help="Maximum number of request rows to verify. Use 0 or less to verify all rows.", + ) + wer: WERVerifierConfig = field( + default_factory=WERVerifierConfig, + help="WER verifier configuration.", + ) + utmos: UTMOSVerifierConfig = field( + default_factory=UTMOSVerifierConfig, + help="UTMOS verifier configuration.", + ) + + @classmethod + def get_type(cls) -> VerificationType: + return VerificationType.AUDIO + + def is_enabled(self) -> bool: + return bool(self.wer.enabled or self.utmos.enabled) diff --git a/veeksha/core/audio_contract.py b/veeksha/core/audio_contract.py new file mode 100644 index 00000000..3f185207 --- /dev/null +++ b/veeksha/core/audio_contract.py @@ -0,0 +1,92 @@ +"""Audio measurement contract: PCM constants and the metric-key vocabulary. + +Pure measurement facts shared by the audio (TTS / realtime) clients and the +audio evaluators. Transport-specific protocol handling lives in the clients. +""" + +from __future__ import annotations + +from enum import StrEnum + +# 16-bit mono PCM at 24 kHz is the shared baseline across TTS dialects. +DEFAULT_AUDIO_SAMPLE_RATE = 24000 +BYTES_PER_SAMPLE = 2 +WAV_HEADER_BYTES = 44 + + +def pcm_bytes_to_duration_ms( + n_bytes: float, sample_rate: int = DEFAULT_AUDIO_SAMPLE_RATE +) -> float: + """Duration in ms of ``n_bytes`` of 16-bit mono PCM at ``sample_rate``.""" + return n_bytes / (sample_rate * BYTES_PER_SAMPLE) * 1000.0 + + +class AudioMetricKey(StrEnum): + TTFC = "ttfc" + END_TO_END_LATENCY = "end_to_end_latency" + GENERATED_AUDIO_DURATION = "generated_audio_duration" + RTF = "rtf" + CHUNK_COUNT = "chunk_count" + RAW_PCM = "raw_pcm" + SAMPLE_RATE = "sample_rate" + PCM_BYTE_COUNT = "pcm_byte_count" + INPUT_CHARS = "input_chars" + INPUT_TOKENS = "input_tokens" + INPUT_TEXT = "input_text" + PROVIDER = "provider" + PROVIDER_MODEL = "provider_model" + PROVIDER_PROTOCOL = "provider_protocol" + SESSION_SIZE = "session_size" + SESSION_DURATION = "session_duration" + + # ----- Realtime input-streaming interactivity keys ----- + # Time convention for all realtime event-time values below: every *_offset_ms + # / timestamp value is a float millisecond offset relative to request start + # (the client's WS-connect initiation), measured with time.monotonic(). + # + # Raw-contract keys are emitted by the websocket client: + TEXT_DELTA_TIMESTAMPS = "text_delta_timestamps" # list[[offset_ms, n_chars]] + AUDIO_CHUNK_TIMESTAMPS = ( + "audio_chunk_timestamps" # list[[offset_ms, n_bytes_decoded_pcm]] + ) + WS_CONNECT_LATENCY_MS = "ws_connect_latency_ms" + SESSION_READY_OFFSET_MS = "session_ready_offset_ms" # nullable + RESPONSE_TRIGGER_OFFSET_MS = "response_trigger_offset_ms" + RESPONSE_CREATED_OFFSET_MS = "response_created_offset_ms" # nullable + INPUT_COMMIT_OFFSET_MS = "input_commit_offset_ms" + AUDIO_DONE_OFFSET_MS = "audio_done_offset_ms" # nullable + RESPONSE_DONE_OFFSET_MS = "response_done_offset_ms" # nullable + + # Stable request-level interactivity keys emitted by the evaluator. + FIRST_INPUT_TO_FIRST_AUDIO_MS = "first_input_to_first_audio_ms" + FIRST_INPUT_TO_FIRST_PLAYABLE_AUDIO_MS = "first_input_to_first_playable_audio_ms" + TRIGGER_TO_FIRST_PLAYABLE_AUDIO_MS = "trigger_to_first_playable_audio_ms" + REQUEST_START_TO_FIRST_AUDIO_MS = "request_start_to_first_audio_ms" + REQUEST_START_TO_FIRST_PLAYABLE_AUDIO_MS = ( + "request_start_to_first_playable_audio_ms" + ) + AUDIO_BEFORE_COMMIT_RATIO = "audio_before_commit_ratio" + DUPLEX_OVERLAP_OBSERVED = "duplex_overlap_observed" + DUPLEX_OVERLAP_MS = "duplex_overlap_ms" + POST_COMMIT_AUDIO_DELIVERY_MS = "post_commit_audio_delivery_ms" + REQUIRED_STARTUP_DELAY_MS = "required_startup_delay_ms" + ZERO_DELAY_STALL_COUNT = "zero_delay_stall_count" + ZERO_DELAY_TOTAL_STALL_MS = "zero_delay_total_stall_ms" + ZERO_DELAY_LONGEST_STALL_MS = "zero_delay_longest_stall_ms" + ZERO_DELAY_STALL_FREE = "zero_delay_stall_free" + + # Diagnostic delivery/finalization metrics. + STREAMING_RTF = "streaming_rtf" + DONE_AFTER_LAST_AUDIO_MS = "done_after_last_audio_ms" + + # Etalon-inspired playable-frame deadline metrics. The untagged user score + # uses AudioChannelPerformanceConfig.fluidity_startup_delay_ms. + USER_AUDIO_FLUIDITY_INDEX = "user_audio_fluidity_index" + TTS_SERVICE_FLUIDITY_INDEX = "tts_service_fluidity_index" + TTS_SERVICE_FLUIDITY_ELIGIBLE = "tts_service_fluidity_eligible" + UNATTRIBUTED_MISSED_DEADLINES = "unattributed_missed_deadlines" + AUDIO_FLUIDITY_TOTAL_DEADLINES = "audio_fluidity_total_deadlines" + AUDIO_FLUIDITY_MISSED_DEADLINES = "audio_fluidity_missed_deadlines" + AUDIO_PLAYABLE_FRAME_COUNT = "audio_playable_frame_count" + AUDIO_FLUIDITY_FRAME_MS = "audio_fluidity_frame_ms" + AUDIO_FLUIDITY_STARTUP_DELAY_MS = "audio_fluidity_startup_delay_ms" diff --git a/veeksha/core/tokenizer.py b/veeksha/core/tokenizer.py index 04a61277..a74925bf 100644 --- a/veeksha/core/tokenizer.py +++ b/veeksha/core/tokenizer.py @@ -11,8 +11,6 @@ TypeVar, ) -from transformers import AutoTokenizer - from veeksha.types import ChannelModality RawContent = TypeVar("RawContent") @@ -70,10 +68,32 @@ def build_hf_tokenizer_handle(tokenizer) -> TokenizerHandle[str]: def build_hf_tokenizer_handle_from_model(model: str) -> TokenizerHandle[str]: """Instantiate a Hugging Face tokenizer from a model name and wrap it.""" + # Import lazily so clients with provider-native tokenization (for example, + # streaming TTS word pacing) do not load the Hugging Face/Rust tokenizer + # stack merely by importing Veeksha's client package. + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) return build_hf_tokenizer_handle(tokenizer) +def build_word_split_tokenizer_provider(model_name: str) -> "TokenizerProvider": + """Build a TokenizerProvider backed by a simple whitespace word-split tokenizer. + + Used by TTS-style clients whose models do not ship a HuggingFace tokenizer; + tokens are approximated as whitespace-delimited words. + """ + handle = TokenizerHandle( + count_tokens=lambda text: len(text.split()), + decode=lambda ids: " ".join(str(i) for i in ids), + encode=lambda text: list(range(len(text.split()))), + ) + return TokenizerProvider( + {ChannelModality.TEXT: handle}, + model_name=model_name, + ) + + def gen_prompt_from_corpus( num_tokens: int, pretokenized_lines: Iterable[Sequence[int]], diff --git a/veeksha/evaluator/accuracy/audio.py b/veeksha/evaluator/accuracy/audio.py new file mode 100644 index 00000000..a96d9208 --- /dev/null +++ b/veeksha/evaluator/accuracy/audio.py @@ -0,0 +1,306 @@ +"""Generated audio quality evaluator.""" + +from __future__ import annotations + +import json +import struct +import threading +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Dict, Optional, Set + +from veeksha.config.evaluator import AudioQualityEvaluatorConfig +from veeksha.core.audio_contract import DEFAULT_AUDIO_SAMPLE_RATE, AudioMetricKey +from veeksha.core.request_content import TextChannelRequestContent +from veeksha.core.seeding import SeedManager +from veeksha.evaluator.accuracy.base import BaseAccuracyEvaluator +from veeksha.evaluator.base import EvaluationResult +from veeksha.logger import init_logger +from veeksha.types import ChannelModality + +logger = init_logger(__name__) + + +@dataclass +class AudioQualityRequestRow: + request_id: int + session_id: int + input_text: str + audio_present: bool + error: Optional[str] = None + + +def _make_wav_header( + data_size: int, sample_rate: int = DEFAULT_AUDIO_SAMPLE_RATE +) -> bytes: + """Build a 44-byte mono 16-bit PCM WAV header.""" + byte_rate = sample_rate * 1 * 16 // 8 + block_align = 1 * 16 // 8 + return struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", + 36 + data_size, + b"WAVE", + b"fmt ", + 16, + 1, + 1, + sample_rate, + byte_rate, + block_align, + 16, + b"data", + data_size, + ) + + +class AudioQualityEvaluator(BaseAccuracyEvaluator): + """Accuracy-family evaluator for generated audio quality. + + It persists an audio-quality-scoped copy of generated audio and reference + text, then runs configured audio verifiers such as WER or UTMOS. + """ + + def __init__( + self, + config: AudioQualityEvaluatorConfig, + seed_manager: Optional[SeedManager] = None, + output_dir: Optional[str] = None, + benchmark_start_time: float = 0.0, + ): + super().__init__( + config=config, + seed_manager=seed_manager, + output_dir=output_dir, + benchmark_start_time=benchmark_start_time, + ) + self.config: AudioQualityEvaluatorConfig = config + self._lock = threading.Lock() + self._registered_request_ids: Set[int] = set() + self._included_request_ids: Optional[Set[int]] = None + self._reference_text_by_request_id: Dict[int, str] = {} + self._rows: list[AudioQualityRequestRow] = [] + self._audio_buffers: Dict[int, bytes] = {} + self._audio_metadata: Dict[int, Dict[str, Any]] = {} + self.num_requests = 0 + self.num_completed_requests = 0 + self.num_errored_requests = 0 + self._verification_summary: Optional[dict[str, Any]] = None + + def register_request( + self, + request_id: int, + session_id: int, + dispatched_at: float, + channels: Dict[ChannelModality, Any], + requested_output: Any = None, + ) -> None: + if not self.should_evaluate_channel(ChannelModality.AUDIO): + return + with self._lock: + self.num_requests += 1 + self._registered_request_ids.add(request_id) + text_content = channels.get(ChannelModality.TEXT) + if isinstance(text_content, TextChannelRequestContent): + self._reference_text_by_request_id[request_id] = text_content.input_text + + def record_request_completed( + self, + request_id: int, + session_id: int, + completed_at: float, + response: Any, + error: Optional[Exception] = None, + ) -> None: + if not self.should_evaluate_channel(ChannelModality.AUDIO): + return + with self._lock: + if ( + self._included_request_ids is not None + and request_id not in self._included_request_ids + ): + return + + input_text = self._reference_text_by_request_id.get(request_id, "") + if error is not None or not getattr(response, "success", True): + self.num_errored_requests += 1 + self._rows.append( + AudioQualityRequestRow( + request_id=request_id, + session_id=session_id, + input_text=input_text, + audio_present=False, + error=str( + error or getattr(response, "error_msg", "request failed") + ), + ) + ) + return + + audio_channel = getattr(response, "channels", {}).get(ChannelModality.AUDIO) + audio_content = ( + getattr(audio_channel, "content", None) if audio_channel else None + ) + if not isinstance(audio_content, bytes) or not audio_content: + self.num_errored_requests += 1 + self._rows.append( + AudioQualityRequestRow( + request_id=request_id, + session_id=session_id, + input_text=input_text, + audio_present=False, + error="missing audio channel content", + ) + ) + return + + metrics = getattr(audio_channel, "metrics", {}) or {} + self.num_completed_requests += 1 + self._rows.append( + AudioQualityRequestRow( + request_id=request_id, + session_id=session_id, + input_text=input_text, + audio_present=True, + ) + ) + if self._should_save_audio_files(): + self._audio_buffers[request_id] = audio_content + self._audio_metadata[request_id] = { + AudioMetricKey.RAW_PCM.value: bool( + metrics.get(AudioMetricKey.RAW_PCM.value, False) + ), + AudioMetricKey.SAMPLE_RATE.value: int( + metrics.get( + AudioMetricKey.SAMPLE_RATE.value, DEFAULT_AUDIO_SAMPLE_RATE + ) + ), + } + + def record_session_completed( + self, + session_id: int, + completed_at: float, + success: bool, + ) -> None: + return + + def finalize(self) -> EvaluationResult: + with self._lock: + metrics = { + "num_requests": self.num_requests, + "num_completed_requests": self.num_completed_requests, + "num_errored_requests": self.num_errored_requests, + "verification_enabled": self.config.verification.is_enabled(), + } + if self._verification_summary is not None: + metrics["verification"] = self._verification_summary + return EvaluationResult( + evaluator_type="audio_quality", + channel=ChannelModality.AUDIO, + metrics=metrics, + ) + + def save(self, output_dir: str) -> None: + audio_quality_dir = Path(output_dir).parent / "audio_quality" + metrics_dir = audio_quality_dir / "metrics" + metrics_dir.mkdir(parents=True, exist_ok=True) + self._save_request_rows(metrics_dir / "request_level_metrics.jsonl") + if self._should_save_audio_files(): + self._save_audio_files(audio_quality_dir / "audio_files") + self._maybe_run_verification(audio_quality_dir) + + def get_completed_request_count(self) -> int: + with self._lock: + return self.num_completed_requests + + def get_session_counts(self) -> tuple[int, int, int]: + with self._lock: + in_progress = max( + 0, + len(self._registered_request_ids) + - self.num_completed_requests + - self.num_errored_requests, + ) + return self.num_completed_requests, self.num_errored_requests, in_progress + + def set_included_requests(self, request_ids: Set[int]) -> None: + with self._lock: + self._included_request_ids = set(request_ids) + + def get_registered_request_ids(self) -> Set[int]: + with self._lock: + return set(self._registered_request_ids) + + def _should_save_audio_files(self) -> bool: + return bool(self.config.save_audio_files) + + def _save_request_rows(self, path: Path) -> None: + with self._lock: + rows = list(self._rows) + with path.open("w", encoding="utf-8") as f: + for row in rows: + f.write(json.dumps(asdict(row)) + "\n") + + def _save_audio_files(self, audio_dir: Path) -> None: + with self._lock: + buffers = dict(self._audio_buffers) + metadata = dict(self._audio_metadata) + if not buffers: + return + audio_dir.mkdir(parents=True, exist_ok=True) + for request_id, audio_data in buffers.items(): + meta = metadata.get(request_id, {}) + raw_pcm = bool(meta.get(AudioMetricKey.RAW_PCM.value, False)) + sample_rate = int( + meta.get(AudioMetricKey.SAMPLE_RATE.value, DEFAULT_AUDIO_SAMPLE_RATE) + ) + wav_path = audio_dir / f"request_{request_id}.wav" + with wav_path.open("wb") as f: + if raw_pcm: + f.write(_make_wav_header(len(audio_data), sample_rate)) + f.write(audio_data) + logger.info("Saved %d audio quality files to %s", len(buffers), audio_dir) + + def _maybe_run_verification(self, audio_quality_dir: Path) -> None: + verification_config = self.config.verification + if not verification_config.is_enabled(): + return + if not self._should_save_audio_files(): + logger.warning( + "Audio quality verification is enabled but audio saving is disabled; " + "skipping because no audio artifacts were saved." + ) + return + + from veeksha.verification.audio import ( + AudioVerificationError, + run_audio_verification, + ) + + try: + summary = run_audio_verification( + config=verification_config, + output_dir=audio_quality_dir, + ) + summary_dict = summary.to_dict() + with self._lock: + self._verification_summary = summary_dict + logger.info( + "Audio quality verification complete: %d transcribed, %d above threshold, " + "%d UTMOS scored, %d errors", + summary.transcribed_requests, + summary.failed_requests, + summary.utmos_evaluated, + summary.error_requests + len(summary.errors), + ) + except AudioVerificationError: + if verification_config.fail_on_threshold: + raise + logger.exception("Audio quality verification failed; continuing") + except Exception: + if verification_config.fail_on_threshold: + raise + logger.exception( + "Audio quality verification failed; continuing because fail_on_threshold=False" + ) diff --git a/veeksha/evaluator/accuracy/base.py b/veeksha/evaluator/accuracy/base.py index 18dcb3d2..c11f4837 100644 --- a/veeksha/evaluator/accuracy/base.py +++ b/veeksha/evaluator/accuracy/base.py @@ -15,19 +15,9 @@ from collections import defaultdict from typing import Any, Dict, Optional, Set, Tuple, cast -# NOTE: `lm_eval` is an external dependency; type checkers in some environments -# may not have it available, so we silence missing-import diagnostics here. -from lm_eval.evaluator_utils import ( # pyright: ignore[reportMissingImports] # type: ignore[import-not-found] - consolidate_group_results, - consolidate_results, - get_subtask_list, - prepare_print_tasks, -) - -from veeksha.config.evaluator import LMEvalAccuracyEvaluatorConfig +from veeksha.config.evaluator import BaseEvaluatorConfig, LMEvalAccuracyEvaluatorConfig from veeksha.core.seeding import SeedManager from veeksha.evaluator.base import BaseEvaluator, EvaluationResult -from veeksha.generator.session.lmeval import LMEvalSessionGenerator from veeksha.logger import init_logger from veeksha.types import ChannelModality, LMEvalOutputType @@ -44,7 +34,7 @@ class BaseAccuracyEvaluator(BaseEvaluator, ABC): def __init__( self, - config: LMEvalAccuracyEvaluatorConfig, + config: BaseEvaluatorConfig, seed_manager: Optional[SeedManager] = None, output_dir: Optional[str] = None, benchmark_start_time: float = 0.0, @@ -73,6 +63,8 @@ def __init__( benchmark_start_time=benchmark_start_time, ) + from veeksha.generator.session.lmeval import LMEvalSessionGenerator + if not isinstance(session_generator, LMEvalSessionGenerator): raise ValueError( "LMEvalAccuracyEvaluator requires LMEvalSessionGenerator via session_generator=..." @@ -277,6 +269,13 @@ def _parse_logprobs(self, ctxlen: int, lp: Any) -> Tuple[float, bool]: raise KeyError("Unsupported logprobs structure for completions response.") def _evaluate_lmeval(self) -> Dict[str, Any]: + from lm_eval.evaluator_utils import ( # pyright: ignore[reportMissingImports] # type: ignore[import-not-found] + consolidate_group_results, + consolidate_results, + get_subtask_list, + prepare_print_tasks, + ) + # Bind generator-owned tasks and compute lm-eval metrics eval_tasks = self.session_generator.eval_tasks limits = self.session_generator.limits diff --git a/veeksha/evaluator/cdf_sketch.py b/veeksha/evaluator/cdf_sketch.py index a6f26ca2..cb9a009d 100644 --- a/veeksha/evaluator/cdf_sketch.py +++ b/veeksha/evaluator/cdf_sketch.py @@ -2,10 +2,10 @@ import numpy as np import pandas as pd -import rekha as rk import wandb from ddsketch import DDSketch +import rekha as rk from veeksha.evaluator.plot_utils import ( apply_axis_scale, format_axis_label, diff --git a/veeksha/evaluator/performance/asr.py b/veeksha/evaluator/performance/asr.py new file mode 100644 index 00000000..047d6b4c --- /dev/null +++ b/veeksha/evaluator/performance/asr.py @@ -0,0 +1,323 @@ +"""ASR-specific scoring helpers for realtime speech-to-text benchmarks.""" + +from __future__ import annotations + +import threading +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, DefaultDict, Dict, List, Optional + +import jiwer + +from veeksha.evaluator.performance.asr_interactivity import compute_interactivity_stats +from veeksha.evaluator.performance.asr_normalizer import EnglishTextNormalizer + +_normalizer = EnglishTextNormalizer() + + +@dataclass(frozen=True) +class WERStats: + """Edit counts and WER percentage for one normalized transcript comparison.""" + + errors: int + reference_words: int + wer: float + + +def compute_wer_stats(reference: str, hypothesis: str) -> WERStats: + """WER counts and percentage using the leaderboard normalizer + jiwer.""" + ref = _normalizer(reference) + hyp = _normalizer(hypothesis) + output = jiwer.process_words(ref, hyp) + errors = output.substitutions + output.deletions + output.insertions + reference_words = output.hits + output.substitutions + output.deletions + if reference_words == 0: + wer = 0.0 if errors == 0 else 100.0 + else: + wer = (errors / reference_words) * 100 + return WERStats(errors=errors, reference_words=reference_words, wer=wer) + + +@dataclass +class WERAggregate: + """Accumulates multiple WER aggregation modes for comparable ASR samples.""" + + sample_count: int = 0 + wer_sum: float = 0.0 + duration_weighted_wer_sum: float = 0.0 + duration_s_sum: float = 0.0 + errors: int = 0 + reference_words: int = 0 + + def add(self, stats: WERStats, duration_s: float) -> None: + self.sample_count += 1 + self.wer_sum += stats.wer + self.errors += stats.errors + self.reference_words += stats.reference_words + if duration_s > 0: + self.duration_weighted_wer_sum += stats.wer * duration_s + self.duration_s_sum += duration_s + + def summary(self, prefix: str) -> Dict[str, Optional[float]]: + sample_mean = ( + self.wer_sum / self.sample_count if self.sample_count > 0 else None + ) + corpus = ( + (self.errors / self.reference_words) * 100 + if self.reference_words > 0 + else ( + 100.0 if self.errors > 0 else (0.0 if self.sample_count > 0 else None) + ) + ) + duration_weighted = ( + self.duration_weighted_wer_sum / self.duration_s_sum + if self.duration_s_sum > 0 + else None + ) + return { + f"{prefix}_sample_count": float(self.sample_count), + f"{prefix}_sample_mean_wer": sample_mean, + f"{prefix}_corpus_wer": corpus, + f"{prefix}_duration_weighted_wer": duration_weighted, + } + + +@dataclass +class ASRScoredSample: + dataset: str + duration_s: float + final_stats: WERStats + partial_stats: Optional[WERStats] + + +@dataclass +class ASRDatasetAggregates: + final: WERAggregate = field(default_factory=WERAggregate) + partial: WERAggregate = field(default_factory=WERAggregate) + + +class ASRMetricAccumulator: + """Builds ASR WER summaries across request-scoped ASR samples. + + Thread-safe: scoring runs concurrently on completion worker threads + outside the evaluator locks, so sample collection guards its own state. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._samples: List[ASRScoredSample] = [] + + def add_clip_sample( + self, + *, + dataset: str, + duration_s: float, + final_stats: WERStats, + partial_stats: Optional[WERStats], + ) -> None: + sample = ASRScoredSample( + dataset=dataset, + duration_s=duration_s, + final_stats=final_stats, + partial_stats=partial_stats, + ) + with self._lock: + self._samples.append(sample) + + def get_summary(self) -> Dict[str, Optional[float]]: + with self._lock: + samples = list(self._samples) + if not samples: + return {} + + overall = ASRDatasetAggregates() + by_dataset: DefaultDict[str, ASRDatasetAggregates] = defaultdict( + ASRDatasetAggregates + ) + + for sample in samples: + overall.final.add(sample.final_stats, sample.duration_s) + by_dataset[sample.dataset].final.add(sample.final_stats, sample.duration_s) + if sample.partial_stats is not None: + overall.partial.add(sample.partial_stats, sample.duration_s) + by_dataset[sample.dataset].partial.add( + sample.partial_stats, sample.duration_s + ) + + summary = {} + summary.update(overall.final.summary("asr_final")) + summary.update(overall.partial.summary("asr_partial")) + for dataset, aggregates in sorted(by_dataset.items()): + slug = _metric_slug(dataset) + summary.update(aggregates.final.summary(f"asr_dataset_{slug}_final")) + summary.update(aggregates.partial.summary(f"asr_dataset_{slug}_partial")) + return summary + + +def _metric_slug(value: str) -> str: + return "".join(ch if ch.isalnum() else "_" for ch in value.lower()).strip("_") + + +@dataclass +class ASRRequestMetrics: + """Request-level ASR fields persisted into request_level_metrics.jsonl.""" + + audio_file: Optional[str] + final_transcript: str + expected_transcript: str + partial_transcript: Optional[str] + reference_word_timestamps: Optional[List[Dict[str, Any]]] + transcript_snapshots: List[Dict[str, Any]] + final_wer: Optional[float] + partial_wer: Optional[float] + dataset: str + source_id: Optional[str] + sample_id: Optional[str] + time_to_first_visible_text: Optional[float] + time_to_first_partial: Optional[float] + time_to_final_transcript: Optional[float] + interactivity: Optional[float] + interactivity_word_count: int + + def to_request_row(self) -> Dict[str, Any]: + return { + "dataset": self.dataset, + "source_id": self.source_id, + "sample_id": self.sample_id, + "audio_file": self.audio_file, + "time_to_first_visible_text": ( + round(self.time_to_first_visible_text, 3) + if self.time_to_first_visible_text is not None + else None + ), + "time_to_first_partial": ( + round(self.time_to_first_partial, 3) + if self.time_to_first_partial is not None + else None + ), + "time_to_final_transcript": ( + round(self.time_to_final_transcript, 3) + if self.time_to_final_transcript is not None + else None + ), + "interactivity": ( + round(self.interactivity, 3) if self.interactivity is not None else None + ), + "interactivity_word_count": self.interactivity_word_count, + "partial_transcript": self.partial_transcript, + "final_transcript": self.final_transcript, + "expected_transcript": self.expected_transcript, + "reference_word_timestamps": self.reference_word_timestamps, + "transcript_snapshots": self.transcript_snapshots, + "partial_wer": ( + round(self.partial_wer, 3) if self.partial_wer is not None else None + ), + "final_wer": ( + round(self.final_wer, 3) if self.final_wer is not None else None + ), + } + + +def score_asr_request( + *, + request_id: int, + channel_metrics: Dict[str, Any], + duration_s: float, + accumulator: ASRMetricAccumulator, +) -> ASRRequestMetrics: + """Validate one STT response, score it, and add it to ASR aggregates.""" + final_transcript = channel_metrics.get("final_transcript") + expected_transcript = channel_metrics.get("expected_transcript") + if final_transcript is None or expected_transcript is None: + raise ValueError( + f"STT response for request {request_id} missing " + f"final_transcript={final_transcript!r} / " + f"expected_transcript={expected_transcript!r}." + ) + + partial_transcript = channel_metrics.get("partial_transcript") + dataset = str(channel_metrics.get("dataset") or "unknown") + audio_file = _optional_str(channel_metrics.get("audio_file")) + source_id = _optional_str(channel_metrics.get("source_id")) + sample_id = _optional_str(channel_metrics.get("sample_id")) + reference_word_timestamps = _optional_list_of_dicts( + channel_metrics.get("reference_word_timestamps") + ) + transcript_snapshots = _list_of_dicts_or_empty( + channel_metrics.get("transcript_snapshots") + ) + time_to_first_visible_text = _optional_float( + channel_metrics.get("time_to_first_visible_text") + ) + time_to_first_partial = _optional_float( + channel_metrics.get("time_to_first_partial") + ) + time_to_final_transcript = _optional_float( + channel_metrics.get("time_to_final_transcript") + ) + interactivity_stats = compute_interactivity_stats(channel_metrics) + interactivity = ( + interactivity_stats.mean_latency_ms if interactivity_stats is not None else None + ) + interactivity_word_count = ( + interactivity_stats.word_count if interactivity_stats is not None else 0 + ) + + final_stats = compute_wer_stats(str(expected_transcript), str(final_transcript)) + final_wer = final_stats.wer + partial_stats = None + partial_wer = None + if partial_transcript: + partial_stats = compute_wer_stats( + str(expected_transcript), str(partial_transcript) + ) + partial_wer = partial_stats.wer + accumulator.add_clip_sample( + dataset=dataset, + duration_s=duration_s, + final_stats=final_stats, + partial_stats=partial_stats, + ) + + return ASRRequestMetrics( + audio_file=audio_file, + final_transcript=str(final_transcript), + expected_transcript=str(expected_transcript), + partial_transcript=partial_transcript, + reference_word_timestamps=reference_word_timestamps, + transcript_snapshots=transcript_snapshots, + final_wer=final_wer, + partial_wer=partial_wer, + dataset=dataset, + source_id=source_id, + sample_id=sample_id, + time_to_first_visible_text=time_to_first_visible_text, + time_to_first_partial=time_to_first_partial, + time_to_final_transcript=time_to_final_transcript, + interactivity=interactivity, + interactivity_word_count=interactivity_word_count, + ) + + +def _optional_float(value: Any) -> Optional[float]: + if value is None: + return None + return float(value) + + +def _optional_list_of_dicts(value: Any) -> Optional[List[Dict[str, Any]]]: + if value is None: + return None + return _list_of_dicts_or_empty(value) + + +def _list_of_dicts_or_empty(value: Any) -> List[Dict[str, Any]]: + if not isinstance(value, list): + return [] + return [dict(item) for item in value if isinstance(item, dict)] + + +def _optional_str(value: Any) -> Optional[str]: + if value is None: + return None + return str(value) diff --git a/veeksha/evaluator/performance/asr_interactivity.py b/veeksha/evaluator/performance/asr_interactivity.py new file mode 100644 index 00000000..ab103de2 --- /dev/null +++ b/veeksha/evaluator/performance/asr_interactivity.py @@ -0,0 +1,203 @@ +"""Interactivity scoring for realtime speech-to-text transcripts.""" + +from __future__ import annotations + +import difflib +import re +from dataclasses import dataclass +from typing import Any, Optional + +from veeksha.evaluator.performance.asr_normalizer import EnglishTextNormalizer + +_normalizer = EnglishTextNormalizer() +_WORD_RE = re.compile(r"[a-z0-9]+(?:'[a-z0-9]+)?") + + +@dataclass(frozen=True) +class ReferenceWord: + word: str + start_ms: float + end_ms: float + + +@dataclass(frozen=True) +class TranscriptSnapshot: + elapsed_ms: float + transcript: str + + +@dataclass(frozen=True) +class InteractivityStats: + mean_latency_ms: float + word_count: int + latencies_ms: list[float] + + +def compute_interactivity_stats( + channel_metrics: dict[str, Any], +) -> Optional[InteractivityStats]: + """Compute word visibility latency for one STT request. + + Missing ``reference_word_timestamps`` means the trace does not support this + metric, so this returns ``None``. Missing or empty snapshots mean the metric + is unavailable for this request. + """ + if channel_metrics.get("reference_word_timestamps") is None: + return None + + reference_words = _expand_reference_words( + _parse_reference_words(channel_metrics["reference_word_timestamps"]) + ) + snapshot_rows = channel_metrics.get("transcript_snapshots") + if snapshot_rows is None: + return None + snapshots = _parse_transcript_snapshots(snapshot_rows) + if not snapshots: + return None + + if not reference_words: + return None + + reference_tokens = [word.word for word in reference_words] + first_seen_ms: list[Optional[float]] = [None] * len(reference_tokens) + + # Intern tokens as small ints: SequenceMatcher only relies on element + # equality/hashing, so a bijective token->id mapping shared between the + # reference and every hypothesis yields identical matching blocks while + # hashing and comparing much faster than strings. + token_ids: dict[str, int] = {} + reference_token_ids = [ + token_ids.setdefault(token, len(token_ids)) for token in reference_tokens + ] + + matcher = difflib.SequenceMatcher(autojunk=False) + matcher.set_seq1(reference_token_ids) + + previous_transcript: Optional[str] = None + previous_hypothesis_ids: Optional[list[int]] = None + previous_blocks: Optional[list[difflib.Match]] = None + + for snapshot in sorted(snapshots, key=lambda item: item.elapsed_ms): + # Streaming partials often repeat the previous transcript (or reduce + # to the same normalized token sequence); the matching blocks depend + # only on the token sequences, so reuse them. + if snapshot.transcript == previous_transcript: + hypothesis_ids = previous_hypothesis_ids + else: + hypothesis_ids = [ + token_ids.setdefault(token, len(token_ids)) + for token in _normalize_words(snapshot.transcript) + ] + previous_transcript = snapshot.transcript + if not hypothesis_ids: + previous_hypothesis_ids = hypothesis_ids + previous_blocks = [] + continue + if hypothesis_ids == previous_hypothesis_ids: + blocks = previous_blocks + else: + matcher.set_seq2(hypothesis_ids) + # Equal-opcode ranges are exactly the matching blocks (the final + # sentinel block has size 0 and marks nothing). + blocks = matcher.get_matching_blocks() + previous_hypothesis_ids = hypothesis_ids + previous_blocks = blocks + + elapsed_ms = snapshot.elapsed_ms + for ref_start, _hyp_start, block_size in blocks: + for ref_index in range(ref_start, ref_start + block_size): + if ( + first_seen_ms[ref_index] is None + and elapsed_ms >= reference_words[ref_index].start_ms + ): + first_seen_ms[ref_index] = elapsed_ms + + latencies = [ + max(0.0, seen_ms - word.end_ms) + for seen_ms, word in zip(first_seen_ms, reference_words) + if seen_ms is not None + ] + if not latencies: + return None + return InteractivityStats( + mean_latency_ms=sum(latencies) / len(latencies), + word_count=len(latencies), + latencies_ms=latencies, + ) + + +def _parse_reference_words(value: Any) -> list[ReferenceWord]: + rows = _expect_list(value, "reference_word_timestamps") + if not rows: + raise ValueError("reference_word_timestamps must not be empty.") + + words: list[ReferenceWord] = [] + for index, row in enumerate(rows): + item = _expect_dict(row, f"reference_word_timestamps[{index}]") + words.append( + ReferenceWord( + word=_expect_str(item, "word", index), + start_ms=_expect_float(item, "start_ms", index), + end_ms=_expect_float(item, "end_ms", index), + ) + ) + return words + + +def _parse_transcript_snapshots(value: Any) -> list[TranscriptSnapshot]: + rows = _expect_list(value, "transcript_snapshots") + snapshots: list[TranscriptSnapshot] = [] + for index, row in enumerate(rows): + item = _expect_dict(row, f"transcript_snapshots[{index}]") + snapshots.append( + TranscriptSnapshot( + elapsed_ms=_expect_float(item, "elapsed_ms", index), + transcript=_expect_str(item, "transcript", index), + ) + ) + return snapshots + + +def _expect_list(value: Any, field: str) -> list[Any]: + if not isinstance(value, list): + raise TypeError(f"{field} must be a list, got {type(value).__name__}.") + return value + + +def _expect_dict(value: Any, field: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise TypeError(f"{field} must be an object, got {type(value).__name__}.") + return value + + +def _expect_str(item: dict[str, Any], field: str, index: int) -> str: + value = item[field] + if not isinstance(value, str): + raise TypeError(f"{field} at index {index} must be str.") + return value + + +def _expect_float(item: dict[str, Any], field: str, index: int) -> float: + value = item[field] + try: + return float(value) + except (TypeError, ValueError) as exc: + raise TypeError(f"{field} at index {index} must be numeric.") from exc + + +def _expand_reference_words(words: list[ReferenceWord]) -> list[ReferenceWord]: + expanded: list[ReferenceWord] = [] + for word in words: + expanded.extend( + ReferenceWord( + word=token, + start_ms=word.start_ms, + end_ms=word.end_ms, + ) + for token in _normalize_words(word.word) + ) + return expanded + + +def _normalize_words(text: str) -> list[str]: + return _WORD_RE.findall(_normalizer(text)) diff --git a/veeksha/evaluator/performance/asr_normalizer/__init__.py b/veeksha/evaluator/performance/asr_normalizer/__init__.py new file mode 100644 index 00000000..6ce0def1 --- /dev/null +++ b/veeksha/evaluator/performance/asr_normalizer/__init__.py @@ -0,0 +1,11 @@ +"""Vendored Open ASR Leaderboard text normalizer. + +Adapted from https://github.com/huggingface/open_asr_leaderboard so transcript +WER is scored the same way the leaderboard scores it. +""" + +from veeksha.evaluator.performance.asr_normalizer.normalizer import ( + EnglishTextNormalizer, +) + +__all__ = ["EnglishTextNormalizer"] diff --git a/veeksha/evaluator/performance/asr_normalizer/english_abbreviations.py b/veeksha/evaluator/performance/asr_normalizer/english_abbreviations.py new file mode 100644 index 00000000..5ad2ea10 --- /dev/null +++ b/veeksha/evaluator/performance/asr_normalizer/english_abbreviations.py @@ -0,0 +1,1745 @@ +# Adapted from https://github.com/huggingface/open_asr_leaderboard +# (normalizer/english_abbreviations.py). British -> American spelling map used +# by EnglishTextNormalizer. Vendored verbatim. +english_spelling_normalizer = { + "accessorise": "accessorize", + "accessorised": "accessorized", + "accessorises": "accessorizes", + "accessorising": "accessorizing", + "acclimatisation": "acclimatization", + "acclimatise": "acclimatize", + "acclimatised": "acclimatized", + "acclimatises": "acclimatizes", + "acclimatising": "acclimatizing", + "accoutrements": "accouterments", + "aeon": "eon", + "aeons": "eons", + "aerogramme": "aerogram", + "aerogrammes": "aerograms", + "aeroplane": "airplane", + "aeroplanes": "airplanes", + "aesthete": "esthete", + "aesthetes": "esthetes", + "aesthetic": "esthetic", + "aesthetically": "esthetically", + "aesthetics": "esthetics", + "aetiology": "etiology", + "ageing": "aging", + "aggrandisement": "aggrandizement", + "agonise": "agonize", + "agonised": "agonized", + "agonises": "agonizes", + "agonising": "agonizing", + "agonisingly": "agonizingly", + "almanack": "almanac", + "almanacks": "almanacs", + "aluminium": "aluminum", + "amortisable": "amortizable", + "amortisation": "amortization", + "amortisations": "amortizations", + "amortise": "amortize", + "amortised": "amortized", + "amortises": "amortizes", + "amortising": "amortizing", + "amphitheatre": "amphitheater", + "amphitheatres": "amphitheaters", + "anaemia": "anemia", + "anaemic": "anemic", + "anaesthesia": "anesthesia", + "anaesthetic": "anesthetic", + "anaesthetics": "anesthetics", + "anaesthetise": "anesthetize", + "anaesthetised": "anesthetized", + "anaesthetises": "anesthetizes", + "anaesthetising": "anesthetizing", + "anaesthetist": "anesthetist", + "anaesthetists": "anesthetists", + "anaesthetize": "anesthetize", + "anaesthetized": "anesthetized", + "anaesthetizes": "anesthetizes", + "anaesthetizing": "anesthetizing", + "analogue": "analog", + "analogues": "analogs", + "analyse": "analyze", + "analysed": "analyzed", + "analyses": "analyzes", + "analysing": "analyzing", + "anglicise": "anglicize", + "anglicised": "anglicized", + "anglicises": "anglicizes", + "anglicising": "anglicizing", + "annualised": "annualized", + "antagonise": "antagonize", + "antagonised": "antagonized", + "antagonises": "antagonizes", + "antagonising": "antagonizing", + "apologise": "apologize", + "apologised": "apologized", + "apologises": "apologizes", + "apologising": "apologizing", + "appal": "appall", + "appals": "appalls", + "appetiser": "appetizer", + "appetisers": "appetizers", + "appetising": "appetizing", + "appetisingly": "appetizingly", + "arbour": "arbor", + "arbours": "arbors", + "archaeologically": "archeologically", + "archaeologist": "archeologist", + "archaeologists": "archeologists", + "archaeology": "archeology", + "archaeological": "archeological", + "ardour": "ardor", + "armour": "armor", + "armoured": "armored", + "armourer": "armorer", + "armourers": "armorers", + "armouries": "armories", + "armoury": "armory", + "artefact": "artifact", + "artefacts": "artifacts", + "authorise": "authorize", + "authorised": "authorized", + "authorises": "authorizes", + "authorising": "authorizing", + "axe": "ax", + "backpedalled": "backpedaled", + "backpedalling": "backpedaling", + "bannister": "banister", + "bannisters": "banisters", + "baptise": "baptize", + "baptised": "baptized", + "baptises": "baptizes", + "baptising": "baptizing", + "bastardise": "bastardize", + "bastardised": "bastardized", + "bastardises": "bastardizes", + "bastardising": "bastardizing", + "battleax": "battleaxe", + "baulk": "balk", + "baulked": "balked", + "baulking": "balking", + "baulks": "balks", + "bedevilled": "bedeviled", + "bedevilling": "bedeviling", + "behaviour": "behavior", + "behavioural": "behavioral", + "behaviourism": "behaviorism", + "behaviourist": "behaviorist", + "behaviourists": "behaviorists", + "behaviours": "behaviors", + "behove": "behoove", + "behoved": "behooved", + "behoves": "behooves", + "bejewelled": "bejeweled", + "belabour": "belabor", + "belaboured": "belabored", + "belabouring": "belaboring", + "belabours": "belabors", + "bevelled": "beveled", + "bevvies": "bevies", + "bevvy": "bevy", + "biassed": "biased", + "biassing": "biasing", + "bingeing": "binging", + "bougainvillaea": "bougainvillea", + "bougainvillaeas": "bougainvilleas", + "bowdlerise": "bowdlerize", + "bowdlerised": "bowdlerized", + "bowdlerises": "bowdlerizes", + "bowdlerising": "bowdlerizing", + "breathalyse": "breathalyze", + "breathalysed": "breathalyzed", + "breathalyser": "breathalyzer", + "breathalysers": "breathalyzers", + "breathalyses": "breathalyzes", + "breathalysing": "breathalyzing", + "brutalise": "brutalize", + "brutalised": "brutalized", + "brutalises": "brutalizes", + "brutalising": "brutalizing", + "busses": "buses", + "bussing": "busing", + "caesarean": "cesarean", + "caesareans": "cesareans", + "calibre": "caliber", + "calibres": "calibers", + "calliper": "caliper", + "callipers": "calipers", + "callisthenics": "calisthenics", + "canalise": "canalize", + "canalised": "canalized", + "canalises": "canalizes", + "canalising": "canalizing", + "cancellation": "cancelation", + "cancellations": "cancelations", + "cancelled": "canceled", + "cancelling": "canceling", + "candour": "candor", + "cannibalise": "cannibalize", + "cannibalised": "cannibalized", + "cannibalises": "cannibalizes", + "cannibalising": "cannibalizing", + "canonise": "canonize", + "canonised": "canonized", + "canonises": "canonizes", + "canonising": "canonizing", + "capitalise": "capitalize", + "capitalised": "capitalized", + "capitalises": "capitalizes", + "capitalising": "capitalizing", + "caramelise": "caramelize", + "caramelised": "caramelized", + "caramelises": "caramelizes", + "caramelising": "caramelizing", + "carbonise": "carbonize", + "carbonised": "carbonized", + "carbonises": "carbonizes", + "carbonising": "carbonizing", + "carolled": "caroled", + "carolling": "caroling", + "catalogue": "catalog", + "catalogued": "cataloged", + "catalogues": "catalogs", + "cataloguing": "cataloging", + "catalyse": "catalyze", + "catalysed": "catalyzed", + "catalyses": "catalyzes", + "catalysing": "catalyzing", + "categorise": "categorize", + "categorised": "categorized", + "categorises": "categorizes", + "categorising": "categorizing", + "cauterise": "cauterize", + "cauterised": "cauterized", + "cauterises": "cauterizes", + "cauterising": "cauterizing", + "cavilled": "caviled", + "cavilling": "caviling", + "centigramme": "centigram", + "centigrammes": "centigrams", + "centilitre": "centiliter", + "centilitres": "centiliters", + "centimetre": "centimeter", + "centimetres": "centimeters", + "centralise": "centralize", + "centralised": "centralized", + "centralises": "centralizes", + "centralising": "centralizing", + "centre": "center", + "centred": "centered", + "centrefold": "centerfold", + "centrefolds": "centerfolds", + "centrepiece": "centerpiece", + "centrepieces": "centerpieces", + "centres": "centers", + "channelled": "channeled", + "channelling": "channeling", + "characterise": "characterize", + "characterised": "characterized", + "characterises": "characterizes", + "characterising": "characterizing", + "cheque": "check", + "chequebook": "checkbook", + "chequebooks": "checkbooks", + "chequered": "checkered", + "cheques": "checks", + "chilli": "chili", + "chimaera": "chimera", + "chimaeras": "chimeras", + "chiselled": "chiseled", + "chiselling": "chiseling", + "circularise": "circularize", + "circularised": "circularized", + "circularises": "circularizes", + "circularising": "circularizing", + "civilise": "civilize", + "civilised": "civilized", + "civilises": "civilizes", + "civilising": "civilizing", + "clamour": "clamor", + "clamoured": "clamored", + "clamouring": "clamoring", + "clamours": "clamors", + "clangour": "clangor", + "clarinettist": "clarinetist", + "clarinettists": "clarinetists", + "collectivise": "collectivize", + "collectivised": "collectivized", + "collectivises": "collectivizes", + "collectivising": "collectivizing", + "colonisation": "colonization", + "colonise": "colonize", + "colonised": "colonized", + "coloniser": "colonizer", + "colonisers": "colonizers", + "colonises": "colonizes", + "colonising": "colonizing", + "colour": "color", + "colourant": "colorant", + "colourants": "colorants", + "coloured": "colored", + "coloureds": "coloreds", + "colourful": "colorful", + "colourfully": "colorfully", + "colouring": "coloring", + "colourize": "colorize", + "colourized": "colorized", + "colourizes": "colorizes", + "colourizing": "colorizing", + "colourless": "colorless", + "colours": "colors", + "commercialise": "commercialize", + "commercialised": "commercialized", + "commercialises": "commercializes", + "commercialising": "commercializing", + "compartmentalise": "compartmentalize", + "compartmentalised": "compartmentalized", + "compartmentalises": "compartmentalizes", + "compartmentalising": "compartmentalizing", + "computerise": "computerize", + "computerised": "computerized", + "computerises": "computerizes", + "computerising": "computerizing", + "conceptualise": "conceptualize", + "conceptualised": "conceptualized", + "conceptualises": "conceptualizes", + "conceptualising": "conceptualizing", + "connexion": "connection", + "connexions": "connections", + "contextualise": "contextualize", + "contextualised": "contextualized", + "contextualises": "contextualizes", + "contextualising": "contextualizing", + "cosier": "cozier", + "cosies": "cozies", + "cosiest": "coziest", + "cosily": "cozily", + "cosiness": "coziness", + "cosy": "cozy", + "councillor": "councilor", + "councillors": "councilors", + "counselled": "counseled", + "counselling": "counseling", + "counsellor": "counselor", + "counsellors": "counselors", + "crenelated": "crenellated", + "criminalise": "criminalize", + "criminalised": "criminalized", + "criminalises": "criminalizes", + "criminalising": "criminalizing", + "criticise": "criticize", + "criticised": "criticized", + "criticises": "criticizes", + "criticising": "criticizing", + "crueller": "crueler", + "cruellest": "cruelest", + "crystallisation": "crystallization", + "crystallise": "crystallize", + "crystallised": "crystallized", + "crystallises": "crystallizes", + "crystallising": "crystallizing", + "cudgelled": "cudgeled", + "cudgelling": "cudgeling", + "customise": "customize", + "customised": "customized", + "customises": "customizes", + "customising": "customizing", + "cypher": "cipher", + "cyphers": "ciphers", + "decentralisation": "decentralization", + "decentralise": "decentralize", + "decentralised": "decentralized", + "decentralises": "decentralizes", + "decentralising": "decentralizing", + "decriminalisation": "decriminalization", + "decriminalise": "decriminalize", + "decriminalised": "decriminalized", + "decriminalises": "decriminalizes", + "decriminalising": "decriminalizing", + "defence": "defense", + "defenceless": "defenseless", + "defences": "defenses", + "dehumanisation": "dehumanization", + "dehumanise": "dehumanize", + "dehumanised": "dehumanized", + "dehumanises": "dehumanizes", + "dehumanising": "dehumanizing", + "demeanour": "demeanor", + "demilitarisation": "demilitarization", + "demilitarise": "demilitarize", + "demilitarised": "demilitarized", + "demilitarises": "demilitarizes", + "demilitarising": "demilitarizing", + "demobilisation": "demobilization", + "demobilise": "demobilize", + "demobilised": "demobilized", + "demobilises": "demobilizes", + "demobilising": "demobilizing", + "democratisation": "democratization", + "democratise": "democratize", + "democratised": "democratized", + "democratises": "democratizes", + "democratising": "democratizing", + "demonise": "demonize", + "demonised": "demonized", + "demonises": "demonizes", + "demonising": "demonizing", + "demoralisation": "demoralization", + "demoralise": "demoralize", + "demoralised": "demoralized", + "demoralises": "demoralizes", + "demoralising": "demoralizing", + "denationalisation": "denationalization", + "denationalise": "denationalize", + "denationalised": "denationalized", + "denationalises": "denationalizes", + "denationalising": "denationalizing", + "deodorise": "deodorize", + "deodorised": "deodorized", + "deodorises": "deodorizes", + "deodorising": "deodorizing", + "depersonalise": "depersonalize", + "depersonalised": "depersonalized", + "depersonalises": "depersonalizes", + "depersonalising": "depersonalizing", + "deputise": "deputize", + "deputised": "deputized", + "deputises": "deputizes", + "deputising": "deputizing", + "desensitisation": "desensitization", + "desensitise": "desensitize", + "desensitised": "desensitized", + "desensitises": "desensitizes", + "desensitising": "desensitizing", + "destabilisation": "destabilization", + "destabilise": "destabilize", + "destabilised": "destabilized", + "destabilises": "destabilizes", + "destabilising": "destabilizing", + "dialled": "dialed", + "dialling": "dialing", + "dialogue": "dialog", + "dialogues": "dialogs", + "diarrhoea": "diarrhea", + "digitise": "digitize", + "digitised": "digitized", + "digitises": "digitizes", + "digitising": "digitizing", + "disc": "disk", + "discolour": "discolor", + "discoloured": "discolored", + "discolouring": "discoloring", + "discolours": "discolors", + "discs": "disks", + "disembowelled": "disemboweled", + "disembowelling": "disemboweling", + "disfavour": "disfavor", + "dishevelled": "disheveled", + "dishonour": "dishonor", + "dishonourable": "dishonorable", + "dishonourably": "dishonorably", + "dishonoured": "dishonored", + "dishonouring": "dishonoring", + "dishonours": "dishonors", + "disorganisation": "disorganization", + "disorganised": "disorganized", + "distil": "distill", + "distils": "distills", + "dramatisation": "dramatization", + "dramatisations": "dramatizations", + "dramatise": "dramatize", + "dramatised": "dramatized", + "dramatises": "dramatizes", + "dramatising": "dramatizing", + "draught": "draft", + "draughtboard": "draftboard", + "draughtboards": "draftboards", + "draughtier": "draftier", + "draughtiest": "draftiest", + "draughts": "drafts", + "draughtsman": "draftsman", + "draughtsmanship": "draftsmanship", + "draughtsmen": "draftsmen", + "draughtswoman": "draftswoman", + "draughtswomen": "draftswomen", + "draughty": "drafty", + "drivelled": "driveled", + "drivelling": "driveling", + "duelled": "dueled", + "duelling": "dueling", + "economise": "economize", + "economised": "economized", + "economises": "economizes", + "economising": "economizing", + "editorialise": "editorialize", + "editorialised": "editorialized", + "editorialises": "editorializes", + "editorialising": "editorializing", + "edoema": "edema", + "empathise": "empathize", + "empathised": "empathized", + "empathises": "empathizes", + "empathising": "empathizing", + "emphasise": "emphasize", + "emphasised": "emphasized", + "emphasises": "emphasizes", + "emphasising": "emphasizing", + "enamelled": "enameled", + "enamelling": "enameling", + "enamoured": "enamored", + "encyclopaedia": "encyclopedia", + "encyclopaedias": "encyclopedias", + "encyclopaedic": "encyclopedic", + "endeavour": "endeavor", + "endeavoured": "endeavored", + "endeavouring": "endeavoring", + "endeavours": "endeavors", + "energise": "energize", + "energised": "energized", + "energises": "energizes", + "energising": "energizing", + "enrol": "enroll", + "enrols": "enrolls", + "enthral": "enthrall", + "enthrals": "enthralls", + "epaulette": "epaulet", + "epaulettes": "epaulets", + "epicentre": "epicenter", + "epicentres": "epicenters", + "epilogue": "epilog", + "epilogues": "epilogs", + "epitomise": "epitomize", + "epitomised": "epitomized", + "epitomises": "epitomizes", + "epitomising": "epitomizing", + "equalisation": "equalization", + "equalise": "equalize", + "equalised": "equalized", + "equaliser": "equalizer", + "equalisers": "equalizers", + "equalises": "equalizes", + "equalising": "equalizing", + "eulogise": "eulogize", + "eulogised": "eulogized", + "eulogises": "eulogizes", + "eulogising": "eulogizing", + "evangelise": "evangelize", + "evangelised": "evangelized", + "evangelises": "evangelizes", + "evangelising": "evangelizing", + "exorcise": "exorcize", + "exorcised": "exorcized", + "exorcises": "exorcizes", + "exorcising": "exorcizing", + "extemporisation": "extemporization", + "extemporise": "extemporize", + "extemporised": "extemporized", + "extemporises": "extemporizes", + "extemporising": "extemporizing", + "externalisation": "externalization", + "externalisations": "externalizations", + "externalise": "externalize", + "externalised": "externalized", + "externalises": "externalizes", + "externalising": "externalizing", + "factorise": "factorize", + "factorised": "factorized", + "factorises": "factorizes", + "factorising": "factorizing", + "faecal": "fecal", + "faeces": "feces", + "familiarisation": "familiarization", + "familiarise": "familiarize", + "familiarised": "familiarized", + "familiarises": "familiarizes", + "familiarising": "familiarizing", + "fantasise": "fantasize", + "fantasised": "fantasized", + "fantasises": "fantasizes", + "fantasising": "fantasizing", + "favour": "favor", + "favourable": "favorable", + "favourably": "favorably", + "favoured": "favored", + "favouring": "favoring", + "favourite": "favorite", + "favourites": "favorites", + "favouritism": "favoritism", + "favours": "favors", + "feminise": "feminize", + "feminised": "feminized", + "feminises": "feminizes", + "feminising": "feminizing", + "fertilisation": "fertilization", + "fertilise": "fertilize", + "fertilised": "fertilized", + "fertiliser": "fertilizer", + "fertilisers": "fertilizers", + "fertilises": "fertilizes", + "fertilising": "fertilizing", + "fervour": "fervor", + "fibre": "fiber", + "fibreglass": "fiberglass", + "fibres": "fibers", + "fictionalisation": "fictionalization", + "fictionalisations": "fictionalizations", + "fictionalise": "fictionalize", + "fictionalised": "fictionalized", + "fictionalises": "fictionalizes", + "fictionalising": "fictionalizing", + "fillet": "filet", + "filleted": "fileted", + "filleting": "fileting", + "fillets": "filets", + "finalisation": "finalization", + "finalise": "finalize", + "finalised": "finalized", + "finalises": "finalizes", + "finalising": "finalizing", + "flautist": "flutist", + "flautists": "flutists", + "flavour": "flavor", + "flavoured": "flavored", + "flavouring": "flavoring", + "flavourings": "flavorings", + "flavourless": "flavorless", + "flavours": "flavors", + "flavoursome": "flavorsome", + "flyer / flier": "flier / flyer", + "foetal": "fetal", + "foetid": "fetid", + "foetus": "fetus", + "foetuses": "fetuses", + "formalisation": "formalization", + "formalise": "formalize", + "formalised": "formalized", + "formalises": "formalizes", + "formalising": "formalizing", + "fossilisation": "fossilization", + "fossilise": "fossilize", + "fossilised": "fossilized", + "fossilises": "fossilizes", + "fossilising": "fossilizing", + "fraternisation": "fraternization", + "fraternise": "fraternize", + "fraternised": "fraternized", + "fraternises": "fraternizes", + "fraternising": "fraternizing", + "fulfil": "fulfill", + "fulfilment": "fulfillment", + "fulfils": "fulfills", + "funnelled": "funneled", + "funnelling": "funneling", + "gage": "gauge", + "gaged": "gauged", + "gages": "gauges", + "gaging": "gauging", + "galvanise": "galvanize", + "galvanised": "galvanized", + "galvanises": "galvanizes", + "galvanising": "galvanizing", + "gambolled": "gamboled", + "gambolling": "gamboling", + "gaol": "jail", + "gaolbird": "jailbird", + "gaolbirds": "jailbirds", + "gaolbreak": "jailbreak", + "gaolbreaks": "jailbreaks", + "gaoled": "jailed", + "gaoler": "jailer", + "gaolers": "jailers", + "gaoling": "jailing", + "gaols": "jails", + "gasses": "gases", + "generalisation": "generalization", + "generalisations": "generalizations", + "generalise": "generalize", + "generalised": "generalized", + "generalises": "generalizes", + "generalising": "generalizing", + "ghettoise": "ghettoize", + "ghettoised": "ghettoized", + "ghettoises": "ghettoizes", + "ghettoising": "ghettoizing", + "gipsies": "gypsies", + "glamor": "glamour", + "glamorise": "glamorize", + "glamorised": "glamorized", + "glamorises": "glamorizes", + "glamorising": "glamorizing", + "globalisation": "globalization", + "globalise": "globalize", + "globalised": "globalized", + "globalises": "globalizes", + "globalising": "globalizing", + "glueing": "gluing", + "goitre": "goiter", + "goitres": "goiters", + "gonorrhoea": "gonorrhea", + "gramme": "gram", + "grammes": "grams", + "gravelled": "graveled", + "grey": "gray", + "greyed": "grayed", + "greying": "graying", + "greyish": "grayish", + "greyness": "grayness", + "greys": "grays", + "grovelled": "groveled", + "grovelling": "groveling", + "groyne": "groin", + "groynes": "groins", + "gruelling": "grueling", + "gruellingly": "gruelingly", + "gryphon": "griffin", + "gryphons": "griffins", + "gynaecological": "gynecological", + "gynaecologist": "gynecologist", + "gynaecologists": "gynecologists", + "gynaecology": "gynecology", + "haematological": "hematological", + "haematologist": "hematologist", + "haematologists": "hematologists", + "haematology": "hematology", + "haemoglobin": "hemoglobin", + "haemophilia": "hemophilia", + "haemophiliac": "hemophiliac", + "haemophiliacs": "hemophiliacs", + "haemorrhage": "hemorrhage", + "haemorrhaged": "hemorrhaged", + "haemorrhages": "hemorrhages", + "haemorrhaging": "hemorrhaging", + "haemorrhoids": "hemorrhoids", + "harbour": "harbor", + "harboured": "harbored", + "harbouring": "harboring", + "harbours": "harbors", + "harmonisation": "harmonization", + "harmonise": "harmonize", + "harmonised": "harmonized", + "harmonises": "harmonizes", + "harmonising": "harmonizing", + "homoeopath": "homeopath", + "homoeopathic": "homeopathic", + "homoeopaths": "homeopaths", + "homoeopathy": "homeopathy", + "homogenise": "homogenize", + "homogenised": "homogenized", + "homogenises": "homogenizes", + "homogenising": "homogenizing", + "honour": "honor", + "honourable": "honorable", + "honourably": "honorably", + "honoured": "honored", + "honouring": "honoring", + "honours": "honors", + "hospitalisation": "hospitalization", + "hospitalise": "hospitalize", + "hospitalised": "hospitalized", + "hospitalises": "hospitalizes", + "hospitalising": "hospitalizing", + "humanise": "humanize", + "humanised": "humanized", + "humanises": "humanizes", + "humanising": "humanizing", + "humour": "humor", + "humoured": "humored", + "humouring": "humoring", + "humourless": "humorless", + "humours": "humors", + "hybridise": "hybridize", + "hybridised": "hybridized", + "hybridises": "hybridizes", + "hybridising": "hybridizing", + "hypnotise": "hypnotize", + "hypnotised": "hypnotized", + "hypnotises": "hypnotizes", + "hypnotising": "hypnotizing", + "hypothesise": "hypothesize", + "hypothesised": "hypothesized", + "hypothesises": "hypothesizes", + "hypothesising": "hypothesizing", + "idealisation": "idealization", + "idealise": "idealize", + "idealised": "idealized", + "idealises": "idealizes", + "idealising": "idealizing", + "idolise": "idolize", + "idolised": "idolized", + "idolises": "idolizes", + "idolising": "idolizing", + "immobilisation": "immobilization", + "immobilise": "immobilize", + "immobilised": "immobilized", + "immobiliser": "immobilizer", + "immobilisers": "immobilizers", + "immobilises": "immobilizes", + "immobilising": "immobilizing", + "immortalise": "immortalize", + "immortalised": "immortalized", + "immortalises": "immortalizes", + "immortalising": "immortalizing", + "immunisation": "immunization", + "immunise": "immunize", + "immunised": "immunized", + "immunises": "immunizes", + "immunising": "immunizing", + "impanelled": "impaneled", + "impanelling": "impaneling", + "imperilled": "imperiled", + "imperilling": "imperiling", + "individualise": "individualize", + "individualised": "individualized", + "individualises": "individualizes", + "individualising": "individualizing", + "industrialise": "industrialize", + "industrialised": "industrialized", + "industrialises": "industrializes", + "industrialising": "industrializing", + "inflexion": "inflection", + "inflexions": "inflections", + "initialise": "initialize", + "initialised": "initialized", + "initialises": "initializes", + "initialising": "initializing", + "initialled": "initialed", + "initialling": "initialing", + "instal": "install", + "instalment": "installment", + "instalments": "installments", + "instals": "installs", + "instil": "instill", + "instils": "instills", + "institutionalisation": "institutionalization", + "institutionalise": "institutionalize", + "institutionalised": "institutionalized", + "institutionalises": "institutionalizes", + "institutionalising": "institutionalizing", + "intellectualise": "intellectualize", + "intellectualised": "intellectualized", + "intellectualises": "intellectualizes", + "intellectualising": "intellectualizing", + "internalisation": "internalization", + "internalise": "internalize", + "internalised": "internalized", + "internalises": "internalizes", + "internalising": "internalizing", + "internationalisation": "internationalization", + "internationalise": "internationalize", + "internationalised": "internationalized", + "internationalises": "internationalizes", + "internationalising": "internationalizing", + "ionisation": "ionization", + "ionise": "ionize", + "ionised": "ionized", + "ioniser": "ionizer", + "ionisers": "ionizers", + "ionises": "ionizes", + "ionising": "ionizing", + "italicise": "italicize", + "italicised": "italicized", + "italicises": "italicizes", + "italicising": "italicizing", + "itemise": "itemize", + "itemised": "itemized", + "itemises": "itemizes", + "itemising": "itemizing", + "jeopardise": "jeopardize", + "jeopardised": "jeopardized", + "jeopardises": "jeopardizes", + "jeopardising": "jeopardizing", + "jewelled": "jeweled", + "jeweller": "jeweler", + "jewellers": "jewelers", + "jewellery": "jewelry", + "judgement": "judgment", + "kilogramme": "kilogram", + "kilogrammes": "kilograms", + "kilometre": "kilometer", + "kilometres": "kilometers", + "labelled": "labeled", + "labelling": "labeling", + "labour": "labor", + "laboured": "labored", + "labourer": "laborer", + "labourers": "laborers", + "labouring": "laboring", + "labours": "labors", + "lacklustre": "lackluster", + "legalisation": "legalization", + "legalise": "legalize", + "legalised": "legalized", + "legalises": "legalizes", + "legalising": "legalizing", + "legitimise": "legitimize", + "legitimised": "legitimized", + "legitimises": "legitimizes", + "legitimising": "legitimizing", + "leukaemia": "leukemia", + "levelled": "leveled", + "leveller": "leveler", + "levellers": "levelers", + "levelling": "leveling", + "libelled": "libeled", + "libelling": "libeling", + "libellous": "libelous", + "liberalisation": "liberalization", + "liberalise": "liberalize", + "liberalised": "liberalized", + "liberalises": "liberalizes", + "liberalising": "liberalizing", + "licence": "license", + "licenced": "licensed", + "licences": "licenses", + "licencing": "licensing", + "likeable": "likable", + "lionisation": "lionization", + "lionise": "lionize", + "lionised": "lionized", + "lionises": "lionizes", + "lionising": "lionizing", + "liquidise": "liquidize", + "liquidised": "liquidized", + "liquidiser": "liquidizer", + "liquidisers": "liquidizers", + "liquidises": "liquidizes", + "liquidising": "liquidizing", + "litre": "liter", + "litres": "liters", + "localise": "localize", + "localised": "localized", + "localises": "localizes", + "localising": "localizing", + "louvre": "louver", + "louvred": "louvered", + "louvres": "louvers", + "lustre": "luster", + "magnetise": "magnetize", + "magnetised": "magnetized", + "magnetises": "magnetizes", + "magnetising": "magnetizing", + "manoeuvrability": "maneuverability", + "manoeuvrable": "maneuverable", + "manoeuvre": "maneuver", + "manoeuvred": "maneuvered", + "manoeuvres": "maneuvers", + "manoeuvring": "maneuvering", + "manoeuvrings": "maneuverings", + "marginalisation": "marginalization", + "marginalise": "marginalize", + "marginalised": "marginalized", + "marginalises": "marginalizes", + "marginalising": "marginalizing", + "marshalled": "marshaled", + "marshalling": "marshaling", + "marvelled": "marveled", + "marvelling": "marveling", + "marvellous": "marvelous", + "marvellously": "marvelously", + "materialisation": "materialization", + "materialise": "materialize", + "materialised": "materialized", + "materialises": "materializes", + "materialising": "materializing", + "maximisation": "maximization", + "maximise": "maximize", + "maximised": "maximized", + "maximises": "maximizes", + "maximising": "maximizing", + "meagre": "meager", + "mechanisation": "mechanization", + "mechanise": "mechanize", + "mechanised": "mechanized", + "mechanises": "mechanizes", + "mechanising": "mechanizing", + "mediaeval": "medieval", + "memorialise": "memorialize", + "memorialised": "memorialized", + "memorialises": "memorializes", + "memorialising": "memorializing", + "memorise": "memorize", + "memorised": "memorized", + "memorises": "memorizes", + "memorising": "memorizing", + "mesmerise": "mesmerize", + "mesmerised": "mesmerized", + "mesmerises": "mesmerizes", + "mesmerising": "mesmerizing", + "metabolise": "metabolize", + "metabolised": "metabolized", + "metabolises": "metabolizes", + "metabolising": "metabolizing", + "metre": "meter", + "metres": "meters", + "mhm": "hmm", + "micrometre": "micrometer", + "micrometres": "micrometers", + "militarise": "militarize", + "militarised": "militarized", + "militarises": "militarizes", + "militarising": "militarizing", + "milligramme": "milligram", + "milligrammes": "milligrams", + "millilitre": "milliliter", + "millilitres": "milliliters", + "millimetre": "millimeter", + "millimetres": "millimeters", + "miniaturisation": "miniaturization", + "miniaturise": "miniaturize", + "miniaturised": "miniaturized", + "miniaturises": "miniaturizes", + "miniaturising": "miniaturizing", + "minibusses": "minibuses", + "minimise": "minimize", + "minimised": "minimized", + "minimises": "minimizes", + "minimising": "minimizing", + "misbehaviour": "misbehavior", + "misdemeanour": "misdemeanor", + "misdemeanours": "misdemeanors", + "misspelt": "misspelled", + "mitre": "miter", + "mitres": "miters", + "mm": "hmm", + "mmm": "hmm", + "mobilisation": "mobilization", + "mobilise": "mobilize", + "mobilised": "mobilized", + "mobilises": "mobilizes", + "mobilising": "mobilizing", + "modelled": "modeled", + "modeller": "modeler", + "modellers": "modelers", + "modelling": "modeling", + "modernise": "modernize", + "modernised": "modernized", + "modernises": "modernizes", + "modernising": "modernizing", + "moisturise": "moisturize", + "moisturised": "moisturized", + "moisturiser": "moisturizer", + "moisturisers": "moisturizers", + "moisturises": "moisturizes", + "moisturising": "moisturizing", + "monologue": "monolog", + "monologues": "monologs", + "monopolisation": "monopolization", + "monopolise": "monopolize", + "monopolised": "monopolized", + "monopolises": "monopolizes", + "monopolising": "monopolizing", + "moralise": "moralize", + "moralised": "moralized", + "moralises": "moralizes", + "moralising": "moralizing", + "motorised": "motorized", + "mould": "mold", + "moulded": "molded", + "moulder": "molder", + "mouldered": "moldered", + "mouldering": "moldering", + "moulders": "molders", + "mouldier": "moldier", + "mouldiest": "moldiest", + "moulding": "molding", + "mouldings": "moldings", + "moulds": "molds", + "mouldy": "moldy", + "moult": "molt", + "moulted": "molted", + "moulting": "molting", + "moults": "molts", + "moustache": "mustache", + "moustached": "mustached", + "moustaches": "mustaches", + "moustachioed": "mustachioed", + "multicoloured": "multicolored", + "nationalisation": "nationalization", + "nationalisations": "nationalizations", + "nationalise": "nationalize", + "nationalised": "nationalized", + "nationalises": "nationalizes", + "nationalising": "nationalizing", + "naturalisation": "naturalization", + "naturalise": "naturalize", + "naturalised": "naturalized", + "naturalises": "naturalizes", + "naturalising": "naturalizing", + "neighbour": "neighbor", + "neighbourhood": "neighborhood", + "neighbourhoods": "neighborhoods", + "neighbouring": "neighboring", + "neighbourliness": "neighborliness", + "neighbourly": "neighborly", + "neighbours": "neighbors", + "neutralisation": "neutralization", + "neutralise": "neutralize", + "neutralised": "neutralized", + "neutralises": "neutralizes", + "neutralising": "neutralizing", + "normalisation": "normalization", + "normalise": "normalize", + "normalised": "normalized", + "normalises": "normalizes", + "normalising": "normalizing", + "odour": "odor", + "odourless": "odorless", + "odours": "odors", + "oesophagus": "esophagus", + "oesophaguses": "esophaguses", + "oestrogen": "estrogen", + "offence": "offense", + "offences": "offenses", + "omelette": "omelet", + "omelettes": "omelets", + "optimise": "optimize", + "optimised": "optimized", + "optimises": "optimizes", + "optimising": "optimizing", + "organisation": "organization", + "organisational": "organizational", + "organisations": "organizations", + "organise": "organize", + "organised": "organized", + "organiser": "organizer", + "organisers": "organizers", + "organises": "organizes", + "organising": "organizing", + "orthopaedic": "orthopedic", + "orthopaedics": "orthopedics", + "ostracise": "ostracize", + "ostracised": "ostracized", + "ostracises": "ostracizes", + "ostracising": "ostracizing", + "outmanoeuvre": "outmaneuver", + "outmanoeuvred": "outmaneuvered", + "outmanoeuvres": "outmaneuvers", + "outmanoeuvring": "outmaneuvering", + "overemphasise": "overemphasize", + "overemphasised": "overemphasized", + "overemphasises": "overemphasizes", + "overemphasising": "overemphasizing", + "oxidisation": "oxidization", + "oxidise": "oxidize", + "oxidised": "oxidized", + "oxidises": "oxidizes", + "oxidising": "oxidizing", + "paederast": "pederast", + "paederasts": "pederasts", + "paediatric": "pediatric", + "paediatrician": "pediatrician", + "paediatricians": "pediatricians", + "paediatrics": "pediatrics", + "paedophile": "pedophile", + "paedophiles": "pedophiles", + "paedophilia": "pedophilia", + "palaeolithic": "paleolithic", + "palaeontologist": "paleontologist", + "palaeontologists": "paleontologists", + "palaeontology": "paleontology", + "panelled": "paneled", + "panelling": "paneling", + "panellist": "panelist", + "panellists": "panelists", + "paralyse": "paralyze", + "paralysed": "paralyzed", + "paralyses": "paralyzes", + "paralysing": "paralyzing", + "parcelled": "parceled", + "parcelling": "parceling", + "parlour": "parlor", + "parlours": "parlors", + "particularise": "particularize", + "particularised": "particularized", + "particularises": "particularizes", + "particularising": "particularizing", + "passivisation": "passivization", + "passivise": "passivize", + "passivised": "passivized", + "passivises": "passivizes", + "passivising": "passivizing", + "pasteurisation": "pasteurization", + "pasteurise": "pasteurize", + "pasteurised": "pasteurized", + "pasteurises": "pasteurizes", + "pasteurising": "pasteurizing", + "patronise": "patronize", + "patronised": "patronized", + "patronises": "patronizes", + "patronising": "patronizing", + "patronisingly": "patronizingly", + "pedalled": "pedaled", + "pedalling": "pedaling", + "pedestrianisation": "pedestrianization", + "pedestrianise": "pedestrianize", + "pedestrianised": "pedestrianized", + "pedestrianises": "pedestrianizes", + "pedestrianising": "pedestrianizing", + "penalise": "penalize", + "penalised": "penalized", + "penalises": "penalizes", + "penalising": "penalizing", + "pencilled": "penciled", + "pencilling": "penciling", + "personalise": "personalize", + "personalised": "personalized", + "personalises": "personalizes", + "personalising": "personalizing", + "pharmacopoeia": "pharmacopeia", + "pharmacopoeias": "pharmacopeias", + "philosophise": "philosophize", + "philosophised": "philosophized", + "philosophises": "philosophizes", + "philosophising": "philosophizing", + "philtre": "filter", + "philtres": "filters", + "phoney": "phony", + "plagiarise": "plagiarize", + "plagiarised": "plagiarized", + "plagiarises": "plagiarizes", + "plagiarising": "plagiarizing", + "plough": "plow", + "ploughed": "plowed", + "ploughing": "plowing", + "ploughman": "plowman", + "ploughmen": "plowmen", + "ploughs": "plows", + "ploughshare": "plowshare", + "ploughshares": "plowshares", + "polarisation": "polarization", + "polarise": "polarize", + "polarised": "polarized", + "polarises": "polarizes", + "polarising": "polarizing", + "politicisation": "politicization", + "politicise": "politicize", + "politicised": "politicized", + "politicises": "politicizes", + "politicising": "politicizing", + "popularisation": "popularization", + "popularise": "popularize", + "popularised": "popularized", + "popularises": "popularizes", + "popularising": "popularizing", + "pouffe": "pouf", + "pouffes": "poufs", + "practise": "practice", + "practised": "practiced", + "practises": "practices", + "practising": "practicing", + "praesidium": "presidium", + "praesidiums": "presidiums", + "pressurisation": "pressurization", + "pressurise": "pressurize", + "pressurised": "pressurized", + "pressurises": "pressurizes", + "pressurising": "pressurizing", + "pretence": "pretense", + "pretences": "pretenses", + "primaeval": "primeval", + "prioritisation": "prioritization", + "prioritise": "prioritize", + "prioritised": "prioritized", + "prioritises": "prioritizes", + "prioritising": "prioritizing", + "privatisation": "privatization", + "privatisations": "privatizations", + "privatise": "privatize", + "privatised": "privatized", + "privatises": "privatizes", + "privatising": "privatizing", + "professionalisation": "professionalization", + "professionalise": "professionalize", + "professionalised": "professionalized", + "professionalises": "professionalizes", + "professionalising": "professionalizing", + "programme": "program", + "programmes": "programs", + "prologue": "prolog", + "prologues": "prologs", + "propagandise": "propagandize", + "propagandised": "propagandized", + "propagandises": "propagandizes", + "propagandising": "propagandizing", + "proselytise": "proselytize", + "proselytised": "proselytized", + "proselytiser": "proselytizer", + "proselytisers": "proselytizers", + "proselytises": "proselytizes", + "proselytising": "proselytizing", + "psychoanalyse": "psychoanalyze", + "psychoanalysed": "psychoanalyzed", + "psychoanalyses": "psychoanalyzes", + "psychoanalysing": "psychoanalyzing", + "publicise": "publicize", + "publicised": "publicized", + "publicises": "publicizes", + "publicising": "publicizing", + "pulverisation": "pulverization", + "pulverise": "pulverize", + "pulverised": "pulverized", + "pulverises": "pulverizes", + "pulverising": "pulverizing", + "pummelled": "pummeled", + "pummelling": "pummeling", + "pyjama": "pajama", + "pyjamas": "pajamas", + "pzazz": "pizzazz", + "quarrelled": "quarreled", + "quarrelling": "quarreling", + "radicalise": "radicalize", + "radicalised": "radicalized", + "radicalises": "radicalizes", + "radicalising": "radicalizing", + "rancour": "rancor", + "randomise": "randomize", + "randomised": "randomized", + "randomises": "randomizes", + "randomising": "randomizing", + "rationalisation": "rationalization", + "rationalisations": "rationalizations", + "rationalise": "rationalize", + "rationalised": "rationalized", + "rationalises": "rationalizes", + "rationalising": "rationalizing", + "ravelled": "raveled", + "ravelling": "raveling", + "realisable": "realizable", + "realisation": "realization", + "realisations": "realizations", + "realise": "realize", + "realised": "realized", + "realises": "realizes", + "realising": "realizing", + "recognisable": "recognizable", + "recognisably": "recognizably", + "recognisance": "recognizance", + "recognise": "recognize", + "recognised": "recognized", + "recognises": "recognizes", + "recognising": "recognizing", + "reconnoitre": "reconnoiter", + "reconnoitred": "reconnoitered", + "reconnoitres": "reconnoiters", + "reconnoitring": "reconnoitering", + "refuelled": "refueled", + "refuelling": "refueling", + "regularisation": "regularization", + "regularise": "regularize", + "regularised": "regularized", + "regularises": "regularizes", + "regularising": "regularizing", + "remodelled": "remodeled", + "remodelling": "remodeling", + "remould": "remold", + "remoulded": "remolded", + "remoulding": "remolding", + "remoulds": "remolds", + "reorganisation": "reorganization", + "reorganisations": "reorganizations", + "reorganise": "reorganize", + "reorganised": "reorganized", + "reorganises": "reorganizes", + "reorganising": "reorganizing", + "revelled": "reveled", + "reveller": "reveler", + "revellers": "revelers", + "revelling": "reveling", + "revitalise": "revitalize", + "revitalised": "revitalized", + "revitalises": "revitalizes", + "revitalising": "revitalizing", + "revolutionise": "revolutionize", + "revolutionised": "revolutionized", + "revolutionises": "revolutionizes", + "revolutionising": "revolutionizing", + "rhapsodise": "rhapsodize", + "rhapsodised": "rhapsodized", + "rhapsodises": "rhapsodizes", + "rhapsodising": "rhapsodizing", + "rigour": "rigor", + "rigours": "rigors", + "ritualised": "ritualized", + "rivalled": "rivaled", + "rivalling": "rivaling", + "romanticise": "romanticize", + "romanticised": "romanticized", + "romanticises": "romanticizes", + "romanticising": "romanticizing", + "rumour": "rumor", + "rumoured": "rumored", + "rumours": "rumors", + "sabre": "saber", + "sabres": "sabers", + "saltpetre": "saltpeter", + "sanitise": "sanitize", + "sanitised": "sanitized", + "sanitises": "sanitizes", + "sanitising": "sanitizing", + "satirise": "satirize", + "satirised": "satirized", + "satirises": "satirizes", + "satirising": "satirizing", + "saviour": "savior", + "saviours": "saviors", + "savour": "savor", + "savoured": "savored", + "savouries": "savories", + "savouring": "savoring", + "savours": "savors", + "savoury": "savory", + "scandalise": "scandalize", + "scandalised": "scandalized", + "scandalises": "scandalizes", + "scandalising": "scandalizing", + "sceptic": "skeptic", + "sceptical": "skeptical", + "sceptically": "skeptically", + "scepticism": "skepticism", + "sceptics": "skeptics", + "sceptre": "scepter", + "sceptres": "scepters", + "scrutinise": "scrutinize", + "scrutinised": "scrutinized", + "scrutinises": "scrutinizes", + "scrutinising": "scrutinizing", + "secularisation": "secularization", + "secularise": "secularize", + "secularised": "secularized", + "secularises": "secularizes", + "secularising": "secularizing", + "sensationalise": "sensationalize", + "sensationalised": "sensationalized", + "sensationalises": "sensationalizes", + "sensationalising": "sensationalizing", + "sensitise": "sensitize", + "sensitised": "sensitized", + "sensitises": "sensitizes", + "sensitising": "sensitizing", + "sentimentalise": "sentimentalize", + "sentimentalised": "sentimentalized", + "sentimentalises": "sentimentalizes", + "sentimentalising": "sentimentalizing", + "sepulchre": "sepulcher", + "sepulchres": "sepulchers", + "serialisation": "serialization", + "serialisations": "serializations", + "serialise": "serialize", + "serialised": "serialized", + "serialises": "serializes", + "serialising": "serializing", + "sermonise": "sermonize", + "sermonised": "sermonized", + "sermonises": "sermonizes", + "sermonising": "sermonizing", + "sheikh": "sheik", + "shovelled": "shoveled", + "shovelling": "shoveling", + "shrivelled": "shriveled", + "shrivelling": "shriveling", + "signalise": "signalize", + "signalised": "signalized", + "signalises": "signalizes", + "signalising": "signalizing", + "signalled": "signaled", + "signalling": "signaling", + "smoulder": "smolder", + "smouldered": "smoldered", + "smouldering": "smoldering", + "smoulders": "smolders", + "snivelled": "sniveled", + "snivelling": "sniveling", + "snorkelled": "snorkeled", + "snorkelling": "snorkeling", + "snowplough": "snowplow", + "snowploughs": "snowplows", + "socialisation": "socialization", + "socialise": "socialize", + "socialised": "socialized", + "socialises": "socializes", + "socialising": "socializing", + "sodomise": "sodomize", + "sodomised": "sodomized", + "sodomises": "sodomizes", + "sodomising": "sodomizing", + "solemnise": "solemnize", + "solemnised": "solemnized", + "solemnises": "solemnizes", + "solemnising": "solemnizing", + "sombre": "somber", + "specialisation": "specialization", + "specialisations": "specializations", + "specialise": "specialize", + "specialised": "specialized", + "specialises": "specializes", + "specialising": "specializing", + "spectre": "specter", + "spectres": "specters", + "spiralled": "spiraled", + "spiralling": "spiraling", + "splendour": "splendor", + "splendours": "splendors", + "squirrelled": "squirreled", + "squirrelling": "squirreling", + "stabilisation": "stabilization", + "stabilise": "stabilize", + "stabilised": "stabilized", + "stabiliser": "stabilizer", + "stabilisers": "stabilizers", + "stabilises": "stabilizes", + "stabilising": "stabilizing", + "standardisation": "standardization", + "standardise": "standardize", + "standardised": "standardized", + "standardises": "standardizes", + "standardising": "standardizing", + "stencilled": "stenciled", + "stencilling": "stenciling", + "sterilisation": "sterilization", + "sterilisations": "sterilizations", + "sterilise": "sterilize", + "sterilised": "sterilized", + "steriliser": "sterilizer", + "sterilisers": "sterilizers", + "sterilises": "sterilizes", + "sterilising": "sterilizing", + "stigmatisation": "stigmatization", + "stigmatise": "stigmatize", + "stigmatised": "stigmatized", + "stigmatises": "stigmatizes", + "stigmatising": "stigmatizing", + "storey": "story", + "storeys": "stories", + "subsidisation": "subsidization", + "subsidise": "subsidize", + "subsidised": "subsidized", + "subsidiser": "subsidizer", + "subsidisers": "subsidizers", + "subsidises": "subsidizes", + "subsidising": "subsidizing", + "succour": "succor", + "succoured": "succored", + "succouring": "succoring", + "succours": "succors", + "sulphate": "sulfate", + "sulphates": "sulfates", + "sulphide": "sulfide", + "sulphides": "sulfides", + "sulphur": "sulfur", + "sulphurous": "sulfurous", + "summarise": "summarize", + "summarised": "summarized", + "summarises": "summarizes", + "summarising": "summarizing", + "swivelled": "swiveled", + "swivelling": "swiveling", + "symbolise": "symbolize", + "symbolised": "symbolized", + "symbolises": "symbolizes", + "symbolising": "symbolizing", + "sympathise": "sympathize", + "sympathised": "sympathized", + "sympathiser": "sympathizer", + "sympathisers": "sympathizers", + "sympathises": "sympathizes", + "sympathising": "sympathizing", + "synchronisation": "synchronization", + "synchronise": "synchronize", + "synchronised": "synchronized", + "synchronises": "synchronizes", + "synchronising": "synchronizing", + "synthesise": "synthesize", + "synthesised": "synthesized", + "synthesiser": "synthesizer", + "synthesisers": "synthesizers", + "synthesises": "synthesizes", + "synthesising": "synthesizing", + "syphon": "siphon", + "syphoned": "siphoned", + "syphoning": "siphoning", + "syphons": "siphons", + "systematisation": "systematization", + "systematise": "systematize", + "systematised": "systematized", + "systematises": "systematizes", + "systematising": "systematizing", + "tantalise": "tantalize", + "tantalised": "tantalized", + "tantalises": "tantalizes", + "tantalising": "tantalizing", + "tantalisingly": "tantalizingly", + "tasselled": "tasseled", + "technicolour": "technicolor", + "temporise": "temporize", + "temporised": "temporized", + "temporises": "temporizes", + "temporising": "temporizing", + "tenderise": "tenderize", + "tenderised": "tenderized", + "tenderises": "tenderizes", + "tenderising": "tenderizing", + "terrorise": "terrorize", + "terrorised": "terrorized", + "terrorises": "terrorizes", + "terrorising": "terrorizing", + "theatre": "theater", + "theatregoer": "theatergoer", + "theatregoers": "theatergoers", + "theatres": "theaters", + "theorise": "theorize", + "theorised": "theorized", + "theorises": "theorizes", + "theorising": "theorizing", + "tonne": "ton", + "tonnes": "tons", + "towelled": "toweled", + "towelling": "toweling", + "toxaemia": "toxemia", + "tranquillise": "tranquilize", + "tranquillised": "tranquilized", + "tranquilliser": "tranquilizer", + "tranquillisers": "tranquilizers", + "tranquillises": "tranquilizes", + "tranquillising": "tranquilizing", + "tranquillity": "tranquility", + "tranquillize": "tranquilize", + "tranquillized": "tranquilized", + "tranquillizer": "tranquilizer", + "tranquillizers": "tranquilizers", + "tranquillizes": "tranquilizes", + "tranquillizing": "tranquilizing", + "tranquilly": "tranquility", + "transistorised": "transistorized", + "traumatise": "traumatize", + "traumatised": "traumatized", + "traumatises": "traumatizes", + "traumatising": "traumatizing", + "travelled": "traveled", + "traveller": "traveler", + "travellers": "travelers", + "travelling": "traveling", + "travelogue": "travelog", + "travelogues": "travelogs", + "trialled": "trialed", + "trialling": "trialing", + "tricolour": "tricolor", + "tricolours": "tricolors", + "trivialise": "trivialize", + "trivialised": "trivialized", + "trivialises": "trivializes", + "trivialising": "trivializing", + "tumour": "tumor", + "tumours": "tumors", + "tunnelled": "tunneled", + "tunnelling": "tunneling", + "tyrannise": "tyrannize", + "tyrannised": "tyrannized", + "tyrannises": "tyrannizes", + "tyrannising": "tyrannizing", + "tyre": "tire", + "tyres": "tires", + "unauthorised": "unauthorized", + "uncivilised": "uncivilized", + "underutilised": "underutilized", + "unequalled": "unequaled", + "unfavourable": "unfavorable", + "unfavourably": "unfavorably", + "unionisation": "unionization", + "unionise": "unionize", + "unionised": "unionized", + "unionises": "unionizes", + "unionising": "unionizing", + "unorganised": "unorganized", + "unravelled": "unraveled", + "unravelling": "unraveling", + "unrecognisable": "unrecognizable", + "unrecognised": "unrecognized", + "unrivalled": "unrivaled", + "unsavoury": "unsavory", + "untrammelled": "untrammeled", + "urbanisation": "urbanization", + "urbanise": "urbanize", + "urbanised": "urbanized", + "urbanises": "urbanizes", + "urbanising": "urbanizing", + "utilisable": "utilizable", + "utilisation": "utilization", + "utilise": "utilize", + "utilised": "utilized", + "utilises": "utilizes", + "utilising": "utilizing", + "valour": "valor", + "vandalise": "vandalize", + "vandalised": "vandalized", + "vandalises": "vandalizes", + "vandalising": "vandalizing", + "vaporisation": "vaporization", + "vaporise": "vaporize", + "vaporised": "vaporized", + "vaporises": "vaporizes", + "vaporising": "vaporizing", + "vapour": "vapor", + "vapours": "vapors", + "verbalise": "verbalize", + "verbalised": "verbalized", + "verbalises": "verbalizes", + "verbalising": "verbalizing", + "victimisation": "victimization", + "victimise": "victimize", + "victimised": "victimized", + "victimises": "victimizes", + "victimising": "victimizing", + "videodisc": "videodisk", + "videodiscs": "videodisks", + "vigour": "vigor", + "visualisation": "visualization", + "visualisations": "visualizations", + "visualise": "visualize", + "visualised": "visualized", + "visualises": "visualizes", + "visualising": "visualizing", + "vocalisation": "vocalization", + "vocalisations": "vocalizations", + "vocalise": "vocalize", + "vocalised": "vocalized", + "vocalises": "vocalizes", + "vocalising": "vocalizing", + "vulcanised": "vulcanized", + "vulgarisation": "vulgarization", + "vulgarise": "vulgarize", + "vulgarised": "vulgarized", + "vulgarises": "vulgarizes", + "vulgarising": "vulgarizing", + "waggon": "wagon", + "waggons": "wagons", + "watercolour": "watercolor", + "watercolours": "watercolors", + "weaselled": "weaseled", + "weaselling": "weaseling", + "westernisation": "westernization", + "westernise": "westernize", + "westernised": "westernized", + "westernises": "westernizes", + "westernising": "westernizing", + "womanise": "womanize", + "womanised": "womanized", + "womaniser": "womanizer", + "womanisers": "womanizers", + "womanises": "womanizes", + "womanising": "womanizing", + "woollen": "woolen", + "woollens": "woolens", + "woollies": "woolies", + "woolly": "wooly", + "worshipped": "worshiped", + "worshipper": "worshiper", + "worshipping": "worshiping", + "yodelled": "yodeled", + "yodelling": "yodeling", + "yoghourt": "yogurt", + "yoghourts": "yogurts", + "yoghurt": "yogurt", + "yoghurts": "yogurts", +} diff --git a/veeksha/evaluator/performance/asr_normalizer/normalizer.py b/veeksha/evaluator/performance/asr_normalizer/normalizer.py new file mode 100644 index 00000000..b20734b1 --- /dev/null +++ b/veeksha/evaluator/performance/asr_normalizer/normalizer.py @@ -0,0 +1,738 @@ +# Adapted from https://github.com/huggingface/open_asr_leaderboard +# (normalizer/normalizer.py), the standard text normalizer used to score the +# Open ASR Leaderboard. Vendored verbatim so WER matches the leaderboard. +# +# Copyright 2022 The OpenAI team and The HuggingFace Team. All rights reserved. +# Most of the code is copy pasted from the original whisper repository +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +import unicodedata +from fractions import Fraction +from typing import Iterator, List, Match, Optional, Union + +import regex + +from .english_abbreviations import english_spelling_normalizer + +# non-ASCII letters that are not separated by "NFKD" normalization +ADDITIONAL_DIACRITICS = { + "Ε“": "oe", + "Ε’": "OE", + "ΓΈ": "o", + "Ø": "O", + "Γ¦": "ae", + "Γ†": "AE", + "ß": "ss", + "ẞ": "SS", + "Δ‘": "d", + "Đ": "D", + "Γ°": "d", + "Ð": "D", + "ΓΎ": "th", + "Þ": "th", + "Ε‚": "l", + "Ł": "L", +} + + +class _SymbolReplacementTable(dict): + """Lazy per-codepoint replacement table for ``str.translate``. + + Computes the same per-character replacement as the original character + loop and memoizes it, so repeated normalizer calls classify each distinct + codepoint only once. + """ + + def __init__(self, keep: str): + super().__init__() + self._keep = keep + + def __missing__(self, codepoint: int) -> str: + char = chr(codepoint) + if char in self._keep: + replacement = char + elif char in ADDITIONAL_DIACRITICS: + replacement = ADDITIONAL_DIACRITICS[char] + elif unicodedata.category(char) == "Mn": + replacement = "" + elif unicodedata.category(char)[0] in "MSP": + replacement = " " + else: + replacement = char + self[codepoint] = replacement + return replacement + + +_symbol_and_diacritic_tables: dict = {} + + +class _SymbolOnlyTable(dict): + """Lazy per-codepoint table replacing 'MSP' categories with a space.""" + + def __missing__(self, codepoint: int) -> str: + char = chr(codepoint) + replacement = " " if unicodedata.category(char)[0] in "MSP" else char + self[codepoint] = replacement + return replacement + + +_symbol_only_table = _SymbolOnlyTable() + + +def remove_symbols_and_diacritics(s: str, keep=""): + """ + Replace any other markers, symbols, and punctuations with a space, and drop any diacritics (category 'Mn' and some + manual mappings) + """ + table = _symbol_and_diacritic_tables.get(keep) + if table is None: + table = _symbol_and_diacritic_tables[keep] = _SymbolReplacementTable(keep) + return unicodedata.normalize("NFKD", s).translate(table) + + +def remove_symbols(s: str): + """ + Replace any other markers, symbols, punctuations with a space, keeping diacritics + """ + return unicodedata.normalize("NFKC", s).translate(_symbol_only_table) + + +class BasicTextNormalizer: + def __init__(self, remove_diacritics: bool = False, split_letters: bool = False): + self.clean = ( + remove_symbols_and_diacritics if remove_diacritics else remove_symbols + ) + self.split_letters = split_letters + + def __call__(self, s: str): + s = s.lower() + s = re.sub(r"[<\[][^>\]]*[>\]]", "", s) # remove words between brackets + s = re.sub(r"\(([^)]+?)\)", "", s) # remove words between parenthesis + s = self.clean(s).lower() + + if self.split_letters: + s = " ".join(regex.findall(r"\X", s, regex.U)) + + s = re.sub( + r"\s+", " ", s + ) # replace any successive whitespace characters with a space + + return s + + +class BasicMultilingualTextNormalizer: + def __init__(self, remove_diacritics: bool = True): + self.clean = ( + remove_symbols_and_diacritics if remove_diacritics else remove_symbols + ) + + def __call__(self, s: str): + s = s.lower() + s = re.sub(r"[<\[][^>\]]*[>\]]", "", s) # remove words between brackets + s = re.sub(r"\(([^)]+?)\)", "", s) # remove words between parenthesis + s = self.clean(s).lower() + + # Remove punctuations and extra spaces + s = re.sub(r"[^\w\s]", "", s) + s = re.sub(r"\s+", " ", s).strip() + + return s + + +_NUMERIC_RE = re.compile(r"^\d+(\.\d+)?$") +_AND_A_HALF_RE = re.compile(r"\band\s+a\s+half\b") +_LETTER_DIGIT_RE = re.compile(r"([a-z])([0-9])") +_DIGIT_LETTER_RE = re.compile(r"([0-9])([a-z])") +_DIGIT_SUFFIX_SPACE_RE = re.compile(r"([0-9])\s+(st|nd|rd|th|s)\b") +_CURRENCY_CENTS_RE = re.compile(r"([€£$])([0-9]+) (?:and )?Β’([0-9]{1,2})\b") +_ZERO_CENTS_RE = re.compile(r"[€£$]0.([0-9]{1,2})\b") +_ONE_READABILITY_RE = re.compile(r"\b1(s?)\b") +_BRACKET_CONTENT_RE = re.compile(r"[<\[][^>\]]*[>\]]") +_PAREN_CONTENT_RE = re.compile(r"\(([^)]+?)\)") +_SPACE_BEFORE_APOSTROPHE_RE = re.compile(r"\s+'") +_DIGIT_COMMA_RE = re.compile(r"(\d),(\d)") +_PERIOD_NON_DIGIT_RE = re.compile(r"\.([^0-9]|$)") +_SYMBOL_NON_DIGIT_RE = re.compile(r"[.$’€£]([^0-9])") +_NON_DIGIT_PERCENT_RE = re.compile(r"([^0-9])%") +_WHITESPACE_RE = re.compile(r"\s+") + + +class EnglishNumberNormalizer: + """ + Convert any spelled-out numbers into arabic numbers, while handling: + + - remove any commas + - keep the suffixes such as: `1960s`, `274th`, `32nd`, etc. + - spell out currency symbols after the number. e.g. `$20 million` -> `20000000 dollars` + - spell out `one` and `ones` + - interpret successive single-digit numbers as nominal: `one oh one` -> `101` + """ + + def __init__(self): + super().__init__() + + self.zeros = {"o", "oh", "zero"} + # fmt: off + self.ones = { + name: i + for i, name in enumerate( + ["one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten", "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen", "seventeen", "eighteen", "nineteen"], + start=1, + ) + } + # fmt: on + self.ones_plural = { + "sixes" if name == "six" else name + "s": (value, "s") + for name, value in self.ones.items() + } + self.ones_ordinal = { + "zeroth": (0, "th"), + "first": (1, "st"), + "second": (2, "nd"), + "third": (3, "rd"), + "fifth": (5, "th"), + "twelfth": (12, "th"), + **{ + name + ("h" if name.endswith("t") else "th"): (value, "th") + for name, value in self.ones.items() + if value > 3 and value != 5 and value != 12 + }, + } + self.ones_suffixed = {**self.ones_plural, **self.ones_ordinal} + + self.tens = { + "twenty": 20, + "thirty": 30, + "forty": 40, + "fifty": 50, + "sixty": 60, + "seventy": 70, + "eighty": 80, + "ninety": 90, + } + self.tens_plural = { + name.replace("y", "ies"): (value, "s") for name, value in self.tens.items() + } + self.tens_ordinal = { + name.replace("y", "ieth"): (value, "th") + for name, value in self.tens.items() + } + self.tens_suffixed = {**self.tens_plural, **self.tens_ordinal} + + self.multipliers = { + "hundred": 100, + "thousand": 1_000, + "million": 1_000_000, + "billion": 1_000_000_000, + "trillion": 1_000_000_000_000, + "quadrillion": 1_000_000_000_000_000, + "quintillion": 1_000_000_000_000_000_000, + "sextillion": 1_000_000_000_000_000_000_000, + "septillion": 1_000_000_000_000_000_000_000_000, + "octillion": 1_000_000_000_000_000_000_000_000_000, + "nonillion": 1_000_000_000_000_000_000_000_000_000_000, + "decillion": 1_000_000_000_000_000_000_000_000_000_000_000, + } + self.multipliers_plural = { + name + "s": (value, "s") for name, value in self.multipliers.items() + } + self.multipliers_ordinal = { + name + "th": (value, "th") for name, value in self.multipliers.items() + } + self.multipliers_suffixed = { + **self.multipliers_plural, + **self.multipliers_ordinal, + } + self.decimals = {*self.ones, *self.tens, *self.zeros} + + self.preceding_prefixers = { + "minus": "-", + "negative": "-", + "plus": "+", + "positive": "+", + } + self.following_prefixers = { + "pound": "Β£", + "pounds": "Β£", + "euro": "€", + "euros": "€", + "dollar": "$", + "dollars": "$", + "cent": "Β’", + "cents": "Β’", + } + self.prefixes = set( + list(self.preceding_prefixers.values()) + + list(self.following_prefixers.values()) + ) + self.suffixers = { + "per": {"cent": "%"}, + "percent": "%", + } + self.specials = {"and", "double", "triple", "point"} + + self.words = { + key + for mapping in [ + self.zeros, + self.ones, + self.ones_suffixed, + self.tens, + self.tens_suffixed, + self.multipliers, + self.multipliers_suffixed, + self.preceding_prefixers, + self.following_prefixers, + self.suffixers, + self.specials, + ] + for key in mapping + } + self.literal_words = {"one", "ones"} + + def process_words(self, words: List[str]) -> Iterator[str]: + prefix: Optional[str] = None + value: Optional[Union[str, int]] = None + skip = False + + def to_fraction(s: str): + try: + return Fraction(s) + except ValueError: + return None + + def output(result: Union[str, int]): + nonlocal prefix, value + result = str(result) + if prefix is not None: + result = prefix + result + value = None + prefix = None + return result + + if len(words) == 0: + return + + vocabulary = self.words + prefixes = self.prefixes + for i, current in enumerate(words): + if skip: + skip = False + continue + + # Fast path: a word with no leading digit/currency sign that is not + # part of the number vocabulary passes through unchanged whenever no + # number is being accumulated. This mirrors the + # `current not in self.words` branch below with `value`/`prefix` + # both unset, where `output(current)` returns `current` verbatim. + first_char = current[0] + if ( + value is None + and prefix is None + and current not in vocabulary + and not first_char.isdigit() + and first_char not in prefixes + ): + yield current + continue + + prev = words[i - 1] if i != 0 else None + next = words[i + 1] if i != len(words) - 1 else None + + next_is_numeric = next is not None and _NUMERIC_RE.match(next) + has_prefix = first_char in prefixes + current_without_prefix = current[1:] if has_prefix else current + if _NUMERIC_RE.match(current_without_prefix): + # arabic numbers (potentially with signs and fractions) + f = to_fraction(current_without_prefix) + if f is None: + raise ValueError("Converting the fraction failed") + + if value is not None: + if isinstance(value, str) and value.endswith("."): + # concatenate decimals / ip address components + value = str(value) + str(current) + continue + else: + yield output(value) + + prefix = current[0] if has_prefix else prefix + if f.denominator == 1: + value = f.numerator # store integers as int + else: + value = current_without_prefix + elif current not in self.words: + # non-numeric words + if value is not None: + yield output(value) + yield output(current) + elif current in self.zeros: + value = str(value or "") + "0" + elif current in self.ones: + ones = self.ones[current] + + if value is None: + value = ones + elif isinstance(value, str) or prev in self.ones: + if ( + prev in self.tens and ones < 10 + ): # replace the last zero with the digit + value = value[:-1] + str(ones) + else: + value = str(value) + str(ones) + elif ones < 10: + if value % 10 == 0: + value += ones + else: + value = str(value) + str(ones) + else: # eleven to nineteen + if value % 100 == 0: + value += ones + else: + value = str(value) + str(ones) + elif current in self.ones_suffixed: + # ordinal or cardinal; yield the number right away + ones, suffix = self.ones_suffixed[current] + if value is None: + yield output(str(ones) + suffix) + elif isinstance(value, str) or prev in self.ones: + if prev in self.tens and ones < 10: + yield output(value[:-1] + str(ones) + suffix) + else: + yield output(str(value) + str(ones) + suffix) + elif ones < 10: + if value % 10 == 0: + yield output(str(value + ones) + suffix) + else: + yield output(str(value) + str(ones) + suffix) + else: # eleven to nineteen + if value % 100 == 0: + yield output(str(value + ones) + suffix) + else: + yield output(str(value) + str(ones) + suffix) + value = None + elif current in self.tens: + tens = self.tens[current] + if value is None: + value = tens + elif isinstance(value, str): + value = str(value) + str(tens) + else: + if value % 100 == 0: + value += tens + else: + value = str(value) + str(tens) + elif current in self.tens_suffixed: + # ordinal or cardinal; yield the number right away + tens, suffix = self.tens_suffixed[current] + if value is None: + yield output(str(tens) + suffix) + elif isinstance(value, str): + yield output(str(value) + str(tens) + suffix) + else: + if value % 100 == 0: + yield output(str(value + tens) + suffix) + else: + yield output(str(value) + str(tens) + suffix) + elif current in self.multipliers: + multiplier = self.multipliers[current] + if value is None: + value = multiplier + elif isinstance(value, str) or value == 0: + f = to_fraction(value) + p = f * multiplier if f is not None else None + if f is not None and p.denominator == 1: + value = p.numerator + else: + yield output(value) + value = multiplier + else: + before = value // 1000 * 1000 + residual = value % 1000 + value = before + residual * multiplier + elif current in self.multipliers_suffixed: + multiplier, suffix = self.multipliers_suffixed[current] + if value is None: + yield output(str(multiplier) + suffix) + elif isinstance(value, str): + f = to_fraction(value) + p = f * multiplier if f is not None else None + if f is not None and p.denominator == 1: + yield output(str(p.numerator) + suffix) + else: + yield output(value) + yield output(str(multiplier) + suffix) + else: # int + before = value // 1000 * 1000 + residual = value % 1000 + value = before + residual * multiplier + yield output(str(value) + suffix) + value = None + elif current in self.preceding_prefixers: + # apply prefix (positive, minus, etc.) if it precedes a number + if value is not None: + yield output(value) + + if next in self.words or next_is_numeric: + prefix = self.preceding_prefixers[current] + else: + yield output(current) + elif current in self.following_prefixers: + # apply prefix (dollars, cents, etc.) only after a number + if value is not None: + prefix = self.following_prefixers[current] + yield output(value) + else: + yield output(current) + elif current in self.suffixers: + # apply suffix symbols (percent -> '%') + if value is not None: + suffix = self.suffixers[current] + if isinstance(suffix, dict): + if next in suffix: + yield output(str(value) + suffix[next]) + skip = True + else: + yield output(value) + yield output(current) + else: + yield output(str(value) + suffix) + else: + yield output(current) + elif current in self.specials: + if next not in self.words and not next_is_numeric: + # apply special handling only if the next word can be numeric + if value is not None: + yield output(value) + yield output(current) + elif current == "and": + # ignore "and" after hundreds, thousands, etc. + if prev not in self.multipliers: + if value is not None: + yield output(value) + yield output(current) + elif current == "double" or current == "triple": + if next in self.ones or next in self.zeros: + repeats = 2 if current == "double" else 3 + ones = self.ones.get(next, 0) + value = str(value or "") + str(ones) * repeats + skip = True + else: + if value is not None: + yield output(value) + yield output(current) + elif current == "point": + if next in self.decimals or next_is_numeric: + value = str(value or "") + "." + else: + # should all have been covered at this point + raise ValueError(f"Unexpected token: {current}") + else: + # all should have been covered at this point + raise ValueError(f"Unexpected token: {current}") + + if value is not None: + yield output(value) + + def preprocess(self, s: str): + # replace " and a half" with " point five" + results = [] + + segments = _AND_A_HALF_RE.split(s) + for i, segment in enumerate(segments): + if len(segment.strip()) == 0: + continue + if i == len(segments) - 1: + results.append(segment) + else: + results.append(segment) + last_word = segment.rsplit(maxsplit=2)[-1] + if last_word in self.decimals or last_word in self.multipliers: + results.append("point five") + else: + results.append("and a half") + + s = " ".join(results) + + # put a space at number/letter boundary + s = _LETTER_DIGIT_RE.sub(r"\1 \2", s) + s = _DIGIT_LETTER_RE.sub(r"\1 \2", s) + + # but remove spaces which could be a suffix + s = _DIGIT_SUFFIX_SPACE_RE.sub(r"\1\2", s) + + return s + + def postprocess(self, s: str): + def combine_cents(m: Match): + try: + currency = m.group(1) + integer = m.group(2) + cents = int(m.group(3)) + return f"{currency}{integer}.{cents:02d}" + except ValueError: + return m.string + + def extract_cents(m: Match): + try: + return f"Β’{int(m.group(1))}" + except ValueError: + return m.string + + # apply currency postprocessing; "$2 and Β’7" -> "$2.07" + if "Β’" in s: + s = _CURRENCY_CENTS_RE.sub(combine_cents, s) + if "€" in s or "Β£" in s or "$" in s: + s = _ZERO_CENTS_RE.sub(extract_cents, s) + + # write "one(s)" instead of "1(s)", just for the readability + if "1" in s: + s = _ONE_READABILITY_RE.sub(r"one\1", s) + + return s + + def __call__(self, s: str): + s = self.preprocess(s) + s = " ".join(word for word in self.process_words(s.split()) if word is not None) + s = self.postprocess(s) + + return s + + +class EnglishSpellingNormalizer: + """ + Applies British-American spelling mappings as listed in [1]. + + [1] https://www.tysto.com/uk-us-spelling-list.html + """ + + def __init__(self, english_spelling_mapping): + self.mapping = english_spelling_mapping + + def __call__(self, s: str): + return " ".join(self.mapping.get(word, word) for word in s.split()) + + +class EnglishTextNormalizer: + def __init__(self, english_spelling_mapping=english_spelling_normalizer): + self.ignore_patterns = r"\b(hmm|mm|mhm|mmm|uh|um)\b" + self.replacers = { + # common contractions + r"\bwon't\b": "will not", + r"\bcan't\b": "can not", + r"\blet's\b": "let us", + r"\bain't\b": "aint", + r"\by'all\b": "you all", + r"\bwanna\b": "want to", + r"\bgotta\b": "got to", + r"\bgonna\b": "going to", + r"\bi'ma\b": "i am going to", + r"\bimma\b": "i am going to", + r"\bwoulda\b": "would have", + r"\bcoulda\b": "could have", + r"\bshoulda\b": "should have", + r"\bma'am\b": "madam", + # contractions in titles/prefixes + r"\bmr\b": "mister ", + r"\bmrs\b": "missus ", + r"\bst\b": "saint ", + r"\bdr\b": "doctor ", + r"\bprof\b": "professor ", + r"\bcapt\b": "captain ", + r"\bgov\b": "governor ", + r"\bald\b": "alderman ", + r"\bgen\b": "general ", + r"\bsen\b": "senator ", + r"\brep\b": "representative ", + r"\bpres\b": "president ", + r"\brev\b": "reverend ", + r"\bhon\b": "honorable ", + r"\basst\b": "assistant ", + r"\bassoc\b": "associate ", + r"\blt\b": "lieutenant ", + r"\bcol\b": "colonel ", + r"\bjr\b": "junior ", + r"\bsr\b": "senior ", + r"\besq\b": "esquire ", + # prefect tenses, ideally it should be any past participles, but it's harder.. + r"'d been\b": " had been", + r"'s been\b": " has been", + r"'d gone\b": " had gone", + r"'s gone\b": " has gone", + r"'d done\b": " had done", # "'s done" is ambiguous + r"'s got\b": " has got", + # general contractions + r"n't\b": " not", + r"'re\b": " are", + r"'s\b": " is", + r"'d\b": " would", + r"'ll\b": " will", + r"'t\b": " not", + r"'ve\b": " have", + r"'m\b": " am", + } + self.standardize_numbers = EnglishNumberNormalizer() + self.standardize_spellings = EnglishSpellingNormalizer(english_spelling_mapping) + + self._ignore_patterns_re = re.compile(self.ignore_patterns) + # Every replacer pattern is a literal decorated only with zero-width + # \b anchors, so any match must contain the bare literal as a + # substring. A C-speed containment check therefore lets us skip the + # regex pass entirely when it could not match (sub with no match is + # the identity). + self._compiled_replacers = [ + (pattern.replace("\\b", ""), re.compile(pattern), replacement) + for pattern, replacement in self.replacers.items() + ] + + def __call__(self, s: str): + s = s.lower() + + if "[" in s or "<" in s: + s = _BRACKET_CONTENT_RE.sub("", s) # remove words between brackets + if "(" in s: + s = _PAREN_CONTENT_RE.sub("", s) # remove words between parenthesis + s = self._ignore_patterns_re.sub("", s) + if "'" in s: + s = _SPACE_BEFORE_APOSTROPHE_RE.sub( + "'", s + ) # standardize when there's a space before an apostrophe + + for literal, pattern, replacement in self._compiled_replacers: + if literal in s: + s = pattern.sub(replacement, s) + + if "," in s: + s = _DIGIT_COMMA_RE.sub(r"\1\2", s) # remove commas between digits + if "." in s: + s = _PERIOD_NON_DIGIT_RE.sub( + r" \1", s + ) # remove periods not followed by numbers + s = remove_symbols_and_diacritics( + s, keep=".%$’€£" + ) # keep some symbols for numerics + + s = self.standardize_numbers(s) + s = self.standardize_spellings(s) + + # now remove prefix/suffix symbols that are not preceded/followed by numbers + s = _SYMBOL_NON_DIGIT_RE.sub(r" \1", s) + s = _NON_DIGIT_PERCENT_RE.sub(r"\1 ", s) + + s = _WHITESPACE_RE.sub( + " ", s + ) # replace any successive whitespace characters with a space + + return s diff --git a/veeksha/evaluator/performance/audio.py b/veeksha/evaluator/performance/audio.py index 09c1a256..f92931d6 100644 --- a/veeksha/evaluator/performance/audio.py +++ b/veeksha/evaluator/performance/audio.py @@ -1,33 +1,151 @@ -from typing import Any, Dict, Optional +"""Performance evaluator for TTS, realtime TTS, LLM-audio, and STT metrics.""" + +import json +import os +import threading +import time +from dataclasses import dataclass +from typing import Any from veeksha.config.evaluator import ( AudioChannelPerformanceConfig, PerformanceEvaluatorConfig, ) +from veeksha.core.audio_contract import ( + DEFAULT_AUDIO_SAMPLE_RATE, + WAV_HEADER_BYTES, + AudioMetricKey, + pcm_bytes_to_duration_ms, +) from veeksha.evaluator.base import EvaluationResult +from veeksha.evaluator.cdf_sketch import CDFSketch +from veeksha.evaluator.performance.asr import ( + ASRMetricAccumulator, + ASRRequestMetrics, + score_asr_request, +) +from veeksha.evaluator.performance.audio_interactivity import ( + InteractivityMetrics, + RequestTiming, + compute_interactivity_metrics, + parse_request_timing, +) from veeksha.logger import init_logger -from veeksha.types import ChannelModality +from veeksha.types import AudioTask, ChannelModality logger = init_logger(__name__) -class AudioPerformanceEvaluator: - """Performance evaluator for audio generation (skeleton). +def _pcm_byte_count(total_bytes: int, *, raw_pcm: bool) -> int: + if raw_pcm: + return total_bytes + return max(total_bytes - WAV_HEADER_BYTES, 0) + + +def _policy_tag(value: float) -> str: + if float(value).is_integer(): + return str(int(value)) + return str(value).replace(".", "p") - TODO: Implement audio-specific metrics - """ + +@dataclass +class AudioRequestMetrics: + """Metrics for a single audio request.""" + + request_id: int + session_id: int + request_dispatched_at: float + client_completed_at: float + ttfc: float + end_to_end_latency: float + generated_audio_duration: float + rtf: float + chunk_count: int + pcm_byte_count: int + input_chars: int = 0 + audio_task: AudioTask | None = None + asr: ASRRequestMetrics | None = None + input_tokens: int = 0 + input_text: str = "" + provider: str = "" + provider_model: str = "" + provider_protocol: str = "" + session_total_requests: int | None = None + interactivity: InteractivityMetrics | None = None + ws_connect_latency_ms: float | None = None + + +class AudioPerformanceEvaluator: + """Evaluate TTS interactivity and STT recognition performance together.""" def __init__( self, config: PerformanceEvaluatorConfig, - channel_config: Optional[AudioChannelPerformanceConfig] = None, + channel_config: AudioChannelPerformanceConfig | None = None, + benchmark_start_time: float = 0.0, ): self.config = config self.channel_config = channel_config or AudioChannelPerformanceConfig() - logger.warning( - "AudioPerformanceEvaluator is a skeleton implementation. " - "Audio metrics are not yet supported." - ) + self.benchmark_start_time = benchmark_start_time + self.lock = threading.Lock() + + self.summaries: dict[str, CDFSketch] = { + AudioMetricKey.TTFC.value: CDFSketch(AudioMetricKey.TTFC.value, unit="ms"), + AudioMetricKey.END_TO_END_LATENCY.value: CDFSketch( + AudioMetricKey.END_TO_END_LATENCY.value, unit="ms" + ), + AudioMetricKey.GENERATED_AUDIO_DURATION.value: CDFSketch( + AudioMetricKey.GENERATED_AUDIO_DURATION.value, unit="ms" + ), + AudioMetricKey.RTF.value: CDFSketch(AudioMetricKey.RTF.value), + AudioMetricKey.CHUNK_COUNT.value: CDFSketch( + AudioMetricKey.CHUNK_COUNT.value + ), + AudioMetricKey.INPUT_TOKENS.value: CDFSketch( + AudioMetricKey.INPUT_TOKENS.value, unit="tokens" + ), + AudioMetricKey.SESSION_SIZE.value: CDFSketch( + AudioMetricKey.SESSION_SIZE.value + ), + AudioMetricKey.SESSION_DURATION.value: CDFSketch( + AudioMetricKey.SESSION_DURATION.value, unit="ms" + ), + } + + self.asr_latency_summaries: dict[str, CDFSketch] = { + "time_to_first_visible_text": CDFSketch( + "time_to_first_visible_text", unit="ms" + ), + "time_to_first_partial": CDFSketch("time_to_first_partial", unit="ms"), + "time_to_final_transcript": CDFSketch( + "time_to_final_transcript", unit="ms" + ), + "interactivity": CDFSketch("interactivity", unit="ms"), + } + self._asr_metrics = ASRMetricAccumulator() + self._stt_request_count = 0 + self._asr_scoring_seconds = 0.0 + + self._pending_requests: dict[int, dict[str, Any]] = {} + self._completed_metrics: list[AudioRequestMetrics] = [] + self._lifecycle_timestamps: list[dict[str, float | None]] = [] + self._request_rows_streamed = 0 + self._request_time_reference = self.benchmark_start_time + + self._interactivity_sketches_ready = False + self._interactive_request_count = 0 + self._tts_service_fluidity_eligible_count = 0 + self._fixed_delay_stall_free_counts: dict[float, int] = {} + self._buffer_target_stall_free_counts: dict[float, int] = {} + self._buffer_target_eligible_counts: dict[float, int] = {} + self._fluidity_policy_sketches: dict[float, CDFSketch] = {} + self._raw_timing_rows: list[dict[str, Any]] = [] + self._raw_timing_rows_streamed = 0 + + self._total_input_chars = 0 + self._total_generated_audio_duration_ms = 0.0 + self._first_dispatch_at: float | None = None + self._last_completion_at: float | None = None def register_request( self, @@ -38,7 +156,13 @@ def register_request( requested_output: Any = None, ) -> None: """Register an audio request that was dispatched.""" - pass # Skeleton + with self.lock: + if self._request_time_reference == 0.0: + self._request_time_reference = dispatched_at + self._pending_requests[request_id] = { + "session_id": session_id, + "dispatched_at": dispatched_at, + } def record_request_completed( self, @@ -47,38 +171,844 @@ def record_request_completed( completed_at: float, response: Any, ) -> None: - """Record that an audio request completed.""" - pass # Skeleton + """Record a completed TTS, LLM-audio, or STT request. + + ASR scoring runs outside the evaluator lock so completion workers can + normalize and align transcripts concurrently. Shared metric mutation + remains under the evaluator lock. + """ + with self.lock: + dispatch_info = self._pending_requests.pop(request_id, None) + request_time_reference = self._request_time_reference + + dispatched_at = ( + dispatch_info["dispatched_at"] + if dispatch_info + else getattr( + response, + "scheduler_dispatched_at", + request_time_reference, + ) + ) + + channel_response = getattr(response, "channels", {}).get(ChannelModality.AUDIO) + if channel_response is None: + logger.debug("Request %d has no AUDIO channel response", request_id) + return + + cm = channel_response.metrics or {} + ttfc = float(cm.get(AudioMetricKey.TTFC.value, 0.0)) + end_to_end_latency = float(cm.get(AudioMetricKey.END_TO_END_LATENCY.value, 0.0)) + chunk_count = int(cm.get(AudioMetricKey.CHUNK_COUNT.value, 0)) + raw_pcm = bool(cm.get(AudioMetricKey.RAW_PCM.value, False)) + sample_rate = int( + cm.get(AudioMetricKey.SAMPLE_RATE.value, DEFAULT_AUDIO_SAMPLE_RATE) + ) + + raw_audio_task = cm.get("audio_task") + try: + audio_task = ( + AudioTask[raw_audio_task.upper()] + if isinstance(raw_audio_task, str) + else AudioTask(raw_audio_task) + ) + except (KeyError, TypeError, ValueError) as error: + expected = ", ".join(str(task) for task in AudioTask) + raise ValueError( + f"AUDIO response for request {request_id} has unknown " + f"audio_task={raw_audio_task!r}; expected one of: {expected}" + ) from error + + if audio_task is AudioTask.STT: + if AudioMetricKey.PCM_BYTE_COUNT.value not in cm: + raise ValueError( + f"STT response for request {request_id} is missing " + f"{AudioMetricKey.PCM_BYTE_COUNT.value}" + ) + total_bytes = int(cm[AudioMetricKey.PCM_BYTE_COUNT.value]) + else: + audio_content = channel_response.content + total_bytes = ( + len(audio_content) + if isinstance(audio_content, (bytes, bytearray, memoryview)) + else 0 + ) + + pcm_byte_count = _pcm_byte_count(total_bytes, raw_pcm=raw_pcm) + generated_audio_duration = pcm_bytes_to_duration_ms(pcm_byte_count, sample_rate) + rtf = ( + end_to_end_latency / generated_audio_duration + if generated_audio_duration > 0 + else float("inf") + ) + input_chars = int(cm.get(AudioMetricKey.INPUT_CHARS.value, 0) or 0) + input_tokens = int(cm.get(AudioMetricKey.INPUT_TOKENS.value, 0) or 0) + input_text = str(cm.get(AudioMetricKey.INPUT_TEXT.value, "")) + provider = str(cm.get(AudioMetricKey.PROVIDER.value, "")) + provider_model = str(cm.get(AudioMetricKey.PROVIDER_MODEL.value, "")) + provider_protocol = str(cm.get(AudioMetricKey.PROVIDER_PROTOCOL.value, "")) + ws_connect_latency = cm.get(AudioMetricKey.WS_CONNECT_LATENCY_MS.value) + ws_connect_latency_ms = ( + float(ws_connect_latency) if ws_connect_latency is not None else None + ) + interactivity: InteractivityMetrics | None = None + raw_timing_row: dict[str, Any] | None = None + if audio_task is AudioTask.TTS and self.channel_config.interactivity_enabled: + timing = parse_request_timing(cm, sample_rate) + if timing is not None: + interactivity = compute_interactivity_metrics( + timing, + startup_delay_ms_values=( + self.channel_config.startup_delay_ms_values + ), + startup_buffer_ms_values=( + self.channel_config.startup_buffer_ms_values + ), + min_reportable_stall_ms=( + self.channel_config.min_reportable_stall_ms + ), + fluidity_frame_ms=self.channel_config.fluidity_frame_ms, + fluidity_startup_delay_ms=( + self.channel_config.fluidity_startup_delay_ms + ), + fluidity_attribution_mode=( + self.channel_config.fluidity_attribution_mode + ), + ) + if self.channel_config.persist_raw_timing: + raw_timing_row = self._build_raw_timing_row( + request_id, session_id, cm, timing + ) + + asr_metrics: ASRRequestMetrics | None = None + asr_scoring_seconds = 0.0 + if audio_task is AudioTask.STT: + scoring_started_at = time.perf_counter() + asr_metrics = score_asr_request( + request_id=request_id, + channel_metrics=cm, + duration_s=generated_audio_duration / 1000.0, + accumulator=self._asr_metrics, + ) + asr_scoring_seconds = time.perf_counter() - scoring_started_at + + with self.lock: + if audio_task is AudioTask.STT: + self._stt_request_count += 1 + self._asr_scoring_seconds += asr_scoring_seconds + if raw_timing_row is not None: + self._raw_timing_rows.append(raw_timing_row) + + metrics = AudioRequestMetrics( + request_id=request_id, + session_id=session_id, + request_dispatched_at=dispatched_at, + client_completed_at=completed_at, + ttfc=ttfc, + end_to_end_latency=end_to_end_latency, + generated_audio_duration=generated_audio_duration, + rtf=rtf, + chunk_count=chunk_count, + pcm_byte_count=pcm_byte_count, + input_chars=input_chars, + input_tokens=input_tokens, + audio_task=audio_task, + asr=asr_metrics, + input_text=input_text, + provider=provider, + provider_model=provider_model, + provider_protocol=provider_protocol, + session_total_requests=getattr( + response, "session_total_requests", None + ), + interactivity=interactivity, + ws_connect_latency_ms=ws_connect_latency_ms, + ) + self._completed_metrics.append(metrics) + + self._total_input_chars += input_chars + self._total_generated_audio_duration_ms += generated_audio_duration + if ( + self._first_dispatch_at is None + or dispatched_at < self._first_dispatch_at + ): + self._first_dispatch_at = dispatched_at + if ( + self._last_completion_at is None + or completed_at > self._last_completion_at + ): + self._last_completion_at = completed_at + + def normalize_ts(timestamp: float | None) -> float | None: + if timestamp is None: + return None + return round(max(0.0, timestamp - self._request_time_reference), 5) + + self._lifecycle_timestamps.append( + { + "scheduler_ready_at": normalize_ts( + getattr(response, "scheduler_ready_at", None) + ), + "scheduler_dispatched_at": normalize_ts( + getattr(response, "scheduler_dispatched_at", None) + ), + "client_picked_up_at": normalize_ts( + getattr(response, "client_picked_up_at", None) + ), + "client_completed_at": normalize_ts( + getattr(response, "client_completed_at", None) + ), + "result_processed_at": normalize_ts( + getattr(response, "result_processed_at", None) + ), + } + ) + + self.summaries[AudioMetricKey.TTFC.value].put(ttfc) + self.summaries[AudioMetricKey.END_TO_END_LATENCY.value].put( + end_to_end_latency + ) + self.summaries[AudioMetricKey.GENERATED_AUDIO_DURATION.value].put( + generated_audio_duration + ) + self.summaries[AudioMetricKey.RTF.value].put(rtf) + self.summaries[AudioMetricKey.CHUNK_COUNT.value].put(chunk_count) + self.summaries[AudioMetricKey.INPUT_TOKENS.value].put(input_tokens) + + if asr_metrics is not None: + asr_values = { + "time_to_first_visible_text": ( + asr_metrics.time_to_first_visible_text + ), + "time_to_first_partial": asr_metrics.time_to_first_partial, + "time_to_final_transcript": (asr_metrics.time_to_final_transcript), + "interactivity": asr_metrics.interactivity, + } + for metric_name, value in asr_values.items(): + if value is not None: + self.asr_latency_summaries[metric_name].put(value) + + if interactivity is not None: + self._record_interactivity(metrics) + + def _ensure_interactivity_sketches(self) -> None: + if self._interactivity_sketches_ready: + return + specs: list[tuple[AudioMetricKey, str | None]] = [ + (AudioMetricKey.FIRST_INPUT_TO_FIRST_AUDIO_MS, "ms"), + (AudioMetricKey.FIRST_INPUT_TO_FIRST_PLAYABLE_AUDIO_MS, "ms"), + (AudioMetricKey.TRIGGER_TO_FIRST_PLAYABLE_AUDIO_MS, "ms"), + (AudioMetricKey.REQUEST_START_TO_FIRST_AUDIO_MS, "ms"), + (AudioMetricKey.REQUEST_START_TO_FIRST_PLAYABLE_AUDIO_MS, "ms"), + (AudioMetricKey.AUDIO_BEFORE_COMMIT_RATIO, None), + (AudioMetricKey.DUPLEX_OVERLAP_OBSERVED, None), + (AudioMetricKey.DUPLEX_OVERLAP_MS, "ms"), + (AudioMetricKey.POST_COMMIT_AUDIO_DELIVERY_MS, "ms"), + (AudioMetricKey.REQUIRED_STARTUP_DELAY_MS, "ms"), + (AudioMetricKey.ZERO_DELAY_STALL_COUNT, None), + (AudioMetricKey.ZERO_DELAY_TOTAL_STALL_MS, "ms"), + (AudioMetricKey.ZERO_DELAY_LONGEST_STALL_MS, "ms"), + (AudioMetricKey.STREAMING_RTF, None), + (AudioMetricKey.DONE_AFTER_LAST_AUDIO_MS, "ms"), + (AudioMetricKey.WS_CONNECT_LATENCY_MS, "ms"), + (AudioMetricKey.USER_AUDIO_FLUIDITY_INDEX, None), + (AudioMetricKey.TTS_SERVICE_FLUIDITY_INDEX, None), + (AudioMetricKey.TTS_SERVICE_FLUIDITY_ELIGIBLE, None), + (AudioMetricKey.UNATTRIBUTED_MISSED_DEADLINES, None), + (AudioMetricKey.AUDIO_FLUIDITY_TOTAL_DEADLINES, None), + (AudioMetricKey.AUDIO_FLUIDITY_MISSED_DEADLINES, None), + (AudioMetricKey.AUDIO_PLAYABLE_FRAME_COUNT, None), + ] + for key, unit in specs: + self.summaries[key.value] = CDFSketch(key.value, unit=unit) + fluidity_delays = list( + dict.fromkeys( + [ + *self.channel_config.startup_delay_ms_values, + self.channel_config.fluidity_startup_delay_ms, + ] + ) + ) + for delay_ms in fluidity_delays: + metric_name = f"user_audio_fluidity_index_d{_policy_tag(delay_ms)}ms" + sketch = CDFSketch(metric_name) + self.summaries[metric_name] = sketch + self._fluidity_policy_sketches[float(delay_ms)] = sketch + self._interactivity_sketches_ready = True + + def _record_interactivity(self, metrics: AudioRequestMetrics) -> None: + interactivity = metrics.interactivity + if interactivity is None: + return + self._ensure_interactivity_sketches() + + def put(key: AudioMetricKey, value: float | None) -> None: + if value is not None: + self.summaries[key.value].put(value) + + zero_delay = interactivity.fixed_delay_playback[0.0] + put( + AudioMetricKey.FIRST_INPUT_TO_FIRST_AUDIO_MS, + interactivity.first_input_to_first_audio_ms, + ) + put( + AudioMetricKey.FIRST_INPUT_TO_FIRST_PLAYABLE_AUDIO_MS, + interactivity.first_input_to_first_playable_audio_ms, + ) + put( + AudioMetricKey.TRIGGER_TO_FIRST_PLAYABLE_AUDIO_MS, + interactivity.trigger_to_first_playable_audio_ms, + ) + put( + AudioMetricKey.REQUEST_START_TO_FIRST_AUDIO_MS, + interactivity.request_start_to_first_audio_ms, + ) + put( + AudioMetricKey.REQUEST_START_TO_FIRST_PLAYABLE_AUDIO_MS, + interactivity.request_start_to_first_playable_audio_ms, + ) + put( + AudioMetricKey.AUDIO_BEFORE_COMMIT_RATIO, + interactivity.audio_before_commit_ratio, + ) + put( + AudioMetricKey.DUPLEX_OVERLAP_OBSERVED, + float(interactivity.duplex_overlap_observed), + ) + put(AudioMetricKey.DUPLEX_OVERLAP_MS, interactivity.duplex_overlap_ms) + put( + AudioMetricKey.POST_COMMIT_AUDIO_DELIVERY_MS, + interactivity.post_commit_audio_delivery_ms, + ) + put( + AudioMetricKey.REQUIRED_STARTUP_DELAY_MS, + interactivity.required_startup_delay_ms, + ) + put(AudioMetricKey.ZERO_DELAY_STALL_COUNT, float(zero_delay.stall_count)) + put(AudioMetricKey.ZERO_DELAY_TOTAL_STALL_MS, zero_delay.total_stall_ms) + put(AudioMetricKey.ZERO_DELAY_LONGEST_STALL_MS, zero_delay.longest_stall_ms) + put(AudioMetricKey.STREAMING_RTF, interactivity.streaming_rtf) + put( + AudioMetricKey.DONE_AFTER_LAST_AUDIO_MS, + interactivity.done_after_last_audio_ms, + ) + put(AudioMetricKey.WS_CONNECT_LATENCY_MS, metrics.ws_connect_latency_ms) + if interactivity.user_audio_fluidity is not None: + put( + AudioMetricKey.USER_AUDIO_FLUIDITY_INDEX, + interactivity.user_audio_fluidity.fluidity_index, + ) + put( + AudioMetricKey.AUDIO_FLUIDITY_TOTAL_DEADLINES, + float(interactivity.user_audio_fluidity.total_deadlines), + ) + put( + AudioMetricKey.AUDIO_FLUIDITY_MISSED_DEADLINES, + float(interactivity.user_audio_fluidity.missed_deadlines), + ) + put( + AudioMetricKey.AUDIO_PLAYABLE_FRAME_COUNT, + float(interactivity.user_audio_fluidity.playable_frame_count), + ) + put( + AudioMetricKey.TTS_SERVICE_FLUIDITY_ELIGIBLE, + float(interactivity.tts_service_fluidity_eligible), + ) + put( + AudioMetricKey.UNATTRIBUTED_MISSED_DEADLINES, + float(interactivity.unattributed_missed_deadlines), + ) + if interactivity.tts_service_fluidity is not None: + put( + AudioMetricKey.TTS_SERVICE_FLUIDITY_INDEX, + interactivity.tts_service_fluidity.fluidity_index, + ) + self._tts_service_fluidity_eligible_count += 1 + for delay_ms, result in interactivity.fluidity_by_startup_delay.items(): + if result is not None: + self._fluidity_policy_sketches[delay_ms].put(result.fluidity_index) + + self._interactive_request_count += 1 + for delay_ms, result in interactivity.fixed_delay_playback.items(): + if result.stall_free: + self._fixed_delay_stall_free_counts[delay_ms] = ( + self._fixed_delay_stall_free_counts.get(delay_ms, 0) + 1 + ) + for target_ms, result in interactivity.buffer_target_playback.items(): + if result is None: + continue + self._buffer_target_eligible_counts[target_ms] = ( + self._buffer_target_eligible_counts.get(target_ms, 0) + 1 + ) + if result.stall_free: + self._buffer_target_stall_free_counts[target_ms] = ( + self._buffer_target_stall_free_counts.get(target_ms, 0) + 1 + ) def record_session_completed( self, session_id: int, session_size: int, - first_dispatch_at: Optional[float], - last_completion_at: Optional[float], + first_dispatch_at: float | None, + last_completion_at: float | None, ) -> None: """Record session-level metrics.""" - pass # Skeleton + with self.lock: + self.summaries[AudioMetricKey.SESSION_SIZE.value].put(session_size) + if first_dispatch_at is not None and last_completion_at is not None: + self.summaries[AudioMetricKey.SESSION_DURATION.value].put( + max(0.0, last_completion_at - first_dispatch_at) + ) - def finalize(self) -> EvaluationResult: - """Finalize evaluation and return results.""" - return EvaluationResult( - evaluator_type="audio_performance", - channel=ChannelModality.AUDIO, - metrics={ - "status": "not_implemented", - "message": "Audio performance evaluation not yet implemented", - }, + def get_summary(self) -> dict[str, Any]: + """Return aggregate audio, ASR, and playback-policy metrics.""" + perf_summary: dict[str, Any] = {} + for cdf_sketch in self.summaries.values(): + perf_summary.update(cdf_sketch.get_summary()) + + if self._stt_request_count > 0: + for cdf_sketch in self.asr_latency_summaries.values(): + if len(cdf_sketch) > 0: + perf_summary.update(cdf_sketch.get_summary()) + perf_summary.update(self._asr_metrics.get_summary()) + + wall_s = 0.0 + if self._first_dispatch_at is not None and self._last_completion_at is not None: + wall_s = max(0.0, self._last_completion_at - self._first_dispatch_at) + perf_summary["chars_per_sec_aggregate"] = ( + self._total_input_chars / wall_s if wall_s > 0 else None + ) + perf_summary["generated_audio_seconds_per_sec_aggregate"] = ( + (self._total_generated_audio_duration_ms / 1000.0) / wall_s + if wall_s > 0 + else None ) - def get_streaming_metrics(self) -> Optional[Dict[str, Any]]: - """Return current metrics for streaming.""" - return None + if self._interactive_request_count > 0: + count = self._interactive_request_count + perf_summary["interactive_requests_count"] = count + perf_summary[AudioMetricKey.AUDIO_FLUIDITY_FRAME_MS.value] = ( + self.channel_config.fluidity_frame_ms + ) + perf_summary[AudioMetricKey.AUDIO_FLUIDITY_STARTUP_DELAY_MS.value] = ( + self.channel_config.fluidity_startup_delay_ms + ) + perf_summary["fluidity_attribution_mode"] = ( + self.channel_config.fluidity_attribution_mode + ) + perf_summary["tts_service_fluidity_eligible_fraction"] = ( + self._tts_service_fluidity_eligible_count / count + ) + for delay_ms in self.channel_config.startup_delay_ms_values: + tag = _policy_tag(delay_ms) + fraction = self._fixed_delay_stall_free_counts.get(delay_ms, 0) / count + perf_summary[f"fixed_delay_stall_free_fraction_d{tag}ms"] = fraction + if delay_ms == 0.0: + perf_summary["zero_delay_stall_free_fraction"] = fraction + for target_ms in self.channel_config.startup_buffer_ms_values: + tag = _policy_tag(target_ms) + eligible = self._buffer_target_eligible_counts.get(target_ms, 0) + perf_summary[f"buffer_target_eligible_count_b{tag}ms"] = eligible + perf_summary[f"buffer_target_stall_free_fraction_b{tag}ms"] = ( + self._buffer_target_stall_free_counts.get(target_ms, 0) / eligible + if eligible > 0 + else None + ) + return perf_summary + + def finalize(self) -> EvaluationResult: + with self.lock: + if self._stt_request_count > 0: + logger.info( + "Evaluator phase 'asr_scoring' took %.2fs total across " + "%d STT requests (concurrent and overlapped with the run)", + self._asr_scoring_seconds, + self._stt_request_count, + ) + return EvaluationResult( + evaluator_type="audio_performance", + channel=ChannelModality.AUDIO, + metrics=self.get_summary(), + ) + + def get_streaming_metrics(self) -> dict[str, Any] | None: + with self.lock: + return self.get_summary() def save(self, output_dir: str) -> None: - """Save evaluation artifacts.""" - pass # Skeleton + with self.lock: + stages = ( + ("audio_request_level_metrics", self._save_request_level_metrics), + ("audio_raw_timing", self._save_raw_timing), + ("audio_cdf_csvs", self._save_cdf_csvs), + ("audio_cdf_plots", self._plot_cdfs), + ( + "audio_stall_free_policy_plot", + self._plot_stall_free_vs_startup_policy, + ), + ) + for stage_name, stage in stages: + stage_started_at = time.perf_counter() + stage(output_dir) + logger.info( + "Evaluator phase '%s' took %.2fs", + stage_name, + time.perf_counter() - stage_started_at, + ) def flush_streaming_outputs(self, output_dir: str) -> None: - """Flush current metrics for streaming.""" - pass # Skeleton + with self.lock: + rows = self._export_request_rows(self._request_rows_streamed) + if rows: + self._append_request_level_rows(output_dir, rows) + self._request_rows_streamed = len(self._completed_metrics) + self._flush_raw_timing_rows(output_dir) + self._save_cdf_csvs(output_dir) + + def _save_request_level_metrics(self, output_dir: str) -> None: + path = os.path.join(output_dir, "request_level_metrics.jsonl") + with open(path, "w") as file: + for row in self._export_request_rows(0): + file.write(json.dumps(row) + "\n") + + def _export_request_rows(self, start_index: int = 0) -> list[dict[str, Any]]: + rows = [] + for index in range(start_index, len(self._completed_metrics)): + metrics = self._completed_metrics[index] + lifecycle = self._lifecycle_timestamps[index] + row = { + "request_id": metrics.request_id, + "session_id": metrics.session_id, + "session_total_requests": metrics.session_total_requests, + "scheduler_ready_at": lifecycle["scheduler_ready_at"], + "scheduler_dispatched_at": round( + max( + 0.0, + metrics.request_dispatched_at - self._request_time_reference, + ), + 5, + ), + "client_picked_up_at": lifecycle["client_picked_up_at"], + "client_completed_at": round( + max( + 0.0, + metrics.client_completed_at - self._request_time_reference, + ), + 5, + ), + "result_processed_at": lifecycle["result_processed_at"], + "audio_task": ( + str(metrics.audio_task) if metrics.audio_task is not None else None + ), + AudioMetricKey.PROVIDER.value: metrics.provider, + AudioMetricKey.PROVIDER_MODEL.value: metrics.provider_model, + AudioMetricKey.PROVIDER_PROTOCOL.value: metrics.provider_protocol, + AudioMetricKey.TTFC.value: round(metrics.ttfc, 3), + AudioMetricKey.END_TO_END_LATENCY.value: round( + metrics.end_to_end_latency, 3 + ), + AudioMetricKey.GENERATED_AUDIO_DURATION.value: round( + metrics.generated_audio_duration, 3 + ), + AudioMetricKey.RTF.value: round(metrics.rtf, 5), + AudioMetricKey.CHUNK_COUNT.value: metrics.chunk_count, + AudioMetricKey.PCM_BYTE_COUNT.value: metrics.pcm_byte_count, + AudioMetricKey.INPUT_CHARS.value: metrics.input_chars, + AudioMetricKey.INPUT_TOKENS.value: metrics.input_tokens, + AudioMetricKey.INPUT_TEXT.value: metrics.input_text, + } + if metrics.asr is not None: + row.update(metrics.asr.to_request_row()) + row.update(self._interactivity_row_fields(metrics)) + rows.append(row) + return rows + + def _interactivity_row_fields(self, metrics: AudioRequestMetrics) -> dict[str, Any]: + interactivity = metrics.interactivity + if interactivity is None: + return {} + + fields: dict[str, Any] = {} + + def add(key: str, value: Any, ndigits: int = 3) -> None: + if value is None: + return + fields[key] = round(value, ndigits) if isinstance(value, float) else value + + zero_delay = interactivity.fixed_delay_playback[0.0] + add( + AudioMetricKey.FIRST_INPUT_TO_FIRST_AUDIO_MS.value, + interactivity.first_input_to_first_audio_ms, + ) + add( + AudioMetricKey.FIRST_INPUT_TO_FIRST_PLAYABLE_AUDIO_MS.value, + interactivity.first_input_to_first_playable_audio_ms, + ) + add( + AudioMetricKey.TRIGGER_TO_FIRST_PLAYABLE_AUDIO_MS.value, + interactivity.trigger_to_first_playable_audio_ms, + ) + add( + AudioMetricKey.REQUEST_START_TO_FIRST_AUDIO_MS.value, + interactivity.request_start_to_first_audio_ms, + ) + add( + AudioMetricKey.REQUEST_START_TO_FIRST_PLAYABLE_AUDIO_MS.value, + interactivity.request_start_to_first_playable_audio_ms, + ) + add( + AudioMetricKey.AUDIO_BEFORE_COMMIT_RATIO.value, + interactivity.audio_before_commit_ratio, + 5, + ) + fields[AudioMetricKey.DUPLEX_OVERLAP_OBSERVED.value] = int( + interactivity.duplex_overlap_observed + ) + add( + AudioMetricKey.DUPLEX_OVERLAP_MS.value, + interactivity.duplex_overlap_ms, + ) + add( + AudioMetricKey.POST_COMMIT_AUDIO_DELIVERY_MS.value, + interactivity.post_commit_audio_delivery_ms, + ) + add( + AudioMetricKey.REQUIRED_STARTUP_DELAY_MS.value, + interactivity.required_startup_delay_ms, + ) + fields[AudioMetricKey.ZERO_DELAY_STALL_COUNT.value] = zero_delay.stall_count + add( + AudioMetricKey.ZERO_DELAY_TOTAL_STALL_MS.value, + zero_delay.total_stall_ms, + ) + add( + AudioMetricKey.ZERO_DELAY_LONGEST_STALL_MS.value, + zero_delay.longest_stall_ms, + ) + fields[AudioMetricKey.ZERO_DELAY_STALL_FREE.value] = int(zero_delay.stall_free) + add(AudioMetricKey.STREAMING_RTF.value, interactivity.streaming_rtf, 5) + add( + AudioMetricKey.DONE_AFTER_LAST_AUDIO_MS.value, + interactivity.done_after_last_audio_ms, + ) + if interactivity.user_audio_fluidity is not None: + add( + AudioMetricKey.AUDIO_FLUIDITY_FRAME_MS.value, + interactivity.user_audio_fluidity.frame_duration_ms, + ) + add( + AudioMetricKey.AUDIO_FLUIDITY_STARTUP_DELAY_MS.value, + interactivity.user_audio_fluidity.startup_delay_ms, + ) + add( + AudioMetricKey.USER_AUDIO_FLUIDITY_INDEX.value, + interactivity.user_audio_fluidity.fluidity_index, + 5, + ) + fields[AudioMetricKey.AUDIO_FLUIDITY_TOTAL_DEADLINES.value] = ( + interactivity.user_audio_fluidity.total_deadlines + ) + fields[AudioMetricKey.AUDIO_FLUIDITY_MISSED_DEADLINES.value] = ( + interactivity.user_audio_fluidity.missed_deadlines + ) + fields[AudioMetricKey.AUDIO_PLAYABLE_FRAME_COUNT.value] = ( + interactivity.user_audio_fluidity.playable_frame_count + ) + fields[AudioMetricKey.TTS_SERVICE_FLUIDITY_ELIGIBLE.value] = int( + interactivity.tts_service_fluidity_eligible + ) + fields[AudioMetricKey.UNATTRIBUTED_MISSED_DEADLINES.value] = ( + interactivity.unattributed_missed_deadlines + ) + if interactivity.tts_service_fluidity is not None: + add( + AudioMetricKey.TTS_SERVICE_FLUIDITY_INDEX.value, + interactivity.tts_service_fluidity.fluidity_index, + 5, + ) + + for delay_ms, result in interactivity.fluidity_by_startup_delay.items(): + if result is None: + continue + add( + f"user_audio_fluidity_index_d{_policy_tag(delay_ms)}ms", + result.fluidity_index, + 5, + ) + + for delay_ms, result in interactivity.fixed_delay_playback.items(): + if delay_ms == 0.0: + continue + prefix = f"fixed_delay_{_policy_tag(delay_ms)}ms" + fields[f"{prefix}_stall_count"] = result.stall_count + add(f"{prefix}_total_stall_ms", result.total_stall_ms) + add(f"{prefix}_longest_stall_ms", result.longest_stall_ms) + fields[f"{prefix}_stall_free"] = int(result.stall_free) + + for target_ms, result in interactivity.buffer_target_playback.items(): + if result is None: + continue + prefix = f"buffer_target_{_policy_tag(target_ms)}ms" + add(f"{prefix}_startup_wait_ms", result.startup_wait_from_first_audio_ms) + fields[f"{prefix}_stall_count"] = result.stall_count + add(f"{prefix}_total_stall_ms", result.total_stall_ms) + add(f"{prefix}_longest_stall_ms", result.longest_stall_ms) + fields[f"{prefix}_stall_free"] = int(result.stall_free) + return fields + + def _append_request_level_rows( + self, output_dir: str, rows: list[dict[str, Any]] + ) -> None: + path = os.path.join(output_dir, "request_level_metrics.jsonl") + with open(path, "a") as file: + for row in rows: + file.write(json.dumps(row) + "\n") + + def _build_raw_timing_row( + self, + request_id: int, + session_id: int, + channel_metrics: dict[str, Any], + timing: RequestTiming, + ) -> dict[str, Any]: + raw_chunks = ( + channel_metrics.get(AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value) or [] + ) + audio_chunks = [ + [round(float(entry[0]), 1), int(entry[1])] for entry in raw_chunks + ] + audio_chunks.sort(key=lambda chunk: chunk[0]) + return { + "request_id": request_id, + "session_id": session_id, + "sample_rate": timing.sample_rate, + "text_deltas": [ + [round(offset, 1), n_chars] for offset, n_chars in timing.text_deltas + ], + "audio_chunks": audio_chunks, + "response_trigger_ms": ( + round(timing.response_trigger_ms, 1) + if timing.response_trigger_ms is not None + else None + ), + "provider": channel_metrics.get(AudioMetricKey.PROVIDER.value), + "provider_model": channel_metrics.get(AudioMetricKey.PROVIDER_MODEL.value), + "provider_protocol": channel_metrics.get( + AudioMetricKey.PROVIDER_PROTOCOL.value + ), + "commit_ms": ( + round(timing.commit_ms, 1) if timing.commit_ms is not None else None + ), + "audio_done_ms": ( + round(timing.audio_done_ms, 1) + if timing.audio_done_ms is not None + else None + ), + "response_done_ms": ( + round(timing.response_done_ms, 1) + if timing.response_done_ms is not None + else None + ), + } + + def _save_raw_timing(self, output_dir: str) -> None: + if not self.channel_config.persist_raw_timing or not self._raw_timing_rows: + return + path = os.path.join(output_dir, "audio_raw_timing.jsonl") + with open(path, "w") as file: + for row in self._raw_timing_rows: + file.write(json.dumps(row) + "\n") + self._raw_timing_rows_streamed = len(self._raw_timing_rows) + + def _flush_raw_timing_rows(self, output_dir: str) -> None: + if not self.channel_config.persist_raw_timing: + return + new_rows = self._raw_timing_rows[self._raw_timing_rows_streamed :] + if not new_rows: + return + path = os.path.join(output_dir, "audio_raw_timing.jsonl") + with open(path, "a") as file: + for row in new_rows: + file.write(json.dumps(row) + "\n") + self._raw_timing_rows_streamed = len(self._raw_timing_rows) + + def _cdf_summaries(self) -> dict[str, CDFSketch]: + summaries = dict(self.summaries) + if self._stt_request_count > 0: + summaries.update( + { + name: sketch + for name, sketch in self.asr_latency_summaries.items() + if len(sketch) > 0 + } + ) + return summaries + + def _save_cdf_csvs(self, output_dir: str) -> None: + for metric_name, cdf_sketch in self._cdf_summaries().items(): + cdf_sketch._to_df().to_csv( + os.path.join(output_dir, f"audio_{metric_name}.csv"), + index=False, + ) + + def _plot_cdfs(self, output_dir: str) -> None: + for metric_name, cdf_sketch in self._cdf_summaries().items(): + try: + cdf_sketch.plot_cdf(output_dir, f"audio_{metric_name}") + except Exception as error: + logger.warning("Failed to plot CDF for %s: %s", metric_name, error) + + def _plot_stall_free_vs_startup_policy(self, output_dir: str) -> None: + if self._interactive_request_count <= 0: + return + try: + import pandas as pd + import rekha as rk + + rows: list[dict[str, Any]] = [] + for delay_ms in sorted(self.channel_config.startup_delay_ms_values): + rows.append( + { + "startup_budget_ms": delay_ms, + "stall_free_fraction": ( + self._fixed_delay_stall_free_counts.get(delay_ms, 0) + / self._interactive_request_count + ), + "policy": "Fixed delay", + } + ) + for target_ms in sorted(self.channel_config.startup_buffer_ms_values): + eligible = self._buffer_target_eligible_counts.get(target_ms, 0) + if eligible == 0: + continue + rows.append( + { + "startup_budget_ms": target_ms, + "stall_free_fraction": ( + self._buffer_target_stall_free_counts.get(target_ms, 0) + / eligible + ), + "policy": "Buffered audio target", + } + ) + dataframe = pd.DataFrame(rows) + figure = rk.line( + dataframe, + x="startup_budget_ms", + y="stall_free_fraction", + color="policy", + markers=True, + labels={ + "startup_budget_ms": "Startup budget (ms)", + "stall_free_fraction": "Stall-free fraction", + "policy": "Playback policy", + }, + ) + figure.save( + os.path.join(output_dir, "audio_stall_free_vs_startup_policy.png"), + transparent=False, + ) + except Exception as error: + logger.warning( + "Failed to plot stall-free fraction by startup policy: %s", error + ) diff --git a/veeksha/evaluator/performance/audio_interactivity.py b/veeksha/evaluator/performance/audio_interactivity.py new file mode 100644 index 00000000..ee3931cc --- /dev/null +++ b/veeksha/evaluator/performance/audio_interactivity.py @@ -0,0 +1,531 @@ +"""Pure computation for realtime TTS interactivity metrics. + +The websocket client records one monotonic timeline per request. This module +replays that timeline locally under multiple playback policies; no additional +server requests are made for a policy sweep. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Optional + +from veeksha.core.audio_contract import ( + DEFAULT_AUDIO_SAMPLE_RATE, + AudioMetricKey, + pcm_bytes_to_duration_ms, +) + + +@dataclass(frozen=True) +class RequestTiming: + """Validated realtime request timing. + + Audio tuples contain ``(receipt_offset_ms, duration_ms, decoded_bytes)``. + Chunks received at the same timestamp are coalesced so playback metrics do + not depend on transport frame fragmentation. + """ + + text_deltas: list[tuple[float, int]] + audio_chunks: list[tuple[float, float, float]] + response_trigger_ms: Optional[float] + commit_ms: Optional[float] + audio_done_ms: Optional[float] + response_done_ms: Optional[float] + sample_rate: int + + +def _optional_float(metrics: Mapping, key: AudioMetricKey) -> Optional[float]: + value = metrics.get(key.value) + if value is None: + return None + try: + return float(value) + except TypeError, ValueError: + return None + + +def parse_request_timing( + metrics: Mapping, default_sample_rate: int = DEFAULT_AUDIO_SAMPLE_RATE +) -> Optional[RequestTiming]: + """Parse the raw client contract, returning ``None`` without usable audio.""" + raw_chunks = metrics.get(AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value) + if not isinstance(raw_chunks, Sequence) or not raw_chunks: + return None + + try: + sample_rate = int( + metrics.get(AudioMetricKey.SAMPLE_RATE.value, default_sample_rate) + or default_sample_rate + ) + except TypeError, ValueError: + sample_rate = default_sample_rate + if sample_rate <= 0: + sample_rate = default_sample_rate + + parsed_chunks: list[tuple[float, float]] = [] + for entry in raw_chunks: + if not isinstance(entry, Sequence) or len(entry) < 2: + continue + try: + offset_ms = float(entry[0]) + n_bytes = float(entry[1]) + except TypeError, ValueError: + continue + if offset_ms < 0 or n_bytes <= 0: + continue + parsed_chunks.append((offset_ms, n_bytes)) + if not parsed_chunks: + return None + parsed_chunks.sort(key=lambda item: item[0]) + + audio_chunks: list[tuple[float, float, float]] = [] + for offset_ms, n_bytes in parsed_chunks: + if audio_chunks and audio_chunks[-1][0] == offset_ms: + _, previous_duration, previous_bytes = audio_chunks[-1] + audio_chunks[-1] = ( + offset_ms, + previous_duration + pcm_bytes_to_duration_ms(n_bytes, sample_rate), + previous_bytes + n_bytes, + ) + else: + audio_chunks.append( + ( + offset_ms, + pcm_bytes_to_duration_ms(n_bytes, sample_rate), + n_bytes, + ) + ) + + text_deltas: list[tuple[float, int]] = [] + raw_text_deltas = metrics.get(AudioMetricKey.TEXT_DELTA_TIMESTAMPS.value) or [] + if isinstance(raw_text_deltas, Sequence): + for entry in raw_text_deltas: + if not isinstance(entry, Sequence) or len(entry) < 2: + continue + try: + offset_ms = float(entry[0]) + n_chars = int(entry[1]) + except TypeError, ValueError: + continue + if offset_ms >= 0 and n_chars >= 0: + text_deltas.append((offset_ms, n_chars)) + text_deltas.sort(key=lambda item: item[0]) + + return RequestTiming( + text_deltas=text_deltas, + audio_chunks=audio_chunks, + response_trigger_ms=_optional_float( + metrics, AudioMetricKey.RESPONSE_TRIGGER_OFFSET_MS + ), + commit_ms=_optional_float(metrics, AudioMetricKey.INPUT_COMMIT_OFFSET_MS), + audio_done_ms=_optional_float(metrics, AudioMetricKey.AUDIO_DONE_OFFSET_MS), + response_done_ms=_optional_float( + metrics, AudioMetricKey.RESPONSE_DONE_OFFSET_MS + ), + sample_rate=sample_rate, + ) + + +@dataclass(frozen=True) +class PlaybackSimResult: + """Playback outcome for one startup policy.""" + + playback_start_ms: float + startup_wait_from_first_audio_ms: float + stall_count: int + total_stall_ms: float + longest_stall_ms: float + stall_free: bool + buffer_margin_min_ms: float + buffer_margin_mean_ms: float + + +@dataclass(frozen=True) +class AudioFluidityResult: + """Etalon-style deadline acceptance for fixed-duration playable audio. + + Network chunks are first converted into complete ``frame_duration_ms`` PCM + frames. A frame becomes available when cumulative received audio crosses + the next frame boundary. This makes the score independent of how a + transport fragments the same playable-audio supply. + + ``startup_delay_ms`` is playback slack after the first complete frame is + available. Request-to-first-playable-frame latency is reported separately, + so this score isolates continuity after audio can begin. + """ + + startup_delay_ms: float + frame_duration_ms: float + playable_frame_count: int + total_deadlines: int + missed_deadlines: int + fluidity_index: float + + +def _playable_frame_arrival_ms( + chunks: Sequence[tuple[float, float, float]], + frame_duration_ms: float, +) -> list[float]: + """Return receipt times for complete fixed-duration PCM frames. + + A partial tail shorter than one frame is deliberately excluded: it cannot + fill another complete playback period and counting it as a full frame would + over-penalize normal end-of-utterance delivery. + """ + if frame_duration_ms <= 0: + raise ValueError("frame_duration_ms must be > 0") + + arrivals: list[float] = [] + cumulative_audio_ms = 0.0 + next_frame_boundary_ms = frame_duration_ms + epsilon = frame_duration_ms * 1e-9 + + for receipt_ms, duration_ms, _ in chunks: + cumulative_audio_ms += duration_ms + while cumulative_audio_ms + epsilon >= next_frame_boundary_ms: + arrivals.append(receipt_ms) + next_frame_boundary_ms += frame_duration_ms + return arrivals + + +def compute_audio_fluidity( + chunks: Sequence[tuple[float, float, float]], + *, + frame_duration_ms: float, + startup_delay_ms: float, +) -> Optional[AudioFluidityResult]: + """Compute the Etalon fluidity-index analogue for playable audio frames. + + The algorithm follows Etalon's deadline/slack/reset semantics. Frames that + arrive early accumulate slack (playback buffer). When an inter-frame gap + exceeds the frame deadline plus available slack, every elapsed playback + deadline is counted as missed and slack is reset. + """ + if startup_delay_ms < 0: + raise ValueError("startup_delay_ms must be >= 0") + + arrivals = _playable_frame_arrival_ms(chunks, frame_duration_ms) + if not arrivals: + return None + + inter_frame_times_ms = [0.0] + inter_frame_times_ms.extend( + current - previous for previous, current in zip(arrivals, arrivals[1:]) + ) + + total_deadlines = 0 + missed_deadlines = 0 + slack_ms = 0.0 + + for index, inter_frame_ms in enumerate(inter_frame_times_ms): + deadline_ms = startup_delay_ms if index == 0 else frame_duration_ms + if inter_frame_ms <= deadline_ms + slack_ms: + slack_ms += deadline_ms - inter_frame_ms + total_deadlines += 1 + continue + + misses = ( + math.floor((inter_frame_ms - slack_ms - deadline_ms) / frame_duration_ms) + + 1 + ) + missed_deadlines += misses + total_deadlines += misses + slack_ms = 0.0 + + return AudioFluidityResult( + startup_delay_ms=startup_delay_ms, + frame_duration_ms=frame_duration_ms, + playable_frame_count=len(arrivals), + total_deadlines=total_deadlines, + missed_deadlines=missed_deadlines, + fluidity_index=(total_deadlines - missed_deadlines) / total_deadlines, + ) + + +def _simulate_from_start( + chunks: Sequence[tuple[float, float, float]], + playback_start_ms: float, + min_reportable_stall_ms: float, +) -> PlaybackSimResult: + """Drain delivered audio from an explicit wall-clock playback start.""" + finish_ms = playback_start_ms + stalls: list[float] = [] + margins: list[float] = [] + + for index, (receipt_ms, duration_ms, _) in enumerate(chunks): + if index > 0: + margins.append(finish_ms - receipt_ms) + gap_ms = receipt_ms - finish_ms + if index > 0 and gap_ms > min_reportable_stall_ms: + stalls.append(gap_ms) + finish_ms = max(finish_ms, receipt_ms) + duration_ms + + first_audio_ms = chunks[0][0] + return PlaybackSimResult( + playback_start_ms=playback_start_ms, + startup_wait_from_first_audio_ms=playback_start_ms - first_audio_ms, + stall_count=len(stalls), + total_stall_ms=float(sum(stalls)), + longest_stall_ms=float(max(stalls, default=0.0)), + stall_free=not stalls, + buffer_margin_min_ms=float(min(margins, default=0.0)), + buffer_margin_mean_ms=(float(sum(margins) / len(margins)) if margins else 0.0), + ) + + +def simulate_fixed_delay( + chunks: Sequence[tuple[float, float, float]], + startup_delay_ms: float, + min_reportable_stall_ms: float, +) -> Optional[PlaybackSimResult]: + """Start playback a fixed delay after the first audio receipt.""" + if not chunks: + return None + return _simulate_from_start( + chunks, + chunks[0][0] + startup_delay_ms, + min_reportable_stall_ms, + ) + + +def simulate_buffer_target( + chunks: Sequence[tuple[float, float, float]], + target_audio_ms: float, + min_reportable_stall_ms: float, + *, + audio_done_ms: Optional[float], + response_done_ms: Optional[float], +) -> Optional[PlaybackSimResult]: + """Start once a target duration is buffered, or at a terminal event. + + Falling back to a terminal event lets short, complete utterances play even + when their total duration is below the configured target. Unterminated + partial streams remain ineligible for a target they never reached. + """ + if not chunks: + return None + + playback_start_ms: Optional[float] = None + buffered_ms = 0.0 + for receipt_ms, duration_ms, _ in chunks: + buffered_ms += duration_ms + if target_audio_ms <= 0 or buffered_ms >= target_audio_ms: + playback_start_ms = receipt_ms + break + + if playback_start_ms is None: + terminal_ms = audio_done_ms + if terminal_ms is None: + terminal_ms = response_done_ms + if terminal_ms is None: + return None + playback_start_ms = max(terminal_ms, chunks[-1][0]) + + return _simulate_from_start(chunks, playback_start_ms, min_reportable_stall_ms) + + +def _required_startup_delay_ms( + chunks: Sequence[tuple[float, float, float]], +) -> float: + """Return the minimum exact fixed delay that prevents every underrun.""" + first_audio_ms = chunks[0][0] + cumulative_previous_ms = 0.0 + worst_deficit_ms = 0.0 + for index in range(1, len(chunks)): + cumulative_previous_ms += chunks[index - 1][1] + deficit_ms = chunks[index][0] - first_audio_ms - cumulative_previous_ms + worst_deficit_ms = max(worst_deficit_ms, deficit_ms) + return max(0.0, worst_deficit_ms) + + +@dataclass(frozen=True) +class InteractivityMetrics: + """Stable and diagnostic metrics derived from one request timeline.""" + + first_input_to_first_audio_ms: Optional[float] + first_input_to_first_playable_audio_ms: Optional[float] + trigger_to_first_playable_audio_ms: Optional[float] + request_start_to_first_audio_ms: float + request_start_to_first_playable_audio_ms: Optional[float] + audio_before_commit_ratio: Optional[float] + duplex_overlap_observed: bool + duplex_overlap_ms: float + post_commit_audio_delivery_ms: Optional[float] + required_startup_delay_ms: float + fixed_delay_playback: dict[float, PlaybackSimResult] + buffer_target_playback: dict[float, Optional[PlaybackSimResult]] + streaming_rtf: Optional[float] + done_after_last_audio_ms: Optional[float] + user_audio_fluidity: Optional[AudioFluidityResult] + tts_service_fluidity: Optional[AudioFluidityResult] + tts_service_fluidity_eligible: bool + unattributed_missed_deadlines: int + fluidity_by_startup_delay: dict[float, Optional[AudioFluidityResult]] + + +def compute_interactivity_metrics( + timing: RequestTiming, + *, + startup_delay_ms_values: Sequence[float], + startup_buffer_ms_values: Sequence[float], + min_reportable_stall_ms: float, + fluidity_frame_ms: float, + fluidity_startup_delay_ms: float, + fluidity_attribution_mode: str, +) -> InteractivityMetrics: + """Derive all playback policies from a single captured request timeline.""" + chunks = timing.audio_chunks + first_audio_ms = chunks[0][0] + last_audio_ms = chunks[-1][0] + total_audio_ms = float(sum(chunk[1] for chunk in chunks)) + total_audio_bytes = float(sum(chunk[2] for chunk in chunks)) + + playable_frame_arrivals = _playable_frame_arrival_ms(chunks, fluidity_frame_ms) + first_playable_audio_ms = ( + playable_frame_arrivals[0] if playable_frame_arrivals else None + ) + + first_input_to_first_audio_ms = ( + first_audio_ms - timing.text_deltas[0][0] if timing.text_deltas else None + ) + first_input_to_first_playable_audio_ms = ( + first_playable_audio_ms - timing.text_deltas[0][0] + if timing.text_deltas and first_playable_audio_ms is not None + else None + ) + trigger_to_first_playable_audio_ms = ( + first_playable_audio_ms - timing.response_trigger_ms + if timing.response_trigger_ms is not None + and first_playable_audio_ms is not None + else None + ) + + audio_before_commit_ratio: Optional[float] = None + post_commit_audio_delivery_ms: Optional[float] = None + duplex_overlap_observed = False + duplex_overlap_ms = 0.0 + if timing.commit_ms is not None: + bytes_before_commit = sum( + n_bytes + for receipt_ms, _, n_bytes in chunks + if receipt_ms <= timing.commit_ms + ) + audio_before_commit_ratio = ( + bytes_before_commit / total_audio_bytes if total_audio_bytes > 0 else None + ) + post_commit_audio_delivery_ms = max(0.0, last_audio_ms - timing.commit_ms) + if ( + first_playable_audio_ms is not None + and first_playable_audio_ms < timing.commit_ms + ): + duplex_overlap_observed = True + duplex_overlap_ms = max( + 0.0, + min(last_audio_ms, timing.commit_ms) - first_playable_audio_ms, + ) + + fixed_delay_playback = { + float(delay_ms): simulate_fixed_delay( + chunks, float(delay_ms), min_reportable_stall_ms + ) + for delay_ms in startup_delay_ms_values + } + # Chunks is non-empty, so fixed-delay simulations are always non-null. + fixed_delay_playback = { + delay_ms: result + for delay_ms, result in fixed_delay_playback.items() + if result is not None + } + + buffer_target_playback = { + float(target_ms): simulate_buffer_target( + chunks, + float(target_ms), + min_reportable_stall_ms, + audio_done_ms=timing.audio_done_ms, + response_done_ms=timing.response_done_ms, + ) + for target_ms in startup_buffer_ms_values + } + + fluidity_delays = list( + dict.fromkeys( + [ + *(float(value) for value in startup_delay_ms_values), + float(fluidity_startup_delay_ms), + ] + ) + ) + fluidity_by_startup_delay = { + delay_ms: compute_audio_fluidity( + chunks, + frame_duration_ms=fluidity_frame_ms, + startup_delay_ms=delay_ms, + ) + for delay_ms in fluidity_delays + } + user_audio_fluidity = fluidity_by_startup_delay[float(fluidity_startup_delay_ms)] + + # A raw playback miss is always valid as a user-experience observation, but + # it is not automatically attributable to the TTS service during duplex + # input: the upstream text source may simply have supplied nothing to + # synthesize. In conservative mode, publish a service score only when all + # text was available before playback could begin. A controlled oversupply + # trace may explicitly opt into service attribution. + complete_input_before_playback = bool( + timing.commit_ms is not None + and first_playable_audio_ms is not None + and timing.commit_ms <= first_playable_audio_ms + ) + tts_service_fluidity_eligible = bool( + user_audio_fluidity is not None + and ( + complete_input_before_playback + or fluidity_attribution_mode == "source_oversupplied" + ) + ) + tts_service_fluidity = ( + user_audio_fluidity if tts_service_fluidity_eligible else None + ) + unattributed_missed_deadlines = ( + 0 + if tts_service_fluidity_eligible or user_audio_fluidity is None + else user_audio_fluidity.missed_deadlines + ) + + streaming_rtf: Optional[float] = None + if len(chunks) >= 2: + delivered_after_first_ms = total_audio_ms - chunks[0][1] + if delivered_after_first_ms > 0: + streaming_rtf = (last_audio_ms - first_audio_ms) / delivered_after_first_ms + + done_after_last_audio_ms = ( + timing.response_done_ms - last_audio_ms + if timing.response_done_ms is not None + else None + ) + + return InteractivityMetrics( + first_input_to_first_audio_ms=first_input_to_first_audio_ms, + first_input_to_first_playable_audio_ms=(first_input_to_first_playable_audio_ms), + trigger_to_first_playable_audio_ms=trigger_to_first_playable_audio_ms, + request_start_to_first_audio_ms=first_audio_ms, + request_start_to_first_playable_audio_ms=first_playable_audio_ms, + audio_before_commit_ratio=audio_before_commit_ratio, + duplex_overlap_observed=duplex_overlap_observed, + duplex_overlap_ms=duplex_overlap_ms, + post_commit_audio_delivery_ms=post_commit_audio_delivery_ms, + required_startup_delay_ms=_required_startup_delay_ms(chunks), + fixed_delay_playback=fixed_delay_playback, + buffer_target_playback=buffer_target_playback, + streaming_rtf=streaming_rtf, + done_after_last_audio_ms=done_after_last_audio_ms, + user_audio_fluidity=user_audio_fluidity, + tts_service_fluidity=tts_service_fluidity, + tts_service_fluidity_eligible=tts_service_fluidity_eligible, + unattributed_missed_deadlines=unattributed_missed_deadlines, + fluidity_by_startup_delay=fluidity_by_startup_delay, + ) diff --git a/veeksha/evaluator/performance/base.py b/veeksha/evaluator/performance/base.py index 75e02aa7..e4d1bbb3 100644 --- a/veeksha/evaluator/performance/base.py +++ b/veeksha/evaluator/performance/base.py @@ -9,11 +9,12 @@ from typing import Any, DefaultDict, Dict, Optional from veeksha.config.evaluator import PerformanceEvaluatorConfig +from veeksha.core.audio_contract import AudioMetricKey from veeksha.core.seeding import SeedManager from veeksha.evaluator.base import BaseEvaluator, EvaluationResult from veeksha.logger import init_logger from veeksha.slo.runner import evaluate_and_save_slos -from veeksha.types import ChannelModality +from veeksha.types import ChannelModality, ClientType logger = init_logger(__name__) @@ -47,11 +48,15 @@ def __init__( seed_manager: Optional[SeedManager] = None, output_dir: Optional[str] = None, benchmark_start_time: float = 0.0, + client_type: Optional[ClientType] = None, ): super().__init__(config, seed_manager) self.config: PerformanceEvaluatorConfig = config self.output_dir = output_dir self.benchmark_start_time = benchmark_start_time + self.client_type = client_type + self._realtime_first_audio_count = 0 + self._realtime_stream_completed_count = 0 self._channel_evaluators: Dict[ChannelModality, Any] = {} @@ -95,12 +100,15 @@ def __init__( for channel in self.config.target_channels: channel_config = self.config.get_channel_config(channel) if channel_config: + evaluator_kwargs: Dict[str, Any] = { + "config": self.config, + "channel_config": channel_config, + "benchmark_start_time": self.benchmark_start_time, + } self._channel_evaluators[channel] = ( ChannelPerformanceEvaluatorRegistry.get( channel, - config=self.config, - channel_config=channel_config, - benchmark_start_time=self.benchmark_start_time, + **evaluator_kwargs, ) ) @@ -208,6 +216,7 @@ def record_request_completed( with self.lock: self.end_time = completed_at + self._record_realtime_tts_outcome(response) # Update session tracking self._update_session_metrics_for_request( session_id=session_id, @@ -232,14 +241,17 @@ def record_request_completed( self.num_completed_requests += 1 - # Delegate to channel evaluators - for channel, evaluator in self._channel_evaluators.items(): - evaluator.record_request_completed( - request_id=request_id, - session_id=session_id, - completed_at=completed_at, - response=response, - ) + # Delegate to channel evaluators outside the aggregate lock: channel + # evaluators do their own locking, and the audio evaluator's ASR + # scoring is CPU-heavy β€” holding the lock here would serialize all + # completion workers behind a single scoring thread. + for channel, evaluator in self._channel_evaluators.items(): + evaluator.record_request_completed( + request_id=request_id, + session_id=session_id, + completed_at=completed_at, + response=response, + ) if self.config.stream_metrics and self._stream_thread: self._stream_has_updates.set() @@ -355,6 +367,39 @@ def record_session_completed( if session_id in self.session_stats: self._finalize_session(session_id, termination) + def _record_realtime_tts_outcome(self, response: Any) -> None: + """Count realtime start/completion outcomes, including failed requests.""" + if self.client_type != ClientType.REALTIME_TTS: + return + channel_response = getattr(response, "channels", {}).get(ChannelModality.AUDIO) + channel_metrics = ( + channel_response.metrics if channel_response is not None else {} + ) or {} + audio_timestamps = channel_metrics.get( + AudioMetricKey.AUDIO_CHUNK_TIMESTAMPS.value + ) + if audio_timestamps or channel_metrics.get(AudioMetricKey.CHUNK_COUNT.value, 0): + self._realtime_first_audio_count += 1 + if ( + channel_metrics.get(AudioMetricKey.RESPONSE_DONE_OFFSET_MS.value) + is not None + ): + self._realtime_stream_completed_count += 1 + + def _get_realtime_tts_summary(self) -> Dict[str, float]: + if self.client_type != ClientType.REALTIME_TTS: + return {} + count = self.num_requests + return { + "realtime_requests_count": float(count), + "first_audio_success_rate": ( + self._realtime_first_audio_count / count if count > 0 else 0.0 + ), + "stream_completion_rate": ( + self._realtime_stream_completed_count / count if count > 0 else 0.0 + ), + } + def get_aggregated_summary(self) -> Dict[str, float]: """Get aggregate summary metrics.""" return { @@ -374,14 +419,19 @@ def get_aggregated_summary(self) -> Dict[str, float]: "Cancelled Sessions": float(self.num_sessions_cancelled), "Incomplete Sessions": float(self.num_sessions_incomplete), "Observed Session Dispatch Rate": self._session_dispatch_rate(), + **self._get_realtime_tts_summary(), } def _build_summary_stats(self) -> Dict[str, Any]: - """Combine aggregate stats with error code frequencies.""" + """Combine aggregate stats, channel-level metrics, and error code frequencies.""" summary_stats: Dict[str, Any] = { **self.get_aggregated_summary(), "error_code_freq": dict(self.error_code_freq), } + for evaluator in self._channel_evaluators.values(): + channel_summary = evaluator.get_summary() + if channel_summary: + summary_stats.update(channel_summary) return summary_stats def finalize(self) -> EvaluationResult: @@ -395,7 +445,13 @@ def finalize(self) -> EvaluationResult: channel_metrics = {} for channel, evaluator in self._channel_evaluators.items(): + finalize_started_at = time.perf_counter() channel_result = evaluator.finalize() + logger.info( + "Evaluator phase '%s_finalize' took %.2fs", + channel, + time.perf_counter() - finalize_started_at, + ) channel_metrics[str(channel)] = channel_result.metrics combined_metrics.update(channel_result.metrics) @@ -411,17 +467,33 @@ def save(self, output_dir: str) -> None: os.makedirs(output_dir, exist_ok=True) # Save aggregate summary + summary_started_at = time.perf_counter() summary = self._build_summary_stats() summary_path = os.path.join(output_dir, "summary_stats.json") with open(summary_path, "w") as f: json.dump(summary, f, indent=2) + logger.info( + "Evaluator phase 'summary_stats' took %.2fs", + time.perf_counter() - summary_started_at, + ) # Delegate to channel evaluators for channel, evaluator in self._channel_evaluators.items(): + save_started_at = time.perf_counter() evaluator.save(output_dir) + logger.info( + "Evaluator phase '%s_save' took %.2fs", + channel, + time.perf_counter() - save_started_at, + ) # request-level metrics are persisted now + slo_started_at = time.perf_counter() evaluate_and_save_slos(slo_configs=self.config.slos, metrics_dir=output_dir) + logger.info( + "Evaluator phase 'slo_evaluation' took %.2fs", + time.perf_counter() - slo_started_at, + ) def get_streaming_metrics(self) -> Optional[Dict[str, Any]]: """Return current metrics for streaming updates.""" diff --git a/veeksha/evaluator/performance/text.py b/veeksha/evaluator/performance/text.py index 68cb341f..6ed6b0f0 100644 --- a/veeksha/evaluator/performance/text.py +++ b/veeksha/evaluator/performance/text.py @@ -1009,7 +1009,6 @@ def _store_ttfc_violin_plots(self, output_dir: str) -> None: try: import rekha as rk # type: ignore[import-not-found] - from veeksha.evaluator.plot_utils import ( apply_axis_scale, format_axis_label, diff --git a/veeksha/evaluator/registry.py b/veeksha/evaluator/registry.py index eec4c5d3..9ae33d27 100644 --- a/veeksha/evaluator/registry.py +++ b/veeksha/evaluator/registry.py @@ -29,6 +29,14 @@ def get_key_from_str(cls, key_str: str) -> EvaluationType: ), ) +EvaluatorRegistry.register( + EvaluationType.AUDIO_QUALITY, + _LazyLoader( + "veeksha.evaluator.accuracy.audio", + "AudioQualityEvaluator", + ), +) + class ChannelPerformanceEvaluatorRegistry(BaseRegistry): """Registry for channel-specific performance evaluators.""" diff --git a/veeksha/generator/session/trace/__init__.py b/veeksha/generator/session/trace/__init__.py index 32e08372..1ba89cfa 100644 --- a/veeksha/generator/session/trace/__init__.py +++ b/veeksha/generator/session/trace/__init__.py @@ -7,9 +7,13 @@ from veeksha.generator.session.trace.request_log import ( RequestLogTraceFlavorGenerator, ) +from veeksha.generator.session.trace.seed_tts_text import ( + SeedTTSTextTraceFlavorGenerator, +) from veeksha.generator.session.trace.shared_prefix import ( SharedPrefixTraceFlavorGenerator, ) +from veeksha.generator.session.trace.sharegpt import ShareGPTTraceFlavorGenerator from veeksha.generator.session.trace.timed_synthetic_session import ( TimedSyntheticSessionTraceFlavorGenerator, ) @@ -22,4 +26,6 @@ "RAGTraceFlavorGenerator", "RequestLogTraceFlavorGenerator", "UntimedContentMultiTurnTraceFlavorGenerator", + "SeedTTSTextTraceFlavorGenerator", + "ShareGPTTraceFlavorGenerator", ] diff --git a/veeksha/generator/session/trace/audio.py b/veeksha/generator/session/trace/audio.py new file mode 100644 index 00000000..b967be9b --- /dev/null +++ b/veeksha/generator/session/trace/audio.py @@ -0,0 +1,242 @@ +"""Audio trace flavor generator for STT benchmarking. + +Reads a JSONL trace file where each line contains: + {"session_id": 0, "audio_file": "/path/to/audio.wav"} + +Each line must also include an ``expected_transcript`` field for WER +evaluation. + +Each row becomes a single-request session with an AUDIO channel +containing the file path in ``AudioChannelRequestContent.input_audio``. +""" + +import os +from typing import Any, List + +import numpy as np +import pandas as pd + +from veeksha.config.generator.session import ( + AudioTraceFlavorConfig, + TraceSessionGeneratorConfig, +) +from veeksha.core.request import Request +from veeksha.core.request_content import AudioChannelRequestContent +from veeksha.core.seeding import SeedManager +from veeksha.core.session import Session +from veeksha.core.tokenizer import TokenizerProvider +from veeksha.generator.session.trace.base_flavor import ( + TraceFlavorGeneratorBase, +) +from veeksha.logger import init_logger +from veeksha.types import ChannelModality + +logger = init_logger(__name__) + + +class AudioTraceFlavorGenerator(TraceFlavorGeneratorBase): + """Trace flavor that feeds audio file paths for STT benchmarking. + + Each row in the JSONL becomes a single-request session whose AUDIO + channel points to the audio file on disk. + """ + + def __init__( + self, + config: TraceSessionGeneratorConfig, + flavor_config: AudioTraceFlavorConfig, + seed_manager: SeedManager, + tokenizer_provider: TokenizerProvider, + ): + self.config = config + self.flavor_config = flavor_config + self.seed_manager = seed_manager + self.tokenizer_provider = tokenizer_provider + self.tokenizer = tokenizer_provider.for_modality(ChannelModality.TEXT) + + if not os.path.exists(config.trace_file): + raise FileNotFoundError(f"Trace file not found: {config.trace_file}") + + self.trace_df = pd.read_json(config.trace_file, lines=True) + self._validate_trace() + + # Ground truth is mandatory; fail at load, not silently per request. + col = self.trace_df["expected_transcript"] + missing = col.isna() | (col.astype(str).str.strip() == "") + if missing.any(): + examples = self.trace_df.loc[missing, "audio_file"].head(3).tolist() + raise ValueError( + f"{int(missing.sum())} audio trace row(s) missing " + f"expected_transcript (e.g. {examples})." + ) + + # Resolve relative paths against audio_dir when provided, otherwise + # against the manifest directory so manifests are portable. + trace_dir = os.path.dirname(os.path.abspath(config.trace_file)) + audio_base = flavor_config.audio_dir or trace_dir + if not os.path.isabs(audio_base): + audio_base = os.path.join(trace_dir, audio_base) + self.trace_df["audio_file"] = self.trace_df["audio_file"].apply( + lambda p: p if os.path.isabs(str(p)) else os.path.join(audio_base, str(p)) + ) + missing_audio = ~self.trace_df["audio_file"].apply(os.path.exists) + if missing_audio.any(): + examples = self.trace_df.loc[missing_audio, "audio_file"].head(3).tolist() + raise FileNotFoundError( + f"{int(missing_audio.sum())} audio trace file(s) missing " + f"(e.g. {examples})." + ) + + logger.info( + "Loaded %d audio sessions from %s", + len(self.trace_df.groupby("session_id")), + config.trace_file, + ) + + # wrapping state + self._num_wraps = 0 + self._session_groups = None + self._current_session_id = 0 + self._current_request_id = 0 + self._rng = seed_manager.random("trace_shuffling") + self._duration_rng = seed_manager.random("audio_target_duration") + + @property + def required_columns(self) -> List[str]: + return ["session_id", "audio_file", "expected_transcript"] + + def prepare_session(self, group: pd.DataFrame) -> Session: + """Convert a single trace row into a single-request Session.""" + session_id = self._next_session_id() + row = group.iloc[0] + audio_file = str(row["audio_file"]) + + channels = { + ChannelModality.AUDIO: AudioChannelRequestContent( + input_audio=audio_file, + ) + } + + metadata = self._row_metadata(row) + if self.flavor_config.target_duration_s is not None: + metadata = self._apply_target_duration( + metadata, self._sample_target_duration_s() + ) + metadata["audio_file"] = audio_file + + request = Request( + id=self._next_request_id(), + channels=channels, + metadata=metadata, + session_context={ + "node_id": 0, + "wait_after_ready": 0.0, + "parent_nodes": [], + "history_parent": None, + }, + ) + graph = self._build_linear_session_graph(1, [0.0]) + return Session( + id=session_id, + session_graph=graph, + requests={0: request}, + ) + + def _row_metadata(self, row: pd.Series) -> dict: + """Pass manifest metadata through to the STT client/evaluator.""" + metadata: dict[str, Any] = {} + for column in self.trace_df.columns: + if column in {"session_id", "audio_file"}: + continue + value = row[column] + if not isinstance(value, (dict, list)) and pd.isna(value): + continue + if isinstance(value, np.generic): + value = value.item() + metadata[column] = value + metadata["expected_transcript"] = str(metadata["expected_transcript"]) + return metadata + + def _sample_target_duration_s(self) -> float: + """Per-session streamed duration: clipped Gaussian around the median. + + Deterministic under the run seed. Normal(M, sigma) re-drawn until + inside [M - S, M + S]; sigma defaults to S/2 (bounds at 2 sigma, + ~4.6% of draws re-sampled). Symmetry keeps the median at M. + """ + target_duration_s = self.flavor_config.target_duration_s + assert target_duration_s is not None + spread_s = self.flavor_config.target_duration_spread_s + if spread_s is None: + return target_duration_s + sigma_s = self.flavor_config.target_duration_sigma_s + if sigma_s is None: + sigma_s = spread_s / 2.0 + # Clipped Gaussian via rejection: symmetric about the target, so the + # median stays at target_duration_s; the clip bounds are hard limits + # (every clip must be at least target + spread long). + low_s = target_duration_s - spread_s + high_s = target_duration_s + spread_s + while True: + duration_s = self._duration_rng.gauss(target_duration_s, sigma_s) + if low_s <= duration_s <= high_s: + return duration_s + + def _apply_target_duration( + self, metadata: dict[str, Any], target_duration_s: float + ) -> dict[str, Any]: + """Limit a trace row to the target streamed audio prefix. + + Trims ``reference_word_timestamps`` and ``expected_transcript`` to the + words that end within the prefix, and records the slice bounds + (``input_audio_start_ms`` / ``input_audio_end_ms``) that the STT + client uses to cut the decoded PCM before streaming. + """ + target_end_ms = target_duration_s * 1000.0 + word_timestamps = metadata.get("reference_word_timestamps") + if not isinstance(word_timestamps, list): + raise ValueError( + "Audio trace target_duration_s requires reference_word_timestamps " + "for transcript trimming." + ) + + trimmed_timestamps: list[dict[str, Any]] = [] + words: list[str] = [] + for index, word_timing in enumerate(word_timestamps): + if not isinstance(word_timing, dict): + raise ValueError( + "reference_word_timestamps entries must be objects; " + f"got {type(word_timing).__name__} at index {index}" + ) + if "word" not in word_timing or "end_ms" not in word_timing: + raise ValueError( + "reference_word_timestamps entries must contain word and end_ms" + ) + end_ms = float(word_timing["end_ms"]) + if end_ms <= target_end_ms: + trimmed_timestamps.append(dict(word_timing)) + words.append(str(word_timing["word"])) + + # Ground truth is mandatory (see __init__); an empty trimmed + # transcript would make WER scoring degenerate. + if not words: + raise ValueError( + f"No reference word ends within target_duration_s=" + f"{target_duration_s}; the trimmed expected_transcript would " + "be empty." + ) + + trimmed_metadata = dict(metadata) + trimmed_metadata["reference_word_timestamps"] = trimmed_timestamps + trimmed_metadata["expected_transcript"] = " ".join(words) + trimmed_metadata["input_audio_start_ms"] = 0.0 + trimmed_metadata["input_audio_end_ms"] = target_end_ms + trimmed_metadata["duration_s"] = target_duration_s + return trimmed_metadata + + def wrap(self) -> pd.DataFrame: + """Wrap trace for new epoch with shuffled session order.""" + df = self.trace_df.copy() + max_sid = int(df["session_id"].max()) if not df.empty else 0 + df["session_id"] = df["session_id"] + max_sid + 1 + return self._shuffle_sessions(df) diff --git a/veeksha/generator/session/trace/registry.py b/veeksha/generator/session/trace/registry.py index 167123fe..e2a9d793 100644 --- a/veeksha/generator/session/trace/registry.py +++ b/veeksha/generator/session/trace/registry.py @@ -44,3 +44,24 @@ def get_key_from_str(cls, key_str: str) -> TraceFlavorType: "UntimedContentMultiTurnTraceFlavorGenerator", ), ) +TraceFlavorGeneratorRegistry.register( + TraceFlavorType.SHAREGPT, + _LazyLoader( + "veeksha.generator.session.trace.sharegpt", + "ShareGPTTraceFlavorGenerator", + ), +) +TraceFlavorGeneratorRegistry.register( + TraceFlavorType.AUDIO, + _LazyLoader( + "veeksha.generator.session.trace.audio", + "AudioTraceFlavorGenerator", + ), +) +TraceFlavorGeneratorRegistry.register( + TraceFlavorType.SEED_TTS_TEXT, + _LazyLoader( + "veeksha.generator.session.trace.seed_tts_text", + "SeedTTSTextTraceFlavorGenerator", + ), +) diff --git a/veeksha/generator/session/trace/seed_tts_text.py b/veeksha/generator/session/trace/seed_tts_text.py new file mode 100644 index 00000000..33c2fce7 --- /dev/null +++ b/veeksha/generator/session/trace/seed_tts_text.py @@ -0,0 +1,277 @@ +"""Seed TTS text dataset trace flavor for TTS benchmarking.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, List + +import pandas as pd + +from veeksha.config.generator.session import ( + SeedTTSTextTraceFlavorConfig, + TraceSessionGeneratorConfig, +) +from veeksha.core.seeding import SeedManager +from veeksha.core.session import Session +from veeksha.core.tokenizer import TokenizerProvider +from veeksha.generator.session.trace.base_flavor import TraceFlavorGeneratorBase +from veeksha.logger import init_logger +from veeksha.types import ChannelModality + +logger = init_logger(__name__) + + +def _load_hf_dataset(flavor_config: SeedTTSTextTraceFlavorConfig) -> Any: + """Load the Seed TTS text dataset from HF or a local compatible copy.""" + try: + import datasets + except ImportError as exc: + raise ImportError( + "SeedTTSTextTraceFlavorGenerator requires the 'datasets' package." + ) from exc + + if flavor_config.local_path: + path = Path(flavor_config.local_path) + if not path.exists(): + raise FileNotFoundError(f"Seed TTS local_path not found: {path}") + + try: + dataset = datasets.load_from_disk(str(path)) + except Exception: + logger.debug( + "Could not load %s with datasets.load_from_disk; trying load_dataset", + path, + exc_info=True, + ) + else: + return _select_split(dataset, flavor_config.split) + + if path.is_file(): + suffix = path.suffix.lower() + if suffix in {".json", ".jsonl"}: + return datasets.load_dataset( + "json", + data_files=str(path), + split=flavor_config.split, + ) + if suffix == ".csv": + return datasets.load_dataset( + "csv", + data_files=str(path), + split=flavor_config.split, + ) + if suffix == ".parquet": + return datasets.load_dataset( + "parquet", + data_files=str(path), + split=flavor_config.split, + ) + raise ValueError( + f"Unsupported Seed TTS local file extension '{suffix}'. " + "Use JSON/JSONL, CSV, Parquet, or a saved HF dataset directory." + ) + + args = [str(path)] + if flavor_config.subset: + args.append(flavor_config.subset) + return datasets.load_dataset(*args, split=flavor_config.split) + + args = [flavor_config.dataset_name] + if flavor_config.subset: + args.append(flavor_config.subset) + return datasets.load_dataset(*args, split=flavor_config.split) + + +def _select_split(dataset: Any, split: str) -> Any: + """Select a split when a local saved dataset returns a DatasetDict.""" + if hasattr(dataset, "to_pandas"): + return dataset + if isinstance(dataset, dict): + if split not in dataset: + raise ValueError( + f"Seed TTS dataset split '{split}' not found. " + f"Available splits: {sorted(dataset.keys())}" + ) + return dataset[split] + return dataset + + +def _dataset_to_dataframe(dataset: Any) -> pd.DataFrame: + if hasattr(dataset, "to_pandas"): + return dataset.to_pandas() + return pd.DataFrame(list(dataset)) + + +def _word_count(text: str) -> int: + return len(text.split()) + + +def _truncate_to_words(text: str, target_words: int) -> str: + words = text.split() + if len(words) <= target_words: + return text + return " ".join(words[:target_words]) + + +class SeedTTSTextTraceFlavorGenerator(TraceFlavorGeneratorBase): + """Trace flavor that emits one text-only TTS session per Seed TTS row.""" + + def __init__( + self, + config: TraceSessionGeneratorConfig, + flavor_config: SeedTTSTextTraceFlavorConfig, + seed_manager: SeedManager, + tokenizer_provider: TokenizerProvider, + ): + self.config = config + self.flavor_config = flavor_config + self.seed_manager = seed_manager + self.tokenizer_provider = tokenizer_provider + self.tokenizer = tokenizer_provider.for_modality(ChannelModality.TEXT) + + raw_dataset = _load_hf_dataset(flavor_config) + raw_df = _dataset_to_dataframe(raw_dataset) + self.trace_df = self._prepare_trace_df(raw_df) + + self._num_wraps = 0 + self._session_groups = None + self._current_session_id = 0 + self._current_request_id = 0 + self._rng = seed_manager.random("seed_tts_text_shuffling") + self._length_rng = seed_manager.random("seed_tts_text_input_length") + + logger.info( + "Loaded %d Seed TTS text rows from %s/%s split=%s", + len(self.trace_df), + flavor_config.local_path or flavor_config.dataset_name, + flavor_config.subset, + flavor_config.split, + ) + + @property + def required_columns(self) -> List[str]: + return [self.flavor_config.text_column] + + def _prepare_trace_df(self, raw_df: pd.DataFrame) -> pd.DataFrame: + text_col = self.flavor_config.text_column + id_col = self.flavor_config.id_column + + if text_col not in raw_df.columns: + raise ValueError( + f"Seed TTS dataset missing text column '{text_col}'. " + f"Available columns: {list(raw_df.columns)}" + ) + + has_id_col = bool(id_col) and id_col in raw_df.columns + rows = [] + skipped_empty = 0 + skipped_short = 0 + min_length = ( + self.flavor_config.min_chars + if self.flavor_config.use_chars + else self.flavor_config.min_tokens + ) + unit = "chars" if self.flavor_config.use_chars else "words" + + for source_index, row in raw_df.iterrows(): + value = row[text_col] + if pd.isna(value): + skipped_empty += 1 + continue + + text = str(value).strip() + if not text: + skipped_empty += 1 + continue + + char_count = len(text) + word_count = _word_count(text) + current_length = char_count if self.flavor_config.use_chars else word_count + if current_length < min_length: + skipped_short += 1 + continue + + source_id = row[id_col] if has_id_col else source_index + rows.append( + { + "session_id": len(rows), + "text": text, + "source_id": str(source_id), + "source_index": ( + int(source_index) + if isinstance(source_index, int) + else str(source_index) + ), + "word_count": word_count, + "char_count": char_count, + } + ) + + if skipped_empty: + logger.info("Skipped %d empty Seed TTS text rows", skipped_empty) + if skipped_short: + logger.info( + "Skipped %d Seed TTS text rows shorter than %d %s", + skipped_short, + min_length, + unit, + ) + + if not rows: + raise ValueError( + f"Seed TTS dataset contains no rows in column '{text_col}' with " + f"at least {min_length} {unit}." + ) + + return pd.DataFrame(rows) + + def prepare_session(self, group: pd.DataFrame) -> Session: + session_id = self._next_session_id() + row = group.iloc[0] + text = str(row["text"]) + + if self.flavor_config.use_chars: + target_chars = self._length_rng.randint( + self.flavor_config.min_chars, + self.flavor_config.max_chars, + ) + if len(text) > target_chars: + text = text[:target_chars] + else: + target_words = self._length_rng.randint( + self.flavor_config.min_tokens, + self.flavor_config.max_tokens, + ) + text = _truncate_to_words(text, target_words) + + target_prompt_tokens = len(self.tokenizer.encode(text)) + + request = self._create_text_request( + node_id=0, + prompt_text=text, + target_output_tokens=1, + wait_after_ready=0.0, + parent_nodes=[], + target_prompt_tokens=target_prompt_tokens, + ) + request.metadata.update( + { + "dataset": self.flavor_config.dataset_name, + "dataset_subset": self.flavor_config.subset, + "dataset_split": self.flavor_config.split, + "dataset_source": self.flavor_config.local_path or "huggingface", + "source_id": row["source_id"], + "source_index": row["source_index"], + "input_words": _word_count(text), + "input_chars": len(text), + } + ) + + graph = self._build_linear_session_graph(1, [0.0]) + return Session(id=session_id, session_graph=graph, requests={0: request}) + + def wrap(self) -> pd.DataFrame: + df = self.trace_df.copy() + max_sid = int(df["session_id"].max()) if not df.empty else 0 + df["session_id"] = df["session_id"] + max_sid + 1 + return self._shuffle_sessions(df) diff --git a/veeksha/generator/session/trace/sharegpt.py b/veeksha/generator/session/trace/sharegpt.py new file mode 100644 index 00000000..150493e2 --- /dev/null +++ b/veeksha/generator/session/trace/sharegpt.py @@ -0,0 +1,255 @@ +"""ShareGPT trace flavor generator for TTS benchmarking. + +Reads ShareGPT-format conversation data and uses assistant turn text +as TTS request input. Each assistant turn becomes its own single-request +session. Input text is truncated to a length sampled uniformly between +min_tokens..max_tokens (token mode, default) or min_chars..max_chars +(char mode, when chars are set on the flavor config). +""" + +import json +import os +from typing import Any, Dict, List + +import pandas as pd + +from veeksha.config.generator.session import ( + ShareGPTTraceFlavorConfig, + TraceSessionGeneratorConfig, +) +from veeksha.core.seeding import SeedManager +from veeksha.core.session import Session +from veeksha.core.tokenizer import TokenizerProvider +from veeksha.generator.session.trace.base_flavor import ( + TraceFlavorGeneratorBase, +) +from veeksha.logger import init_logger + +logger = init_logger(__name__) + + +def _load_sharegpt_file(path: str) -> List[Dict[str, Any]]: + """Load a ShareGPT file (JSON array or JSONL).""" + with open(path, "r", encoding="utf-8") as f: + content = f.read().strip() + + # Try JSON array first + if content.startswith("["): + return json.loads(content) + + # Fall back to JSONL + conversations = [] + for line in content.splitlines(): + line = line.strip() + if line: + conversations.append(json.loads(line)) + return conversations + + +def _alpha_ratio(text: str) -> float: + """Return ratio of alphabetic characters to total non-space characters.""" + non_space = [c for c in text if not c.isspace()] + if not non_space: + return 0.0 + return sum(1 for c in non_space if c.isalpha()) / len(non_space) + + +def _flatten_to_dataframe( + conversations: List[Dict[str, Any]], + assistant_role: str, + tokenizer: Any, + min_length: int, + use_chars: bool, + min_alpha_ratio: float = 0.5, +) -> pd.DataFrame: + """Convert ShareGPT conversations to a flat DataFrame of assistant turns. + + Each assistant turn gets its own session_id (1 request per session). + Turns shorter than min_length are skipped, where the unit is chars + (use_chars=True) or tokens (use_chars=False). + Turns with alpha ratio below min_alpha_ratio are skipped (filters out + number sequences, code snippets, etc.). + Returns DataFrame with columns: session_id, text, token_count, char_count + """ + unit = "chars" if use_chars else "tokens" + rows = [] + session_id = 0 + skipped_length = 0 + skipped_alpha = 0 + for conv in conversations: + turns = conv.get("conversations", []) + for turn in turns: + role = turn.get("from", "") + if role == assistant_role: + text = turn.get("value", "").strip() + if not text: + continue + # Filter out junk text (number sequences, code, etc.) + if min_alpha_ratio > 0 and _alpha_ratio(text) < min_alpha_ratio: + skipped_alpha += 1 + continue + char_count = len(text) + if use_chars: + # Skip token encode in char mode -- it's the hot path filter. + if char_count < min_length: + skipped_length += 1 + continue + token_count = -1 + else: + token_count = len(tokenizer.encode(text)) + if token_count < min_length: + skipped_length += 1 + continue + rows.append( + { + "session_id": session_id, + "text": text, + "token_count": token_count, + "char_count": char_count, + } + ) + session_id += 1 + + if skipped_length: + logger.info( + "Skipped %d assistant turns shorter than %d %s", + skipped_length, + min_length, + unit, + ) + if skipped_alpha: + logger.info( + "Skipped %d assistant turns with alpha ratio below %.2f " + "(number sequences, code, etc.)", + skipped_alpha, + min_alpha_ratio, + ) + + if not rows: + min_key = "min_chars" if use_chars else "min_tokens" + raise ValueError( + f"No assistant turns found with role='{assistant_role}' " + f"and >= {min_length} {unit} and alpha_ratio >= {min_alpha_ratio}. " + f"Check 'assistant_role', '{min_key}', and 'min_alpha_ratio' config." + ) + + return pd.DataFrame(rows) + + +class ShareGPTTraceFlavorGenerator(TraceFlavorGeneratorBase): + """Trace flavor that reads ShareGPT conversations for TTS benchmarking. + + Each assistant turn becomes a single-request session. Input text is + truncated to a length sampled uniformly between min_tokens..max_tokens + (token mode) or min_chars..max_chars (char mode, when chars are set). + Turns shorter than the minimum are skipped entirely. + """ + + def __init__( + self, + config: TraceSessionGeneratorConfig, + flavor_config: ShareGPTTraceFlavorConfig, + seed_manager: SeedManager, + tokenizer_provider: TokenizerProvider, + ): + # We override the base __init__ because ShareGPT uses a different + # file format (JSON array / JSONL with "conversations" field) rather + # than the standard JSONL with session_id/input_length columns. + self.config = config + self.flavor_config = flavor_config + self.seed_manager = seed_manager + self.tokenizer_provider = tokenizer_provider + from veeksha.types import ChannelModality + + self.tokenizer = tokenizer_provider.for_modality(ChannelModality.TEXT) + + if not os.path.exists(config.trace_file): + raise FileNotFoundError(f"Trace file not found: {config.trace_file}") + + raw_conversations = _load_sharegpt_file(config.trace_file) + logger.info( + "Loaded %d conversations from %s", + len(raw_conversations), + config.trace_file, + ) + + use_chars = flavor_config.use_chars + min_length = flavor_config.min_chars if use_chars else flavor_config.min_tokens + self.trace_df = _flatten_to_dataframe( + raw_conversations, + flavor_config.assistant_role, + self.tokenizer, + min_length, + use_chars, + flavor_config.min_alpha_ratio, + ) + logger.info( + "Extracted %d assistant turns (1 session each, min_%s=%d)", + len(self.trace_df), + "chars" if use_chars else "tokens", + min_length, + ) + + # wrapping state (same as base) + self._num_wraps = 0 + self._session_groups = None + self._current_session_id = 0 + self._current_request_id = 0 + self._rng = seed_manager.random("trace_shuffling") + self._length_rng = seed_manager.random("sharegpt_input_length") + + @property + def required_columns(self) -> List[str]: + return ["session_id", "text"] + + def _truncate_to_tokens(self, text: str, target_tokens: int) -> str: + """Truncate text to exactly target_tokens by encoding and decoding.""" + token_ids = self.tokenizer.encode(text) + if len(token_ids) <= target_tokens: + return text + truncated_ids = token_ids[:target_tokens] + return self.tokenizer.decode(truncated_ids) + + def prepare_session(self, group: pd.DataFrame) -> Session: + """Convert a single assistant turn into a single-request Session.""" + session_id = self._next_session_id() + row = group.iloc[0] + text = str(row["text"]) + + if self.flavor_config.use_chars: + target_chars = self._length_rng.randint( + self.flavor_config.min_chars, + self.flavor_config.max_chars, + ) + if len(text) > target_chars: + text = text[:target_chars] + actual_tokens = len(self.tokenizer.encode(text)) + else: + target_tokens = self._length_rng.randint( + self.flavor_config.min_tokens, + self.flavor_config.max_tokens, + ) + text = self._truncate_to_tokens(text, target_tokens) + actual_tokens = len(self.tokenizer.encode(text)) + + request = self._create_text_request( + node_id=0, + prompt_text=text, + target_output_tokens=1, # TTS doesn't use output tokens + wait_after_ready=0.0, + parent_nodes=[], + target_prompt_tokens=actual_tokens, + ) + graph = self._build_linear_session_graph(1, [0.0]) + return Session( + id=session_id, + session_graph=graph, + requests={0: request}, + ) + + def wrap(self) -> pd.DataFrame: + """Wrap trace for new epoch with shuffled session order.""" + df = self.trace_df.copy() + max_sid = int(df["session_id"].max()) if not df.empty else 0 + df["session_id"] = df["session_id"] + max_sid + 1 + return self._shuffle_sessions(df) diff --git a/veeksha/launcher.py b/veeksha/launcher.py new file mode 100644 index 00000000..56ab5b50 --- /dev/null +++ b/veeksha/launcher.py @@ -0,0 +1,91 @@ +"""Public CLI/API boundary for orchestrated Veeksha launcher sweeps. + +Configuration lives in ``veeksha.config.launcher`` and lifecycle orchestration +lives in ``veeksha.orchestration``. This module is intentionally thin, matching +the public-module shape used by ``veeksha.capacity_search``. +""" + +from __future__ import annotations + +import argparse +import logging +import resource +import sys +from pathlib import Path + +from veeksha.config.launcher import LauncherConfig, LauncherConfigError, RetryConfig +from veeksha.orchestration.launcher import BenchmarkAttemptResult, LauncherOrchestrator + +logger = logging.getLogger(__name__) + + +def _raise_open_file_limit() -> None: + """Raise the soft open-file limit for high-concurrency launcher sweeps.""" + try: + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + except (OSError, ValueError) as exc: + logger.warning("Could not read RLIMIT_NOFILE: %s", exc) + return + + desired = 1 << 20 + if hard != resource.RLIM_INFINITY: + desired = min(desired, hard) + if desired <= soft: + return + + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (desired, hard)) + logger.info("Raised RLIMIT_NOFILE soft limit from %s to %s", soft, desired) + except (OSError, ValueError) as exc: + logger.warning("Could not raise RLIMIT_NOFILE to %s: %s", desired, exc) + + +def run_launcher(config: LauncherConfig) -> None: + """Run an already-parsed launcher config.""" + LauncherOrchestrator(config).run() + + +def run_launcher_file(path: str | Path) -> None: + """Load a launcher YAML file and run it.""" + run_launcher(LauncherConfig.from_file(path)) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="veeksha launcher", + description="Launch an engine and run an orchestrated Veeksha sweep.", + ) + parser.add_argument( + "--config", required=True, help="Path to Veeksha launcher YAML config." + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + _raise_open_file_limit() + parser = build_parser() + args = parser.parse_args(argv) + try: + run_launcher_file(args.config) + except LauncherConfigError as exc: + parser.error(str(exc)) + except KeyboardInterrupt: + return 130 + return 0 + + +__all__ = [ + "BenchmarkAttemptResult", + "LauncherConfig", + "LauncherConfigError", + "LauncherOrchestrator", + "RetryConfig", + "build_parser", + "main", + "run_launcher", + "run_launcher_file", +] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/veeksha/orchestration/__init__.py b/veeksha/orchestration/__init__.py index 190c6b43..9eb94401 100644 --- a/veeksha/orchestration/__init__.py +++ b/veeksha/orchestration/__init__.py @@ -2,18 +2,32 @@ create_server_manager, managed_server, ) +from veeksha.orchestration.launcher import LauncherOrchestrator +from veeksha.orchestration.managed_engines import ( + BaseEngineRunner, + EngineError, + EngineRestartLimitExceeded, + SglangOmniDockerRunner, + VajraSubprocessRunner, + VllmOmniDockerRunner, + create_engine_runner, +) from veeksha.orchestration.registry import ServerManagerRegistry from veeksha.orchestration.resource_manager import ResourceManager from veeksha.orchestration.server_manager import BaseServerManager -from veeksha.orchestration.sglang_server import SGLangServerManager -from veeksha.orchestration.vllm_server import VLLMServerManager __all__ = [ "BaseServerManager", - "VLLMServerManager", - "SGLangServerManager", "create_server_manager", "ServerManagerRegistry", "managed_server", "ResourceManager", + "LauncherOrchestrator", + "BaseEngineRunner", + "EngineError", + "EngineRestartLimitExceeded", + "VajraSubprocessRunner", + "VllmOmniDockerRunner", + "SglangOmniDockerRunner", + "create_engine_runner", ] diff --git a/veeksha/orchestration/benchmark_orchestrator.py b/veeksha/orchestration/benchmark_orchestrator.py index fb976fda..365f127e 100644 --- a/veeksha/orchestration/benchmark_orchestrator.py +++ b/veeksha/orchestration/benchmark_orchestrator.py @@ -7,14 +7,14 @@ from typing import Any, Dict, Generator from veeksha.config.server import BaseServerConfig +from veeksha.orchestration.managed_engines import BaseEngineRunner from veeksha.orchestration.registry import ServerManagerRegistry -from veeksha.orchestration.server_manager import BaseServerManager def create_server_manager( config: BaseServerConfig, output_dir: str, -) -> BaseServerManager: +) -> BaseEngineRunner: """Create appropriate server manager based on config type.""" return ServerManagerRegistry.get( config.get_type(), @@ -42,6 +42,7 @@ def managed_server( Yields: Dictionary with server info: + - endpoint: Endpoint config - api_base: API base URL - api_key: API key - server_manager: Server manager instance @@ -49,28 +50,19 @@ def managed_server( server_manager = create_server_manager(config, output_dir=output_dir) try: - launch_result = server_manager.launch() - if isinstance(launch_result, tuple): - success, error = launch_result - else: - success, error = bool(launch_result), None - if not success: - raise RuntimeError(f"Failed to launch server: {error}") + server_manager.start() + endpoint = server_manager.get_endpoint() - if not server_manager.wait_for_ready(): - raise RuntimeError("Server failed to become ready") - - api_base = config.get_api_base_url() - - # Set environment variables for clients - os.environ["OPENAI_API_KEY"] = config.api_key or "EMPTY" - os.environ["OPENAI_API_BASE"] = api_base + # Set environment variables for clients. + os.environ["OPENAI_API_KEY"] = endpoint.api_key or "EMPTY" + os.environ["OPENAI_API_BASE"] = endpoint.api_base yield { - "api_base": api_base, - "api_key": config.api_key, + "endpoint": endpoint, + "api_base": endpoint.api_base, + "api_key": endpoint.api_key, "server_manager": server_manager, } finally: - server_manager.shutdown() + server_manager.stop() diff --git a/veeksha/orchestration/launcher.py b/veeksha/orchestration/launcher.py new file mode 100644 index 00000000..fe3d4ecb --- /dev/null +++ b/veeksha/orchestration/launcher.py @@ -0,0 +1,452 @@ +"""Orchestrator that runs a sweep under an optional managed server.""" + +from __future__ import annotations + +import json +import os +import subprocess +import time +from contextlib import ExitStack +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable, Optional + +import yaml + +from veeksha.config.endpoint import EndpointConfig +from veeksha.config.launcher import LauncherConfig +from veeksha.orchestration.launcher_events import ( + attempt_log_name, + console_message, + run_event_payload, + sweep_plan_payload, +) +from veeksha.orchestration.launcher_progress import ( + BenchmarkProgressReader, + BenchmarkRequestProgress, + LauncherProgressReporter, + request_progress_payload, +) +from veeksha.orchestration.managed_engines import ( + BaseEngineRunner, + EngineRestartLimitExceeded, + create_engine_runner, +) +from veeksha.orchestration.processes import ProcessTerminator +from veeksha.sweeps import planner as sweep_planner + +_BENCHMARK_PROGRESS_INTERVAL_SECONDS = 30.0 + + +@dataclass(frozen=True) +class BenchmarkAttemptResult: + success: bool + reason: str + returncode: Optional[int] = None + request_progress: Optional[BenchmarkRequestProgress] = None + + +class LauncherOrchestrator: + def __init__( + self, + config: LauncherConfig, + *, + engine_runner: Optional[BaseEngineRunner] = None, + popen_factory: Callable[..., subprocess.Popen] = subprocess.Popen, + sleep_fn: Callable[[float], None] = time.sleep, + monotonic_fn: Callable[[], float] = time.monotonic, + terminator: Optional[ProcessTerminator] = None, + progress_reporter: Optional[LauncherProgressReporter] = None, + ): + self.config = config + self.output_dir = Path(config.output_dir) + self.server_output_dir = self.output_dir / "server" + self.benchmark_log_dir = self.output_dir / "benchmark_logs" + self.generated_config_dir = self.output_dir / "generated_configs" + self._engine = engine_runner + self._popen_factory = popen_factory + self._sleep = sleep_fn + self._monotonic = monotonic_fn + self._terminator = terminator or ProcessTerminator() + self._event_file = None + self._progress = progress_reporter or LauncherProgressReporter() + + def run(self) -> None: + self.output_dir.mkdir(parents=True, exist_ok=True) + self.benchmark_log_dir.mkdir(parents=True, exist_ok=True) + self._persist_launcher_config() + + events_path = self.output_dir / "launcher_events.jsonl" + with ExitStack() as stack: + self._event_file = stack.enter_context( + events_path.open("a", encoding="utf-8") + ) + stack.callback(self._clear_event_file) + stack.callback(self._progress.close) + + self._record_event( + "launcher_start", + output_dir=str(self.output_dir), + events_file=str(events_path), + benchmark_log_dir=str(self.benchmark_log_dir), + ) + + engine = self._resolve_engine() + endpoint = self.config.endpoint + if engine is not None: + endpoint = engine.get_endpoint() + self._validate_endpoint(endpoint) + plan = sweep_planner.build_sweep_plan_from_config( + self.config.sweep, + endpoint=endpoint, + tmp_parent=self.generated_config_dir, + ) + self._record_event( + "sweep_plan_ready", + **sweep_plan_payload(plan, self.generated_config_dir), + ) + + if engine is not None: + self._record_event("engine_start", **self._engine_event_payload(engine)) + engine.start() + stack.callback(self._stop_engine, engine) + self._record_event("engine_ready", **self._engine_event_payload(engine)) + elif endpoint is not None: + self._record_event( + "endpoint_external", **self._endpoint_event_payload(endpoint) + ) + else: + self._record_event("engine_unmanaged") + for run in plan.runs: + self._run_descriptor(engine, plan.spec, run) + self._record_event("sweep_complete", runs=len(plan.runs)) + + def _clear_event_file(self) -> None: + self._event_file = None + + def _stop_engine(self, engine: BaseEngineRunner) -> None: + self._record_event("engine_stop", **self._engine_event_payload(engine)) + engine.stop() + self._record_event("engine_stopped", **self._engine_event_payload(engine)) + + def _resolve_engine(self) -> Optional[BaseEngineRunner]: + if self._engine is not None: + return self._engine + if self.config.server is None: + return None + self.server_output_dir.mkdir(parents=True, exist_ok=True) + return create_engine_runner(self.config.server, self.server_output_dir) + + def _validate_endpoint(self, endpoint: Optional[EndpointConfig]) -> None: + if endpoint is None: + return + if endpoint.engine_type != self.config.sweep.engine: + raise RuntimeError( + "endpoint engine_type must match sweep.engine " + f"(endpoint.engine_type={endpoint.engine_type}, " + f"sweep.engine={self.config.sweep.engine})" + ) + + def _engine_event_payload(self, engine: BaseEngineRunner) -> dict[str, object]: + endpoint = engine.get_endpoint() + payload = { + "runner": engine.__class__.__name__, + "engine_type": endpoint.engine_type, + "model": endpoint.model, + "api_base": endpoint.api_base, + "health_url": endpoint.health_url or engine.config.health_check_url, + "engine_log_dir": str(engine.output_dir), + } + container_name = getattr(engine, "container_name", None) + if isinstance(container_name, str): + payload["container"] = container_name + return payload + + def _endpoint_event_payload(self, endpoint: EndpointConfig) -> dict[str, object]: + return { + "engine_type": endpoint.engine_type, + "model": endpoint.model, + "api_base": endpoint.api_base, + "health_url": endpoint.health_url, + "host": endpoint.host, + "port": endpoint.port, + } + + def _run_descriptor( + self, + engine: Optional[BaseEngineRunner], + spec: sweep_planner.SweepSpec, + run: sweep_planner.SweepRunDescriptor, + ) -> None: + if engine is not None: + engine.reset_restart_budget() + max_attempts = self.config.retry.max_attempts_per_run + attempt = 1 + while attempt <= max_attempts: + if engine is not None and not self._ensure_engine_ready(engine): + self._record_run_exhausted( + spec, + run, + f"sweep run {run.run_index}/{run.run_count} aborted: " + "engine restart budget exhausted", + ) + return + stdout_path, stderr_path = self._attempt_log_paths(run.run_index, attempt) + start_payload = run_event_payload( + spec, run, completed_runs=run.run_index - 1 + ) + start_payload.update( + attempt=attempt, + command=run.command, + config=str(run.run_config), + output_dir=run.output_dir, + stdout_log=str(stdout_path), + stderr_log=str(stderr_path), + timeout_seconds=run.timeout_seconds, + **request_progress_payload( + BenchmarkRequestProgress( + completed_requests=0, + total_requests=( + run.max_sessions if run.max_sessions > 0 else None + ), + ) + ), + ) + self._record_event("benchmark_attempt_start", **start_payload) + result = self._run_benchmark_attempt( + engine, spec, run, attempt, stdout_path, stderr_path + ) + if result.success: + success_payload = run_event_payload( + spec, run, completed_runs=run.run_index + ) + success_payload.update(attempt=attempt, returncode=result.returncode) + if result.request_progress is not None: + success_payload.update( + request_progress_payload(result.request_progress) + ) + self._record_event("benchmark_attempt_success", **success_payload) + if run.run_index < run.run_count and self._cooldown_seconds > 0: + cooldown_payload = run_event_payload( + spec, run, completed_runs=run.run_index + ) + cooldown_payload.update( + seconds=self._cooldown_seconds, + next_run_index=run.run_index + 1, + ) + self._record_event("cooldown_start", **cooldown_payload) + self._sleep(self._cooldown_seconds) + return + + failure_payload = run_event_payload( + spec, run, completed_runs=run.run_index - 1 + ) + failure_payload.update( + attempt=attempt, + reason=result.reason, + returncode=result.returncode, + ) + if result.request_progress is not None: + failure_payload.update( + request_progress_payload(result.request_progress) + ) + self._record_event("benchmark_attempt_failed", **failure_payload) + if attempt >= max_attempts: + self._record_run_exhausted( + spec, + run, + f"sweep run {run.run_index}/{run.run_count} failed after " + f"{max_attempts} attempts: {result.reason}", + ) + return + + if engine is not None and ( + result.reason.startswith("engine_") + or self.config.retry.restart_engine_before_retry + ): + if not self._restart_engine(engine, result.reason): + self._record_run_exhausted( + spec, + run, + f"sweep run {run.run_index}/{run.run_count} aborted: " + "engine restart budget exhausted", + ) + return + attempt += 1 + + def _record_run_exhausted( + self, + spec: sweep_planner.SweepSpec, + run: sweep_planner.SweepRunDescriptor, + message: str, + ) -> None: + payload = run_event_payload(spec, run, completed_runs=run.run_index - 1) + payload.update(message=message) + self._record_event("benchmark_attempts_exhausted", **payload) + if self.config.retry.fail_sweep_after_exhausted_retries: + raise RuntimeError(message) + + @property + def _cooldown_seconds(self) -> float: + return float(self.config.sweep.cooldown_seconds or 0) + + def _ensure_engine_ready(self, engine: BaseEngineRunner) -> bool: + """Ensure the engine is alive and healthy, restarting it if needed. + + Returns False only when a required restart hit the restart budget, so + the caller can fail the run gracefully instead of crashing the sweep. + """ + if not engine.is_alive(): + return self._restart_engine(engine, "engine_not_alive_before_run") + if not engine.health_check(): + return self._restart_engine(engine, "engine_unhealthy_before_run") + return True + + def _restart_engine(self, engine: BaseEngineRunner, reason: str) -> bool: + """Restart the engine, emitting events. Return False if the budget is exhausted.""" + self._record_event( + "engine_restart", reason=reason, **self._engine_event_payload(engine) + ) + try: + engine.restart() + except EngineRestartLimitExceeded as exc: + self._record_event( + "engine_restart_exhausted", + reason=reason, + message=str(exc), + **self._engine_event_payload(engine), + ) + return False + self._record_event("engine_ready", **self._engine_event_payload(engine)) + return True + + def _run_benchmark_attempt( + self, + engine: Optional[BaseEngineRunner], + spec: sweep_planner.SweepSpec, + run: sweep_planner.SweepRunDescriptor, + attempt: int, + stdout_path: Path, + stderr_path: Path, + ) -> BenchmarkAttemptResult: + progress_path = self.benchmark_log_dir / attempt_log_name( + run.run_index, attempt, "progress.json" + ) + with ( + stdout_path.open("w", encoding="utf-8") as stdout_file, + stderr_path.open("w", encoding="utf-8") as stderr_file, + ): + process = self._popen_factory( + run.command, + cwd=str(sweep_planner.REPO_ROOT), + stdout=stdout_file, + stderr=stderr_file, + text=True, + start_new_session=True, + env={**os.environ, "VEEKSHA_PROGRESS_FILE": str(progress_path)}, + ) + started_at = self._monotonic() + progress_reader = BenchmarkProgressReader(progress_path, run.max_sessions) + if engine is None: + last_health_check = 0.0 + else: + last_health_check = started_at - engine.config.health_check_interval + last_progress_report = started_at + while True: + returncode = process.poll() + if returncode is not None: + return BenchmarkAttemptResult( + success=returncode == 0, + reason="completed" if returncode == 0 else "benchmark_failed", + returncode=returncode, + request_progress=progress_reader.read(), + ) + + if engine is not None and not engine.is_alive(): + self._terminate_benchmark_process(process) + return BenchmarkAttemptResult( + success=False, + reason="engine_exited", + request_progress=progress_reader.read(), + ) + + poll_interval = 1.0 + now = self._monotonic() + request_progress = progress_reader.read() + self._progress.update_attempt_requests( + request_progress.completed_requests, request_progress.total_requests + ) + if engine is not None: + poll_interval = min(1.0, engine.config.health_check_interval) + if now - last_health_check >= engine.config.health_check_interval: + last_health_check = now + if not engine.health_check(): + self._terminate_benchmark_process(process) + return BenchmarkAttemptResult( + success=False, + reason="engine_unhealthy", + request_progress=progress_reader.read(), + ) + + if now - last_progress_report >= _BENCHMARK_PROGRESS_INTERVAL_SECONDS: + last_progress_report = now + self._record_benchmark_progress( + spec, + run, + attempt, + elapsed_seconds=now - started_at, + request_progress=request_progress, + ) + + self._sleep(poll_interval) + + def _terminate_benchmark_process(self, process: subprocess.Popen) -> None: + self._terminator.terminate(process) + + def _record_benchmark_progress( + self, + spec: sweep_planner.SweepSpec, + run: sweep_planner.SweepRunDescriptor, + attempt: int, + *, + elapsed_seconds: float, + request_progress: BenchmarkRequestProgress, + ) -> None: + payload = run_event_payload(spec, run, completed_runs=run.run_index - 1) + payload.update( + attempt=attempt, + elapsed_seconds=round(elapsed_seconds, 1), + timeout_seconds=run.timeout_seconds, + **request_progress_payload(request_progress), + ) + self._record_event("benchmark_attempt_progress", **payload) + + def _attempt_log_paths(self, run_index: int, attempt: int) -> tuple[Path, Path]: + return ( + self.benchmark_log_dir / attempt_log_name(run_index, attempt, "stdout.log"), + self.benchmark_log_dir / attempt_log_name(run_index, attempt, "stderr.log"), + ) + + def _persist_launcher_config(self) -> None: + path = self.output_dir / "launcher_config.yml" + with path.open("w", encoding="utf-8") as f: + yaml.safe_dump(self.config.to_dict(), f, sort_keys=False) + + def _record_event(self, event: str, **payload) -> None: + if self._event_file is None: + return + record = { + "timestamp": datetime.now(timezone.utc).isoformat(), + "event": event, + **payload, + } + self._event_file.write(json.dumps(record, sort_keys=True) + "\n") + self._event_file.flush() + self._progress.handle_event(event, payload) + self._print_event(event, payload) + + def _print_event(self, event: str, payload: dict) -> None: + message = console_message(event, payload) + if message is not None: + self._progress.write(message) diff --git a/veeksha/orchestration/launcher_events.py b/veeksha/orchestration/launcher_events.py new file mode 100644 index 00000000..fd09735e --- /dev/null +++ b/veeksha/orchestration/launcher_events.py @@ -0,0 +1,210 @@ +"""Launcher event payloads and console messages for orchestration.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Optional + +from veeksha.orchestration.launcher_progress import ( + positive_int_or_none, + progress_percentage, +) + + +def sweep_plan_payload(plan: Any, generated_config_dir: Path) -> dict: + payload = { + "runs": len(plan.runs), + "generated_config_dir": str(generated_config_dir), + "sweep_type": plan.spec.sweep_type, + "engine": plan.spec.engine, + "model": plan.spec.model, + "trace": plan.runs[0].trace_source if plan.runs else "unknown", + "base_config": str(plan.base_config_path), + "concurrencies": sorted({run.concurrency for run in plan.runs}), + "input_sizes": sorted( + {run.input_size for run in plan.runs if run.input_size is not None} + ), + } + if getattr(plan, "endpoint", None) is not None: + payload["endpoint"] = plan.endpoint.to_dict() + return payload + + +def run_event_payload(spec: Any, run: Any, *, completed_runs: int) -> dict: + completed_runs = min(max(completed_runs, 0), run.run_count) + return { + "sweep_type": spec.sweep_type, + "engine": spec.engine, + "model": spec.model, + "trace": run.trace_source, + "run_index": run.run_index, + "run_count": run.run_count, + "run_name": run.run_name, + "concurrency": run.concurrency, + "input_size": run.input_size, + "completed_runs": completed_runs, + "remaining_runs": max(run.run_count - completed_runs, 0), + "sweep_progress_pct": progress_percentage(completed_runs, run.run_count), + } + + +def console_message(event: str, payload: dict) -> Optional[str]: + if event == "launcher_start": + return ( + f"output={payload['output_dir']} events={payload['events_file']} " + f"benchmark_logs={payload['benchmark_log_dir']}" + ) + if event == "sweep_plan_ready": + dimensions = _plan_dimensions(payload) + if dimensions: + dimensions = f" {dimensions}" + return ( + f"sweep plan ready: {_sweep_details(payload)} runs={payload['runs']}" + f"{dimensions} generated_configs={payload['generated_config_dir']} " + f"base_config={payload['base_config']}" + ) + if event == "engine_unmanaged": + return "no server or endpoint configured; running sweep against config api_base" + if event == "endpoint_external": + return f"using external endpoint: {_endpoint_details(payload)}" + if event == "engine_start": + return f"starting engine: {_engine_details(payload)}" + if event == "engine_ready": + return f"engine ready: api_base={payload['api_base']}" + if event == "engine_restart": + return ( + f"restarting engine: reason={payload['reason']} {_engine_details(payload)}" + ) + if event == "engine_restart_exhausted": + return ( + f"engine restart budget exhausted: reason={payload['reason']} " + f"{_engine_details(payload)}" + ) + if event == "engine_stop": + return f"stopping engine: {_engine_details(payload)}" + if event == "engine_stopped": + return "engine stopped" + if event == "benchmark_attempt_start": + return ( + f"run {payload['run_index']}/{payload['run_count']} attempt " + f"{payload['attempt']} starting: {_run_details(payload)} " + f"sweep_progress={_sweep_progress(payload)} config={payload['config']} " + f"stdout={payload['stdout_log']} stderr={payload['stderr_log']}" + ) + if event == "benchmark_attempt_progress": + return ( + f"run {payload['run_index']}/{payload['run_count']} attempt " + f"{payload['attempt']} running: " + f"requests_processed={_request_progress(payload)} " + f"sweep_progress={_sweep_progress(payload)}" + ) + if event == "benchmark_attempt_success": + return ( + f"run {payload['run_index']}/{payload['run_count']} attempt " + f"{payload['attempt']} succeeded rc={payload['returncode']}; " + f"requests_processed={_request_progress(payload)} " + f"sweep_progress={_sweep_progress(payload)}" + ) + if event == "benchmark_attempt_failed": + return ( + f"run {payload['run_index']}/{payload['run_count']} attempt " + f"{payload['attempt']} failed reason={payload['reason']} " + f"rc={payload.get('returncode')}; " + f"requests_processed={_request_progress(payload)} " + f"sweep_progress={_sweep_progress(payload)}" + ) + if event == "benchmark_attempts_exhausted": + return f"{payload['message']}; sweep_progress={_sweep_progress(payload)}" + if event == "cooldown_start": + return ( + f"cooldown before next run: {payload['seconds']}s after " + f"run {payload['run_index']}/{payload['run_count']}; " + f"sweep_progress={_sweep_progress(payload)} " + f"next_run={payload['next_run_index']}/{payload['run_count']}" + ) + if event == "sweep_complete": + return f"sweep complete: {payload['runs']} run(s)" + return None + + +def attempt_log_name(run_index: int, attempt: int, suffix: str) -> str: + return f"run_{run_index:03d}_attempt_{attempt:02d}_{suffix}" + + +def _sweep_details(payload: dict) -> str: + return ( + f"sweep={payload['sweep_type']} engine={payload['engine']} " + f"model={payload['model']} trace={payload['trace']}" + ) + + +def _plan_dimensions(payload: dict) -> str: + parts = [] + if payload.get("concurrencies"): + parts.append(f"concurrencies={_format_values(payload['concurrencies'])}") + if payload.get("input_sizes"): + parts.append(f"input_sizes={_format_values(payload['input_sizes'])}") + return " ".join(parts) + + +def _run_details(payload: dict) -> str: + parts = [ + _sweep_details(payload), + f"run={payload['run_name']}", + f"concurrency={payload['concurrency']}", + ] + if payload.get("input_size") is not None: + parts.append(f"input_size={payload['input_size']}") + return " ".join(parts) + + +def _sweep_progress(payload: dict) -> str: + return ( + f"{payload['completed_runs']}/{payload['run_count']} complete " + f"({float(payload['sweep_progress_pct']):.1f}%)" + ) + + +def _request_progress(payload: dict) -> str: + completed_requests = int(payload["requests_completed"]) + request_total = positive_int_or_none(payload.get("request_total")) + if request_total is None: + return str(completed_requests) + request_pct = payload.get("request_progress_pct") + if request_pct is None: + request_pct = progress_percentage(completed_requests, request_total) + return f"{completed_requests}/{request_total} ({float(request_pct):.1f}%)" + + +def _format_values(value: object) -> str: + if isinstance(value, (list, tuple)): + return ",".join(str(item) for item in value) + return str(value) + + +def _engine_details(payload: dict) -> str: + parts = [ + f"runner={payload['runner']}", + f"engine_type={payload.get('engine_type', 'unknown')}", + f"model={payload.get('model', 'unknown')}", + f"api_base={payload['api_base']}", + f"health={payload['health_url']}", + f"logs={payload['engine_log_dir']}", + ] + if payload.get("container"): + parts.append(f"container={payload['container']}") + return " ".join(parts) + + +def _endpoint_details(payload: dict) -> str: + parts = [ + f"engine_type={payload.get('engine_type', 'unknown')}", + f"model={payload.get('model', 'unknown')}", + f"api_base={payload['api_base']}", + ] + if payload.get("health_url") is not None: + parts.append(f"health={payload['health_url']}") + if payload.get("port"): + parts.append(f"host={payload.get('host', 'unknown')}") + parts.append(f"port={payload['port']}") + return " ".join(parts) diff --git a/veeksha/orchestration/launcher_progress.py b/veeksha/orchestration/launcher_progress.py new file mode 100644 index 00000000..598461f7 --- /dev/null +++ b/veeksha/orchestration/launcher_progress.py @@ -0,0 +1,199 @@ +"""Progress parsing and rendering for orchestrated launcher runs.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +from tqdm import tqdm + + +@dataclass(frozen=True) +class BenchmarkRequestProgress: + completed_requests: int + total_requests: Optional[int] = None + + +class BenchmarkProgressReader: + """Read benchmark progress from the JSON file the benchmark publishes. + + The launcher points the benchmark subprocess at this file via the + ``VEEKSHA_PROGRESS_FILE`` environment variable; the benchmark writes + ``{"completed": int, "total": int | null}`` to it atomically. This is a + dedicated machine-readable channel, so there is no console scraping and no + coupling to the benchmark's progress-bar formatting. + """ + + def __init__(self, progress_path: Path, fallback_total_requests: int): + self._progress_path = progress_path + total_requests = positive_int_or_none(fallback_total_requests) + self._progress = BenchmarkRequestProgress( + completed_requests=0, total_requests=total_requests + ) + + def read(self) -> BenchmarkRequestProgress: + try: + raw = self._progress_path.read_text(encoding="utf-8") + data = json.loads(raw) + except (FileNotFoundError, OSError, ValueError): + # Not written yet, transiently unreadable, or mid-write: keep last value. + return self._progress + + completed = data.get("completed") + if not isinstance(completed, int) or isinstance(completed, bool): + return self._progress + + total = positive_int_or_none(data.get("total")) + if total is None: + total = self._progress.total_requests + self._progress = BenchmarkRequestProgress( + completed_requests=max(completed, 0), total_requests=total + ) + return self._progress + + +def request_progress_payload(progress: BenchmarkRequestProgress) -> dict: + payload = { + "requests_completed": progress.completed_requests, + "request_total": progress.total_requests, + } + if progress.total_requests is None: + payload["request_progress_pct"] = None + else: + payload["request_progress_pct"] = progress_percentage( + progress.completed_requests, progress.total_requests + ) + return payload + + +def positive_int_or_none(value: object) -> Optional[int]: + if not isinstance(value, int) or value <= 0: + return None + return value + + +def progress_percentage(numerator: float, denominator: float) -> float: + if denominator <= 0: + return 100.0 + return round(min(max(numerator / denominator * 100, 0.0), 100.0), 1) + + +class LauncherProgressReporter: + """TQDM-backed console progress for orchestrated launcher runs.""" + + def __init__(self) -> None: + self._sweep_bar: Optional[Any] = None + self._attempt_bar: Optional[Any] = None + + def handle_event(self, event: str, payload: dict) -> None: + if event == "sweep_plan_ready": + self._start_sweep(payload) + return + if event == "benchmark_attempt_start": + self._start_attempt(payload) + return + if event == "benchmark_attempt_progress": + self.update_attempt_requests( + int(payload["requests_completed"]), + positive_int_or_none(payload.get("request_total")), + ) + return + if event == "benchmark_attempt_success": + self.update_attempt_requests( + int(payload["requests_completed"]), + positive_int_or_none(payload.get("request_total")), + ) + self._close_attempt() + self._set_sweep_completed(int(payload["completed_runs"])) + return + if event == "benchmark_attempt_failed": + self.update_attempt_requests( + int(payload["requests_completed"]), + positive_int_or_none(payload.get("request_total")), + ) + self._close_attempt() + return + if event == "benchmark_attempts_exhausted": + self._close_attempt() + self._set_sweep_completed(int(payload["run_index"])) + return + if event == "sweep_complete": + self._finish_sweep() + + def update_attempt_requests( + self, completed_requests: int, total_requests: Optional[int] + ) -> None: + if self._attempt_bar is None: + return + bounded_completed = max(completed_requests, 0) + if total_requests is not None: + self._attempt_bar.total = total_requests + bounded_completed = min(bounded_completed, total_requests) + self._attempt_bar.n = bounded_completed + self._attempt_bar.set_postfix_str( + f"requests={bounded_completed}/{total_requests}" + ) + else: + previous_completed = int(self._attempt_bar.n or 0) + if bounded_completed >= previous_completed: + self._attempt_bar.update(bounded_completed - previous_completed) + else: + self._attempt_bar.n = bounded_completed + self._attempt_bar.set_postfix_str(f"requests={bounded_completed}") + self._attempt_bar.refresh() + + def write(self, message: str) -> None: + tqdm.write(f"[veeksha-launcher] {message}") + + def close(self) -> None: + self._close_attempt() + if self._sweep_bar is not None: + self._sweep_bar.close() + self._sweep_bar = None + + def _start_sweep(self, payload: dict) -> None: + self.close() + total_runs = int(payload["runs"]) + self._sweep_bar = tqdm( + total=total_runs, + desc=f"sweep {payload['sweep_type']} {payload['engine']} {payload['model']}", + unit="run", + dynamic_ncols=True, + position=0, + ) + self._set_sweep_completed(0) + + def _start_attempt(self, payload: dict) -> None: + self._close_attempt() + total_requests = positive_int_or_none(payload.get("request_total")) + self._attempt_bar = tqdm( + total=total_requests, + desc=( + f"run {payload['run_index']}/{payload['run_count']} " + f"attempt {payload['attempt']}" + ), + unit="req", + dynamic_ncols=True, + position=1, + leave=False, + ) + self.update_attempt_requests(0, total_requests) + + def _close_attempt(self) -> None: + if self._attempt_bar is not None: + self._attempt_bar.close() + self._attempt_bar = None + + def _set_sweep_completed(self, completed_runs: int) -> None: + if self._sweep_bar is None: + return + total = int(self._sweep_bar.total or 0) + self._sweep_bar.n = min(max(completed_runs, 0), total) + self._sweep_bar.set_postfix_str(f"complete={self._sweep_bar.n}/{total}") + self._sweep_bar.refresh() + + def _finish_sweep(self) -> None: + if self._sweep_bar is not None: + self._set_sweep_completed(int(self._sweep_bar.total or 0)) diff --git a/veeksha/orchestration/managed_engines.py b/veeksha/orchestration/managed_engines.py new file mode 100644 index 00000000..6ca50bfc --- /dev/null +++ b/veeksha/orchestration/managed_engines.py @@ -0,0 +1,595 @@ +"""Managed server lifecycle runners for orchestrated Veeksha sweeps.""" + +from __future__ import annotations + +import json +import os +import shlex +import subprocess +import time +import uuid +from abc import ABC, abstractmethod +from datetime import datetime, timezone +from pathlib import Path +from typing import IO, Optional + +import requests + +from veeksha.config.server import ( + ManagedServerConfig, + SglangServerConfig, + VajraServerConfig, + VllmServerConfig, +) +from veeksha.orchestration.processes import ProcessTerminator +from veeksha.orchestration.server_manager import BaseServerManager + +_ENGINE_DETAILS_FILENAME = "engine_details.json" + + +class EngineError(RuntimeError): + """Base engine lifecycle error.""" + + +class EngineRestartLimitExceeded(EngineError): + """Raised when the configured restart budget is exhausted.""" + + +class BaseEngineRunner(BaseServerManager, ABC): + def __init__( + self, + config: ManagedServerConfig, + output_dir: str | Path, + *, + terminator: Optional[ProcessTerminator] = None, + ): + super().__init__(config, output_dir=str(output_dir)) + self.output_dir = Path(output_dir) + self._delete_log_file_on_cleanup = False + self._terminator = terminator or ProcessTerminator() + + @property + def is_running(self) -> bool: + raise EngineError( + "managed server lifecycle uses is_alive(); is_running is invalid" + ) + + def get_api_base(self) -> str: + return self.get_endpoint().api_base + + def get_endpoint(self): + return self.config.get_endpoint() + + def launch(self) -> tuple[bool, Optional[str]]: + raise EngineError("managed server lifecycle uses start(); launch is invalid") + + def shutdown(self, force: bool = False) -> bool: + raise EngineError("managed server lifecycle uses stop(); shutdown is invalid") + + def get_server_logs(self, lines: int = 50) -> tuple[str, str]: + return self.tail_logs(lines), "" + + def _build_launch_command(self) -> list[str]: + return [] + + def _ensure_managed_port_available(self) -> None: + if self._is_port_in_use(): + raise EngineError( + f"port {self.config.port} on host {self.config.host!r} is already in use; " + "refusing to attach managed server health checks to an existing listener" + ) + + @abstractmethod + def start(self) -> None: + """Start the engine and wait until ready.""" + + @abstractmethod + def stop(self) -> None: + """Stop the engine if it is running.""" + + @abstractmethod + def is_alive(self) -> bool: + """Return whether the underlying process/container is still running.""" + + @abstractmethod + def tail_logs(self, lines: int = 80) -> str: + """Return recent engine logs for diagnostics.""" + + def health_check(self) -> bool: + try: + response = requests.get(self.config.health_check_url, timeout=5) + return response.status_code == 200 + except requests.RequestException: + return False + + def wait_for_ready(self, timeout: Optional[int] = None) -> bool: + try: + self._wait_for_ready_or_raise(timeout=timeout) + except (EngineError, TimeoutError): + return False + return True + + def _wait_for_ready_or_raise(self, timeout: Optional[float] = None) -> None: + startup_timeout = self.config.startup_timeout if timeout is None else timeout + start = time.monotonic() + while time.monotonic() - start < startup_timeout: + if not self.is_alive(): + raise EngineError( + "engine exited before becoming ready\n" + self.tail_logs() + ) + if self.health_check(): + return + time.sleep(self.config.health_check_interval) + raise TimeoutError( + f"engine did not become ready within {startup_timeout}s\n" + + self.tail_logs() + ) + + def restart(self) -> None: + if self.restart_count >= self.config.max_restarts: + raise EngineRestartLimitExceeded( + f"engine restart budget exhausted ({self.config.max_restarts})" + ) + self.restart_count += 1 + self.stop() + self.start() + + def reset_restart_budget(self) -> None: + """Reset the restart counter so ``max_restarts`` applies per sweep run.""" + self.restart_count = 0 + + def _write_engine_details(self, payload: dict[str, object]) -> Path: + self.output_dir.mkdir(parents=True, exist_ok=True) + path = self.output_dir / _ENGINE_DETAILS_FILENAME + record = { + "recorded_at_utc": datetime.now(timezone.utc).isoformat(), + **payload, + } + path.write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return path + + +class VajraSubprocessRunner(BaseEngineRunner): + def __init__( + self, + config: VajraServerConfig, + output_dir: str | Path, + *, + terminator: Optional[ProcessTerminator] = None, + ): + super().__init__(config, output_dir, terminator=terminator) + self.config = config + self._process: Optional[subprocess.Popen] = None + self._stdout_file: Optional[IO[str]] = None + self._stderr_file: Optional[IO[str]] = None + self._stdout_path: Optional[Path] = None + self._stderr_path: Optional[Path] = None + + def start(self) -> None: + if self.is_alive(): + return + self._ensure_managed_port_available() + allocation_success, allocation_error = self._ensure_gpu_allocation() + if not allocation_success: + raise EngineError(allocation_error or "failed to allocate GPUs") + + try: + self.output_dir.mkdir(parents=True, exist_ok=True) + self.start_count += 1 + self._write_engine_details(self._build_git_details()) + self._stdout_path = self.output_dir / f"vajra_stdout_{self.start_count}.log" + self._stderr_path = self.output_dir / f"vajra_stderr_{self.start_count}.log" + self._stdout_file = self._stdout_path.open("w", encoding="utf-8") + self._stderr_file = self._stderr_path.open("w", encoding="utf-8") + self._process = subprocess.Popen( + list(self.config.command), + cwd=str(self._setup_dir()), + stdout=self._stdout_file, + stderr=self._stderr_file, + start_new_session=True, + env=self._build_env(), + text=True, + ) + self._wait_for_ready_or_raise() + except BaseException: + self.stop() + raise + + def stop(self) -> None: + try: + if self._process is not None: + self._terminator.terminate(self._process) + finally: + self._close_logs() + self._process = None + self._release_allocated_resources() + + def is_alive(self) -> bool: + return self._process is not None and self._process.poll() is None + + def tail_logs(self, lines: int = 80) -> str: + chunks = [] + for label, path in ( + ("stderr", self._stderr_path), + ("stdout", self._stdout_path), + ): + if path is None or not path.exists(): + continue + try: + content = path.read_text( + encoding="utf-8", errors="replace" + ).splitlines() + chunks.append(f"[{label}]\n" + "\n".join(content[-lines:])) + except OSError as exc: + chunks.append(f"[{label}] failed to read {path}: {exc}") + return "\n".join(chunks) + + def _build_env(self) -> dict[str, str]: + env = os.environ.copy() + env.update(self.config.env or {}) + setup_dir = str(self._setup_dir()) + pythonpath = env.get("PYTHONPATH") + paths = [ + path + for path in (pythonpath.split(os.pathsep) if pythonpath else []) + if path and path != setup_dir + ] + env["PYTHONPATH"] = os.pathsep.join([setup_dir, *paths]) + if self.config.gpu_ids is not None: + env["CUDA_VISIBLE_DEVICES"] = ",".join( + str(gpu) for gpu in self.config.gpu_ids + ) + return env + + def _setup_dir(self) -> Path: + if not self.config.setup_dir: + raise EngineError("vajra server.setup_dir is required") + return Path(self.config.setup_dir).expanduser() + + def _close_logs(self) -> None: + for file_obj in (self._stdout_file, self._stderr_file): + if file_obj is not None and not file_obj.closed: + file_obj.close() + self._stdout_file = None + self._stderr_file = None + + def _build_git_details(self) -> dict[str, object]: + setup_dir = self._setup_dir() + result = subprocess.run( + ["git", "-C", str(setup_dir), "log", "-1", "--format=%H"], + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise EngineError( + "failed to read latest Vajra git commit with git log " + f"in {setup_dir}: {result.stderr.strip() or result.stdout.strip()}" + ) + git_commit_id = result.stdout.strip() + if not git_commit_id: + raise EngineError(f"git log returned no commit ID in {setup_dir}") + return { + "engine_type": self.config.type, + "source": "git", + "git_commit_id": git_commit_id, + "git_source_dir": str(setup_dir), + "git_command": "git log -1 --format=%H", + } + + +class VllmOmniDockerRunner(BaseEngineRunner): + def __init__( + self, + config: VllmServerConfig | SglangServerConfig, + output_dir: str | Path, + *, + terminator: Optional[ProcessTerminator] = None, + ): + super().__init__(config, output_dir, terminator=terminator) + self.config = config + self._container_name = _unique_container_name(config.container_name_prefix) + self._container_id: Optional[str] = None + self._log_process: Optional[subprocess.Popen] = None + self._log_file: Optional[IO[str]] = None + + @property + def container_name(self) -> str: + return self._container_name + + def start(self) -> None: + if self.is_alive(): + return + self._ensure_managed_port_available() + allocation_success, allocation_error = self._ensure_gpu_allocation() + if not allocation_success: + raise EngineError(allocation_error or "failed to allocate GPUs") + + try: + self.output_dir.mkdir(parents=True, exist_ok=True) + self.start_count += 1 + cmd = self._build_docker_run_cmd() + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise EngineError( + "docker run failed " + f"(rc={result.returncode})\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + self._container_id = result.stdout.strip() + self._write_engine_details(self._build_docker_details()) + self._start_log_streamer() + self._wait_for_ready_or_raise() + except BaseException: + self.stop() + raise + + def stop(self) -> None: + try: + self._stop_log_streamer() + if self._container_id is None and not self.is_alive(): + return + try: + subprocess.run( + ["docker", "stop", self._container_name], + capture_output=True, + timeout=60, + ) + except subprocess.TimeoutExpired: + pass + subprocess.run( + ["docker", "rm", "-f", self._container_name], capture_output=True + ) + self._container_id = None + finally: + self._release_allocated_resources() + + def is_alive(self) -> bool: + result = subprocess.run( + [ + "docker", + "inspect", + "-f", + "{{.State.Running}}", + self._container_name, + ], + capture_output=True, + text=True, + ) + return result.stdout.strip() == "true" + + def tail_logs(self, lines: int = 80) -> str: + result = subprocess.run( + [ + "docker", + "logs", + "--tail", + str(lines), + self._container_name, + ], + capture_output=True, + text=True, + ) + return f"{result.stdout}\n{result.stderr}".strip() + + def _build_docker_run_cmd(self) -> list[str]: + cmd = [ + "docker", + "run", + "-d", + "--name", + self._container_name, + ] + if self.config.docker_runtime: + cmd.extend(["--runtime", self.config.docker_runtime]) + cmd.extend(["--gpus", self._docker_gpu_arg()]) + if self.config.ipc_mode: + cmd.append(f"--ipc={self.config.ipc_mode}") + cmd.extend(self._extra_run_args()) + cmd.extend( + [ + "-p", + f"{self.config.port}:{self.config.resolved_container_port}", + "--init", + ] + ) + for volume in [ + *self.config.volumes, + *self._extra_volumes(), + self._deploy_config_volume(), + ]: + cmd.extend(["-v", volume]) + config_env = self.config.env or {} + for key, value in config_env.items(): + cmd.extend(["-e", f"{key}={value}"]) + for key in self.config.pass_env: + if key not in config_env: + if key not in os.environ: + raise EngineError( + f"required environment variable is not set: {key}" + ) + cmd.extend(["-e", key]) + cmd.append(self.config.image) + cmd.extend(self._build_server_cmd()) + return cmd + + def _docker_gpu_arg(self) -> str: + if self.config.docker_gpus: + return self.config.docker_gpus + if self.config.gpu_ids is not None: + return "device=" + ",".join(str(gpu) for gpu in self.config.gpu_ids) + return "all" + + def _deploy_config_volume(self) -> str: + deploy_config = Path(self.config.deploy_config) + if not deploy_config.is_file(): + raise EngineError(f"deploy config not found: {deploy_config}") + return f"{deploy_config}:{self.config.resolved_container_deploy_config}:ro" + + def _build_server_cmd(self) -> list[str]: + # Base command is vLLM-specific; SglangOmniDockerRunner overrides this. + assert isinstance(self.config, VllmServerConfig) + serve = [ + "vllm", + "serve", + self.config.hf_model, + "--omni", + "--port", + str(self.config.resolved_container_port), + "--deploy-config", + self.config.resolved_container_deploy_config, + *self.config.engine_args, + ] + if not self.config.uses_bootstrap: + return serve + serve_line = "exec " + " ".join(shlex.quote(part) for part in serve) + script = self.config.resolved_bootstrap + "\n" + serve_line + return ["bash", "-lc", script] + + def _extra_run_args(self) -> list[str]: + """Extra ``docker run`` flags inserted before the port mapping.""" + return [] + + def _extra_volumes(self) -> list[str]: + """Extra ``-v`` mounts beyond ``volumes`` and the deploy config.""" + return [] + + def _log_prefix(self) -> str: + return "vllm_docker" + + def _start_log_streamer(self) -> None: + self._log_file = ( + self.output_dir / f"{self._log_prefix()}_{self.start_count}.log" + ).open("w", encoding="utf-8") + self._log_process = subprocess.Popen( + ["docker", "logs", "-f", self._container_name], + stdout=self._log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + text=True, + ) + + def _stop_log_streamer(self) -> None: + if self._log_process is not None: + self._terminator.terminate(self._log_process) + self._log_process = None + if self._log_file is not None and not self._log_file.closed: + self._log_file.close() + self._log_file = None + + def _build_docker_details(self) -> dict[str, object]: + if self._container_id is None: + raise EngineError("cannot record Docker details before container start") + return { + "engine_type": self.config.type, + "source": "docker", + "docker_image": self.config.image, + "docker_image_hash": self._docker_inspect_value( + "{{.Image}}", self._container_id + ), + "container_id": self._container_id, + "container_name": self._container_name, + } + + def _docker_inspect_value(self, format_expr: str, target: str) -> str: + result = subprocess.run( + [ + "docker", + "inspect", + "-f", + format_expr, + target, + ], + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise EngineError( + "docker inspect failed " + f"(target={target}, rc={result.returncode}): {result.stderr.strip()}" + ) + value = result.stdout.strip() + if not value: + raise EngineError(f"docker inspect returned no value for {target}") + return value + + +class SglangOmniDockerRunner(VllmOmniDockerRunner): + """Runs sglang-omni's OpenAI-compatible server in a Docker container. + + Differences from the vLLM Omni runner: + + - The container command is ``sgl-omni serve --model-path ... --config ...`` + (``--config``, not ``--deploy-config``; no ``--omni`` flag). + - The stock ``frankleeeee/sglang-omni:dev`` image only ships prerequisites, + so the command is wrapped in a ``bash`` bootstrap that installs + sglang-omni from the mounted source checkout before serving (configurable; + set ``engine.bootstrap`` to an empty string for a pre-baked image). + - Adds ``--shm-size`` (sglang-omni uses a shared-memory relay) and mounts the + source checkout. Health is checked at ``/health``. + """ + + def __init__( + self, + config: SglangServerConfig, + output_dir: str | Path, + *, + terminator: Optional[ProcessTerminator] = None, + ): + super().__init__(config, output_dir, terminator=terminator) + self.config = config + + def _log_prefix(self) -> str: + return "sglang_docker" + + def _extra_run_args(self) -> list[str]: + args: list[str] = [] + if self.config.shm_size: + args.extend(["--shm-size", self.config.shm_size]) + args.extend(self.config.docker_run_args) + return args + + def _extra_volumes(self) -> list[str]: + if self.config.uses_bootstrap and self.config.source_dir: + return [f"{self.config.source_dir}:{self.config.container_source_dir}"] + return [] + + def _build_server_cmd(self) -> list[str]: + serve = [ + "sgl-omni", + "serve", + "--model-path", + self.config.model_path, + "--config", + self.config.resolved_container_deploy_config, + "--host", + "0.0.0.0", + "--port", + str(self.config.resolved_container_port), + ] + if self.config.model_name: + serve.extend(["--model-name", self.config.model_name]) + serve.extend(self.config.engine_args) + if not self.config.uses_bootstrap: + return serve + serve_line = "exec " + " ".join(shlex.quote(part) for part in serve) + script = self.config.resolved_bootstrap + "\n" + serve_line + return ["bash", "-lc", script] + + +def _unique_container_name(prefix: str) -> str: + return f"{prefix}-{os.getpid()}-{uuid.uuid4().hex[:8]}" + + +def create_engine_runner( + config: ManagedServerConfig, output_dir: str | Path +) -> BaseEngineRunner: + from veeksha.orchestration.registry import ServerManagerRegistry + + manager = ServerManagerRegistry.get( + config.get_type(), config=config, output_dir=str(output_dir) + ) + if not isinstance(manager, BaseEngineRunner): + raise EngineError(f"unsupported engine manager: {type(manager).__name__}") + return manager diff --git a/veeksha/orchestration/processes.py b/veeksha/orchestration/processes.py new file mode 100644 index 00000000..cd317b12 --- /dev/null +++ b/veeksha/orchestration/processes.py @@ -0,0 +1,58 @@ +"""Subprocess lifecycle helpers for Veeksha orchestration.""" + +from __future__ import annotations + +import os +import signal +import subprocess +from typing import Optional + + +class ProcessTerminator: + """Terminate a subprocess and its process group when possible.""" + + def __init__(self, terminate_timeout: float = 30.0, kill_timeout: float = 10.0): + self.terminate_timeout = terminate_timeout + self.kill_timeout = kill_timeout + + def terminate(self, process: subprocess.Popen) -> None: + if process.poll() is not None: + return + + pgid = self._process_group_id(process) + if pgid is not None: + self._signal_group(pgid, signal.SIGTERM) + else: + process.terminate() + + try: + process.wait(timeout=self.terminate_timeout) + return + except subprocess.TimeoutExpired: + pass + + if pgid is not None: + self._signal_group(pgid, signal.SIGKILL) + else: + process.kill() + process.wait(timeout=self.kill_timeout) + + @staticmethod + def _process_group_id(process: subprocess.Popen) -> Optional[int]: + pid = getattr(process, "pid", None) + if pid is None: + return None + try: + pgid = os.getpgid(pid) + except ProcessLookupError: + return None + if pgid == os.getpgrp(): + return None + return pgid + + @staticmethod + def _signal_group(pgid: int, sig: signal.Signals) -> None: + try: + os.killpg(pgid, sig) + except ProcessLookupError: + return diff --git a/veeksha/orchestration/registry.py b/veeksha/orchestration/registry.py index 96b2e5fd..435a668a 100644 --- a/veeksha/orchestration/registry.py +++ b/veeksha/orchestration/registry.py @@ -12,15 +12,21 @@ def get_key_from_str(cls, key_str: str) -> ServerType: ServerManagerRegistry.register( ServerType.VLLM, _LazyLoader( - "veeksha.orchestration.vllm_server", - "VLLMServerManager", + "veeksha.orchestration.managed_engines", + "VllmOmniDockerRunner", + ), +) +ServerManagerRegistry.register( + ServerType.VAJRA, + _LazyLoader( + "veeksha.orchestration.managed_engines", + "VajraSubprocessRunner", ), ) - ServerManagerRegistry.register( ServerType.SGLANG, _LazyLoader( - "veeksha.orchestration.sglang_server", - "SGLangServerManager", + "veeksha.orchestration.managed_engines", + "SglangOmniDockerRunner", ), ) diff --git a/veeksha/orchestration/server_manager.py b/veeksha/orchestration/server_manager.py index 44e913bf..668d8ba4 100644 --- a/veeksha/orchestration/server_manager.py +++ b/veeksha/orchestration/server_manager.py @@ -47,6 +47,8 @@ def __init__(self, config: BaseServerConfig, output_dir: Optional[str] = None): self._delete_log_file_on_cleanup = True self.resource_manager = ResourceManager() self._allocated_job_id: Optional[str] = None # Track allocated resources + self.restart_count = 0 + self.start_count = 0 @property def is_running(self) -> bool: @@ -65,6 +67,66 @@ def _build_launch_command(self) -> list[str]: List of command arguments """ + def _should_auto_allocate_gpus(self) -> bool: + if self.config.gpu_ids is not None: + return False + # Docker configs may ask Docker itself for a device set/count, e.g. + # ``--gpus all`` or ``--gpus 2``. Treat that as an explicit override. + if getattr(self.config, "docker_gpus", None): + return False + return True + + def _ensure_gpu_allocation(self) -> tuple[bool, Optional[str]]: + if not self._should_auto_allocate_gpus(): + return True, None + + num_gpus = self.config.get_num_gpus() + job_id = f"server_{self.config.host}_{self.config.port}_{int(time.time())}" + resource_mapping = self.resource_manager.wait_for_resources( + num_gpus=num_gpus, + timeout=300, + job_id=job_id, + contiguous=self.config.require_contiguous_gpus, + ) + + if resource_mapping is None: + logger.error(f"Failed to allocate {num_gpus} GPUs for server") + return False, f"Failed to allocate {num_gpus} GPUs for server" + + self._allocated_job_id = job_id + gpu_ids = [gpu_id for _, gpu_id in resource_mapping] + self.config = replace(self.config, gpu_ids=gpu_ids) + return True, None + + def _release_allocated_resources(self) -> None: + if self._allocated_job_id is None: + return + try: + self.resource_manager.release_resources(self._allocated_job_id) + except Exception as e: + logger.error(f"Error releasing resources: {e}") + finally: + self._allocated_job_id = None + + def start(self) -> None: + success, error = self.launch() + if not success: + raise RuntimeError(error or "failed to launch server") + + def stop(self) -> None: + self.shutdown() + + def restart(self) -> None: + max_restarts = getattr(self.config, "max_restarts", None) + if max_restarts is not None and self.restart_count >= max_restarts: + raise RuntimeError(f"server restart budget exhausted ({max_restarts})") + self.restart_count += 1 + self.stop() + self.start() + + def reset_restart_budget(self) -> None: + self.restart_count = 0 + def _create_log_file(self) -> IO[str]: """Create a log file for the server process.""" timestamp = time.strftime("%Y%m%d-%H%M%S") @@ -118,27 +180,9 @@ def launch(self) -> tuple[bool, Optional[str]]: return False, error_msg try: - # auto-allocate if not specified - if self.config.gpu_ids is None: - num_gpus = self.config.get_num_gpus() - - job_id = ( - f"server_{self.config.host}_{self.config.port}_{int(time.time())}" - ) - resource_mapping = self.resource_manager.wait_for_resources( - num_gpus=num_gpus, - timeout=300, # 5 minute timeout - job_id=job_id, - contiguous=self.config.require_contiguous_gpus, - ) - - if resource_mapping is None: - logger.error(f"Failed to allocate {num_gpus} GPUs for server") - return False, f"Failed to allocate {num_gpus} GPUs for server" - - self._allocated_job_id = job_id - gpu_ids = [gpu_id for _, gpu_id in resource_mapping] - self.config = replace(self.config, gpu_ids=gpu_ids) + allocation_success, allocation_error = self._ensure_gpu_allocation() + if not allocation_success: + return False, allocation_error command = self._build_launch_command() logger.info(f"Launching server with command: {' '.join(command)}") @@ -174,15 +218,13 @@ def launch(self) -> tuple[bool, Optional[str]]: text=True, ) + self.start_count += 1 logger.info(f"Server process started with PID: {self.process.pid}") self._is_running = True return True, None except Exception as e: - # release GPUs - if self._allocated_job_id is not None: - self.resource_manager.release_resources(self._allocated_job_id) - self._allocated_job_id = None + self._release_allocated_resources() if self._log_file is not None: self._log_file.close() self._log_file = None @@ -360,13 +402,7 @@ def shutdown(self, force: bool = False) -> bool: # reset state and clean up resources, even with exceptions self._is_running = False - if self._allocated_job_id is not None: - try: - self.resource_manager.release_resources(self._allocated_job_id) - except Exception as e: - logger.error(f"Error releasing resources: {e}") - finally: - self._allocated_job_id = None + self._release_allocated_resources() if self._log_file: try: diff --git a/veeksha/orchestration/sglang_server.py b/veeksha/orchestration/sglang_server.py deleted file mode 100644 index 9be38f97..00000000 --- a/veeksha/orchestration/sglang_server.py +++ /dev/null @@ -1,58 +0,0 @@ -from typing import List - -from veeksha.config.server import SglangServerConfig -from veeksha.logger import init_logger -from veeksha.orchestration.server_manager import BaseServerManager - -logger = init_logger(__name__) - - -class SGLangServerManager(BaseServerManager): - """Manager for SGLang inference servers.""" - - def __init__(self, config: SglangServerConfig, output_dir: str): - super().__init__(config, output_dir=output_dir) - - def _build_launch_command(self) -> List[str]: - command = [ - "python", - "-m", - "sglang.launch_server", - "--model-path", - self.config.model, - "--host", - self.config.host, - "--port", - str(self.config.port), - ] - - if self.config.tensor_parallel_size > 1: - command.extend( - ["--tensor-parallel-size", str(self.config.tensor_parallel_size)] - ) - - if self.config.dtype: - command.extend(["--dtype", self.config.dtype]) - - if self.config.max_model_len is not None: - command.extend(["--context-length", str(self.config.max_model_len)]) - - additional_args_dict = self.get_additional_args_dict() - for key, value in additional_args_dict.items(): - if isinstance(value, bool) and value: - command.append(f"--{key}") - elif value is None: - continue - elif isinstance(value, (list, tuple)): - command.append(f"--{key}") - command.append(f"[{', '.join([str(v) for v in value])}]") - else: - command.extend([f"--{key}", str(value)]) - - command.extend(self._parse_additional_sglang_args()) - - return command - - def _parse_additional_sglang_args(self) -> List[str]: - args = [] - return args diff --git a/veeksha/orchestration/vllm_server.py b/veeksha/orchestration/vllm_server.py deleted file mode 100644 index e9cdb184..00000000 --- a/veeksha/orchestration/vllm_server.py +++ /dev/null @@ -1,74 +0,0 @@ -from typing import List - -from veeksha.config.server import VllmServerConfig -from veeksha.logger import init_logger -from veeksha.orchestration.server_manager import BaseServerManager - -logger = init_logger(__name__) - - -class VLLMServerManager(BaseServerManager): - """Manager for vLLM inference servers.""" - - def __init__(self, config: VllmServerConfig, output_dir: str): - super().__init__(config, output_dir=output_dir) - - def _build_launch_command(self) -> List[str]: - command = [ - "python", - "-m", - "vllm.entrypoints.openai.api_server", - "--model", - self.config.model, - "--host", - self.config.host, - "--port", - str(self.config.port), - "--api-key", - self.config.api_key, - ] - - if self.config.tensor_parallel_size > 1: - command.extend( - ["--tensor-parallel-size", str(self.config.tensor_parallel_size)] - ) - - if self.config.dtype: - command.extend(["--dtype", self.config.dtype]) - - if self.config.max_model_len is not None: - command.extend(["--max-model-len", str(self.config.max_model_len)]) - - additional_args_dict = self.get_additional_args_dict() - special_keys = {"rope_scaling"} - for key, value in additional_args_dict.items(): - if key in special_keys: - continue - if isinstance(value, bool): - if value: - command.append(f"--{key}") - continue - if value is None: - continue - elif isinstance(value, (list, tuple)): - command.append(f"--{key}") - command.extend([str(v) for v in value]) - else: - command.extend([f"--{key}", str(value)]) - - command.extend(self._parse_additional_vllm_args()) - - return command - - def _parse_additional_vllm_args(self) -> List[str]: - args = [] - additional_args_dict = self.get_additional_args_dict() - - if "rope_scaling" in additional_args_dict: - import json - - rope_config = additional_args_dict["rope_scaling"] - rope_json = json.dumps(rope_config) - args.extend(["--rope-scaling", rope_json]) - - return args diff --git a/veeksha/slo/metrics.py b/veeksha/slo/metrics.py index dda3027f..99ab7a50 100644 --- a/veeksha/slo/metrics.py +++ b/veeksha/slo/metrics.py @@ -7,8 +7,14 @@ def lower_is_better(metric: str) -> bool: """Return whether lower values are better for the given metric.""" - # Today we only support latency-like metrics. - return True + higher_is_better = { + "audio_before_commit_ratio", + "duplex_overlap_observed", + "tts_service_fluidity_eligible", + "tts_service_fluidity_index", + "user_audio_fluidity_index", + } + return metric not in higher_is_better def extract_metric_values( @@ -27,8 +33,6 @@ def extract_metric_values( key = metric if metric == "e2e": key = "end_to_end_latency" - elif metric == "tpot": - key = "tpot" raw = request_level_metrics.get(key, []) if not isinstance(raw, list) or not raw: @@ -69,5 +73,5 @@ def _coerce_float(value: Any) -> Optional[float]: return None try: return float(value) - except (TypeError, ValueError): + except TypeError, ValueError: return None diff --git a/veeksha/slo/slo.py b/veeksha/slo/slo.py index 123d8b76..a8f2f63b 100644 --- a/veeksha/slo/slo.py +++ b/veeksha/slo/slo.py @@ -5,7 +5,7 @@ from veeksha.config.slo import BaseSloConfig, ConstantSloConfig from veeksha.logger import init_logger -from veeksha.slo.metrics import extract_metric_values +from veeksha.slo.metrics import extract_metric_values, lower_is_better logger = init_logger(__name__) @@ -56,17 +56,20 @@ def evaluate(self, request_level_metrics: Dict[str, Any]) -> Tuple[bool, float]: metric_value = _percentile(values, self.config.percentile) threshold = self.get_threshold() - # Currently, supported SLO metrics are latency-like (lower is better). if math.isclose(metric_value, threshold, rel_tol=1e-9, abs_tol=1e-12): met = True - else: + elif lower_is_better(self.config.metric): met = metric_value < threshold + else: + met = metric_value > threshold return met, metric_value def __str__(self) -> str: return ( f"ConstantSlo(metric={self.config.metric}, " - f"p{self.config.percentile*100:.0f} <= {self.config.value})" + f"p{self.config.percentile*100:.0f} " + f"{'<=' if lower_is_better(self.config.metric) else '>='} " + f"{self.config.value})" ) def get_slo_metric_key(self) -> str: diff --git a/veeksha/sweeps/__init__.py b/veeksha/sweeps/__init__.py new file mode 100644 index 00000000..586acb8e --- /dev/null +++ b/veeksha/sweeps/__init__.py @@ -0,0 +1,18 @@ +"""Sweep planning APIs.""" + +from veeksha.sweeps.config import SweepConfig, SweepConfigError +from veeksha.sweeps.planner import ( + SweepPlan, + SweepRunDescriptor, + build_sweep_plan_from_config, +) +from veeksha.sweeps.specs import SweepSpec + +__all__ = [ + "SweepConfig", + "SweepConfigError", + "SweepPlan", + "SweepRunDescriptor", + "SweepSpec", + "build_sweep_plan_from_config", +] diff --git a/veeksha/sweeps/config.py b/veeksha/sweeps/config.py new file mode 100644 index 00000000..0da42f95 --- /dev/null +++ b/veeksha/sweeps/config.py @@ -0,0 +1,139 @@ +"""Typed configuration model for sweep planning.""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping, Optional, Tuple + +TRACE_SHAREGPT = "sharegpt" +TRACE_SEED_TTS = "seed_tts" +TRACE_ALIASES = { + "sharegpt": TRACE_SHAREGPT, + "share_gpt": TRACE_SHAREGPT, + "seed_tts": TRACE_SEED_TTS, + "seedtts": TRACE_SEED_TTS, + "seed_tts_text": TRACE_SEED_TTS, +} + +DEFAULT_TIMEOUT_SECONDS = 600 +DEFAULT_COOLDOWN_SECONDS = 120 +DEFAULT_MAX_SESSIONS = 20000 +DEFAULT_MIN_TOKENS = 20 +DEFAULT_MAX_TOKENS = 150 + + +class SweepConfigError(ValueError): + """Raised when a programmatic sweep config is invalid.""" + + +@dataclass(frozen=True) +class SweepConfig: + """Typed sweep configuration shared by CLI and programmatic callers.""" + + sweep_type: str + engine: str + model: str + trace: str = TRACE_SHAREGPT + base_config: Optional[str] = None + output_dir: Optional[str] = None + min_tokens: Optional[int] = None + max_tokens: Optional[int] = None + min_chars: Optional[int] = None + max_chars: Optional[int] = None + concurrencies: Optional[Tuple[int, ...]] = None + concurrency: Optional[int] = None + sizes: Optional[Tuple[int, ...]] = None + range_start: Optional[int] = None + range_end: Optional[int] = None + step: Optional[int] = None + timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS + cooldown_seconds: int = DEFAULT_COOLDOWN_SECONDS + max_sessions: int = DEFAULT_MAX_SESSIONS + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any]) -> "SweepConfig": + unknown = sorted(set(raw) - {field.name for field in fields(cls)}) + if unknown: + raise SweepConfigError( + f"unsupported sweep config keys: {', '.join(unknown)}" + ) + + trace = _optional_str(raw, "trace") + if trace is None: + trace = TRACE_SHAREGPT + + timeout_seconds = _optional_int(raw, "timeout_seconds") + if timeout_seconds is None: + timeout_seconds = DEFAULT_TIMEOUT_SECONDS + cooldown_seconds = _optional_int(raw, "cooldown_seconds") + if cooldown_seconds is None: + cooldown_seconds = DEFAULT_COOLDOWN_SECONDS + max_sessions = _optional_int(raw, "max_sessions") + if max_sessions is None: + max_sessions = DEFAULT_MAX_SESSIONS + + return cls( + sweep_type=_required_str(raw, "sweep_type"), + engine=_required_str(raw, "engine"), + model=_required_str(raw, "model"), + trace=_normalize_trace(trace), + base_config=_optional_str(raw, "base_config"), + output_dir=_optional_str(raw, "output_dir"), + min_tokens=_optional_int(raw, "min_tokens"), + max_tokens=_optional_int(raw, "max_tokens"), + min_chars=_optional_int(raw, "min_chars"), + max_chars=_optional_int(raw, "max_chars"), + concurrencies=_optional_int_tuple(raw, "concurrencies"), + concurrency=_optional_int(raw, "concurrency"), + sizes=_optional_int_tuple(raw, "sizes"), + range_start=_optional_int(raw, "range_start"), + range_end=_optional_int(raw, "range_end"), + step=_optional_int(raw, "step"), + timeout_seconds=timeout_seconds, + cooldown_seconds=cooldown_seconds, + max_sessions=max_sessions, + ) + + +def _normalize_trace(raw: str) -> str: + key = raw.strip().lower().replace("-", "_") + if key not in TRACE_ALIASES: + supported = ", ".join(sorted(TRACE_ALIASES)) + raise SweepConfigError(f"trace must be one of: {supported}") + return TRACE_ALIASES[key] + + +def _required_str(raw: Mapping[str, Any], key: str) -> str: + value = raw.get(key) + if not isinstance(value, str) or not value: + raise SweepConfigError(f"sweep.{key} must be a non-empty string") + return value + + +def _optional_str(raw: Mapping[str, Any], key: str) -> Optional[str]: + value = raw.get(key) + if value is None: + return None + if not isinstance(value, str): + raise SweepConfigError(f"sweep.{key} must be a string") + return value + + +def _optional_int(raw: Mapping[str, Any], key: str) -> Optional[int]: + value = raw.get(key) + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise SweepConfigError(f"sweep.{key} must be an integer") + return value + + +def _optional_int_tuple(raw: Mapping[str, Any], key: str) -> Optional[Tuple[int, ...]]: + value = raw.get(key) + if value is None: + return None + if not isinstance(value, (list, tuple)) or not all( + isinstance(item, int) and not isinstance(item, bool) for item in value + ): + raise SweepConfigError(f"sweep.{key} must be a list of integers") + return tuple(value) diff --git a/veeksha/sweeps/planner.py b/veeksha/sweeps/planner.py new file mode 100755 index 00000000..e3f4f2a1 --- /dev/null +++ b/veeksha/sweeps/planner.py @@ -0,0 +1,599 @@ +#!/usr/bin/env python3 +"""Run Veeksha benchmark sweeps from one launcher. + +Examples +-------- + # Concurrency sweep + python scripts/sweep.py --sweep-type concurrency --engine vajra --model qwen-tts + + # Input-size sweep at fixed concurrency + python scripts/sweep.py --sweep-type input --engine vllm --model qwen-tts \\ + --concurrency 16 --sizes 20,60,100,200,500 + + # Preview generated configs and benchmark commands without running them + python scripts/sweep.py --sweep-type concurrency --engine vllm \\ + --model qwen3-omni --concurrencies 1,2 --dry-run + + # Use the Seed TTS text dataset instead of the default ShareGPT trace + python scripts/sweep.py --sweep-type concurrency --engine vajra \\ + --model qwen-tts --trace seed_tts + # STT concurrency sweep (uses the audio trace in its base config) + python scripts/sweep.py --sweep-type concurrency --engine vajra \\ + --model voxtral --cooldown-seconds 30 + + # Override output_dir for generated benchmark configs + python scripts/sweep.py --sweep-type concurrency --engine vajra \\ + --model qwen-tts --output-dir 'benchmark_output/{run_name}' +""" + +from __future__ import annotations + +import argparse +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass, replace +from pathlib import Path +from typing import List, Optional, Tuple + +from veeksha.config.endpoint import EndpointConfig +from veeksha.sweeps.config import ( + DEFAULT_COOLDOWN_SECONDS, + DEFAULT_MAX_SESSIONS, + DEFAULT_MAX_TOKENS, + DEFAULT_MIN_TOKENS, + DEFAULT_TIMEOUT_SECONDS, + TRACE_ALIASES, + TRACE_SHAREGPT, + SweepConfig, + SweepConfigError, +) +from veeksha.sweeps.specs import ( + CONCURRENCY_SWEEP, + INPUT_SWEEP, + MODEL_ALIASES, + REPO_ROOT, + SPECS, + SweepSpec, + supported_combinations, +) +from veeksha.sweeps.utils import ( + apply_endpoint, + apply_trace_source, + benchmark_command, + build_run_config, + format_template, + input_sizes, + load_config, + write_config, +) + + +@dataclass(frozen=True) +class SweepRunDescriptor: + run_index: int + run_count: int + concurrency: int + input_size: Optional[int] + run_config: Path + run_name: str + output_dir: Optional[str] + command: List[str] + timeout_seconds: int + max_sessions: int + trace_source: str + + +@dataclass(frozen=True) +class SweepPlan: + spec: SweepSpec + tmp_parent: Path + runs: Tuple[SweepRunDescriptor, ...] + base_config_path: Path + endpoint: Optional[EndpointConfig] = None + + def cleanup(self) -> None: + for child in self.tmp_parent.iterdir(): + child.unlink() + self.tmp_parent.rmdir() + + +def _parse_csv_ints(raw: str, *, name: str) -> Tuple[int, ...]: + values: List[int] = [] + for item in raw.split(","): + item = item.strip() + if not item: + continue + try: + value = int(item) + except ValueError as exc: + raise argparse.ArgumentTypeError( + f"{name} must be comma-separated integers" + ) from exc + if value <= 0: + raise argparse.ArgumentTypeError(f"{name} values must be positive") + values.append(value) + if not values: + raise argparse.ArgumentTypeError(f"{name} cannot be empty") + return tuple(values) + + +def _parse_concurrencies(raw: str) -> Tuple[int, ...]: + return _parse_csv_ints(raw, name="concurrencies") + + +def _parse_sizes(raw: str) -> Tuple[int, ...]: + return _parse_csv_ints(raw, name="sizes") + + +def _parse_trace(raw: str) -> str: + key = raw.strip().lower().replace("-", "_") + if key not in TRACE_ALIASES: + supported = ", ".join(sorted(TRACE_ALIASES)) + raise argparse.ArgumentTypeError(f"trace must be one of: {supported}") + return TRACE_ALIASES[key] + + +def _normalize_model(raw: str) -> str: + key = raw.strip().lower() + if key not in MODEL_ALIASES: + raise SweepConfigError( + f"Unsupported model: {raw}\nSupported combinations:\n{supported_combinations()}" + ) + return MODEL_ALIASES[key] + + +def _resolve_base_config_path(sweep_config: SweepConfig, spec: SweepSpec) -> Path: + if sweep_config.base_config is None: + return spec.config_path + path = Path(sweep_config.base_config) + if not path.is_absolute(): + path = REPO_ROOT / path + if not path.is_file(): + raise SweepConfigError(f"sweep.base_config does not exist: {path}") + return path + + +def _build_sweep_plan( + sweep_config: SweepConfig, + spec: SweepSpec, + *, + endpoint: Optional[EndpointConfig] = None, + tmp_parent: Optional[Path] = None, +) -> SweepPlan: + base_config_path = _resolve_base_config_path(sweep_config, spec) + base_config = load_config(base_config_path) + if not spec.audio_input: + apply_trace_source(base_config, sweep_config) + trace_source = "audio" if spec.audio_input else sweep_config.trace + if tmp_parent is None: + tmp_parent = Path(tempfile.mkdtemp(prefix=f"{spec.temp_prefix}.")) + else: + tmp_parent.mkdir(parents=True, exist_ok=True) + + if spec.sweep_type == CONCURRENCY_SWEEP: + concurrencies = sweep_config.concurrencies or spec.default_concurrencies + work_items = [(concurrency, None) for concurrency in concurrencies] + else: + concurrency = sweep_config.concurrency or spec.default_concurrency + if concurrency is None: + raise ValueError("Input sweep requires --concurrency") + work_items = [(concurrency, size) for size in input_sizes(sweep_config, spec)] + + runs: List[SweepRunDescriptor] = [] + run_count = len(work_items) + for index, (concurrency, input_size) in enumerate(work_items, start=1): + config_name = format_template( + spec.run_config_template, + concurrency=concurrency, + input_size=input_size, + ) + run_config = tmp_parent / config_name + run_name = format_template( + spec.run_name_template, + concurrency=concurrency, + input_size=input_size, + ) + run_cfg = build_run_config( + base_config, + spec, + concurrency=concurrency, + input_size=input_size, + timeout_seconds=sweep_config.timeout_seconds, + max_sessions=sweep_config.max_sessions, + output_dir_template=sweep_config.output_dir, + min_tokens=sweep_config.min_tokens, + max_tokens=sweep_config.max_tokens, + min_chars=sweep_config.min_chars, + max_chars=sweep_config.max_chars, + ) + apply_endpoint(run_cfg, endpoint) + write_config(run_config, run_cfg) + runs.append( + SweepRunDescriptor( + run_index=index, + run_count=run_count, + concurrency=concurrency, + input_size=input_size, + run_config=run_config, + run_name=run_name, + output_dir=run_cfg.get("output_dir"), + command=benchmark_command(run_config), + timeout_seconds=sweep_config.timeout_seconds, + max_sessions=sweep_config.max_sessions, + trace_source=trace_source, + ) + ) + + return SweepPlan( + spec=spec, + tmp_parent=tmp_parent, + runs=tuple(runs), + base_config_path=base_config_path, + endpoint=endpoint, + ) + + +def _print_run_header( + *, + spec: SweepSpec, + run_index: int, + run_count: int, + concurrency: int, + input_size: Optional[int], + run_config: Path, + run_name: str, + timeout_seconds: int, + max_sessions: int, + dry_run: bool, + trace_source: str, + output_dir: Optional[str], +) -> None: + mode = "DRY RUN " if dry_run else "" + target = f"concurrency={concurrency}" + if input_size is not None: + target += f" input_chars={input_size}" + print("=" * 62) + print( + f" {mode}Running {spec.engine}/{spec.model} {spec.sweep_type} sweep: " + f"{target} ({run_index}/{run_count})" + ) + print(f" config : {run_config}") + print(f" base config : {spec.config_path}") + print(f" trace : {trace_source}") + if output_dir is not None: + print(f" output_dir : {output_dir}") + print(f" wandb : {run_name}") + if spec.write_runtime_limits: + print(f" timeout : {timeout_seconds}s") + print(f" max_sessions : {max_sessions}") + print("=" * 62) + + +def _run_sweep(sweep_config: SweepConfig, spec: SweepSpec, *, dry_run: bool) -> int: + plan = _build_sweep_plan(sweep_config, spec) + try: + for run in plan.runs: + _print_run_header( + spec=spec, + run_index=run.run_index, + run_count=run.run_count, + concurrency=run.concurrency, + input_size=run.input_size, + run_config=run.run_config, + run_name=run.run_name, + timeout_seconds=run.timeout_seconds, + max_sessions=run.max_sessions, + dry_run=dry_run, + trace_source=run.trace_source, + output_dir=run.output_dir, + ) + print(" ".join(run.command)) + if not dry_run: + subprocess.run(run.command, cwd=REPO_ROOT, check=True) + if run.run_index < run.run_count and sweep_config.cooldown_seconds > 0: + print( + f"-- cooldown {sweep_config.cooldown_seconds}s before next run --" + ) + time.sleep(sweep_config.cooldown_seconds) + finally: + if dry_run: + print(f"Dry-run configs left in: {plan.tmp_parent}") + else: + plan.cleanup() + + if spec.sweep_type == CONCURRENCY_SWEEP: + values = " ".join(str(run.concurrency) for run in plan.runs) + print(f"Sweep complete for concurrencies: {values}") + else: + sizes = " ".join( + str(run.input_size) for run in plan.runs if run.input_size is not None + ) + concurrency = plan.runs[0].concurrency if plan.runs else "-" + print(f"Input-size sweep complete (concurrency={concurrency}, sizes={sizes})") + return 0 + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--sweep-type", choices=(CONCURRENCY_SWEEP, INPUT_SWEEP), required=True + ) + parser.add_argument("--engine", choices=("vajra", "vllm", "sglang"), required=True) + parser.add_argument( + "--model", + choices=tuple(sorted(MODEL_ALIASES)), + required=True, + help="Model/workload alias. qwen3-tts and tts map to qwen-tts.", + ) + parser.add_argument( + "--trace", + type=_parse_trace, + default=TRACE_SHAREGPT, + help="Trace text source. Supported: sharegpt, seed_tts.", + ) + parser.add_argument( + "--output-dir", + help=( + "Override benchmark output_dir in generated configs. Supports " + "{concurrency}, {input_size}, {run_name}, {date_tag}, {engine}, " + "{model}, and {sweep_type}." + ), + ) + parser.add_argument( + "--min-tokens", + type=int, + help="Minimum input word count for concurrency sweeps. Defaults to 20.", + ) + parser.add_argument( + "--max-tokens", + type=int, + help="Maximum input word count for concurrency sweeps. Defaults to 150.", + ) + parser.add_argument( + "--min-chars", + type=int, + help="Minimum input char count for concurrency sweeps.", + ) + parser.add_argument( + "--max-chars", + type=int, + help="Maximum input char count for concurrency sweeps.", + ) + parser.add_argument( + "--concurrencies", + type=_parse_concurrencies, + help="Comma-separated concurrency values for concurrency sweeps.", + ) + parser.add_argument( + "--concurrency", + type=int, + help="Fixed concurrency for input sweeps.", + ) + parser.add_argument( + "--sizes", + type=_parse_sizes, + help="Comma-separated input sizes for input sweeps.", + ) + parser.add_argument("--range-start", type=int, help="Input sweep range start.") + parser.add_argument("--range-end", type=int, help="Input sweep range end.") + parser.add_argument("--step", type=int, help="Input sweep range step.") + parser.add_argument("--timeout-seconds", type=int, default=DEFAULT_TIMEOUT_SECONDS) + parser.add_argument( + "--cooldown-seconds", type=int, default=DEFAULT_COOLDOWN_SECONDS + ) + parser.add_argument("--max-sessions", type=int, default=DEFAULT_MAX_SESSIONS) + parser.add_argument( + "--dry-run", + action="store_true", + help="Write temporary configs and print commands without running benchmarks.", + ) + return parser + + +def _sweep_config_from_args(args: argparse.Namespace) -> SweepConfig: + return SweepConfig( + sweep_type=args.sweep_type, + engine=args.engine, + model=args.model, + trace=args.trace, + output_dir=args.output_dir, + min_tokens=args.min_tokens, + max_tokens=args.max_tokens, + min_chars=args.min_chars, + max_chars=args.max_chars, + concurrencies=args.concurrencies, + concurrency=args.concurrency, + sizes=args.sizes, + range_start=args.range_start, + range_end=args.range_end, + step=args.step, + timeout_seconds=args.timeout_seconds, + cooldown_seconds=args.cooldown_seconds, + max_sessions=args.max_sessions, + ) + + +def _validate_length_config(sweep_config: SweepConfig) -> SweepConfig: + char_pair = (sweep_config.min_chars is not None, sweep_config.max_chars is not None) + token_pair = ( + sweep_config.min_tokens is not None, + sweep_config.max_tokens is not None, + ) + has_chars = any(char_pair) + has_tokens = any(token_pair) + + if sweep_config.sweep_type == INPUT_SWEEP: + if has_chars or has_tokens: + raise SweepConfigError( + "input sweeps derive min_chars/max_chars from sizes; do not pass " + "min_tokens/max_tokens or min_chars/max_chars" + ) + return sweep_config + + if has_chars and has_tokens: + raise SweepConfigError( + "min_chars/max_chars and min_tokens/max_tokens are mutually exclusive" + ) + if has_chars and not all(char_pair): + raise SweepConfigError("min_chars and max_chars must be specified together") + if has_tokens and not all(token_pair): + raise SweepConfigError("min_tokens and max_tokens must be specified together") + + if has_chars: + min_chars = sweep_config.min_chars + max_chars = sweep_config.max_chars + if min_chars is None or max_chars is None: + raise SweepConfigError("min_chars and max_chars must be specified together") + if min_chars <= 0 or max_chars <= 0: + raise SweepConfigError("min_chars and max_chars must be positive") + if min_chars > max_chars: + raise SweepConfigError("min_chars must be <= max_chars") + return sweep_config + + if has_tokens: + min_tokens = sweep_config.min_tokens + max_tokens = sweep_config.max_tokens + if min_tokens is None or max_tokens is None: + raise SweepConfigError( + "min_tokens and max_tokens must be specified together" + ) + if min_tokens <= 0 or max_tokens <= 0: + raise SweepConfigError("min_tokens and max_tokens must be positive") + if min_tokens > max_tokens: + raise SweepConfigError("min_tokens must be <= max_tokens") + return sweep_config + + return replace( + sweep_config, + min_tokens=DEFAULT_MIN_TOKENS, + max_tokens=DEFAULT_MAX_TOKENS, + ) + + +def validate_sweep_config(sweep_config: SweepConfig) -> SweepConfig: + """Validate and normalize a typed sweep config.""" + if sweep_config.timeout_seconds <= 0: + raise SweepConfigError("timeout_seconds must be positive") + if sweep_config.cooldown_seconds < 0: + raise SweepConfigError("cooldown_seconds must be non-negative") + if sweep_config.max_sessions <= 0: + raise SweepConfigError("max_sessions must be positive") + if sweep_config.concurrency is not None and sweep_config.concurrency <= 0: + raise SweepConfigError("concurrency must be positive") + + sweep_config = _validate_length_config(sweep_config) + + if sweep_config.sweep_type == CONCURRENCY_SWEEP: + disallowed = [] + for name, value in ( + ("concurrency", sweep_config.concurrency), + ("sizes", sweep_config.sizes), + ("range_start", sweep_config.range_start), + ("range_end", sweep_config.range_end), + ("step", sweep_config.step), + ): + if value is not None: + disallowed.append(name) + if disallowed: + raise SweepConfigError( + "concurrency sweeps do not accept input-sweep options: " + + ", ".join(disallowed) + ) + elif sweep_config.sweep_type == INPUT_SWEEP: + if sweep_config.concurrencies is not None: + raise SweepConfigError("input sweeps use concurrency, not concurrencies") + else: + raise SweepConfigError( + f"sweep_type must be one of: {CONCURRENCY_SWEEP}, {INPUT_SWEEP}" + ) + + return sweep_config + + +def build_parser() -> argparse.ArgumentParser: + """Build the sweep CLI parser.""" + return _build_parser() + + +def validate_sweep_args( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + """Validate CLI arguments through the typed sweep config boundary.""" + try: + validate_sweep_config(_sweep_config_from_args(args)) + except SweepConfigError as exc: + parser.error(str(exc)) + + +def resolve_sweep_config(config: SweepConfig) -> Tuple[SweepConfig, SweepSpec]: + """Resolve a typed sweep config to a normalized config and sweep spec.""" + config = validate_sweep_config(config) + model = _normalize_model(config.model) + config = replace(config, model=model) + spec = SPECS.get((config.sweep_type, config.engine, model)) + if spec is None: + raise SweepConfigError( + "Unsupported sweep combination:\n" + f" {config.sweep_type} {config.engine} {model}\n" + f"Supported combinations:\n{supported_combinations()}" + ) + return config, spec + + +def resolve_sweep_spec( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> SweepSpec: + """Resolve CLI arguments to a sweep specification.""" + try: + _, spec = resolve_sweep_config(_sweep_config_from_args(args)) + except SweepConfigError as exc: + parser.error(str(exc)) + return spec + + +def build_sweep_plan( + sweep_config: SweepConfig, + spec: SweepSpec, + *, + endpoint: Optional[EndpointConfig] = None, + tmp_parent: Optional[Path] = None, +) -> SweepPlan: + """Build a concrete set of benchmark runs for a resolved sweep spec.""" + return _build_sweep_plan( + sweep_config, spec, endpoint=endpoint, tmp_parent=tmp_parent + ) + + +def build_sweep_plan_from_config( + config: SweepConfig, + *, + endpoint: Optional[EndpointConfig] = None, + tmp_parent: Optional[Path] = None, +) -> SweepPlan: + """Build a concrete sweep plan from a typed programmatic config.""" + config, spec = resolve_sweep_config(config) + return build_sweep_plan(config, spec, endpoint=endpoint, tmp_parent=tmp_parent) + + +def run_sweep(sweep_config: SweepConfig, spec: SweepSpec, *, dry_run: bool) -> int: + """Execute a resolved sweep plan.""" + return _run_sweep(sweep_config, spec, dry_run=dry_run) + + +def main() -> int: + parser = build_parser() + args = parser.parse_args() + try: + sweep_config, spec = resolve_sweep_config(_sweep_config_from_args(args)) + except SweepConfigError as exc: + parser.error(str(exc)) + return run_sweep(sweep_config, spec, dry_run=args.dry_run) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/veeksha/sweeps/specs.py b/veeksha/sweeps/specs.py new file mode 100644 index 00000000..69f089a6 --- /dev/null +++ b/veeksha/sweeps/specs.py @@ -0,0 +1,278 @@ +"""Static sweep specifications.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Optional, Tuple + +REPO_ROOT = Path(__file__).resolve().parents[2] +CONFIG_DIR = REPO_ROOT / "configs" + +CONCURRENCY_SWEEP = "concurrency" +INPUT_SWEEP = "input" + +MODEL_ALIASES = { + "qwen-tts": "qwen-tts", + "qwen3-tts": "qwen-tts", + "tts": "qwen-tts", + "qwen3-omni": "qwen3-omni", + "omni": "qwen3-omni", + "higgs-audio-v2": "higgs-audio-v2", + "higgs-audio-v3": "higgs-audio-v3", + "higgs-v3": "higgs-audio-v3", + "higgs-audio": "higgs-audio-v2", + "higgs": "higgs-audio-v2", + "fish-speech": "fish-speech", + "fish-speech-s2-pro": "fish-speech", + "fishaudio-s2-pro": "fish-speech", + "s2-pro": "fish-speech", + "vibe-voice": "vibe-voice", + "vibe": "vibe-voice", + "voxtral": "voxtral", + "voxtral-mini": "voxtral", + "stt": "voxtral", +} + + +@dataclass(frozen=True) +class SweepSpec: + sweep_type: str + engine: str + model: str + config_name: str + temp_prefix: str + run_config_template: str + run_name_template: str + default_concurrencies: Tuple[int, ...] = () + default_concurrency: Optional[int] = None + default_range_start: Optional[int] = None + default_range_end: Optional[int] = None + default_step: Optional[int] = None + write_runtime_limits: bool = True + disable_audio_for_input: bool = False + + # Audio-input workloads own their trace and do not use text length bounds. + audio_input: bool = False + + @property + def config_path(self) -> Path: + return CONFIG_DIR / self.config_name + + +SPECS: Dict[Tuple[str, str, str], SweepSpec] = { + (CONCURRENCY_SWEEP, "vajra", "qwen-tts"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vajra", + model="qwen-tts", + config_name="vajra.yaml", + temp_prefix="vajra_qwen_tts_sweep", + run_config_template="vajra_qwen_c{concurrency}.yaml", + run_name_template="vajra_qwen3tts_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64), + ), + (CONCURRENCY_SWEEP, "vllm", "qwen-tts"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vllm", + model="qwen-tts", + config_name="tts_vllm_omni.yaml", + temp_prefix="tts_vllm_omni_sweep", + run_config_template="tts_vllm_omni_c{concurrency}.yaml", + run_name_template="tts_vllm_omni_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "sglang", "qwen-tts"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="sglang", + model="qwen-tts", + config_name="tts_sglang_omni.yaml", + temp_prefix="tts_sglang_omni_sweep", + run_config_template="tts_sglang_omni_c{concurrency}.yaml", + run_name_template="tts_sglang_omni_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "sglang", "higgs-audio-v3"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="sglang", + model="higgs-audio-v3", + config_name="sglang_higgs_audio_v3.yaml", + temp_prefix="sglang_higgs_audio_v3_sweep", + run_config_template="sglang_higgs_audio_v3_c{concurrency}.yaml", + run_name_template="sglang_higgs_audio_v3_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16), + ), + (CONCURRENCY_SWEEP, "vllm", "higgs-audio-v2"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vllm", + model="higgs-audio-v2", + config_name="higgs_audio_v2_vllm_omni.yaml", + temp_prefix="higgs_audio_v2_vllm_omni_sweep", + run_config_template="higgs_audio_v2_vllm_omni_c{concurrency}.yaml", + run_name_template="higgs_audio_v2_vllm_omni_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "vllm", "fish-speech"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vllm", + model="fish-speech", + config_name="fish_speech_vllm_omni.yaml", + temp_prefix="fish_speech_vllm_omni_sweep", + run_config_template="fish_speech_vllm_omni_c{concurrency}.yaml", + run_name_template="fish_speech_vllm_omni_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "vajra", "fish-speech"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vajra", + model="fish-speech", + config_name="vajra_fish_speech_s2_pro.yaml", + temp_prefix="vajra_fish_speech_s2_pro_sweep", + run_config_template="vajra_fish_speech_s2_pro_c{concurrency}.yaml", + run_name_template="vajra_fish_speech_s2_pro_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "vajra", "higgs-audio-v2"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vajra", + model="higgs-audio-v2", + config_name="vajra_higgs_audio_v2.yaml", + temp_prefix="vajra_higgs_audio_v2_sweep", + run_config_template="vajra_higgs_audio_v2_c{concurrency}.yaml", + run_name_template="vajra_higgs_audio_v2_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "vajra", "higgs-audio-v3"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vajra", + model="higgs-audio-v3", + config_name="vajra_higgs_audio_v3.yaml", + temp_prefix="vajra_higgs_audio_v3_sweep", + run_config_template="vajra_higgs_audio_v3_c{concurrency}.yaml", + run_name_template="vajra_higgs_audio_v3_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64), + ), + (CONCURRENCY_SWEEP, "vajra", "qwen3-omni"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vajra", + model="qwen3-omni", + config_name="vajra_qwen.yaml", + temp_prefix="vajra_qwen3_omni_sweep", + run_config_template="vajra_qwen3_omni_c{concurrency}.yaml", + run_name_template="vajra_qwen3_omni_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "vllm", "qwen3-omni"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vllm", + model="qwen3-omni", + config_name="qwen3_omni.yaml", + temp_prefix="vllm_qwen3_omni_sweep", + run_config_template="vllm_qwen3_omni_c{concurrency}.yaml", + run_name_template="vllm_qwen3_omni_c_{concurrency}_10_minutes", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64, 128), + ), + (CONCURRENCY_SWEEP, "vajra", "vibe-voice"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vajra", + model="vibe-voice", + config_name="vajra_vibe_voice.yaml", + temp_prefix="vajra_vibe_voice_0_5_sweep", + run_config_template="vajra_vibe_voice_c{concurrency}.yaml", + run_name_template="vj_vibe_voice_0.5_{date_tag}_c={concurrency}", + default_concurrencies=(1, 2, 4, 8), + write_runtime_limits=False, + ), + (CONCURRENCY_SWEEP, "vajra", "voxtral"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vajra", + model="voxtral", + config_name="stt_vajra.yaml", + temp_prefix="stt_vajra_voxtral_sweep", + run_config_template="stt_vajra_voxtral_c{concurrency}.yaml", + run_name_template="stt_vajra_voxtral_c_{concurrency}", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64), + audio_input=True, + ), + (CONCURRENCY_SWEEP, "vllm", "voxtral"): SweepSpec( + sweep_type=CONCURRENCY_SWEEP, + engine="vllm", + model="voxtral", + config_name="stt_vllm_realtime.yaml", + temp_prefix="stt_vllm_voxtral_sweep", + run_config_template="stt_vllm_voxtral_c{concurrency}.yaml", + run_name_template="stt_vllm_voxtral_c_{concurrency}", + default_concurrencies=(1, 2, 4, 8, 16, 32, 64), + audio_input=True, + ), + (INPUT_SWEEP, "vajra", "qwen-tts"): SweepSpec( + sweep_type=INPUT_SWEEP, + engine="vajra", + model="qwen-tts", + config_name="vajra.yaml", + temp_prefix="vajra_qwen_tts_inputsweep", + run_config_template="vajra_qwen_c{concurrency}_chars{input_size}.yaml", + run_name_template=( + "vajra_qwen3tts_c_{concurrency}_chars_{input_size}_10_minutes" + ), + default_concurrency=64, + default_range_start=380, + default_range_end=500, + default_step=40, + disable_audio_for_input=True, + ), + (INPUT_SWEEP, "vllm", "qwen-tts"): SweepSpec( + sweep_type=INPUT_SWEEP, + engine="vllm", + model="qwen-tts", + config_name="tts_vllm_omni.yaml", + temp_prefix="vllm_qwen_tts_inputsweep", + run_config_template="vllm_qwen_tts_c{concurrency}_chars{input_size}.yaml", + run_name_template=( + "vllm_qwen3tts_c_{concurrency}_chars_{input_size}_10_minutes" + ), + default_concurrency=16, + default_range_start=180, + default_range_end=500, + default_step=40, + disable_audio_for_input=True, + ), + (INPUT_SWEEP, "vllm", "higgs-audio-v2"): SweepSpec( + sweep_type=INPUT_SWEEP, + engine="vllm", + model="higgs-audio-v2", + config_name="higgs_audio_v2_vllm_omni.yaml", + temp_prefix="vllm_higgs_audio_v2_inputsweep", + run_config_template="vllm_higgs_audio_v2_c{concurrency}_chars{input_size}.yaml", + run_name_template=( + "vllm_higgs_audio_v2_c_{concurrency}_chars_{input_size}_10_minutes" + ), + default_concurrency=16, + default_range_start=180, + default_range_end=500, + default_step=40, + disable_audio_for_input=True, + ), + (INPUT_SWEEP, "vllm", "qwen3-omni"): SweepSpec( + sweep_type=INPUT_SWEEP, + engine="vllm", + model="qwen3-omni", + config_name="qwen3_omni.yaml", + temp_prefix="vllm_qwen3_omni_inputsweep", + run_config_template="vllm_qwen3_omni_c{concurrency}_chars{input_size}.yaml", + run_name_template=( + "vllm_qwen3_omni_c_{concurrency}_chars_{input_size}_10_minutes" + ), + default_concurrency=16, + default_range_start=20, + default_range_end=500, + default_step=40, + disable_audio_for_input=True, + ), +} + + +def supported_combinations() -> str: + rows = sorted(SPECS) + return "\n".join( + f" {kind:11s} {engine:8s} {model}" for kind, engine, model in rows + ) diff --git a/veeksha/sweeps/utils.py b/veeksha/sweeps/utils.py new file mode 100644 index 00000000..b5f187fd --- /dev/null +++ b/veeksha/sweeps/utils.py @@ -0,0 +1,264 @@ +"""YAML and template helpers for sweep planning.""" + +from __future__ import annotations + +import copy +import sys +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Optional, Sequence, Tuple + +import yaml + +from veeksha.config.endpoint import EndpointConfig +from veeksha.config.utils import serialize_config_value +from veeksha.sweeps.config import TRACE_SEED_TTS, TRACE_SHAREGPT, SweepConfig +from veeksha.sweeps.specs import INPUT_SWEEP, SweepSpec + + +def load_config(path: Path) -> Dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(f"Base config not found: {path}") + with path.open("r", encoding="utf-8") as f: + data = yaml.safe_load(f) + if not isinstance(data, dict): + raise ValueError(f"Base config must be a YAML mapping: {path}") + return data + + +def write_config(path: Path, config: Dict[str, Any]) -> None: + with path.open("w", encoding="utf-8") as f: + yaml.safe_dump( + serialize_config_value(config), f, sort_keys=False, width=1_000_000 + ) + + +def mapping_at(config: Dict[str, Any], path: Sequence[str]) -> Dict[str, Any]: + node: Any = config + for key in path: + if not isinstance(node, dict) or key not in node: + dotted = ".".join(path) + raise KeyError(f"Missing YAML mapping: {dotted}") + node = node[key] + if not isinstance(node, dict): + dotted = ".".join(path) + raise TypeError(f"Expected YAML mapping at {dotted}") + return node + + +def set_required(config: Dict[str, Any], path: Sequence[str], value: Any) -> None: + parent = mapping_at(config, path[:-1]) + key = path[-1] + if key not in parent: + raise KeyError(f"Missing YAML key: {'.'.join(path)}") + parent[key] = value + + +def apply_trace_source(config: Dict[str, Any], sweep_config: SweepConfig) -> None: + session_generator = mapping_at(config, ("session_generator",)) + if session_generator.get("type") != "trace": + raise ValueError( + "--trace can only be used with trace session generator configs" + ) + + flavor = session_generator.get("flavor") + if not isinstance(flavor, dict): + raise TypeError("Expected YAML mapping at session_generator.flavor") + + if sweep_config.trace == TRACE_SHAREGPT: + flavor["type"] = "sharegpt" + if "assistant_role" not in flavor: + flavor["assistant_role"] = "gpt" + return + + if sweep_config.trace != TRACE_SEED_TTS: + raise ValueError(f"Unsupported trace source: {sweep_config.trace}") + + seed_flavor = {"type": "seed_tts_text"} + for key in ("min_tokens", "max_tokens", "min_chars", "max_chars"): + if key in flavor: + seed_flavor[key] = flavor[key] + + session_generator["trace_file"] = "" + session_generator["flavor"] = seed_flavor + + +def format_template( + template: str, *, concurrency: int, input_size: Optional[int] = None +) -> str: + return template.format( + concurrency=concurrency, + input_size=input_size, + date_tag=_date_tag(), + ) + + +def input_sizes(sweep_config: SweepConfig, spec: SweepSpec) -> Tuple[int, ...]: + if sweep_config.sizes: + return sweep_config.sizes + + start = ( + sweep_config.range_start + if sweep_config.range_start is not None + else spec.default_range_start + ) + end = ( + sweep_config.range_end + if sweep_config.range_end is not None + else spec.default_range_end + ) + step = sweep_config.step if sweep_config.step is not None else spec.default_step + if start is None or end is None or step is None: + raise ValueError("Input sweep requires --sizes or range defaults") + if start <= 0 or end <= 0 or step <= 0: + raise ValueError("--range-start, --range-end, and --step must be positive") + if start > end: + raise ValueError("--range-start must be <= --range-end") + return tuple(range(start, end + 1, step)) + + +def build_run_config( + base_config: Dict[str, Any], + spec: SweepSpec, + *, + concurrency: int, + input_size: Optional[int], + timeout_seconds: int, + max_sessions: int, + output_dir_template: Optional[str], + min_tokens: Optional[int], + max_tokens: Optional[int], + min_chars: Optional[int], + max_chars: Optional[int], +) -> Dict[str, Any]: + config = copy.deepcopy(base_config) + + set_required( + config, + ("traffic_scheduler", "target_concurrent_sessions"), + concurrency, + ) + set_required(config, ("runtime", "num_client_threads"), concurrency) + if spec.write_runtime_limits: + set_required(config, ("runtime", "benchmark_timeout"), timeout_seconds) + set_required(config, ("runtime", "max_sessions"), max_sessions) + + run_name = format_template( + spec.run_name_template, + concurrency=concurrency, + input_size=input_size, + ) + set_required(config, ("wandb", "run_name"), run_name) + + if output_dir_template: + config["output_dir"] = _format_output_dir( + output_dir_template, + spec, + concurrency=concurrency, + input_size=input_size, + run_name=run_name, + ) + + if spec.sweep_type == INPUT_SWEEP: + if input_size is None: + raise ValueError("input_size is required for input sweeps") + flavor = mapping_at(config, ("session_generator", "flavor")) + _clear_length_bounds(flavor) + flavor["min_chars"] = input_size + flavor["max_chars"] = input_size + if spec.disable_audio_for_input: + _disable_audio_saving(config) + elif not spec.audio_input: + # STT audio traces have no text length bounds to rewrite. + _apply_length_bounds( + config, + min_tokens=min_tokens, + max_tokens=max_tokens, + min_chars=min_chars, + max_chars=max_chars, + ) + + return config + + +def benchmark_command(config_path: Path) -> list[str]: + return [ + sys.executable, + "-Xgil=0", + "-m", + "veeksha.benchmark", + "--config", + str(config_path), + ] + + +def apply_endpoint(config: Dict[str, Any], endpoint: Optional[EndpointConfig]) -> None: + if endpoint is None: + return + config["endpoint"] = endpoint.to_dict() + client = mapping_at(config, ("client",)) + endpoint.apply_to_client_mapping(client) + + +def _date_tag() -> str: + return datetime.now().strftime("%m_%d") + + +def _format_output_dir( + template: str, + spec: SweepSpec, + *, + concurrency: int, + input_size: Optional[int], + run_name: str, +) -> str: + try: + return template.format( + concurrency=concurrency, + input_size=input_size, + run_name=run_name, + date_tag=_date_tag(), + engine=spec.engine, + model=spec.model, + sweep_type=spec.sweep_type, + ) + except KeyError as exc: + raise ValueError(f"Unknown --output-dir template field: {exc.args[0]}") from exc + + +def _disable_audio_saving(config: Dict[str, Any]) -> None: + evaluators = config.get("evaluators") + if not isinstance(evaluators, list): + return + for evaluator in evaluators: + if not isinstance(evaluator, dict): + continue + if evaluator.get("type") == "audio_quality" and "save_audio_files" in evaluator: + evaluator["save_audio_files"] = False + + +def _clear_length_bounds(flavor: Dict[str, Any]) -> None: + for key in ("min_tokens", "max_tokens", "min_chars", "max_chars"): + flavor.pop(key, None) + + +def _apply_length_bounds( + config: Dict[str, Any], + *, + min_tokens: Optional[int], + max_tokens: Optional[int], + min_chars: Optional[int], + max_chars: Optional[int], +) -> None: + flavor = mapping_at(config, ("session_generator", "flavor")) + _clear_length_bounds(flavor) + + if min_chars is not None and max_chars is not None: + flavor["min_chars"] = min_chars + flavor["max_chars"] = max_chars + return + + if min_tokens is None or max_tokens is None: + raise ValueError("min_tokens and max_tokens must be set together") + flavor["min_tokens"] = min_tokens + flavor["max_tokens"] = max_tokens diff --git a/veeksha/types/__init__.py b/veeksha/types/__init__.py index 3fa16ff6..e0d81d16 100644 --- a/veeksha/types/__init__.py +++ b/veeksha/types/__init__.py @@ -36,6 +36,9 @@ class TraceFlavorType(BaseIntEnum): RAG = 3 REQUEST_LOG = 4 UNTIMED_CONTENT_MULTI_TURN = 5 + SHAREGPT = 6 + SEED_TTS_TEXT = 7 + AUDIO = 8 class ChannelModality(BaseIntEnum): @@ -45,6 +48,14 @@ class ChannelModality(BaseIntEnum): VIDEO = 4 +class AudioTask(BaseIntEnum): + """Workload represented by an AUDIO response.""" + + TTS = 1 + STT = 2 + LLM_AUDIO = 3 + + class SessionGraphType(BaseIntEnum): LINEAR = 1 SINGLE_REQUEST = 2 @@ -55,6 +66,7 @@ class SessionGraphType(BaseIntEnum): class EvaluationType(BaseIntEnum): PERFORMANCE = 1 ACCURACY_LMEVAL = 2 + AUDIO_QUALITY = 3 # ----- Client ----- @@ -62,12 +74,26 @@ class ClientType(BaseIntEnum): OPENAI_CHAT_COMPLETIONS = 1 OPENAI_COMPLETIONS = 2 OPENAI_ROUTER = 3 + TTS = 4 + REALTIME_TTS = 5 + STT = 6 + ELEVENLABS_STREAMING_TTS = 7 + DEEPGRAM_FLUX_STREAMING_TTS = 8 + ELEVENLABS_HTTP_TTS = 9 + DEEPGRAM_FLUX_HTTP_TTS = 10 + DEEPGRAM_AURA_STREAMING_TTS = 11 # ----- Server ----- class ServerType(BaseIntEnum): VLLM = 1 SGLANG = 2 + VAJRA = 3 + + +# ----- Verification ----- +class VerificationType(BaseIntEnum): + AUDIO = 1 # ----- SLOs ----- diff --git a/veeksha/verification/__init__.py b/veeksha/verification/__init__.py new file mode 100644 index 00000000..648e9670 --- /dev/null +++ b/veeksha/verification/__init__.py @@ -0,0 +1 @@ +"""Post-run verification helpers.""" diff --git a/veeksha/verification/audio.py b/veeksha/verification/audio.py new file mode 100644 index 00000000..94dfce0a --- /dev/null +++ b/veeksha/verification/audio.py @@ -0,0 +1,545 @@ +"""Post-run audio verification with composable verifiers.""" + +from __future__ import annotations + +import json +import math +import statistics +import string +import threading +from abc import ABC, abstractmethod +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Callable, Iterable, Optional + +import numpy as np + +from veeksha.config.verification import AudioVerificationConfig +from veeksha.core.audio_contract import AudioMetricKey +from veeksha.logger import init_logger + +logger = init_logger(__name__) + + +TranscribeFn = Callable[[Path], str] +_utmos_lock = threading.Lock() +_utmos_jit_model: Any | None = None +_utmos_jit_key: tuple[str, str, str] | None = None +_utmos_jit_load_failed_keys: set[tuple[str, str, str]] = set() + + +class AudioVerificationError(RuntimeError): + """Raised when strict audio verification fails.""" + + +@dataclass(frozen=True) +class AudioVerifierContext: + request_id: int + reference_text: str + audio_path: Path + has_audio: bool + + +class AudioOutputVerifier(ABC): + """Base class for a verifier that evaluates persisted audio.""" + + name: str + + @abstractmethod + def verify(self, context: AudioVerifierContext) -> dict[str, Any]: + """Return row fields produced by this verifier.""" + + +class WERVerifier(AudioOutputVerifier): + name = "wer" + + def __init__(self, config: AudioVerificationConfig, transcribe_audio: TranscribeFn): + self.config = config + self.transcribe_audio = transcribe_audio + + def verify(self, context: AudioVerifierContext) -> dict[str, Any]: + if not context.reference_text: + return { + "wer_error": f"Missing {AudioMetricKey.INPUT_TEXT.value} in request-level metrics" + } + if not context.has_audio: + return {"wer_error": f"Missing audio file: {context.audio_path}"} + try: + transcript = self.transcribe_audio(context.audio_path) + wer = compute_wer(context.reference_text, transcript) + return { + "transcript": transcript, + "wer": wer, + "passed": wer <= self.config.wer.threshold, + } + except Exception as exc: + return {"wer_error": f"Transcription failed: {exc}"} + + +class UTMOSVerifier(AudioOutputVerifier): + name = "utmos" + + def __init__(self, config: AudioVerificationConfig): + self.config = config + + def verify(self, context: AudioVerifierContext) -> dict[str, Any]: + if not context.has_audio: + return {"utmos_error": f"Missing audio file: {context.audio_path}"} + try: + utmos = _utmos_predict_audio_path(context.audio_path, self.config) + if utmos is None: + return {"utmos_error": "UTMOS unavailable or returned no finite score"} + return {"utmos": utmos} + except Exception as exc: + return {"utmos_error": f"UTMOS failed: {exc}"} + + +class LocalWhisperTranscriber: + """In-process faster-whisper transcriber for WER verification.""" + + def __init__(self, config: AudioVerificationConfig): + self.config = config + whisper_config = config.wer.whisper + try: + from faster_whisper import WhisperModel + except ImportError as exc: + raise AudioVerificationError( + "WER verification requires faster-whisper. Install the audio-verification extra." + ) from exc + + self.model = WhisperModel( + whisper_config.model, + device=whisper_config.device, + compute_type=whisper_config.compute_type, + ) + + def transcribe(self, audio_path: Path) -> str: + whisper_config = self.config.wer.whisper + segments, _ = self.model.transcribe( + str(audio_path), + language=whisper_config.language, + task="transcribe", + beam_size=whisper_config.beam_size, + ) + return " ".join(segment.text.strip() for segment in segments).strip() + + +@dataclass +class AudioVerificationRow: + request_id: int + reference_text: str + transcript: str + wer: Optional[float] + passed: Optional[bool] + audio_path: str + utmos: Optional[float] = None + error: Optional[str] = None + wer_error: Optional[str] = None + utmos_error: Optional[str] = None + + +@dataclass +class AudioVerificationSummary: + total_requests: int + evaluated_requests: int + skipped_requests: int + max_requests: int + transcribed_requests: int + passed_requests: int + failed_requests: int + error_requests: int + wer_avg: Optional[float] + wer_p50: Optional[float] + wer_p90: Optional[float] + wer_p99: Optional[float] + wer_max: Optional[float] + wer_threshold: float + fail_on_threshold: bool + wer_enabled: bool + utmos_enabled: bool + utmos_evaluated: int + utmos_mean: Optional[float] + utmos_median: Optional[float] + utmos_failed: int + errors: list[str] + + def to_dict(self) -> dict: + return asdict(self) + + +def normalize_text(text: str) -> str: + """Normalize English text with the Seed-TTS WER punctuation protocol.""" + normalized = text + for char in string.punctuation: + if char == "'": + continue + normalized = normalized.replace(char, "") + return normalized.lower() + + +def _jiwer_wer(reference: str, hypothesis: str) -> float: + try: + from jiwer import compute_measures + + return float(compute_measures(reference, hypothesis)["wer"]) + except ImportError: + import jiwer + + return float(jiwer.process_words(reference, hypothesis).wer) + + +def compute_wer(reference: str, hypothesis: str) -> float: + """Compute English WER using Seed-TTS normalization.""" + reference_n = normalize_text(reference) + hypothesis_n = normalize_text(hypothesis) + reference_words = reference_n.split() + hypothesis_words = hypothesis_n.split() + if not reference_words: + return 0.0 if not hypothesis_words else 1.0 + + try: + return _jiwer_wer(reference_n, hypothesis_n) + except ImportError: + return _edit_distance(reference_words, hypothesis_words) / len(reference_words) + + +def _audio_path_to_f32_16k(path: Path) -> np.ndarray: + import scipy.signal + import soundfile as sf + + data, sample_rate = sf.read(path, dtype="float32", always_2d=True) + if data.size == 0: + return np.zeros(0, dtype=np.float32) + + mono = np.mean(data, axis=1).astype(np.float32) + if int(sample_rate) == 16000: + return mono + + target_len = max(1, int(len(mono) * 16000 / int(sample_rate))) + return scipy.signal.resample(mono, target_len).astype(np.float32) + + +def _utmos_key(config: AudioVerificationConfig) -> tuple[str, str, str]: + return (config.utmos.hf_repo, config.utmos.jit_file, config.utmos.device) + + +def _resolve_utmos_device(config: AudioVerificationConfig) -> str: + import torch + + device = config.utmos.device.strip() + if device.startswith("cuda") and not torch.cuda.is_available(): + raise AudioVerificationError( + f"UTMOS requested {device}, but CUDA is unavailable." + ) + return device + + +def _ensure_utmos_jit_model(config: AudioVerificationConfig) -> Any | None: + global _utmos_jit_key, _utmos_jit_model + + key = _utmos_key(config) + with _utmos_lock: + if key in _utmos_jit_load_failed_keys: + return None + if _utmos_jit_model is not None and _utmos_jit_key == key: + return _utmos_jit_model + + try: + import torch + from huggingface_hub import hf_hub_download + + path = hf_hub_download( + repo_id=config.utmos.hf_repo, + filename=config.utmos.jit_file, + repo_type="model", + ) + model = torch.jit.load(path, map_location=_resolve_utmos_device(config)) + model.eval() + except Exception as exc: + logger.warning( + "UTMOS JIT unavailable; install torch, scipy, soundfile, and " + "huggingface_hub, then check GPU/model access: %s", + exc, + ) + _utmos_jit_load_failed_keys.add(key) + return None + + _utmos_jit_model = model + _utmos_jit_key = key + return _utmos_jit_model + + +def _utmos_predict_f32_16k( + wav_f32: np.ndarray, config: AudioVerificationConfig +) -> float | None: + import torch + + if len(wav_f32) == 0: + return None + + model = _ensure_utmos_jit_model(config) + if model is None: + return None + + try: + model_device = next(model.parameters()).device + except StopIteration: + try: + model_device = next(model.buffers()).device + except StopIteration: + model_device = torch.device(config.utmos.device) + + wav = np.ascontiguousarray(wav_f32, dtype=np.float32) + model_input = ( + torch.from_numpy(wav).unsqueeze(0).to(device=model_device, dtype=torch.float32) + ) + with torch.no_grad(): + output = model(model_input) + value = float(output.reshape(-1)[0].item()) + if not math.isfinite(value): + return None + return value + + +def _utmos_predict_audio_path( + audio_path: Path, config: AudioVerificationConfig +) -> float | None: + wav_16k = _audio_path_to_f32_16k(audio_path) + return _utmos_predict_f32_16k(wav_16k, config) + + +def build_audio_verifiers( + config: AudioVerificationConfig, + transcribe_audio: Optional[TranscribeFn] = None, +) -> list[AudioOutputVerifier]: + """Build the configured audio output verifiers.""" + verifiers: list[AudioOutputVerifier] = [] + if config.wer.enabled: + if transcribe_audio is None: + raise AudioVerificationError( + "WER verification requires a transcription function" + ) + verifiers.append(WERVerifier(config, transcribe_audio)) + if config.utmos.enabled: + verifiers.append(UTMOSVerifier(config)) + return verifiers + + +def verify_audio_outputs( + output_dir: str | Path, + config: AudioVerificationConfig, + transcribe_audio: Optional[TranscribeFn] = None, +) -> AudioVerificationSummary: + """Verify saved audio files and persist JSON artifacts.""" + verifiers = build_audio_verifiers(config, transcribe_audio) + output_path = Path(output_dir) + verification_dir = output_path / "verification" + verification_dir.mkdir(parents=True, exist_ok=True) + + request_metrics_path = output_path / "metrics" / "request_level_metrics.jsonl" + audio_dir = output_path / "audio_files" + rows: list[AudioVerificationRow] = [] + errors: list[str] = [] + total_requests = 0 + skipped_requests = 0 + + if not request_metrics_path.exists(): + errors.append(f"Missing request metrics file: {request_metrics_path}") + else: + metric_rows = list(_load_jsonl(request_metrics_path)) + total_requests = len(metric_rows) + if config.max_requests > 0: + skipped_requests = max(total_requests - config.max_requests, 0) + metric_rows = metric_rows[: config.max_requests] + if skipped_requests: + logger.info( + "Audio verification limited to first %d of %d request rows", + config.max_requests, + total_requests, + ) + + for metric_row in metric_rows: + request_id = metric_row.get("request_id") + if request_id is None: + errors.append("Skipping request row without request_id") + continue + + audio_path = audio_dir / f"request_{request_id}.wav" + context = AudioVerifierContext( + request_id=int(request_id), + reference_text=str( + metric_row.get(AudioMetricKey.INPUT_TEXT.value) or "" + ), + audio_path=audio_path, + has_audio=audio_path.exists(), + ) + row_data: dict[str, Any] = { + "transcript": "", + "wer": None, + "passed": None, + "utmos": None, + "wer_error": None, + "utmos_error": None, + } + for verifier in verifiers: + row_data.update(verifier.verify(context)) + + row_error = _combine_errors(row_data["wer_error"], row_data["utmos_error"]) + rows.append( + AudioVerificationRow( + request_id=context.request_id, + reference_text=context.reference_text, + transcript=row_data["transcript"], + wer=row_data["wer"], + passed=row_data["passed"], + audio_path=str(context.audio_path), + utmos=row_data["utmos"], + error=row_error, + wer_error=row_data["wer_error"], + utmos_error=row_data["utmos_error"], + ) + ) + + _save_rows(verification_dir / "audio_verification.jsonl", rows) + summary = _build_summary(rows, config, errors, total_requests, skipped_requests) + _save_summary(verification_dir / "audio_summary.json", summary) + + if summary.failed_requests: + logger.warning( + "Audio verification found %d requests above WER threshold %.4f", + summary.failed_requests, + config.wer.threshold, + ) + if summary.error_requests or summary.errors: + logger.warning( + "Audio verification completed with %d request errors and %d run errors", + summary.error_requests, + len(summary.errors), + ) + + if config.fail_on_threshold: + failures = [] + if summary.failed_requests: + failures.append( + f"{summary.failed_requests} audio requests exceeded WER threshold " + f"{config.wer.threshold}" + ) + if summary.error_requests: + failures.append( + f"{summary.error_requests} audio requests could not be verified" + ) + if summary.errors: + failures.append(f"{len(summary.errors)} run-level verification errors") + if failures: + raise AudioVerificationError("; ".join(failures)) + + return summary + + +def run_audio_verification( + config: AudioVerificationConfig, + output_dir: str | Path, +) -> AudioVerificationSummary: + """Run configured post-run audio verification metrics.""" + verification_dir = Path(output_dir) / "verification" + verification_dir.mkdir(parents=True, exist_ok=True) + + transcriber = LocalWhisperTranscriber(config) if config.wer.enabled else None + return verify_audio_outputs( + output_dir=output_dir, + config=config, + transcribe_audio=transcriber.transcribe if transcriber is not None else None, + ) + + +def _load_jsonl(path: Path) -> Iterable[dict]: + with path.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + yield json.loads(line) + + +def _save_rows(path: Path, rows: list[AudioVerificationRow]) -> None: + with path.open("w", encoding="utf-8") as f: + for row in rows: + f.write(json.dumps(asdict(row)) + "\n") + + +def _save_summary(path: Path, summary: AudioVerificationSummary) -> None: + with path.open("w", encoding="utf-8") as f: + json.dump(summary.to_dict(), f, indent=2) + + +def _build_summary( + rows: list[AudioVerificationRow], + config: AudioVerificationConfig, + errors: list[str], + total_requests: int, + skipped_requests: int, +) -> AudioVerificationSummary: + wers = [row.wer for row in rows if row.wer is not None] + utmos_values = [row.utmos for row in rows if row.utmos is not None] + passed_requests = sum(1 for row in rows if row.passed is True) + failed_requests = sum(1 for row in rows if row.passed is False) + error_requests = sum(1 for row in rows if row.error is not None) + utmos_failed = sum(1 for row in rows if row.utmos_error is not None) + return AudioVerificationSummary( + total_requests=total_requests, + evaluated_requests=len(rows), + skipped_requests=skipped_requests, + max_requests=config.max_requests, + transcribed_requests=len(wers), + passed_requests=passed_requests, + failed_requests=failed_requests, + error_requests=error_requests, + wer_avg=(sum(wers) / len(wers)) if wers else None, + wer_p50=_percentile(wers, 0.50), + wer_p90=_percentile(wers, 0.90), + wer_p99=_percentile(wers, 0.99), + wer_max=max(wers) if wers else None, + wer_threshold=config.wer.threshold, + fail_on_threshold=config.fail_on_threshold, + wer_enabled=config.wer.enabled, + utmos_enabled=config.utmos.enabled, + utmos_evaluated=len(utmos_values), + utmos_mean=statistics.fmean(utmos_values) if utmos_values else None, + utmos_median=statistics.median(utmos_values) if utmos_values else None, + utmos_failed=utmos_failed, + errors=errors, + ) + + +def _combine_errors(*errors: Optional[str]) -> Optional[str]: + parts = [error for error in errors if error] + return "; ".join(parts) if parts else None + + +def _percentile(values: list[float], percentile: float) -> Optional[float]: + if not values: + return None + ordered = sorted(values) + if len(ordered) == 1: + return ordered[0] + rank = (len(ordered) - 1) * percentile + lower_idx = int(rank) + upper_idx = min(lower_idx + 1, len(ordered) - 1) + fraction = rank - lower_idx + return ordered[lower_idx] + (ordered[upper_idx] - ordered[lower_idx]) * fraction + + +def _edit_distance(reference_words: list[str], hypothesis_words: list[str]) -> int: + previous = list(range(len(hypothesis_words) + 1)) + for i, reference_word in enumerate(reference_words, start=1): + current = [i] + for j, hypothesis_word in enumerate(hypothesis_words, start=1): + substitution_cost = 0 if reference_word == hypothesis_word else 1 + current.append( + min( + previous[j] + 1, + current[j - 1] + 1, + previous[j - 1] + substitution_cost, + ) + ) + previous = current + return previous[-1]