Output:
Chunk ID: 43dfe8
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 5321
Output:
production/worker.py:3:Proxies HTTP requests to ASR model server, reports metrics for autoscaling.
production/worker.py:60:def metrics_workload_calculator(_payload):
production/worker.py:95:            route="/metrics",
production/worker.py:98:            workload_calculator=metrics_workload_calculator,
production/server.py:40:from starlette.websockets import WebSocketState
production/server.py:66:ENABLE_VAD = os.environ.get("ENABLE_VAD", "true").lower() == "true"
production/server.py:120:        self.metrics = {
production/server.py:129:            "vad_filtered_chunks": 0,
production/server.py:133:            "ws_connections_total": 0,
production/server.py:134:            "ws_connections_active": 0,
production/server.py:135:            "ws_connections_rejected": 0,
production/server.py:215:        input_case: str = "lower_cased"
production/server.py:247:        chunk_size_in_secs: Optional[float] = None
production/server.py:289:        calculate_wer: bool = True
production/server.py:291:        metrics: MetricsConfig = field(default_factory=MetricsConfig)
production/server.py:334:    chunk_size = int(0.080 * 16000 * 2)  # 80ms chunks, 16kHz, 16-bit
production/server.py:341:        for i in range(0, len(audio_bytes), chunk_size):
production/server.py:342:            chunk = audio_bytes[i:i + chunk_size]
production/server.py:343:            state.pipeline.append_streaming_audio(stream_id, chunk)
production/server.py:387:    # Use pipeline's native chunk size if available; fallback to 80ms
production/server.py:388:    frame_secs = getattr(state.pipeline, "chunk_size_in_secs", 0.08) or 0.08
production/server.py:395:        # Feed audio in frame-sized chunks
production/server.py:397:            chunk = pcm_bytes[i:i + frame_bytes]
production/server.py:398:            state.pipeline.append_streaming_audio(stream_id, chunk)
production/server.py:459:    Returns (audio_bytes, language, response_format, chunk_duration_ms).
production/server.py:471:            raise HTTPException(status_code=400, detail="audio_base64 required")
production/server.py:475:            raise HTTPException(status_code=400, detail=f"invalid audio_base64: {e}")
production/server.py:478:        chunk_duration_ms = int(body.get("chunk_duration_ms", 160) or 160)
production/server.py:479:        return audio_data, language, response_format, chunk_duration_ms
production/server.py:485:        raise HTTPException(status_code=400, detail="file required")
production/server.py:494:    chunk_duration_ms = int(form.get("chunk_duration_ms", 160) or 160)
production/server.py:495:    return audio_data, language, response_format, chunk_duration_ms
production/server.py:537:def _build_metrics_payload():
production/server.py:538:    """Shared metrics payload for GET/POST /metrics."""
production/server.py:539:    m = state.metrics
production/server.py:570:            "vad_filtered_chunks": m["vad_filtered_chunks"],
production/server.py:582:        "ws": {
production/server.py:583:            "connections_total": m["ws_connections_total"],
production/server.py:584:            "connections_active": m["ws_connections_active"],
production/server.py:585:            "connections_rejected": m["ws_connections_rejected"],
production/server.py:597:@app.get("/metrics")
production/server.py:598:async def metrics_get():
production/server.py:600:    Detailed metrics endpoint for monitoring/alerting.
production/server.py:603:    return _build_metrics_payload()
production/server.py:606:@app.post("/metrics")
production/server.py:607:async def metrics_post(_request: Request):
production/server.py:609:    return _build_metrics_payload()
production/server.py:619:    state.metrics["http_requests_total"] += 1
production/server.py:621:        state.metrics["http_requests_err"] += 1
production/server.py:622:        raise HTTPException(status_code=503, detail="Model not loaded yet")
production/server.py:625:        audio_data, language, _response_format, _chunk_duration_ms = await _parse_audio_request(request)
production/server.py:632:        state.metrics["http_requests_ok"] += 1
production/server.py:641:        state.metrics["http_requests_err"] += 1
production/server.py:643:        raise HTTPException(status_code=500, detail=str(e))
production/server.py:655:    state.metrics["http_requests_total"] += 1
production/server.py:658:        state.metrics["http_requests_err"] += 1
production/server.py:659:        raise HTTPException(status_code=503, detail="Model not loaded yet")
production/server.py:665:        audio_data, language, response_format, _chunk_duration_ms = await _parse_audio_request(request)
production/server.py:667:            raise HTTPException(status_code=400, detail="Empty audio file")
production/server.py:687:        state.metrics["http_requests_ok"] += 1
production/server.py:694:        state.metrics["http_requests_err"] += 1
production/server.py:696:        raise HTTPException(status_code=500, detail=str(e))
production/server.py:708:    state.metrics["http_requests_total"] += 1
production/server.py:711:        state.metrics["http_requests_err"] += 1
production/server.py:712:        raise HTTPException(status_code=503, detail="Model not loaded yet")
production/server.py:715:        audio_data, language, response_format, chunk_duration_ms = await _parse_audio_request(request)
production/server.py:717:            raise HTTPException(status_code=400, detail="Empty audio file")
production/server.py:730:            chunk_size = int(chunk_duration_ms / 1000.0 * 16000 * 2)  # bytes per chunk
production/server.py:739:                for i in range(0, len(pcm_bytes), chunk_size):
production/server.py:740:                    chunk = pcm_bytes[i:i + chunk_size]
production/server.py:741:                    audio_buffer += chunk
production/server.py:811:        state.metrics["http_requests_ok"] += 1
production/server.py:821:        state.metrics["http_requests_err"] += 1
production/server.py:823:        raise HTTPException(status_code=500, detail=str(e))
production/server.py:830:@app.websocket("/ws")
production/server.py:831:async def websocket_streaming(ws: WebSocket):
production/server.py:836:      Client -> Server: binary audio chunks (raw PCM 16-bit 16kHz mono)
production/server.py:843:    state.metrics["ws_connections_total"] += 1
production/server.py:848:        state.metrics["ws_connections_rejected"] += 1
production/server.py:849:        await ws.accept()
production/server.py:850:        await ws.close(code=1013, reason=f"Overloaded: {active}/{MAX_ACTIVE_STREAMS}")
production/server.py:853:    await ws.accept()
production/server.py:854:    state.metrics["ws_connections_active"] += 1
production/server.py:869:        state.metrics["ws_connections_active"] -= 1
production/server.py:872:        if ws.application_state is WebSocketState.CONNECTED:
production/server.py:873:            await ws.close()
production/server.py:877:    await ws.send_text(json.dumps({"type": "Ready", "id": stream_id}))
production/server.py:880:        """Receive audio chunks from client, feed to pipeline"""
production/server.py:887:                message = await ws.receive()
production/server.py:907:                            state.metrics["vad_filtered_chunks"] += 1
production/server.py:950:                    await ws.send_text("END")
production/server.py:958:                    await ws.send_text(json_str)
production/server.py:974:        state.metrics["ws_connections_active"] -= 1
production/server.py:976:        if ws.application_state is WebSocketState.CONNECTED:
production/server.py:977:            await ws.close()
production/server.py:1087:    metrics_report_interval = 100
production/server.py:1128:                    state.metrics["total_batches"] += 1
production/server.py:1129:                    state.metrics["total_outputs"] += len(outputs)
production/server.py:1130:                    state.metrics["total_gpu_ms"] += gpu_ms
production/server.py:1131:                    state.metrics["batch_sizes"].append(len(outputs))
production/server.py:1132:                    state.metrics["gpu_times_ms"].append(gpu_ms)
production/server.py:1151:                            state.metrics["empty_cycles"] += 1
production/server.py:1163:                state.metrics["total_route_ms"] += (time.perf_counter() - route_start) * 1000
production/server.py:1168:                state.metrics["cycle_times_ms"].append(cycle_ms)
production/server.py:1170:                # Periodic metrics log
production/server.py:1171:                if state.metrics["total_batches"] % metrics_report_interval == 0:
production/server.py:1172:                    _log_metrics(active)
production/server.py:1181:def _log_metrics(active_clients: int):
production/server.py:1182:    """Structured metrics log"""
production/server.py:1183:    m = state.metrics
production/server.py:1203:        f"vad_skip={m['vad_filtered_chunks']}"
production/test_stress_full.py:10:- Time-series snapshots from /metrics and /get_endpoint_workers/
production/test_stress_full.py:39:    import jiwer
production/test_stress_full.py:41:    jiwer = None
production/test_stress_full.py:94:    queue_full = "all worker queues full" in status.lower()
production/test_stress_full.py:118:    text = text.lower().strip()
production/test_stress_full.py:125:def compute_wer_cer(reference: str, hypothesis: str) -> Tuple[Optional[float], Optional[float]]:
production/test_stress_full.py:126:    if not jiwer:
production/test_stress_full.py:133:        return float(jiwer.wer(ref, hyp)), float(jiwer.cer(ref, hyp))
production/test_stress_full.py:172:def load_fleurs_samples(max_samples: int, min_duration: float, max_duration: float, seed: int) -> List[AudioSample]:
production/test_stress_full.py:175:    ds = load_dataset("google/fleurs", "hi_in", split="test")
production/test_stress_full.py:176:    rows = list(ds)
production/test_stress_full.py:178:    rnd.shuffle(rows)
production/test_stress_full.py:181:    for row in rows:
production/test_stress_full.py:228:    wer: float
production/test_stress_full.py:258:    wer_pct: float
production/test_stress_full.py:263:async def fetch_server_metrics_via_endpoint(
production/test_stress_full.py:276:            f"{worker_url}/metrics",
production/test_stress_full.py:282:            return {"error": f"metrics status={r.status}"}
production/test_stress_full.py:306:async def metrics_sampler(
production/test_stress_full.py:320:            sm = await fetch_server_metrics_via_endpoint(session, endpoint_name, endpoint_key, route_cost)
production/test_stress_full.py:321:            server_samples.append({"t": ts, "metrics": sm})
production/test_stress_full.py:418:    wer_val, cer_val = compute_wer_cer(ref, hypothesis)
production/test_stress_full.py:437:        wer=float("nan") if wer_val is None else wer_val,
production/test_stress_full.py:469:                "chunk_duration_ms": 160,
production/test_stress_full.py:560:    wer_val, cer_val = compute_wer_cer(ref, hypothesis)
production/test_stress_full.py:579:        wer=float("nan") if wer_val is None else wer_val,
production/test_stress_full.py:587:        m = s.get("metrics", {})
production/test_stress_full.py:628:    wer_vals = [r.wer for r in ok_recs if r.reference and not math.isnan(r.wer)]
production/test_stress_full.py:666:        wer_pct=round(100.0 * safe_mean(wer_vals), 3) if wer_vals else 0.0,
production/test_stress_full.py:681:    metrics_interval: float,
production/test_stress_full.py:717:            metrics_sampler(
production/test_stress_full.py:719:                metrics_interval,
production/test_stress_full.py:783:        f"ft_p95={summary.first_token_ms_p95:.0f}ms wer={summary.wer_pct:.2f}%"
production/test_stress_full.py:800:    details_csv_path = out_dir / f"{out_prefix}_details.csv"
production/test_stress_full.py:827:        "wer_pct",
production/test_stress_full.py:840:    with details_csv_path.open("w", newline="") as f:
production/test_stress_full.py:857:            "wer",
production/test_stress_full.py:878:    return str(csv_path), str(details_csv_path), str(json_path)
production/test_stress_full.py:892:    p.add_argument("--metrics-interval", type=float, default=2.0)
production/test_stress_full.py:895:    p.add_argument("--audio-source", choices=["sine", "fleurs"], default="fleurs")
production/test_stress_full.py:898:    p.add_argument("--sample-pool", type=int, default=120, help="Pool size for fleurs samples")
production/test_stress_full.py:930:        samples = load_fleurs_samples(
production/test_stress_full.py:959:                metrics_interval=args.metrics_interval,
production/test_stress_full.py:971:    csv_path, details_path, json_path = write_outputs(results, prefix)
production/test_stress_full.py:982:            f"tput={s['throughput_x']:>6.2f}x wer={s['wer_pct']:>6.2f}%"
production/test_stress_full.py:986:    print(f"Details CSV : {details_path}")
production/test_stress.py:5:Captures ALL metrics needed for scaling law decisions:
production/test_stress.py:11:- GPU utilization (via /metrics endpoint)
production/test_stress.py:16:    python test_stress.py --url http://ip:8000 --mode ws --levels 1,10,50,100
production/test_stress.py:94:    """Collects and reports all metrics for a stress test level"""
production/test_stress.py:101:        # Per-stream metrics
production/test_stress.py:108:        # Server metrics snapshots
production/test_stress.py:109:        self.server_metrics_before = None
production/test_stress.py:110:        self.server_metrics_after = None
production/test_stress.py:138:async def fetch_server_metrics(url=None, endpoint_name=None, endpoint_key=None, route_cost=10.0):
production/test_stress.py:139:    """Fetch /metrics either directly or through serverless PyWorker."""
production/test_stress.py:155:                    f"{worker_url}/metrics",
production/test_stress.py:169:            async with s.get(f"{url}/metrics", timeout=aiohttp.ClientTimeout(total=5)) as r:
production/test_stress.py:194:async def run_ws_stream(url, pcm_bytes, stream_idx, metrics, duration_s,
production/test_stress.py:195:                        endpoint_name=None, endpoint_key=None, route_cost=10.0, ws_port=None):
production/test_stress.py:198:        if not ws_port:
production/test_stress.py:199:            metrics.errors.append({"stream": stream_idx, "error": "ws_port required for endpoint mode"})
production/test_stress.py:210:                metrics.errors.append({"stream": stream_idx, "error": f"route failed: {route}"})
production/test_stress.py:213:            scheme = "wss" if parsed.scheme == "https" else "ws"
production/test_stress.py:214:            ws_url = f"{scheme}://{parsed.hostname}:{ws_port}/ws"
production/test_stress.py:216:            metrics.errors.append({"stream": stream_idx, "error": str(e)})
production/test_stress.py:219:        ws_url = url.replace("http://", "ws://").replace("https://", "wss://") + "/ws"
production/test_stress.py:220:    chunk_size = int(0.160 * 16000 * 2)  # 160ms chunks
production/test_stress.py:225:            async with s.ws_connect(ws_url, timeout=30) as ws:
production/test_stress.py:227:                metrics.connect_times_ms.append((t_connected - t_start) * 1000)
production/test_stress.py:230:                ready = await asyncio.wait_for(ws.receive(), timeout=10)
production/test_stress.py:236:                    for i in range(0, len(pcm_bytes), chunk_size):
production/test_stress.py:237:                        await ws.send_bytes(pcm_bytes[i:i + chunk_size])
production/test_stress.py:239:                    await ws.send_str("END")
production/test_stress.py:245:                            m = await asyncio.wait_for(ws.receive(), timeout=30)
production/test_stress.py:262:                    metrics.first_token_times_ms.append((first_token - t_start) * 1000)
production/test_stress.py:263:                metrics.total_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:264:                metrics.messages_per_stream.append(len(msgs))
production/test_stress.py:271:                    metrics.texts.append(last_text)
production/test_stress.py:274:        metrics.errors.append({"stream": stream_idx, "error": str(e)})
production/test_stress.py:277:async def run_http_batch(url, wav_bytes, stream_idx, metrics,
production/test_stress.py:293:                    metrics.errors.append({"stream": stream_idx, "error": f"route failed: {route}"})
production/test_stress.py:304:                    metrics.connect_times_ms.append(0)
production/test_stress.py:307:                        metrics.total_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:308:                        metrics.first_token_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:309:                        metrics.messages_per_stream.append(1)
production/test_stress.py:310:                        metrics.texts.append(d.get("text", ""))
production/test_stress.py:320:                            metrics.errors.append({"stream": stream_idx, "error": f"route failed: {route2}"})
production/test_stress.py:332:                                metrics.total_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:333:                                metrics.first_token_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:334:                                metrics.messages_per_stream.append(1)
production/test_stress.py:335:                                metrics.texts.append(d.get("text", ""))
production/test_stress.py:337:                                metrics.errors.append({"stream": stream_idx, "error": f"HTTP {r2.status}"})
production/test_stress.py:339:                        metrics.errors.append({"stream": stream_idx, "error": f"HTTP {r.status}"})
production/test_stress.py:352:                metrics.connect_times_ms.append(0)  # N/A for HTTP
production/test_stress.py:355:                    metrics.total_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:356:                    metrics.first_token_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:357:                    metrics.messages_per_stream.append(1)
production/test_stress.py:358:                    metrics.texts.append(d.get("text", ""))
production/test_stress.py:360:                    metrics.errors.append({"stream": stream_idx, "error": f"HTTP {r.status}"})
production/test_stress.py:362:        metrics.errors.append({"stream": stream_idx, "error": str(e)})
production/test_stress.py:365:async def run_sse_stream(url, wav_bytes, stream_idx, metrics,
production/test_stress.py:380:                    metrics.errors.append({"stream": stream_idx, "error": f"route failed: {route}"})
production/test_stress.py:383:                payload = {"audio_base64": audio_b64, "language": "hi", "chunk_duration_ms": 160}
production/test_stress.py:391:                        metrics.errors.append({"stream": stream_idx, "error": f"SSE {r.status}"})
production/test_stress.py:406:                        metrics.first_token_times_ms.append((first_token - t_start) * 1000)
production/test_stress.py:407:                    metrics.total_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:408:                    metrics.messages_per_stream.append(msgs)
production/test_stress.py:419:                    metrics.errors.append({"stream": stream_idx, "error": f"SSE {r.status}"})
production/test_stress.py:434:                    metrics.first_token_times_ms.append((first_token - t_start) * 1000)
production/test_stress.py:435:                metrics.total_times_ms.append((t_done - t_start) * 1000)
production/test_stress.py:436:                metrics.messages_per_stream.append(msgs)
production/test_stress.py:438:        metrics.errors.append({"stream": stream_idx, "error": str(e)})
production/test_stress.py:442:                           endpoint_name=None, endpoint_key=None, route_cost=10.0, ws_port=None):
production/test_stress.py:449:    metrics = StressMetrics(concurrency, mode, audio_duration_s)
production/test_stress.py:451:    # Capture server metrics before
production/test_stress.py:453:        metrics.server_metrics_before = await fetch_server_metrics(
production/test_stress.py:459:        metrics.server_metrics_before = await fetch_server_metrics(url=url)
production/test_stress.py:461:    metrics.start_time = time.perf_counter()
production/test_stress.py:466:        if mode == "ws":
production/test_stress.py:467:            tasks.append(run_ws_stream(url, pcm, i, metrics, audio_duration_s,
production/test_stress.py:468:                                       endpoint_name, endpoint_key, route_cost, ws_port))
production/test_stress.py:470:            tasks.append(run_http_batch(url, wav, i, metrics, endpoint_name, endpoint_key, route_cost))
production/test_stress.py:472:            tasks.append(run_sse_stream(url, wav, i, metrics, endpoint_name, endpoint_key, route_cost))
production/test_stress.py:475:    metrics.end_time = time.perf_counter()
production/test_stress.py:477:    # Brief pause then capture server metrics after
production/test_stress.py:480:        metrics.server_metrics_after = await fetch_server_metrics(
production/test_stress.py:486:        metrics.server_metrics_after = await fetch_server_metrics(url=url)
production/test_stress.py:489:    s = metrics.summary()


Output:
Chunk ID: cb3384
Wall time: 0.0508 seconds
Process exited with code 0
Original token count: 2704
Output:
#!/usr/bin/env python3
"""
Full-fidelity stress + quality test for Hindi ASR over Vast serverless (HTTP/SSE).

Features:
- Customer-like concurrent clients with multiple requests per client
- Route timing + queue-full detection
- First-token/total latency + server processing time
- Optional FLEURS Hindi audio (real speech) with WER/CER checks
- Time-series snapshots from /metrics and /get_endpoint_workers/
- CSV + JSON outputs suitable for scaling-law fitting
"""

import argparse
import asyncio
import base64
import csv
import io
import json
import math
import os
import random
import re
import struct
import sys
import time
import wave
from collections import Counter
from dataclasses import dataclass, asdict
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import aiohttp
import numpy as np
import requests

try:
    import jiwer
except Exception:
    jiwer = None

VAST_API = "https://console.vast.ai/api/v0"
VAST_RUN = "https://run.vast.ai"

def _env_or_dotenv(name: str) -> Optional[str]:
    if os.environ.get(name):
        return os.environ.get(name)
    env_path = Path(__file__).resolve().parent.parent / ".env"
    if env_path.exists():
        for line in env_path.read_text().splitlines():
            if line.startswith(name + "="):
                return line.split("=", 1)[1].strip()
    return None

def get_endpoint_info(endpoint_name: str, account_key: Optional[str] = None) -> Dict[str, Any]:
    key = account_key or _env_or_dotenv("VAST_KEY")
    if not key:
        raise RuntimeError("VAST_KEY not set; pass --endpoint-key and --endpoint-id")
    r = requests.get(
        f"{VAST_API}/endptjobs/",
        headers={"Accept": "application/json", "Authorization": f"Bearer {key}"},
        timeout=30,
    )
    data = r.json()
    results = data.get("results", []) if isinstance(data, dict) else data
    for ep in results or []:
        if ep.get("endpoint_name") == endpoint_name:
            return ep
    raise RuntimeError(f"endpoint '{endpoint_name}' not found")

async def get_route(session: aiohttp.ClientSession, endpoint_name: str, endpoint_key: str, cost: float) -> Tuple[Optional[Dict[str, Any]], float, Optional[str], bool]:
    t0 = time.perf_counter()
    try:
        async with session.post(
            f"{VAST_RUN}/route/",
            json={"endpoint": endpoint_name, "api_key": endpoint_key, "cost": cost},
            timeout=aiohttp.ClientTimeout(total=15),
        ) as r:
            data = await r.json(content_type=None)
    except Exception as e:
        return None, (time.perf_counter() - t0) * 1000.0, str(e), False

    dt = (time.perf_counter() - t0) * 1000.0
    if isinstance(data, dict) and "url" in data:
        return data, dt, None, False

    status = ""
    if isinstance(data, dict):
        status = str(data.get("status", ""))
    queue_full = "all worker queues full" in status.lower()
    return None, dt, f"route failed: {data}", queue_full

def wrap_auth(route: Dict[str, Any], payload: Dict[str, Any]) -> Dict[str, Any]:
    auth = {k: route[k] for k in ["signature", "cost", "endpoint", "reqnum", "url", "request_idx"] if k in route}
    return {"auth_data": auth, "payload": payload}

def pct(vals: List[float], p: int) -> float:
    if not vals:
        return 0.0
    s = sorted(vals)
    idx = min(int(len(s) * p / 100), len(s) - 1)
    return float(s[idx])

def safe_mean(vals: List[float]) -> float:
    return float(sum(vals) / len(vals)) if vals else 0.0

def normalize_text(text: str) -> str:
    if not text:
        return ""
    text = text.lower().strip()
    text = re.sub(r"<[^>]+>", " ", text)
    text = re.sub(r"[^\u0900-\u097fa-z0-9\s]", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text

def compute_wer_cer(reference: str, hypothesis: str) -> Tuple[Optional[float], Optional[float]]:
    if not jiwer:
        return None, None
    ref = normalize_text(reference)
    hyp = normalize_text(hypothesis)
    if not ref:
        return None, None
    try:
        return float(jiwer.wer(ref, hyp)), float(jiwer.cer(ref, hyp))
    except Exception:
        return None, None

@dataclass
class AudioSample:
    sample_id: str
    duration_s: float
    wav_bytes: bytes
    reference: Optional[str] = None

def _pcm_to_wav_bytes(pcm_int16: np.ndarray, sample_rate: int = 16000) -> bytes:
    bio = io.BytesIO()
    with wave.open(bio, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(pcm_int16.astype(np.int16).tobytes())
    return bio.getvalue()

def make_sine_samples(durations: List[float]) -> List[AudioSample]:
    samples: List[AudioSample] = []
    for i, d in enumerate(durations):
        n = int(16000 * d)
        pcm = np.array(
            [
                min(32767, max(-32768, int(16000 * math.sin(2 * math.pi * 440 * k / 16000))))
                for k in range(n)
            ],
            dtype=np.int16,
        )
        wav = _pcm_to_wav_bytes(pcm)
        samples.append(AudioSample(sample_id=f"sine_{i}_{d:.2f}s", duration_s=float(d), wav_bytes=wav, reference=None))
    return samples

def load_fleurs_samples(max_samples: int, min_duration: float, max_duration: float, seed: int) -> List[AudioSample]:
    from datasets import load_dataset

    ds = load_dataset("google/fleurs", "hi_in", split="test")
    rows = list(ds)
    rnd = random.Random(seed)
    rnd.shuffle(rows)

    out: List[AudioSample] = []
    for row in rows:
        audio = row["audio"]
        arr = np.asarray(audio["array"])  # float32 [-1, 1]
        sr = int(audio["sampling_rate"])
        if sr != 16000:
            # FLEURS hi_in is normally 16kHz, but keep explicit guard.
            continue
        duration = len(arr) / float(sr)
        if duration < min_duration or duration > max_duration:
            continue
        pcm = np.clip(arr, -1.0, 1.0)
        pcm = (pcm * 32767.0).astype(np.int16)
        wav = _pcm_to_wav_bytes(pcm, sample_rate=sr)
        out.append(
            AudioSample(
                sample_id=str(row["id"]),
                duration_s=round(duration, 3),
                wav_bytes=wav,
                reference=row.get("transcription"),
            )
        )
        if len(out) >= max_samples:
            break

    if not out:
        raise RuntimeError("No FLEURS samples matched filters")
    return out

@dataclass
class RequestRecord:
    client_idx: int
    req_idx: int
    mode: str
    sample_id: str
    duration_s: float
    route_ms: float
    connect_ms: float
    first_token_ms: float
    total_ms: float
    server_processing_ms: float
    status_code: int
    ok: bool
    queue_full: bool
    error: str
    reference: str
    hypothesis: str
    wer: float
    cer: float

@dataclass
class LevelResult:
    mode: str
    concurrency: int
    requests_per_client: int
    wall_time_s: float
    ok: int
    errors: int
    queue_full_errors: int
    throughput_x: float
    route_ms_p50: float
    route_ms_p95: float
    route_ms_p99: float
    first_token_ms_p50: float
    first_token_ms_p95: float
    first_token_ms_p99: float
    total_ms_p50: float
    total_ms_p95: float
    total_ms_p99: float
    server_proc_ms_p50: float
    server_proc_ms_p95: float
    server_proc_ms_p99: float
    vram_used_gb_max: float
    worker_cur_load_max: float
    worker_cur_load_p95: float
    worker_cur_perf_avg: float
    wer_pct: float
    cer_pct: float
    status_counts: Dict[str, int]

async def fetch_server_metrics_via_endpoint(
    session: aiohttp.ClientSession,
    endpoint_name: str,
    endpoint_key: str,
    route_cost: float,
) -> Dict[str, Any]:
    route, _, err, _ = await get_route(session, endpoint_name, endpoint_key, route_cost)
    if not route:
        return {"error": err}
    worker_url = route["url"].rstrip("/")
    req = wrap_auth(route, {})
    try:
        async with session.post(
            f"{worker_url}/metrics",
            json=req,
            timeout=aiohttp.ClientTimeout(total=12),
        ) as r:
            if r.status == 200:
                return await r.json(content_type=None)
            return {"error": f"metrics status={r.status}"}
    except Exception as e:
        return {"error": str(e)}

async def fetch_endpoint_workers(
    session: aiohttp.ClientSession,
    endpoint_id: int,
    endpoint_key: str,
) -> List[Dict[str, Any]]:
    try:
        async with session.post(
            f"{VAST_RUN}/get_endpoint_workers/",
            json={"id": endpoint_id, "api_key": endpoint_key},
            timeout=aiohttp.ClientTimeout(total=12),
        ) as r:
            data = await r.json(content_type=None)
            if isinstance(data, dict) and "workers" in data:
                return data["workers"]
            return data if isinstance(data, list) else []
    except Exception:
        return []

async def metrics_sampler(
    stop_evt: asyncio.Event,
    interval_s: float,
    session: aiohttp.ClientSession,
    endpoint_name: Optional[str],
    endpoint_key: Optional[str],
    endpoint_id: Optional[int],
    route_cost: float,
    server_samples: List[Dict[str, Any]],
    worker_samples: List[Dict[str, Any]],
):
    while not stop_evt.is_set():
        ts = time.time()
        if endpoint_name and endpoint_key:
            sm = await fetch_server_metrics_via_endpoint(session, endpoint_name, endpoint_key, route_cost)
            server_samples.append({"t": ts, "metrics": sm})
            if endpoint_id is not None:
                workers = await fetch_endpoint_workers(session, endpoint_id, endpoint_key)
                worker_samples.append({"t": ts, "workers": workers})
        await asyncio.sleep(interval_s)

async def run_http_request(
    session: aiohttp.ClientSession,
    sample: AudioSample,
    endpoint_name: Optional[str],
    endpoint_key: Optional[str],
    route_cost: float,
    direct_url: Optional[str],
) -> RequestRecord:
    t0 = time.perf_counter()
    route_ms = 0.0
    queue_full = False
    connect_ms = 0.0
    status_code = 0
    first_token_ms = 0.0
    server_proc_ms = 0.0
    hypothesis = ""
    err = ""

    try:
        if endpoint_name and endpoint_key:
            route, route_ms, rerr, qf = await get_route(session, endpoint_name, endpoint_key, route_cost)
            queue_full = qf
            if not route:
                raise RuntimeError(rerr or "route failed")
            worker_url = route["url"].rstrip("/")
            payload = {
                "audio_base64": base64.b64encode(sample.wav_bytes).decode("ascii"),
                "language": "hi",
                "audio_duration": sample.duration_s,
            }
            req = wrap_auth(route, payload)
            t_conn = time.perf_counter()
            async with session.post(


Output:
Chunk ID: 4bb8d8
Wall time: 0.0515 seconds
Process exited with code 0
Original token count: 3719
Output:
            async with session.post(
                f"{worker_url}/v1/audio/transcriptions",
                json=req,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as r:
                connect_ms = (time.perf_counter() - t_conn) * 1000.0
                status_code = int(r.status)
                if r.status in (404, 422):
                    route2, route_ms2, rerr2, qf2 = await get_route(session, endpoint_name, endpoint_key, route_cost)
                    route_ms += route_ms2
                    queue_full = queue_full or qf2
                    if not route2:
                        raise RuntimeError(rerr2 or "route fallback failed")
                    req2 = wrap_auth(route2, payload)
                    worker_url2 = route2["url"].rstrip("/")
                    async with session.post(
                        f"{worker_url2}/v1/audio/transcriptions/json",
                        json=req2,
                        timeout=aiohttp.ClientTimeout(total=180),
                    ) as r2:
                        status_code = int(r2.status)
                        data = await r2.json(content_type=None)
                else:
                    data = await r.json(content_type=None)
        else:
            if not direct_url:
                raise RuntimeError("direct_url required when endpoint is not set")
            form = aiohttp.FormData()
            form.add_field("file", sample.wav_bytes, filename="sample.wav", content_type="audio/wav")
            form.add_field("language", "hi")
            t_conn = time.perf_counter()
            async with session.post(
                f"{direct_url.rstrip('/')}/v1/audio/transcriptions",
                data=form,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as r:
                connect_ms = (time.perf_counter() - t_conn) * 1000.0
                status_code = int(r.status)
                data = await r.json(content_type=None)

        if status_code != 200:
            err = f"HTTP {status_code}: {str(data)[:400]}"
        else:
            hypothesis = str(data.get("text", "") or "")
            if "processing_time" in data:
                try:
                    server_proc_ms = float(data["processing_time"]) * 1000.0
                except Exception:
                    server_proc_ms = 0.0

    except Exception as e:
        err = str(e)

    t1 = time.perf_counter()
    total_ms = (t1 - t0) * 1000.0
    first_token_ms = total_ms  # HTTP batch has no token stream

    ref = sample.reference or ""
    wer_val, cer_val = compute_wer_cer(ref, hypothesis)

    return RequestRecord(
        client_idx=-1,
        req_idx=-1,
        mode="http",
        sample_id=sample.sample_id,
        duration_s=sample.duration_s,
        route_ms=route_ms,
        connect_ms=connect_ms,
        first_token_ms=first_token_ms,
        total_ms=total_ms,
        server_processing_ms=server_proc_ms,
        status_code=status_code,
        ok=(err == "" and status_code == 200),
        queue_full=queue_full,
        error=err,
        reference=ref,
        hypothesis=hypothesis,
        wer=float("nan") if wer_val is None else wer_val,
        cer=float("nan") if cer_val is None else cer_val,
    )

async def run_sse_request(
    session: aiohttp.ClientSession,
    sample: AudioSample,
    endpoint_name: Optional[str],
    endpoint_key: Optional[str],
    route_cost: float,
    direct_url: Optional[str],
) -> RequestRecord:
    t0 = time.perf_counter()
    route_ms = 0.0
    queue_full = False
    connect_ms = 0.0
    status_code = 0
    first_token_ms = 0.0
    hypothesis = ""
    err = ""

    try:
        if endpoint_name and endpoint_key:
            route, route_ms, rerr, qf = await get_route(session, endpoint_name, endpoint_key, route_cost)
            queue_full = qf
            if not route:
                raise RuntimeError(rerr or "route failed")
            worker_url = route["url"].rstrip("/")
            payload = {
                "audio_base64": base64.b64encode(sample.wav_bytes).decode("ascii"),
                "language": "hi",
                "chunk_duration_ms": 160,
                "audio_duration": sample.duration_s,
            }
            req = wrap_auth(route, payload)
            t_conn = time.perf_counter()
            async with session.post(
                f"{worker_url}/v1/audio/transcriptions/stream",
                json=req,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as r:
                connect_ms = (time.perf_counter() - t_conn) * 1000.0
                status_code = int(r.status)
                if status_code != 200:
                    body = await r.text()
                    raise RuntimeError(f"SSE {status_code}: {body[:200]}")

                got_first = False
                saw_any_event = False
                saw_done = False
                last_text = ""
                async for line in r.content:
                    l = line.decode(errors="ignore").strip()
                    if not l.startswith("data: "):
                        continue
                    try:
                        d = json.loads(l[6:])
                    except Exception:
                        continue
                    saw_any_event = True
                    txt = str(d.get("text", "") or "")
                    if txt:
                        last_text = txt
                        if not got_first:
                            first_token_ms = (time.perf_counter() - t0) * 1000.0
                            got_first = True
                    if d.get("done"):
                        saw_done = True
                        break
                hypothesis = last_text
                if not saw_any_event or not saw_done:
                    raise RuntimeError("sse stream incomplete: no terminal done event")
        else:
            if not direct_url:
                raise RuntimeError("direct_url required when endpoint is not set")
            form = aiohttp.FormData()
            form.add_field("file", sample.wav_bytes, filename="sample.wav", content_type="audio/wav")
            t_conn = time.perf_counter()
            async with session.post(
                f"{direct_url.rstrip('/')}/v1/audio/transcriptions/stream",
                data=form,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as r:
                connect_ms = (time.perf_counter() - t_conn) * 1000.0
                status_code = int(r.status)
                if status_code != 200:
                    body = await r.text()
                    raise RuntimeError(f"SSE {status_code}: {body[:200]}")
                got_first = False
                saw_any_event = False
                saw_done = False
                last_text = ""
                async for line in r.content:
                    l = line.decode(errors="ignore").strip()
                    if not l.startswith("data: "):
                        continue
                    try:
                        d = json.loads(l[6:])
                    except Exception:
                        continue
                    saw_any_event = True
                    txt = str(d.get("text", "") or "")
                    if txt:
                        last_text = txt
                        if not got_first:
                            first_token_ms = (time.perf_counter() - t0) * 1000.0
                            got_first = True
                    if d.get("done"):
                        saw_done = True
                        break
                hypothesis = last_text
                if not saw_any_event or not saw_done:
                    raise RuntimeError("sse stream incomplete: no terminal done event")
    except Exception as e:
        err = str(e)

    t1 = time.perf_counter()
    total_ms = (t1 - t0) * 1000.0
    if first_token_ms <= 0:
        first_token_ms = total_ms

    ref = sample.reference or ""
    wer_val, cer_val = compute_wer_cer(ref, hypothesis)

    return RequestRecord(
        client_idx=-1,
        req_idx=-1,
        mode="sse",
        sample_id=sample.sample_id,
        duration_s=sample.duration_s,
        route_ms=route_ms,
        connect_ms=connect_ms,
        first_token_ms=first_token_ms,
        total_ms=total_ms,
        server_processing_ms=0.0,
        status_code=status_code,
        ok=(err == "" and status_code == 200),
        queue_full=queue_full,
        error=err,
        reference=ref,
        hypothesis=hypothesis,
        wer=float("nan") if wer_val is None else wer_val,
        cer=float("nan") if cer_val is None else cer_val,
    )

def _series_vram_max(server_samples: List[Dict[str, Any]]) -> float:
    vals: List[float] = []
    for s in server_samples:
        m = s.get("metrics", {})
        if isinstance(m, dict) and "vram_used_gb" in m:
            try:
                vals.append(float(m.get("vram_used_gb", 0.0)))
            except Exception:
                pass
    return max(vals) if vals else 0.0

def _worker_stats(worker_samples: List[Dict[str, Any]]) -> Tuple[float, float, float]:
    loads: List[float] = []
    perfs: List[float] = []
    for s in worker_samples:
        workers = s.get("workers") or []
        for w in workers:
            try:
                loads.append(float(w.get("cur_load", 0.0)))
                perfs.append(float(w.get("cur_perf", 0.0)))
            except Exception:
                pass
    return (max(loads) if loads else 0.0, pct(loads, 95), safe_mean(perfs))

def summarize_level(
    mode: str,
    concurrency: int,
    requests_per_client: int,
    wall_time_s: float,
    records: List[RequestRecord],
    server_samples: List[Dict[str, Any]],
    worker_samples: List[Dict[str, Any]],
) -> LevelResult:
    ok_recs = [r for r in records if r.ok]
    route_ms = [r.route_ms for r in records if r.route_ms > 0]
    first_ms = [r.first_token_ms for r in ok_recs]
    total_ms = [r.total_ms for r in ok_recs]
    proc_ms = [r.server_processing_ms for r in ok_recs if r.server_processing_ms > 0]

    total_audio_ok = sum(r.duration_s for r in ok_recs)
    throughput_x = (total_audio_ok / wall_time_s) if wall_time_s > 0 else 0.0

    wer_vals = [r.wer for r in ok_recs if r.reference and not math.isnan(r.wer)]
    cer_vals = [r.cer for r in ok_recs if r.reference and not math.isnan(r.cer)]

    status_counts: Counter[str] = Counter()
    for r in records:
        k = str(r.status_code if r.status_code else "exception")
        status_counts[k] += 1

    qf = sum(1 for r in records if r.queue_full)

    vram_max = _series_vram_max(server_samples)
    worker_load_max, worker_load_p95, worker_perf_avg = _worker_stats(worker_samples)

    return LevelResult(
        mode=mode,
        concurrency=concurrency,
        requests_per_client=requests_per_client,
        wall_time_s=round(wall_time_s, 3),
        ok=len(ok_recs),
        errors=len(records) - len(ok_recs),
        queue_full_errors=qf,
        throughput_x=round(throughput_x, 3),
        route_ms_p50=round(pct(route_ms, 50), 2),
        route_ms_p95=round(pct(route_ms, 95), 2),
        route_ms_p99=round(pct(route_ms, 99), 2),
        first_token_ms_p50=round(pct(first_ms, 50), 2),
        first_token_ms_p95=round(pct(first_ms, 95), 2),
        first_token_ms_p99=round(pct(first_ms, 99), 2),
        total_ms_p50=round(pct(total_ms, 50), 2),
        total_ms_p95=round(pct(total_ms, 95), 2),
        total_ms_p99=round(pct(total_ms, 99), 2),
        server_proc_ms_p50=round(pct(proc_ms, 50), 2),
        server_proc_ms_p95=round(pct(proc_ms, 95), 2),
        server_proc_ms_p99=round(pct(proc_ms, 99), 2),
        vram_used_gb_max=round(vram_max, 2),
        worker_cur_load_max=round(worker_load_max, 3),
        worker_cur_load_p95=round(worker_load_p95, 3),
        worker_cur_perf_avg=round(worker_perf_avg, 4),
        wer_pct=round(100.0 * safe_mean(wer_vals), 3) if wer_vals else 0.0,
        cer_pct=round(100.0 * safe_mean(cer_vals), 3) if cer_vals else 0.0,
        status_counts=dict(status_counts),
    )

async def run_level(
    mode: str,
    concurrency: int,
    requests_per_client: int,
    samples: List[AudioSample],
    endpoint_name: Optional[str],
    endpoint_key: Optional[str],
    endpoint_id: Optional[int],
    route_cost: float,
    metrics_interval: float,
    direct_url: Optional[str],
    shuffle_samples: bool,
    seed: int,
    inter_request_delay: float,
):
    print("\n" + "=" * 90)
    print(f"Level mode={mode} c={concurrency} reqs/client={requests_per_client}")
    print("=" * 90)

    pool = list(samples)
    if shuffle_samples:
        rnd = random.Random(seed + concurrency)
        rnd.shuffle(pool)

    sample_idx = 0
    sample_lock = asyncio.Lock()

    async def pick_sample() -> AudioSample:
        nonlocal sample_idx
        async with sample_lock:
            s = pool[sample_idx % len(pool)]
            sample_idx += 1
            return s

    timeout = aiohttp.ClientTimeout(total=220)
    connector = aiohttp.TCPConnector(limit=max(200, concurrency * 8), ttl_dns_cache=300)

    records: List[RequestRecord] = []
    server_samples: List[Dict[str, Any]] = []
    worker_samples: List[Dict[str, Any]] = []

    stop_evt = asyncio.Event()

    async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
        sampler_task = asyncio.create_task(
            metrics_sampler(
                stop_evt,
                metrics_interval,
                session,
                endpoint_name,
                endpoint_key,
                endpoint_id,
                route_cost,
                server_samples,
                worker_samples,
            )
        )

        t0 = time.perf_counter()

        async def client_worker(client_idx: int):
            for req_idx in range(requests_per_client):
                sample = await pick_sample()
                if mode == "http":
                    rec = await run_http_request(
                        session,
                        sample,
                        endpoint_name,
                        endpoint_key,
                        route_cost,
                        direct_url,
                    )
                else:
                    rec = await run_sse_request(
                        session,
                        sample,
                        endpoint_name,
                        endpoint_key,
                        route_cost,
                        direct_url,
                    )
                rec.client_idx = client_idx
                rec.req_idx = req_idx
                records.append(rec)

                if inter_request_delay > 0 and req_idx + 1 < requests_per_client:
                    await asyncio.sleep(inter_request_delay)

        await asyncio.gather(*(client_worker(i) for i in range(concurrency)))


# Collaboration Mode: Default

You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active.

Your active mode changes only when new developer instructions with a different `<collaboration_mode>...</collaboration_mode>` change it; user requests or tool descriptions do not change mode by themselves. Known mode names are Default and Plan.

## request_user_input availability

The `request_user_input` tool is unavailable in Default mode. If you call it while in Default mode, it will return an error.

If a decision is necessary and cannot be discovered from local context, ask the user directly. However, in Default mode you should strongly prefer executing the user's request rather than stopping to ask questions.

Summary: auto

## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) When `SKILL.md` references relative paths (e.g., `scripts/foo.py`), resolve them relative to the skill directory listed above first, and only consider other paths if needed.
  3) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  4) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  5) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


**Planning apples-to-apples WebSocket benchmark**


I found that `test_stress_full.py` already gives strong HTTP/SSE metrics plus WER, but it does not include the WS apples-to-apples drift test. I’m now inspecting `server.py` to identify why concurrency is far below NVIDIA’s reported WS benchmark conditions.


Output:
Chunk ID: 78d432
Wall time: 0.0506 seconds
Process exited with code 0
Original token count: 2107
Output:
"""
Production Hindi ASR Server - Standalone FastAPI
Supports: HTTP batch, HTTP streaming (SSE), WebSocket streaming
Designed for Vast.ai serverless deployment with PyWorker integration.

Architecture:
  - NeMo cache-aware RNNT streaming pipeline (Nemotron-0.6B fine-tuned Hindi)
  - Event-driven batching (50ms max wait, immediate dispatch when batch fills)
  - CPU-based Silero VAD gating (filters silence before GPU, ~3x effective capacity)
  - Per-stream state management with fairness scheduling
  - Backpressure: rejects new connections when at capacity (fast-fail > silent degrade)

Performance (A100-SXM4-80GB, from prior benchmarks):
  - Safe zone: 1-50 concurrent streams (WER 22-31%, FT <211ms)
  - Optimal: 50-100 streams (WER 31-33%, FT <402ms)
  - Edge: 100-140 streams (WER 33-37%)
  - VRAM: 10.4GB constant (compute-bound, not memory-bound)
"""

import os
import sys
import time
import json
import wave
import asyncio
import logging
import tempfile
import signal
import base64
from io import BytesIO
from dataclasses import dataclass, field
from typing import Optional, List
from collections import deque

import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.websockets import WebSocketState

# === Logging setup ===
LOG_FILE = os.environ.get("MODEL_LOG", "/var/log/model/server.log")
os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(LOG_FILE),
        logging.StreamHandler(sys.stdout),
    ],
)
logger = logging.getLogger("asr_server")

# === Configuration ===

MODEL_PATH = os.environ.get("MODEL_PATH", "/model/final_model.nemo")
HF_MODEL_REPO = os.environ.get("HF_MODEL_REPO", "BayAreaBoys/nemotron-hindi")
HF_MODEL_FILE = os.environ.get("HF_MODEL_FILE", "final_model.nemo")
SERVER_PORT = int(os.environ.get("SERVER_PORT", "8000"))

# Production limits
MAX_ACTIVE_STREAMS = int(os.environ.get("MAX_ACTIVE_STREAMS", "100"))
MAX_BATCH_WAIT_MS = float(os.environ.get("MAX_BATCH_WAIT_MS", "50"))
ENABLE_VAD = os.environ.get("ENABLE_VAD", "true").lower() == "true"
VAD_THRESHOLD = float(os.environ.get("VAD_THRESHOLD", "0.5"))

@dataclass
class PipelineConfig:
    """Pipeline configuration — matches NeMo's expected OmegaConf structure"""
    model_name: str = ""
    device: str = "cuda"
    device_id: int = 0
    compute_dtype: str = "bfloat16"
    use_amp: bool = True
    # Decoding
    decoding_strategy: str = "greedy_batch"
    use_cuda_graph_decoder: bool = False
    max_symbols: int = 10
    # Streaming
    sample_rate: int = 16000
    batch_size: int = 512
    att_context_size: list = field(default_factory=lambda: [70, 6])
    use_cache: bool = True
    use_feat_cache: bool = True
    num_slots: int = 1024
    # Endpointing
    stop_history_eou: int = 800
    residue_tokens_at_end: int = 2

# === Global server state ===
class ASRServerState:
    """
    Centralized server state. Single instance per process.
    WHY: Avoids global variable soup, enables clean shutdown.
    """
    def __init__(self):
        self.pipeline = None
        self.cfg = None
        self.vad = None
        self.start_time = None

        # Stream management
        self.client_queues: dict[int, asyncio.Queue] = {}
        self.client_queues_lock = asyncio.Lock()
        self.stream_audio_timestamps: dict[int, float] = {}
        self.stream_timestamps_lock = asyncio.Lock()
        self.streams_pending_cleanup: set[int] = set()
        self.streams_cleanup_lock = asyncio.Lock()

        # Event-driven batching
        self._batch_ready_event = asyncio.Event()
        self.inference_task = None
        self.inference_running = False

        # Metrics (rolling window for percentile calculation)
        self.metrics = {
            "total_batches": 0,
            "total_outputs": 0,
            "total_gpu_ms": 0.0,
            "total_route_ms": 0.0,
            "batch_sizes": deque(maxlen=500),
            "gpu_times_ms": deque(maxlen=500),
            "cycle_times_ms": deque(maxlen=500),
            "empty_cycles": 0,
            "vad_filtered_chunks": 0,
            "http_requests_total": 0,
            "http_requests_ok": 0,
            "http_requests_err": 0,
            "ws_connections_total": 0,
            "ws_connections_active": 0,
            "ws_connections_rejected": 0,
        }
        self.model_loaded = False

state = ASRServerState()

# === FastAPI App ===
app = FastAPI(
    title="Hindi ASR Production Server",
    description="Cache-aware streaming ASR for Hindi (Nemotron-0.6B fine-tuned)",
    version="1.0.0",
)

# ============================================================
# Model Loading
# ============================================================

def download_model_if_needed() -> str:
    """Download model from HuggingFace if not found locally. Returns local path."""
    if os.path.exists(MODEL_PATH):
        logger.info(f"Model found at {MODEL_PATH}")
        return MODEL_PATH

    logger.info(f"Model not found at {MODEL_PATH}, downloading from {HF_MODEL_REPO}...")
    from huggingface_hub import hf_hub_download
    hf_token = os.environ.get("HF_TOKEN")
    model_path = hf_hub_download(
        repo_id=HF_MODEL_REPO,
        filename=HF_MODEL_FILE,
        token=hf_token,
        cache_dir="/model",
    )
    logger.info(f"Model downloaded to {model_path}")
    return model_path

def build_pipeline_config(model_path: str):
    """Build OmegaConf config for NeMo pipeline"""
    from omegaconf import OmegaConf
    from dataclasses import dataclass, field
    from typing import Optional, List

    # Full config structure matching CacheAwarePipelineConfig from nemotron_asr
    @dataclass
    class BoostingTreeConfig:
        model_path: Optional[str] = None
        key_phrases_file: Optional[str] = None
        key_phrases_list: Optional[List[str]] = None
        source_lang: str = "en"

    @dataclass
    class GreedyDecodingConfig:
        use_cuda_graph_decoder: bool = False
        max_symbols: int = 10
        ngram_lm_model: Optional[str] = None
        ngram_lm_alpha: float = 0.0
        boosting_tree: BoostingTreeConfig = field(default_factory=BoostingTreeConfig)
        boosting_tree_alpha: float = 0.0

    @dataclass
    class DecodingConfig:
        strategy: str = "greedy_batch"
        preserve_alignments: bool = False
        fused_batch_size: int = -1
        greedy: GreedyDecodingConfig = field(default_factory=GreedyDecodingConfig)

    @dataclass
    class ASRConfig:
        model_name: str = model_path
        device: str = "cuda"
        device_id: int = 0
        compute_dtype: str = "bfloat16"
        use_amp: bool = True
        decoding: DecodingConfig = field(default_factory=DecodingConfig)

    @dataclass
    class ITNConfig:
        input_case: str = "lower_cased"
        whitelist: Optional[str] = None
        overwrite_cache: bool = False
        max_number_of_permutations_per_split: int = 729
        left_padding_size: int = 4
        batch_size: int = 32
        n_jobs: int = 16

    @dataclass
    class ConfidenceConfig:
        exclude_blank: bool = True
        aggregation: str = "mean"
        method_cfg: dict = field(default_factory=lambda: {
            "name": "entropy",
            "entropy_type": "tsallis",
            "alpha": 0.5,
            "entropy_norm": "exp",
        })

    @dataclass
    class EndpointingConfig:
        stop_history_eou: int = 800
        residue_tokens_at_end: int = 2

    @dataclass
    class StreamingConfig:
        sample_rate: int = 16000
        batch_size: int = 512
        word_boundary_tolerance: int = 4
        att_context_size: List[int] = field(default_factory=lambda: [70, 6])
        use_cache: bool = True
        use_feat_cache: bool = True
        chunk_size_in_secs: Optional[float] = None
        request_type: str = "frame"
        num_slots: int = 1024
        exhaustive_batching: bool = False
        batching_delay_secs: float = 0.300

    @dataclass
    class MetricsASRConfig:
        gt_text_attr_name: str = "text"
        clean_groundtruth_text: bool = False
        langid: str = "hi"
        use_cer: bool = True
        ignore_capitalization: bool = True
        ignore_punctuation: bool = True


Output:
Chunk ID: abfc8d
Wall time: 0.0516 seconds
Process exited with code 0
Original token count: 2927
Output:
        ignore_punctuation: bool = True
        strip_punc_space: bool = False

    @dataclass
    class MetricsConfig:
        asr: MetricsASRConfig = field(default_factory=MetricsASRConfig)

    @dataclass
    class CacheAwarePipelineConfig:
        asr: ASRConfig = field(default_factory=ASRConfig)
        itn: ITNConfig = field(default_factory=ITNConfig)
        nmt: Optional[dict] = None
        confidence: ConfidenceConfig = field(default_factory=ConfidenceConfig)
        endpointing: EndpointingConfig = field(default_factory=EndpointingConfig)
        streaming: StreamingConfig = field(default_factory=StreamingConfig)
        matmul_precision: str = "high"
        log_level: int = 20
        pipeline_type: str = "cache_aware"
        asr_decoding_type: str = "rnnt"
        audio_file: Optional[str] = None
        output_filename: Optional[str] = None
        output_dir: Optional[str] = None
        enable_pnc: bool = False
        enable_itn: bool = False
        enable_nmt: bool = False
        asr_output_granularity: str = "segment"
        cache_dir: Optional[str] = None
        lang: Optional[str] = "hi"
        return_tail_result: bool = False
        calculate_wer: bool = True
        calculate_bleu: bool = True
        metrics: MetricsConfig = field(default_factory=MetricsConfig)

    config = CacheAwarePipelineConfig()
    config.asr.model_name = model_path
    return OmegaConf.structured(config)

async def load_model():
    """Load model and build pipeline. Called at startup."""
    from omegaconf import OmegaConf

    # Silence chatty NeMo logs
    logging.getLogger("nemo_logger").setLevel(logging.WARNING)

    logger.info("=== Starting model load ===")
    model_path = download_model_if_needed()
    state.cfg = build_pipeline_config(model_path)

    logger.info(f"Building pipeline: model={model_path}")
    logger.info(f"  att_context_size={state.cfg.streaming.att_context_size}")
    logger.info(f"  batch_size={state.cfg.streaming.batch_size}")

    # Build pipeline (reuses nemotron_asr pipeline code)
    from pipelines import StreamingPipelineBuilder
    state.pipeline = StreamingPipelineBuilder.build_pipeline(state.cfg)

    # Initialize streaming support
    state.pipeline.init_streaming_request_generator()
    logger.info("Streaming request generator initialized")

    # Warmup GPU (3 iterations)
    logger.info("Warming up GPU (3 iterations)...")
    from asr_utils import preprocess_audio
    from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions

    WARMUP_URL = "https://github.com/voxserv/audio_quality_testing_samples/raw/refs/heads/master/mono_44100/156550__acclivity__a-dream-within-a-dream.wav"
    try:
        audio_bytes = preprocess_audio(WARMUP_URL, target_sample_rate=16000)
    except Exception as e:
        # If warmup audio fails to download, generate silence
        logger.warning(f"Warmup audio download failed: {e}, using silence")
        audio_bytes = np.zeros(16000 * 3, dtype=np.int16).tobytes()  # 3s silence

    chunk_size = int(0.080 * 16000 * 2)  # 80ms chunks, 16kHz, 16-bit

    for warmup_iter in range(1, 4):
        options = ASRRequestOptions()
        options.asr_output_granularity = "word"
        stream_id = state.pipeline.open_streaming_session(options=options)

        for i in range(0, len(audio_bytes), chunk_size):
            chunk = audio_bytes[i:i + chunk_size]
            state.pipeline.append_streaming_audio(stream_id, chunk)
            await state.pipeline.process_streaming_batch()

        state.pipeline.close_streaming_session(stream_id)
        try:
            await state.pipeline.process_streaming_batch()
        except Exception:
            pass

        # Cleanup warmup stream
        if stream_id in state.pipeline._state_pool:
            del state.pipeline._state_pool[stream_id]
        if stream_id in state.pipeline._streaming_request_generator.streams:
            state.pipeline._streaming_request_generator.streams.pop(stream_id, None)

        logger.info(f"  Warmup {warmup_iter}/3 complete")

    # Initialize VAD
    if ENABLE_VAD:
        from vad import SileroVAD
        state.vad = SileroVAD(threshold=VAD_THRESHOLD, sample_rate=16000)
        logger.info(f"VAD enabled (threshold={VAD_THRESHOLD})")
    else:
        logger.info("VAD disabled")

    state.start_time = time.perf_counter()
    state.model_loaded = True

    # CRITICAL: This log line triggers Vast.ai PyWorker readiness detection
    logger.info("Application startup complete.")
    print("Application startup complete.", flush=True)

async def transcribe_pcm_bytes_streaming(pcm_bytes: bytes) -> str:
    """
    Transcribe PCM bytes using the streaming pipeline (no .run() required).
    WHY: CacheAwareRNNTStreamingPipeline doesn't implement .run().
    """
    from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions

    options = ASRRequestOptions()
    options.asr_output_granularity = "word"
    stream_id = state.pipeline.open_streaming_session(options=options)

    # Use pipeline's native chunk size if available; fallback to 80ms
    frame_secs = getattr(state.pipeline, "chunk_size_in_secs", 0.08) or 0.08
    frame_bytes = max(2560, int(frame_secs * 16000 * 2))  # 16kHz, 16-bit

    full_segments: list[str] = []
    last_partial = ""

    try:
        # Feed audio in frame-sized chunks
        for i in range(0, len(pcm_bytes), frame_bytes):
            chunk = pcm_bytes[i:i + frame_bytes]
            state.pipeline.append_streaming_audio(stream_id, chunk)
            try:
                requests, outputs = await state.pipeline.process_streaming_batch()
            except Exception as e:
                if type(e).__name__ == "NotEnoughDataException":
                    outputs = []
                else:
                    raise
            for output in outputs:
                if output.stream_id != stream_id:
                    continue
                if output.final_transcript:
                    full_segments.append(output.final_transcript.strip())
                    # Final transcript commits buffered tokens; clear stale partial.
                    last_partial = ""
                if output.partial_transcript:
                    last_partial = output.partial_transcript.strip()

        # Close and drain remaining outputs
        state.pipeline.close_streaming_session(stream_id)
        for _ in range(10):
            try:
                requests, outputs = await state.pipeline.process_streaming_batch()
            except Exception as e:
                if type(e).__name__ == "NotEnoughDataException":
                    outputs = []
                else:
                    break
            if not outputs:
                break
            for output in outputs:
                if output.stream_id != stream_id:
                    continue
                if output.final_transcript:
                    full_segments.append(output.final_transcript.strip())
                    last_partial = ""
                if output.partial_transcript:
                    last_partial = output.partial_transcript.strip()

    finally:
        # Cleanup stream state
        if stream_id in state.pipeline._state_pool:
            del state.pipeline._state_pool[stream_id]
        if stream_id in state.pipeline._streaming_request_generator.streams:
            state.pipeline._streaming_request_generator.streams.pop(stream_id, None)

    # Keep tail partial to avoid truncating long utterances with no final EOU flush.
    combined = " ".join([x for x in [*full_segments, last_partial] if x]).strip()
    final_only = " ".join([x for x in full_segments if x]).strip()
    if len(combined) >= len(final_only):
        return combined
    return final_only or last_partial

# ============================================================
# HTTP Endpoints
# ============================================================

async def _parse_audio_request(request: Request):
    """
    Parse audio payload from either JSON (audio_base64) or multipart form-data (file).
    Returns (audio_bytes, language, response_format, chunk_duration_ms).
    """
    content_type = request.headers.get("content-type", "")

    if content_type.startswith("application/json"):
        body = await request.json()
        audio_b64 = (
            body.get("audio_base64")
            or body.get("audio_b64")
            or body.get("audio")
        )
        if not audio_b64:
            raise HTTPException(status_code=400, detail="audio_base64 required")
        try:
            audio_data = base64.b64decode(audio_b64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"invalid audio_base64: {e}")
        language = body.get("language", "hi")
        response_format = body.get("response_format", "json")
        chunk_duration_ms = int(body.get("chunk_duration_ms", 160) or 160)
        return audio_data, language, response_format, chunk_duration_ms

    # Fallback: multipart form-data
    form = await request.form()
    file_obj = form.get("file")
    if file_obj is None:
        raise HTTPException(status_code=400, detail="file required")

    if hasattr(file_obj, "read"):
        audio_data = await file_obj.read()
    else:
        audio_data = file_obj

    language = form.get("language", "hi")
    response_format = form.get("response_format", "json")
    chunk_duration_ms = int(form.get("chunk_duration_ms", 160) or 160)
    return audio_data, language, response_format, chunk_duration_ms

@app.on_event("startup")
async def startup():
    """Load model on server startup"""
    await load_model()

@app.get("/health")
async def health_check():
    """
    Health check for load balancers and PyWorker.
    Returns 503 when overloaded so /route/ skips this worker.
    """
    try:
        vram_used = torch.cuda.memory_allocated() / 1e9
        vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9
        gpu_name = torch.cuda.get_device_name(0)
    except Exception:
        vram_used = vram_total = 0
        gpu_name = "unknown"

    active = len(state.client_queues)
    status = "healthy" if active < MAX_ACTIVE_STREAMS else "overloaded"
    uptime = round(time.perf_counter() - state.start_time, 0) if state.start_time else 0

    return JSONResponse(
        status_code=200 if status == "healthy" else 503,
        content={
            "status": status,
            "model_loaded": state.model_loaded,
            "active_streams": active,
            "max_streams": MAX_ACTIVE_STREAMS,
            "gpu": gpu_name,
            "vram_used_gb": round(vram_used, 2),
            "vram_total_gb": round(vram_total, 0),
            "inference_running": state.inference_running,
            "uptime_s": uptime,
        },
    )

def _build_metrics_payload():
    """Shared metrics payload for GET/POST /metrics."""
    m = state.metrics

    def pct(arr, p):
        if not arr:
            return 0
        s = sorted(arr)
        idx = min(int(len(s) * p / 100), len(s) - 1)
        return round(s[idx], 2)

    cuda_available = False
    device_count = 0
    gpu_name = "unknown"
    vram_used = vram_total = vram_reserved = 0
    try:
        cuda_available = torch.cuda.is_available()
        device_count = torch.cuda.device_count() if cuda_available else 0
        if cuda_available and device_count > 0:
            gpu_name = torch.cuda.get_device_name(0)
            vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9
            vram_used = torch.cuda.memory_allocated() / 1e9
            vram_reserved = torch.cuda.memory_reserved() / 1e9
    except Exception:


Output:
Chunk ID: 1ed8cb
Wall time: 0.0508 seconds
Process exited with code 0
Original token count: 4223
Output:
    except Exception:
        pass

    return {
        "active_streams": len(state.client_queues),
        "max_streams": MAX_ACTIVE_STREAMS,
        "inference": {
            "total_batches": m["total_batches"],
            "total_outputs": m["total_outputs"],
            "empty_cycles": m["empty_cycles"],
            "vad_filtered_chunks": m["vad_filtered_chunks"],
            "gpu_ms_p50": pct(m["gpu_times_ms"], 50),
            "gpu_ms_p99": pct(m["gpu_times_ms"], 99),
            "cycle_ms_p50": pct(m["cycle_times_ms"], 50),
            "cycle_ms_p99": pct(m["cycle_times_ms"], 99),
            "batch_size_avg": round(sum(m["batch_sizes"]) / max(len(m["batch_sizes"]), 1), 1),
        },
        "http": {
            "requests_total": m["http_requests_total"],
            "requests_ok": m["http_requests_ok"],
            "requests_err": m["http_requests_err"],
        },
        "ws": {
            "connections_total": m["ws_connections_total"],
            "connections_active": m["ws_connections_active"],
            "connections_rejected": m["ws_connections_rejected"],
        },
        "cuda_available": cuda_available,
        "cuda_device_count": device_count,
        "gpu_name": gpu_name,
        "vram_used_gb": round(vram_used, 2),
        "vram_reserved_gb": round(vram_reserved, 2),
        "vram_total_gb": round(vram_total, 0),
        "uptime_s": round(time.perf_counter() - state.start_time, 0) if state.start_time else 0,
    }

@app.get("/metrics")
async def metrics_get():
    """
    Detailed metrics endpoint for monitoring/alerting.
    Returns structured JSON with rolling-window percentiles.
    """
    return _build_metrics_payload()

@app.post("/metrics")
async def metrics_post(_request: Request):
    """POST variant for serverless PyWorker proxy."""
    return _build_metrics_payload()

@app.post("/v1/audio/transcriptions/json")
async def transcribe_json(request: Request):
    """
    JSON-based transcription endpoint for PyWorker benchmarks.
    Accepts: {"audio_base64": "...", "language": "hi"}
    WHY: PyWorker sends benchmark payloads as JSON, not multipart.
    """
    state.metrics["http_requests_total"] += 1
    if not state.model_loaded:
        state.metrics["http_requests_err"] += 1
        raise HTTPException(status_code=503, detail="Model not loaded yet")
    t_start = time.perf_counter()
    try:
        audio_data, language, _response_format, _chunk_duration_ms = await _parse_audio_request(request)
        from asr_utils import preprocess_audio
        pcm_bytes = preprocess_audio(audio_data, target_sample_rate=16000)
        audio_duration = len(pcm_bytes) / (16000 * 2)
        t_infer = time.perf_counter()
        text = await transcribe_pcm_bytes_streaming(pcm_bytes)
        t_done = time.perf_counter()
        state.metrics["http_requests_ok"] += 1
        return {
            "text": text,
            "duration": round(audio_duration, 2),
            "processing_time": round(t_done - t_infer, 3),
        }
    except HTTPException:
        raise
    except Exception as e:
        state.metrics["http_requests_err"] += 1
        logger.error(f"JSON transcribe error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/audio/transcriptions")
async def transcribe_batch(request: Request):
    """
    OpenAI-compatible batch transcription endpoint.
    Accepts audio file, returns full transcription.

    WHY HTTP batch: For non-real-time use cases (file transcription, post-processing).
    Uses pipeline.run() directly — simpler, no streaming overhead.
    """
    state.metrics["http_requests_total"] += 1

    if not state.model_loaded:
        state.metrics["http_requests_err"] += 1
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    t_start = time.perf_counter()

    try:
        # Read uploaded audio (multipart or JSON base64)
        audio_data, language, response_format, _chunk_duration_ms = await _parse_audio_request(request)
        if not audio_data:
            raise HTTPException(status_code=400, detail="Empty audio file")

        # Preprocess audio to 16kHz mono PCM
        from asr_utils import preprocess_audio
        pcm_bytes = preprocess_audio(audio_data, target_sample_rate=16000)
        audio_duration = len(pcm_bytes) / (16000 * 2)  # seconds

        # Run streaming inference (batch)
        t_infer = time.perf_counter()
        text = await transcribe_pcm_bytes_streaming(pcm_bytes)
        t_done = time.perf_counter()

        result = {
            "text": text,
            "language": language,
            "duration": round(audio_duration, 2),
            "processing_time": round(t_done - t_infer, 3),
            "total_time": round(t_done - t_start, 3),
        }

        state.metrics["http_requests_ok"] += 1
        logger.info(f"HTTP batch: {audio_duration:.1f}s audio -> {t_done - t_infer:.2f}s inference")
        return result

    except HTTPException:
        raise
    except Exception as e:
        state.metrics["http_requests_err"] += 1
        logger.error(f"HTTP batch error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/audio/transcriptions/stream")
async def transcribe_stream(request: Request):
    """
    HTTP streaming transcription via Server-Sent Events (SSE).
    Client uploads audio file, gets progressive transcription events.

    WHY SSE: Works through HTTP proxies (including Vast.ai PyWorker).
    Compatible with standard HTTP clients — no WebSocket library needed.
    """
    state.metrics["http_requests_total"] += 1

    if not state.model_loaded:
        state.metrics["http_requests_err"] += 1
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    try:
        audio_data, language, response_format, chunk_duration_ms = await _parse_audio_request(request)
        if not audio_data:
            raise HTTPException(status_code=400, detail="Empty audio file")

        from asr_utils import preprocess_audio
        pcm_bytes = preprocess_audio(audio_data, target_sample_rate=16000)

        async def sse_generator():
            """Stream transcription results as SSE events"""
            from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions

            options = ASRRequestOptions()
            options.asr_output_granularity = "word"
            stream_id = state.pipeline.open_streaming_session(options=options)

            chunk_size = int(chunk_duration_ms / 1000.0 * 16000 * 2)  # bytes per chunk
            att_ctx = state.cfg.streaming.att_context_size[1] + 1
            frame_bytes = att_ctx * 1280 * 2  # samples per frame * 2 bytes

            audio_buffer = b""
            full_segments: list[str] = []
            last_partial = ""

            try:
                for i in range(0, len(pcm_bytes), chunk_size):
                    chunk = pcm_bytes[i:i + chunk_size]
                    audio_buffer += chunk

                    # Feed to pipeline when we have enough for a frame
                    if len(audio_buffer) >= frame_bytes:
                        state.pipeline.append_streaming_audio(stream_id, audio_buffer)
                        audio_buffer = b""

                        # Process batch
                        requests, outputs = await state.pipeline.process_streaming_batch()
                        for output in outputs:
                            if output.stream_id == stream_id:
                                text = output.final_transcript or output.partial_transcript or ""
                                is_final = bool(output.final_transcript)
                                if text.strip():
                                    if is_final:
                                        full_segments.append(text.strip())
                                        last_partial = ""
                                    else:
                                        last_partial = text.strip()
                                    event_data = json.dumps({
                                        "text": text,
                                        "is_final": is_final,
                                        "timestamp": time.time(),
                                    })
                                    yield f"data: {event_data}\n\n"

                    # Simulate real-time pacing (optional, for more realistic SSE)
                    await asyncio.sleep(0)

                # Flush remaining audio
                if audio_buffer:
                    state.pipeline.append_streaming_audio(stream_id, audio_buffer)

                # Close stream and drain remaining outputs
                state.pipeline.close_streaming_session(stream_id)
                for _ in range(10):
                    try:
                        requests, outputs = await state.pipeline.process_streaming_batch()
                        if not outputs:
                            break
                        for output in outputs:
                            if output.stream_id == stream_id:
                                text = output.final_transcript or output.partial_transcript or ""
                                is_final = bool(output.final_transcript)
                                if text.strip():
                                    if is_final:
                                        full_segments.append(text.strip())
                                        last_partial = ""
                                    else:
                                        last_partial = text.strip()
                                    event_data = json.dumps({
                                        "text": text,
                                        "is_final": is_final,
                                        "timestamp": time.time(),
                                    })
                                    yield f"data: {event_data}\n\n"
                    except Exception:
                        break

                # Final event with complete transcription
                combined = " ".join([x for x in [*full_segments, last_partial] if x]).strip()
                yield f"data: {json.dumps({'text': combined, 'is_final': True, 'done': True})}\n\n"

            finally:
                # Cleanup
                if stream_id in state.pipeline._state_pool:
                    del state.pipeline._state_pool[stream_id]
                if stream_id in state.pipeline._streaming_request_generator.streams:
                    state.pipeline._streaming_request_generator.streams.pop(stream_id, None)

        state.metrics["http_requests_ok"] += 1
        return StreamingResponse(
            sse_generator(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    except HTTPException:
        raise
    except Exception as e:
        state.metrics["http_requests_err"] += 1
        logger.error(f"HTTP stream error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))

# ============================================================
# WebSocket Streaming
# ============================================================

@app.websocket("/ws")
async def websocket_streaming(ws: WebSocket):
    """
    Production WebSocket endpoint for real-time streaming ASR.

    Protocol:
      Client -> Server: binary audio chunks (raw PCM 16-bit 16kHz mono)
      Client -> Server: text "END" to signal end of audio
      Server -> Client: JSON {"text": "...", "is_final": bool, "timestamp": float}
      Server -> Client: text "END" to confirm stream complete

    Backpressure: Rejects with 1013 when at capacity.
    """
    state.metrics["ws_connections_total"] += 1

    # Backpressure: fast-fail when overloaded
    active = len(state.client_queues)
    if active >= MAX_ACTIVE_STREAMS:
        state.metrics["ws_connections_rejected"] += 1
        await ws.accept()
        await ws.close(code=1013, reason=f"Overloaded: {active}/{MAX_ACTIVE_STREAMS}")
        return

    await ws.accept()
    state.metrics["ws_connections_active"] += 1

    stream_id = None
    try:
        # Start inference loop if needed
        start_inference_loop_if_needed()

        # Open streaming session
        stream_id = state.pipeline.open_streaming_session()
        transcription_queue = await register_client(stream_id)
        elapsed = time.perf_counter() - state.start_time
        logger.info(f"[+{elapsed:.1f}s] WS client {stream_id} connected ({len(state.client_queues)} total)")

    except Exception as e:
        logger.error(f"WS connection setup error: {e}", exc_info=True)
        state.metrics["ws_connections_active"] -= 1
        if stream_id is not None:
            await unregister_client(stream_id)
        if ws.application_state is WebSocketState.CONNECTED:
            await ws.close()
        return

    # Send ready message
    await ws.send_text(json.dumps({"type": "Ready", "id": stream_id}))

    async def recv_loop():
        """Receive audio chunks from client, feed to pipeline"""
        att_ctx = state.cfg.streaming.att_context_size[1] + 1
        frame_threshold = att_ctx * 1280 * 2  # bytes
        audio_buffer = b""

        try:
            while True:
                message = await ws.receive()

                # Text message = control signal
                if "text" in message:
                    text_msg = message["text"]
                    if text_msg == "END":
                        await transcription_queue.put({"type": "marker"})
                        break

                # Binary message = audio data
                if "bytes" in message:
                    raw_bytes = message["bytes"]
                    if isinstance(raw_bytes, str):
                        raw_bytes = raw_bytes.encode("latin-1")
                    audio_buffer += raw_bytes

                    # Feed to pipeline when buffer exceeds frame threshold
                    if len(audio_buffer) >= frame_threshold:
                        # VAD gate: skip silence before GPU
                        if state.vad is not None and not state.vad.is_speech(audio_buffer, stream_id):
                            state.metrics["vad_filtered_chunks"] += 1
                            audio_buffer = b""
                            continue

                        state.pipeline.append_streaming_audio(stream_id, audio_buffer)
                        audio_buffer = b""

                        # Signal inference loop
                        state._batch_ready_event.set()
                        await asyncio.sleep(0)

        except WebSocketDisconnect:
            pass
        except Exception as e:
            logger.error(f"WS recv error stream {stream_id}: {e}")
        finally:
            # Flush remaining buffer
            if audio_buffer:
                try:
                    state.pipeline.append_streaming_audio(stream_id, audio_buffer)
                except Exception:
                    pass

            # Mark stream as ended
            try:
                state.pipeline._streaming_request_generator.streams[stream_id].mark_end()
            except Exception:
                pass

            state._batch_ready_event.set()
            await asyncio.sleep(0.1)
            await transcription_queue.put(None)  # Shutdown signal

    async def send_loop():
        """Send transcription results to client"""
        step_num = 0
        try:
            while True:
                output_dict = await transcription_queue.get()
                if output_dict is None:
                    break

                if output_dict.get("type") == "marker":
                    await ws.send_text("END")
                    break

                json_str = output_dict.get("json_str")
                if json_str:
                    if step_num == 0:
                        elapsed = time.perf_counter() - state.start_time
                        logger.info(f"[+{elapsed:.1f}s] Stream {stream_id}: first token sent")
                    await ws.send_text(json_str)
                    step_num += 1

        except Exception as e:
            logger.error(f"WS send error stream {stream_id}: {e}")

    tasks = []
    try:
        tasks = [
            asyncio.create_task(recv_loop()),
            asyncio.create_task(send_loop()),
        ]
        await asyncio.gather(*tasks, return_exceptions=True)
    except Exception as e:
        logger.error(f"WS error stream {stream_id}: {e}")
    finally:
        state.metrics["ws_connections_active"] -= 1
        await unregister_client(stream_id)
        if ws.application_state is WebSocketState.CONNECTED:
            await ws.close()
        for task in tasks:
            if not task.done():
                task.cancel()


Summary: auto


Output:
Chunk ID: ec219f
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2436
Output:
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass

# ============================================================
# Stream Management + Inference Loop
# ============================================================

async def register_client(stream_id: int) -> asyncio.Queue:
    """Register a streaming client, return its output queue"""
    queue = asyncio.Queue()
    async with state.client_queues_lock:
        state.client_queues[stream_id] = queue
    async with state.stream_timestamps_lock:
        state.stream_audio_timestamps[stream_id] = 0.0
    return queue

async def unregister_client(stream_id: int):
    """Remove client from tracking, mark for cleanup"""
    try:
        async with state.client_queues_lock:
            state.client_queues.pop(stream_id, None)
        async with state.stream_timestamps_lock:
            state.stream_audio_timestamps.pop(stream_id, None)
        if state.vad is not None:
            state.vad.reset_stream(stream_id)
        if hasattr(state.pipeline, "_streaming_request_generator"):
            state.pipeline._streaming_request_generator.close_stream(stream_id)
        async with state.streams_cleanup_lock:
            state.streams_pending_cleanup.add(stream_id)
    except Exception as e:
        logger.error(f"Unregister error stream {stream_id}: {e}")

async def route_outputs(outputs_with_timestamps):
    """Route inference outputs to per-client queues"""
    current_time = time.time()

    for output_dict in outputs_with_timestamps:
        output = output_dict["output"]
        audio_ts = output_dict["audio_timestamp"]

        text = output.current_step_transcript if output.partial_transcript else output.final_transcript
        is_final = not bool(output.partial_transcript)

        result = {
            "text": text or "",
            "is_final": is_final,
            "timestamp": current_time,
            "audio_timestamp": audio_ts,
        }

        if len(output.final_segments) > 0:
            result["segment_text"] = output.final_segments[0].text
            result["segment_start"] = output.final_segments[0].start
            result["segment_end"] = output.final_segments[0].end

        output_dict["json_str"] = json.dumps(result)

    # Snapshot queues without holding lock during I/O
    async with state.client_queues_lock:
        queues_snapshot = dict(state.client_queues)

    for output_dict in outputs_with_timestamps:
        output = output_dict["output"]
        queue = queues_snapshot.get(output.stream_id)
        if queue:
            try:
                await queue.put(output_dict)
            except Exception as e:
                logger.error(f"Route error stream {output.stream_id}: {e}")

async def cleanup_finished_streams():
    """Clean up streams fully drained from the request generator"""
    async with state.streams_cleanup_lock:
        if not state.streams_pending_cleanup:
            return
        active_ids = set(state.pipeline._streaming_request_generator.streams.keys())
        to_cleanup = state.streams_pending_cleanup - active_ids
        for sid in to_cleanup:
            try:
                if sid in state.pipeline._state_pool:
                    del state.pipeline._state_pool[sid]
                if hasattr(state.pipeline, "context_manager") and state.pipeline.context_manager is not None:
                    try:
                        state.pipeline.context_manager.reset_slots([sid], [True])
                    except Exception:
                        pass
                state.streams_pending_cleanup.discard(sid)
            except Exception as e:
                logger.error(f"Cleanup error stream {sid}: {e}")

async def centralized_inference_loop():
    """
    Production inference loop with event-driven batching.

    WHY event-driven: fires immediately when data arrives (high load) or
    after MAX_BATCH_WAIT_MS timeout (low load). 6x better than fixed 300ms sleep.
    """
    state.inference_running = True
    max_wait = MAX_BATCH_WAIT_MS / 1000.0
    metrics_report_interval = 100

    logger.info(f"Inference loop started (max_wait={MAX_BATCH_WAIT_MS}ms)")

    while state.inference_running:
        try:
            active = len(state.client_queues)

            # Idle if no clients
            if active == 0:
                try:
                    await asyncio.wait_for(state._batch_ready_event.wait(), timeout=0.1)
                    state._batch_ready_event.clear()
                except asyncio.TimeoutError:
                    pass
                continue

            # Event-driven wait
            cycle_start = time.perf_counter()
            try:
                await asyncio.wait_for(state._batch_ready_event.wait(), timeout=max_wait)
            except asyncio.TimeoutError:
                pass
            state._batch_ready_event.clear()

            # Process ALL ready frames (exhaustive)
            all_outputs = []
            async with state.stream_timestamps_lock:
                timestamps = dict(state.stream_audio_timestamps)

            batches_this_cycle = 0
            while True:
                try:
                    batch_start = time.perf_counter()
                    requests, outputs = await state.pipeline.process_streaming_batch()
                    gpu_ms = (time.perf_counter() - batch_start) * 1000

                    if not outputs:
                        break

                    batches_this_cycle += 1
                    state.metrics["total_batches"] += 1
                    state.metrics["total_outputs"] += len(outputs)
                    state.metrics["total_gpu_ms"] += gpu_ms
                    state.metrics["batch_sizes"].append(len(outputs))
                    state.metrics["gpu_times_ms"].append(gpu_ms)

                    # Update timestamps
                    for req in requests:
                        if hasattr(req, "length"):
                            sid = req.stream_id
                            timestamps[sid] = timestamps.get(sid, 0.0) + req.length / 16000.0

                    for output in outputs:
                        all_outputs.append({
                            "output": output,
                            "audio_timestamp": timestamps.get(output.stream_id, 0.0),
                        })

                    await asyncio.sleep(0)

                except Exception as e:
                    if type(e).__name__ == "NotEnoughDataException":
                        if batches_this_cycle == 0:
                            state.metrics["empty_cycles"] += 1
                        break
                    raise

            if all_outputs:
                async with state.stream_timestamps_lock:
                    for sid, ts in timestamps.items():
                        if sid in state.stream_audio_timestamps:
                            state.stream_audio_timestamps[sid] = ts

                route_start = time.perf_counter()
                await route_outputs(all_outputs)
                state.metrics["total_route_ms"] += (time.perf_counter() - route_start) * 1000

                await cleanup_finished_streams()

                cycle_ms = (time.perf_counter() - cycle_start) * 1000
                state.metrics["cycle_times_ms"].append(cycle_ms)

                # Periodic metrics log
                if state.metrics["total_batches"] % metrics_report_interval == 0:
                    _log_metrics(active)

        except Exception as e:
            logger.error(f"Inference loop error: {e}", exc_info=True)
            await asyncio.sleep(0.1)

    logger.warning("Inference loop stopped")

def _log_metrics(active_clients: int):
    """Structured metrics log"""
    m = state.metrics

    def pct(arr, p):
        if not arr:
            return 0
        s = sorted(arr)
        return s[min(int(len(s) * p / 100), len(s) - 1)]

    try:
        vram_used = torch.cuda.memory_allocated() / 1e9
        vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9
    except Exception:
        vram_used = vram_total = 0

    elapsed = time.perf_counter() - state.start_time
    logger.info(
        f"[+{elapsed:.0f}s] METRICS | streams={active_clients} "
        f"batches={m['total_batches']} outputs={m['total_outputs']} "
        f"GPU: p50={pct(m['gpu_times_ms'], 50):.1f}ms p99={pct(m['gpu_times_ms'], 99):.1f}ms | "
        f"VRAM: {vram_used:.1f}/{vram_total:.0f}GB | "
        f"vad_skip={m['vad_filtered_chunks']}"
    )

def start_inference_loop_if_needed():
    """Start centralized inference loop on first client connection"""
    if state.inference_task is None:
        logger.info("Starting inference loop (first client)")
        state.inference_task = asyncio.create_task(centralized_inference_loop())
    elif state.inference_task.done():
        logger.warning("Inference task was done, restarting")
        state.inference_task = asyncio.create_task(centralized_inference_loop())

# ============================================================
# Main
# ============================================================

if __name__ == "__main__":
    logger.info(f"Starting ASR server on port {SERVER_PORT}")
    logger.info(f"  Model: {MODEL_PATH}")
    logger.info(f"  Max streams: {MAX_ACTIVE_STREAMS}")
    logger.info(f"  Batch wait: {MAX_BATCH_WAIT_MS}ms")
    logger.info(f"  VAD: {'enabled' if ENABLE_VAD else 'disabled'}")

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=SERVER_PORT,
        log_level="info",
        access_log=False,  # Reduce noise, we have structured logging
        timeout_keep_alive=300,  # 5min for long WebSocket sessions
    )


**Identifying concurrency bottleneck**


I found a key reason for the mismatch: our HTTP/SSE path is not equivalent to NVIDIA’s benchmark method. Their numbers are real-time WS cache-aware streaming on H100; our current production results came from request-per-file HTTP/SSE, which stresses routing + request overhead differently and cannot be compared directly. I’ll now check existing WS test tooling and run a clean WS benchmark on the dedicated bench endpoint.


rg --files | rg 'ws|websocket|stress|bench|drift'


python3 production/deploy.py workers | sed -n '1,220p'