Output:
Chunk ID: fa5c3a
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 2730
Output:

    def observed_inflight_max(self) -> Optional[int]:
        values = [r.request_inflight for r in self.successes if isinstance(r.request_inflight, int)]
        if not values:
            return None
        return max(values)

    def error_summary(self) -> Dict[str, int]:
        errors: Dict[str, int] = {}
        for r in self.failures:
            key = r.error or f"HTTP {r.status_code}"
            errors[key[:100]] = errors.get(key[:100], 0) + 1
        return errors

def _coerce_number(raw: object) -> Optional[float]:
    if raw is None:
        return None
    v = str(raw).strip()
    if v == "":
        return None
    try:
        if "." in v:
            return float(v)
        return float(int(v))
    except ValueError:
        try:
            return float(v)
        except ValueError:
            return None

def _coerce_bool(raw: object) -> Optional[bool]:
    if raw is None:
        return None
    if isinstance(raw, bool):
        return raw
    v = str(raw).strip().lower()
    if v in {"true", "1", "yes", "y", "on"}:
        return True
    if v in {"false", "0", "no", "n", "off"}:
        return False
    return None

def _normalize_time_ms(raw: Optional[str]) -> float:
    value = _coerce_number(raw)
    if value is None:
        return 0.0
    return float(value)

def _parse_perf_headers(headers: Any) -> Dict[str, Any]:
    """Parse server timing headers into numeric metric fields."""
    metrics: Dict[str, Any] = {}

    # Explicit compact payload.
    detail_payload = headers.get("x-perf-details") if hasattr(headers, "get") else None
    if detail_payload:
        try:
            parsed = json.loads(detail_payload)
            if isinstance(parsed, dict):
                metrics.update(parsed)
        except Exception:
            pass

    header_map = {
        "x-llm-token-total": "llm_token_total",
        "x-llm-batch-count": "llm_batch_count",
        "x-llm-batch-wall-ms": "llm_batch_wall_ms_total",
        "x-llm-batch-gpu-ms": "llm_batch_gpu_ms_total",
        "x-llm-batch-wall-ms-min": "llm_batch_wall_ms_min",
        "x-llm-batch-wall-ms-max": "llm_batch_wall_ms_max",
        "x-llm-batch-wall-ms-p50": "llm_batch_wall_ms_p50",
        "x-llm-batch-gpu-ms-min": "llm_batch_gpu_ms_min",
        "x-llm-batch-gpu-ms-max": "llm_batch_gpu_ms_max",
        "x-llm-batch-gpu-ms-p50": "llm_batch_gpu_ms_p50",
        "x-llm-decode-wall-ms": "llm_decode_wall_ms_total",
        "x-llm-decode-gpu-ms": "llm_decode_gpu_ms_total",
        "x-llm-decode-wall-ms-min": "llm_decode_wall_ms_min",
        "x-llm-decode-wall-ms-max": "llm_decode_wall_ms_max",
        "x-llm-decode-wall-ms-p50": "llm_decode_wall_ms_p50",
        "x-llm-decode-gpu-ms-min": "llm_decode_gpu_ms_min",
        "x-llm-decode-gpu-ms-max": "llm_decode_gpu_ms_max",
        "x-llm-decode-gpu-ms-p50": "llm_decode_gpu_ms_p50",
        "x-llm-decode-calls": "decode_calls",
        "x-llm-decode-cpu-ms": "llm_decode_cpu_ms",
        "x-llm-parse-ms": "llm_parse_ms",
        "x-llm-parse-avg-ms": "llm_parse_ms_avg",
        "x-llm-tokens-per-batch-min": "tokens_per_batch_min",
        "x-llm-tokens-per-batch-max": "tokens_per_batch_max",
        "x-llm-tokens-per-batch-p50": "tokens_per_batch_p50",
        "x-generation-ms": "generation_ms",
        "x-request-inflight": "request_inflight",
        "x-llm-text-chunked": "text_chunked",
        "x-llm-chunks-processed": "chunks_processed",
        "x-llm-sr-applied": "sr_applied",
        "x-llm-time-per-token-ms": "llm_time_per_token_ms",
        "x-llm-time-per-batch-wall-ms": "llm_time_per_batch_wall_ms",
        "x-llm-time-per-batch-gpu-ms": "llm_time_per_batch_gpu_ms",
    }

    bool_keys = {"text_chunked", "sr_applied"}
    int_keys = {
        "llm_token_total",
        "llm_batch_count",
        "decode_calls",
        "chunks_processed",
        "request_inflight",
    }

    if hasattr(headers, "items"):
        for raw_key, raw_val in headers.items():
            lk = str(raw_key).lower()
            target_key = header_map.get(lk)
            if target_key is None:
                continue
            if target_key in bool_keys:
                b = _coerce_bool(raw_val)
                if b is not None:
                    metrics[target_key] = b
                continue

            v = _coerce_number(raw_val)
            if v is None:
                continue
            if target_key in int_keys and float(v).is_integer():
                metrics[target_key] = int(v)
            else:
                metrics[target_key] = int(v) if float(v).is_integer() else float(v)

    if "llm_token_total" in metrics and "generation_ms" in metrics and metrics["llm_token_total"]:
        try:
            metrics["llm_time_per_token_ms"] = float(metrics["generation_ms"]) / float(metrics["llm_token_total"])
        except Exception:
            pass

    if "llm_batch_count" in metrics and "generation_ms" in metrics and metrics["llm_batch_count"]:
        try:
            metrics["llm_time_per_batch_wall_ms"] = float(metrics["generation_ms"]) / float(metrics["llm_batch_count"])
        except Exception:
            pass

    if "llm_batch_count" in metrics and "llm_batch_gpu_ms_total" in metrics and metrics["llm_batch_count"]:
        try:
            metrics["llm_time_per_batch_gpu_ms"] = float(metrics["llm_batch_gpu_ms_total"]) / float(metrics["llm_batch_count"])
        except Exception:
            pass

    return metrics

async def make_request(
    client: httpx.AsyncClient,
    base_url: str,
    text: str,
    speaker: str,
    chunking: bool,
    max_tokens: int,
    stream: bool = False,
    timeout: float = 120.0,
) -> RequestResult:
    url = f"{base_url}{GENERATE_PATH}"
    payload = {
        "text": text,
        "speaker": speaker,
        "stream": stream,
        "chunking": chunking,
        "max_tokens": max_tokens,
        "format": "wav",
    }

    start = time.time()
    ttfb_time = 0.0

    try:
        if stream:
            async with client.stream(
                "POST",
                url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=timeout,
            ) as response:
                first = True
                total_bytes = 0
                async for chunk in response.aiter_bytes():
                    if first:
                        ttfb_time = (time.time() - start) * 1000
                        first = False
                    total_bytes += len(chunk)

                latency = (time.time() - start) * 1000
                timing = _parse_perf_headers(response.headers)

                if response.status_code == 200:
                    pcm_bytes = max(0, total_bytes - 44)
                    audio_seconds = pcm_bytes / (16000 * 2)
                    return RequestResult(
                        success=True,
                        status_code=200,
                        latency_ms=latency,
                        ttfb_ms=ttfb_time,
                        audio_bytes=total_bytes,
                        audio_seconds=audio_seconds,
                        request_id=response.headers.get("x-request-id", ""),
                        speaker=speaker,
                        text_length=len(text),
                        stream=True,
                        timing=timing,
                        request_inflight=timing.get("request_inflight"),
                        server_ttfb_ms=_normalize_time_ms(response.headers.get("x-ttfb-ms")),
                    )

                return RequestResult(
                    success=False,
                    status_code=response.status_code,
                    latency_ms=latency,
                    ttfb_ms=ttfb_time,
                    error=f"HTTP {response.status_code}",
                    speaker=speaker,
                    text_length=len(text),
                    stream=True,
                    timing=timing,
                    request_id=response.headers.get("x-request-id", ""),
                    request_inflight=timing.get("request_inflight"),
                )
        else:
            response = await client.post(
                url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=timeout,
            )
            latency = (time.time() - start) * 1000
            timing = _parse_perf_headers(response.headers)

            server_ttfb = _normalize_time_ms(response.headers.get("x-ttfb-ms"))
            audio_seconds = 0.0
            audio_sec_header = response.headers.get("x-audio-seconds")
            if audio_sec_header:
                parsed = _coerce_number(audio_sec_header)
                if parsed is not None:
                    audio_seconds = parsed
            elif response.status_code == 200:
                pcm_bytes = max(0, len(response.content) - 44)
                audio_seconds = pcm_bytes / (16000 * 2)

            if response.status_code == 200:
                return RequestResult(
                    success=True,
                    status_code=200,
                    latency_ms=latency,
                    ttfb_ms=latency,
                    server_ttfb_ms=server_ttfb,
                    audio_bytes=len(response.content),
                    audio_seconds=audio_seconds,
                    request_id=response.headers.get("x-request-id", ""),
                    speaker=speaker,
                    text_length=len(text),
                    stream=False,
                    timing=timing,
                    request_inflight=timing.get("request_inflight"),
                )

            return RequestResult(
                success=False,
                status_code=response.status_code,
                latency_ms=latency,
                ttfb_ms=latency,
                server_ttfb_ms=server_ttfb,
                audio_bytes=len(response.content),
                error=(response.text[:200] if response.text else f"HTTP {response.status_code}"),
                speaker=speaker,
                text_length=len(text),
                stream=False,
                timing=timing,
                request_id=response.headers.get("x-request-id", ""),
                request_inflight=timing.get("request_inflight"),
            )

    except httpx.TimeoutException:
        return RequestResult(
            success=False,
            status_code=0,
            latency_ms=(time.time() - start) * 1000,
            error="TIMEOUT",
            speaker=speaker,
            text_length=len(text),
            stream=stream,
        )
    except Exception as e:
        return RequestResult(
            success=False,
            status_code=0,
            latency_ms=(time.time() - start) * 1000,
            error=str(e)[:200],
            speaker=speaker,
            text_length=len(text),
            stream=stream,
        )


# 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.


sed -n '560,980p' scripts/stress_test_optimized_local.py


Output:
Chunk ID: 60b63a
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2926
Output:
        )

async def run_level(
    base_url: str,
    num_requests: int,
    concurrency: int,
    text_category: str,
    chunking: bool,
    stream: bool,
    max_tokens: int,
    gpu_monitor: Optional[GPUMonitor] = None,
) -> LevelResult:
    mode = "streaming" if stream else "non-streaming"
    level = LevelResult(
        level_name=f"{concurrency}c-{text_category}-{mode}",
        concurrency=concurrency,
        total_requests=num_requests,
        mode=mode,
        text_category=text_category,
        chunking=chunking,
    )

    semaphore = asyncio.Semaphore(concurrency)

    if text_category == "mixed":
        texts = list(TEST_TEXTS.values())
    else:
        texts = [TEST_TEXTS.get(text_category, TEST_TEXTS["short"])]

    async def limited_request(i: int, client: httpx.AsyncClient) -> RequestResult:
        async with semaphore:
            text = texts[i % len(texts)]
            speaker = SPEAKERS[i % len(SPEAKERS)]
            return await make_request(
                client=client,
                base_url=base_url,
                text=text,
                speaker=speaker,
                chunking=chunking,
                max_tokens=max_tokens,
                stream=stream,
            )

    if gpu_monitor:
        gpu_monitor.snapshots.clear()

    start_time = time.time()
    limits = httpx.Limits(
        max_connections=concurrency + 10,
        max_keepalive_connections=concurrency,
    )
    async with httpx.AsyncClient(limits=limits) as client:
        tasks = [limited_request(i, client) for i in range(num_requests)]
        responses = await asyncio.gather(*tasks, return_exceptions=True)

    level.wall_time_seconds = time.time() - start_time

    for item in responses:
        if isinstance(item, RequestResult):
            level.results.append(item)
        else:
            level.results.append(
                RequestResult(
                    success=False,
                    status_code=0,
                    latency_ms=0,
                    error=str(item)[:200],
                    stream=stream,
                )
            )

    if gpu_monitor:
        level.gpu_summary = gpu_monitor.summary()

    return level

def print_level_report(level: LevelResult):
    print(f"\n{'=' * 72}")
    print(f"  {level.level_name.upper()}")
    print(f"  {level.total_requests} requests, {level.concurrency} concurrent, {level.mode}, chunking={level.chunking}")
    print(f"{'=' * 72}")

    print(f"  Success:     {len(level.successes)}/{level.total_requests} ({level.success_rate:.0%})")
    print(f"  Wall time:   {level.wall_time_seconds:.2f}s")
    print(f"  Throughput:  {level.throughput_rps:.2f} req/s")
    print(f"  Audio total: {level.audio_duration_total:.2f}s")
    print(f"  Eff. RTF:    {level.effective_rtf:.3f}")

    if level.observed_inflight_max() is not None:
        print(f"  Peak inflight: {level.observed_inflight_max()}")

    latency = level.latency_stats
    if latency:
        print("\n  Latency (ms):")
        print(
            f"    min={latency['min']:.0f} avg={latency['avg']:.0f} "
            f"p50={latency['p50']:.0f} p95={latency['p95']:.0f} p99={latency['p99']:.0f} max={latency['max']:.0f}"
        )

    ttfb = level.ttfb_stats
    if ttfb and not level.mode.startswith("stream"):
        print("\n  Server TTFB (ms):")
        print(
            f"    min={ttfb['min']:.0f} avg={ttfb['avg']:.0f} "
            f"p50={ttfb['p50']:.0f} p95={ttfb['p95']:.0f} max={ttfb['max']:.0f}"
        )

    if level.mode == "non-streaming":
        llm = level.timing_stats("generation_ms")
        if llm:
            print("\n  Generation (ms):")
            print(f"    min={llm['min']:.0f} avg={llm['avg']:.0f} p50={llm['p50']:.0f} p95={llm['p95']:.0f} max={llm['max']:.0f}")
        wall = level.timing_stats("llm_batch_wall_ms_total")
        if wall:
            print("\n  LLM batch wall (ms):")
            print(f"    min={wall['min']:.0f} avg={wall['avg']:.0f} p50={wall['p50']:.0f} p95={wall['p95']:.0f} max={wall['max']:.0f}")
        gpu = level.timing_stats("llm_batch_gpu_ms_total")
        if gpu:
            print("\n  LLM batch GPU (ms):")
            print(f"    min={gpu['min']:.0f} avg={gpu['avg']:.0f} p50={gpu['p50']:.0f} p95={gpu['p95']:.0f} max={gpu['max']:.0f}")
        parse = level.timing_stats("llm_parse_ms")
        if parse:
            print("\n  Parse (ms):")
            print(f"    avg={parse['avg']:.0f} p95={parse['p95']:.0f}")
        decode = level.timing_stats("llm_decode_wall_ms_total")
        if decode:
            print("\n  Decode (ms):")
            print(f"    avg={decode['avg']:.0f} p95={decode['p95']:.0f}")

        token_ms, token_count = level.timing_rate("generation_ms", "llm_token_total")
        if token_count:
            print(f"  Token timing: avg {token_ms / token_count * 1000:.3f} us/token")

    if level.gpu_summary and "memory_mb" in level.gpu_summary:
        gpu_summary = level.gpu_summary
        mem = gpu_summary["memory_mb"]
        util = gpu_summary["gpu_util_pct"]
        temp = gpu_summary["temperature_c"]
        print("\n  GPU:")
        print(f"    Memory: {mem['avg']:.0f}MB avg, {mem['max']:.0f}MB peak / {mem['total']:.0f}MB")
        print(f"    Util:   {util['avg']:.0f}% avg, {util['max']:.0f}% peak")
        print(f"    Temp:   {temp['avg']:.0f}C avg, {temp['max']:.0f}C peak")

    errors = level.error_summary()
    if errors:
        print("\n  Errors:")
        for key, count in sorted(errors.items(), key=lambda kv: -kv[1]):
            print(f"    [{count}x] {key}")

def print_summary_table(levels: List[LevelResult]):
    print(f"\n{'=' * 110}")
    print("  STRESS BENCHMARK SUMMARY")
    print(f"{'=' * 110}")
    header = f"{'Level':<26} {'OK%':>5} {'RPS':>7} {'p50':>7} {'p95':>7} {'GPU%':>5} {'GPUmem':>8} {'RTF':>6} {'Inflight':>8}"
    print(header)
    print("-" * 110)

    for level in levels:
        lat = level.latency_stats
        p50 = f"{lat.get('p50', 0):.0f}" if lat else "---"
        p95 = f"{lat.get('p95', 0):.0f}" if lat else "---"
        gpu_util = level.gpu_summary.get("gpu_util_pct", {}).get("avg", 0)
        gpu_mem = level.gpu_summary.get("memory_mb", {}).get("max", 0)
        inflight = level.observed_inflight_max() or 0
        print(
            f"{level.level_name:<26} "
            f"{level.success_rate:>4.0%} "
            f"{level.throughput_rps:>7.2f} "
            f"{p50:>7} "
            f"{p95:>7} "
            f"{gpu_util:>4.0f}% "
            f"{gpu_mem:>7.0f}MB "
            f"{level.effective_rtf:>6.3f} "
            f"{inflight:>8d}"
        )

    print("=" * 110)

def save_results(levels: List[LevelResult], path: str):
    payload = []
    for level in levels:
        payload.append({
            "level": level.level_name,
            "concurrency": level.concurrency,
            "mode": level.mode,
            "chunking": level.chunking,
            "text_category": level.text_category,
            "success_rate": level.success_rate,
            "throughput_rps": level.throughput_rps,
            "wall_time_s": level.wall_time_seconds,
            "latency_ms": level.latency_stats,
            "ttfb_ms": level.ttfb_stats,
            "effective_rtf": level.effective_rtf,
            "audio_total_seconds": level.audio_duration_total,
            "observed_inflight_max": level.observed_inflight_max(),
            "errors": level.error_summary(),
            "gpu": level.gpu_summary,
            "timing_averages": {
                "generation_ms": level.timing_stats("generation_ms"),
                "llm_batch_wall_ms_total": level.timing_stats("llm_batch_wall_ms_total"),
                "llm_batch_gpu_ms_total": level.timing_stats("llm_batch_gpu_ms_total"),
                "llm_parse_ms": level.timing_stats("llm_parse_ms"),
                "llm_decode_wall_ms_total": level.timing_stats("llm_decode_wall_ms_total"),
            },
            "requests": [
                {
                    "request_id": r.request_id,
                    "success": r.success,
                    "status_code": r.status_code,
                    "latency_ms": r.latency_ms,
                    "ttfb_ms": r.ttfb_ms,
                    "server_ttfb_ms": r.server_ttfb_ms,
                    "audio_bytes": r.audio_bytes,
                    "audio_seconds": r.audio_seconds,
                    "request_inflight": r.request_inflight,
                    "error": r.error,
                    "timing": r.timing,
                }
                for r in level.results
            ],
        })

    with open(path, "w") as fd:
        json.dump(payload, fd, indent=2)

async def run_stress_test(args):
    base_url = args.url

    try:
        async with httpx.AsyncClient() as client:
            response = await client.get(f"{base_url}/v1/tts/health", timeout=10)
            health = response.json()
            print(f"  Health: {health.get('status')}")
            print(f"  Version: {health.get('model_version')}")
            print(f"  GPU: {health.get('gpu_available')}")
    except Exception as e:
        print(f"ERROR: server not reachable at {base_url}: {e}")
        sys.exit(1)

    levels = [int(v) for v in args.levels.split(",") if v.strip()]
    gpu_monitor = GPUMonitor(interval_seconds=0.5)
    gpu_monitor.start()

    mode_label = "streaming" if args.stream else "non-streaming"
    results: List[LevelResult] = []

    print(f"\nWarming up (1 request, non-streaming)...")
    warmup = await run_level(
        base_url=base_url,
        num_requests=1,
        concurrency=1,
        text_category=args.text,
        chunking=args.chunking,
        stream=False,
        max_tokens=args.max_tokens,
        gpu_monitor=gpu_monitor,
    )
    print(f"  warmup ok: {warmup.success_rate:.0%}")

    for conc in levels:
        num_requests = max(conc * 3, 10)
        num_requests = min(num_requests, 200)
        level = await run_level(
            base_url=base_url,
            num_requests=num_requests,
            concurrency=conc,
            text_category=args.text,
            chunking=args.chunking,
            stream=args.stream,
            max_tokens=args.max_tokens,
            gpu_monitor=gpu_monitor,
        )
        print_level_report(level)
        results.append(level)
        if conc != levels[-1]:
            print("  cooldown 3s...")
            await asyncio.sleep(3)

    gpu_monitor.stop()
    print_summary_table(results)

    if args.output:
        save_results(results, args.output)
        print(f"\nResults written to {args.output}")

def main():
    parser = argparse.ArgumentParser(description="Detailed local optimized-path TTS benchmark")
    parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (default: {DEFAULT_URL})")
    parser.add_argument("--levels", default="1,5,10,20,50", help="Comma-separated concurrency levels")
    parser.add_argument("--text", default="short", choices=["short", "medium", "long", "mixed"], help="Input text category")
    parser.add_argument("--chunking", action="store_true", help="Enable server-side chunking", default=True)
    parser.add_argument("--no-chunking", action="store_false", dest="chunking", help="Disable chunking")
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=256,
        help="max_tokens passed to /v1/tts/generate (default: 256)",
    )
    parser.add_argument("--stream", action="store_true", help="Use streaming endpoint")
    parser.add_argument("--output", default="stress_test_optimized_local.json", help="Output JSON path")

    args = parser.parse_args()
    asyncio.run(run_stress_test(args))

if __name__ == "__main__":
    main()


Summary: auto


I confirmed the stress script already collects rich request timing headers and writes per-request JSON. I’m checking the server response headers next so I can guarantee we capture the exact GPU/LLM fields you need (and wire any missing ones into the report).


sed -n '1,340p' veena3modal/api/fastapi_app.py


Output:
Chunk ID: 153650
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2937
Output:
"""
FastAPI ASGI app factory for Modal deployment.

Phase 1: expose `/v1/tts/generate` and `/v1/tts/health` while reusing existing
Veena3 inference/streaming code from the Django project.

Phase 2+: extract shared code into `veena3modal/shared/` and remove Django coupling.
"""

# NOTE: Intentionally NOT using `from __future__ import annotations` here
# to ensure FastAPI recognizes Request type annotations properly

import time
import uuid
import json
from typing import Dict, Optional

# Structured logging (JSON formatted, request-scoped)
from veena3modal.shared.logging import (
    get_logger,
    set_request_context,
    clear_request_context,
    log_request_received,
    log_first_audio_emitted,
    log_request_completed,
    log_request_failed,
)

# Prometheus metrics
from veena3modal.shared.metrics import (
    record_request_received as metrics_request_received,
    record_request_completed as metrics_request_completed,
    record_request_failed as metrics_request_failed,
    record_ttfb,
    record_rtf,
    record_audio_duration,
    record_chunks_sent,
    get_metrics_text,
    get_content_type as get_metrics_content_type,
)

# Auth and rate limiting
from veena3modal.api.auth import (
    get_api_validator,
    extract_api_key,
    hash_api_key,
)
from veena3modal.api.rate_limiter import get_rate_limiter
from veena3modal.api.error_handlers import (
    ErrorCode,
    create_error_response,
    get_error_status,
    is_gpu_fault,
    handle_gpu_fault,
    FeatureFlags,
)

logger = get_logger(__name__)

# Sentence storage (fire-and-forget, non-blocking)
# Imported at module level for singleton access
_sentence_store = None

def _get_sentence_store():
    """Lazy load sentence store singleton."""
    global _sentence_store
    if _sentence_store is None:
        from veena3modal.services.sentence_store import get_sentence_store
        _sentence_store = get_sentence_store()
    return _sentence_store

def _build_perf_headers(metrics: Dict[str, object]) -> Dict[str, str]:
    """Build compact timing headers for benchmark tooling."""
    perf: Dict[str, str] = {}
    if not metrics:
        return perf

    def _pick(*keys: str):
        for key in keys:
            if key in metrics and metrics[key] is not None:
                return metrics[key]
        return None

    def _add(name: str, value: object):
        if value is None:
            return
        if isinstance(value, bool):
            perf[name] = str(value).lower()
        else:
            perf[name] = str(value)

    llm_token_total = _pick("llm_token_total", "token_total")
    llm_batch_count = _pick("llm_batch_count", "batch_count")
    generation_ms = _pick("generation_ms") or _pick("total_ms")
    llm_batch_wall_total = _pick("llm_batch_wall_ms_total", "llm_wall_ms", "total_ms")
    llm_batch_gpu_total = _pick("llm_batch_gpu_ms_total", "llm_gpu_batch_ms_total")
    llm_decode_wall_total = _pick("llm_decode_wall_ms_total")
    llm_decode_gpu_total = _pick("llm_decode_gpu_ms_total")
    llm_parse_ms = _pick("llm_parse_ms")
    decode_cpu_ms = _pick("llm_decode_cpu_ms")
    parse_ms = _pick("parse_ms")

    _add("X-LLM-Token-Total", llm_token_total)
    _add("X-LLM-Batch-Count", llm_batch_count)
    _add("X-LLM-Batch-Wall-Ms", llm_batch_wall_total)
    _add("X-LLM-Batch-GPU-Ms", llm_batch_gpu_total)
    _add("X-LLM-Batch-Wall-Ms-Min", _pick("llm_batch_wall_ms_min"))
    _add("X-LLM-Batch-Wall-Ms-Max", _pick("llm_batch_wall_ms_max"))
    _add("X-LLM-Batch-Wall-Ms-P50", _pick("llm_batch_wall_ms_p50"))
    _add("X-LLM-Batch-GPU-Ms-Min", _pick("llm_batch_gpu_ms_min"))
    _add("X-LLM-Batch-GPU-Ms-Max", _pick("llm_batch_gpu_ms_max"))
    _add("X-LLM-Batch-GPU-Ms-P50", _pick("llm_batch_gpu_ms_p50"))
    _add("X-LLM-Decode-Wall-Ms", llm_decode_wall_total)
    _add("X-LLM-Decode-GPU-Ms", llm_decode_gpu_total)
    _add("X-LLM-Decode-Wall-Ms-Min", _pick("llm_decode_wall_ms_min"))
    _add("X-LLM-Decode-Wall-Ms-Max", _pick("llm_decode_wall_ms_max"))
    _add("X-LLM-Decode-Wall-Ms-P50", _pick("llm_decode_wall_ms_p50"))
    _add("X-LLM-Decode-GPU-Ms-Min", _pick("llm_decode_gpu_ms_min"))
    _add("X-LLM-Decode-GPU-Ms-Max", _pick("llm_decode_gpu_ms_max"))
    _add("X-LLM-Decode-GPU-Ms-P50", _pick("llm_decode_gpu_ms_p50"))
    _add("X-LLM-Decode-Calls", _pick("decode_calls"))
    _add("X-LLM-Decode-CPU-Ms", decode_cpu_ms)
    _add("X-LLM-Parse-Ms", llm_parse_ms or parse_ms)
    _add("X-LLM-Parse-Avg-Ms", _pick("llm_parse_ms_avg"))
    _add("X-LLM-Tokens-Per-Batch-Min", _pick("tokens_per_batch_min"))
    _add("X-LLM-Tokens-Per-Batch-Max", _pick("tokens_per_batch_max"))
    _add("X-LLM-Tokens-Per-Batch-P50", _pick("tokens_per_batch_p50"))
    _add("X-Generation-Ms", generation_ms)
    _add("X-LLM-Generation-Ms", _pick("llm_wall_ms", generation_ms))
    _add("X-Request-Inflight", _pick("request_inflight", "active_requests"))
    _add("X-LLM-Text-Chunked", _pick("text_chunked"))
    _add("X-LLM-Chunks-Processed", _pick("chunks_processed"))
    _add("X-LLM-SR-Applied", _pick("sr_applied"))

    if llm_token_total and generation_ms:
        try:
            perf["X-LLM-Time-Per-Token-Ms"] = f"{float(generation_ms) / float(llm_token_total):.3f}"
        except Exception:
            pass

    if llm_batch_count:
        try:
            perf["X-LLM-Time-Per-Batch-Wall-Ms"] = f"{float(generation_ms or 0) / float(llm_batch_count):.3f}"
        except Exception:
            pass
        try:
            perf["X-LLM-Time-Per-Batch-GPU-Ms"] = f"{float(llm_batch_gpu_total or 0) / float(llm_batch_count):.3f}"
        except Exception:
            pass

    # Include compact JSON snapshot for downstream benchmark tooling.
    detail_keys = [
        "request_id",
        "speaker",
        "text_length",
        "llm_token_total",
        "llm_batch_count",
        "llm_batch_wall_ms_total",
        "llm_batch_gpu_ms_total",
        "llm_batch_wall_ms_min",
        "llm_batch_wall_ms_max",
        "llm_batch_wall_ms_p50",
        "llm_batch_gpu_ms_min",
        "llm_batch_gpu_ms_max",
        "llm_batch_gpu_ms_p50",
        "llm_decode_cpu_ms",
        "llm_decode_wall_ms_total",
        "llm_decode_gpu_ms_total",
        "llm_decode_wall_ms_min",
        "llm_decode_wall_ms_max",
        "llm_decode_wall_ms_p50",
        "llm_decode_gpu_ms_min",
        "llm_decode_gpu_ms_max",
        "llm_decode_gpu_ms_p50",
        "decode_calls",
        "decode_cpu_ms",
        "parse_ms",
        "llm_parse_ms",
        "llm_parse_ms_avg",
        "llm_batch_tokens",
        "tokens_per_batch_min",
        "tokens_per_batch_max",
        "tokens_per_batch_p50",
        "generation_ms",
        "audio_duration_seconds",
        "request_inflight",
        "active_requests",
        "chunks_processed",
        "text_chunked",
        "sr_applied",
    ]
    perf["X-Perf-Details"] = json.dumps(
        {k: metrics[k] for k in detail_keys if k in metrics},
        separators=(",", ":"),
    )

    return perf

# Version info (will be set dynamically after model load)
_APP_VERSION = "0.1.0"
_MODEL_VERSION: Optional[str] = None
_startup_time: Optional[float] = None

def get_model_version() -> str:
    """Return the current model version, or 'not_loaded' if not initialized."""
    return _MODEL_VERSION or "not_loaded"

def set_model_version(version: str) -> None:
    """Set the model version after loading."""
    global _MODEL_VERSION
    _MODEL_VERSION = version

def create_app():
    """
    Factory function to create FastAPI app.

    Import inside factory so local tooling can import this module without FastAPI installed.
    """
    from fastapi import FastAPI, HTTPException, Request
    from fastapi.responses import JSONResponse, Response

    global _startup_time
    _startup_time = time.time()

    app = FastAPI(
        title="Veena3 TTS (Modal)",
        version=_APP_VERSION,
        description="High-quality multilingual Text-to-Speech API with true streaming support.",
    )

    @app.get("/v1/tts/health")
    async def tts_health():
        """
        Health check endpoint.

        Returns:
            - status: "healthy" | "degraded" | "unhealthy"
            - model_loaded: whether TTS model is initialized
            - model_version: version string of loaded model
            - uptime_seconds: time since app startup
            - gpu_available: whether GPU is detected (best-effort check)
        """
        # Check model status via runtime
        from veena3modal.services import tts_runtime
        model_loaded = tts_runtime.is_initialized()

        # Use runtime's model version if available
        runtime = tts_runtime.get_runtime()
        model_version = runtime.model_version if runtime else get_model_version()

        # Check GPU availability (best-effort, non-blocking)
        gpu_available = False
        try:
            import torch
            gpu_available = torch.cuda.is_available()
        except ImportError:
            pass  # torch not installed, skip GPU check

        # Compute uptime
        uptime = time.time() - _startup_time if _startup_time else 0

        # Determine overall status
        if model_loaded and gpu_available:
            status = "healthy"
        elif model_loaded or gpu_available:
            status = "degraded"
        else:
            status = "degraded"

        return JSONResponse(
            content={
                "status": status,
                "model_loaded": model_loaded,
                "model_version": model_version,
                "uptime_seconds": round(uptime, 2),
                "gpu_available": gpu_available,
                "app_version": _APP_VERSION,
            },
            headers={
                "X-Model-Version": model_version,
                "X-App-Version": _APP_VERSION,
            },
        )

    @app.get("/v1/tts/metrics")
    async def tts_metrics():
        """
        Prometheus metrics endpoint.

        Returns metrics in Prometheus text format for scraping.
        """
        return Response(
            content=get_metrics_text(),
            media_type=get_metrics_content_type(),
        )

    @app.post("/v1/tts/generate")
    async def tts_generate(request: Request):
        """
        Generate speech from text.

        Supports streaming (M4) and non-streaming (M3) modes.
        Currently implements non-streaming WAV generation only.
        """
        from veena3modal.api.schemas import TTSGenerateRequest, AudioFormat
        from veena3modal.services import tts_runtime

        request_id = str(uuid.uuid4())
        start_time = time.time()

        # === AUTH CHECK ===
        if FeatureFlags.is_auth_enabled():
            # Extract API key from headers
            headers_dict = dict(request.headers)
            api_key = extract_api_key(headers_dict)

            if not api_key:
                return JSONResponse(
                    status_code=401,
                    content=create_error_response(
                        code=ErrorCode.INVALID_API_KEY,
                        message="API key required. Use 'Authorization: Bearer <key>' or 'X-API-Key: <key>' header.",
                        request_id=request_id,
                    ),
                    headers={"X-Request-ID": request_id},
                )

            # Validate API key
            validator = get_api_validator()
            key_hash = hash_api_key(api_key)
            auth_result = validator.validate(key_hash)

            if not auth_result.is_valid:
                status_code = 401 if auth_result.error_code == "INVALID_API_KEY" else 403
                return JSONResponse(


sed -n '340,760p' veena3modal/api/fastapi_app.py


Output:
Chunk ID: dce3e1
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 4507
Output:
                return JSONResponse(
                    status_code=status_code,
                    content=create_error_response(
                        code=ErrorCode(auth_result.error_code),
                        message=auth_result.error_message or "Authentication failed",
                        request_id=request_id,
                    ),
                    headers={"X-Request-ID": request_id},
                )

            # === RATE LIMIT CHECK ===
            if FeatureFlags.is_rate_limiting_enabled():
                rate_limiter = get_rate_limiter()
                allowed, remaining, reset_after = rate_limiter.check(key_hash)

                if not allowed:
                    headers = rate_limiter.get_headers(allowed, remaining, reset_after)
                    headers["X-Request-ID"] = request_id
                    return JSONResponse(
                        status_code=429,
                        content=create_error_response(
                            code=ErrorCode.RATE_LIMIT_EXCEEDED,
                            message=f"Rate limit exceeded. Retry after {int(reset_after) + 1} seconds.",
                            request_id=request_id,
                        ),
                        headers=headers,
                    )

        # Parse and validate request body
        try:
            body = await request.json()
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}")

        try:
            req = TTSGenerateRequest(**body)
        except Exception as e:
            error_msg = str(e)
            # Extract validation error details
            return JSONResponse(
                status_code=400,
                content={
                    "error": {
                        "code": "VALIDATION_ERROR",
                        "message": error_msg,
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Check if runtime is initialized
        if not tts_runtime.is_initialized():
            return JSONResponse(
                status_code=503,
                content={
                    "error": {
                        "code": "MODEL_NOT_LOADED",
                        "message": "TTS model not initialized. Please wait for warmup.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Get resolved speaker
        speaker = req.get_resolved_speaker()

        # Get normalized text (apply text normalization if enabled)
        normalizer_func = None
        if req.normalize:
            try:
                from veena3modal.processing.text_normalizer import normalize_text
                normalizer_func = lambda t: normalize_text(t, verbose=req.normalize_verbose)
            except ImportError:
                logger.warning("TextNormalizer not available, skipping text normalization")

        text = req.get_normalized_text(normalizer_func=normalizer_func)

        # Set request context for structured logging
        set_request_context(request_id=request_id)

        # Log request received (lifecycle event + metrics)
        log_request_received(
            request_id=request_id,
            text_length=len(text),
            speaker=speaker,
            stream=req.stream,
            format=req.format,
        )
        metrics_request_received(speaker=speaker, stream=req.stream, format=req.format)

        # Streaming mode (M4)
        if req.stream:
            return await _handle_streaming_request(req, request_id, text, speaker, start_time)

        # Format encoding will be applied after generation (non-streaming only)
        # For streaming, only WAV is supported
        target_format = req.format

        try:
            # Get output sample rate - use_enum_values=True means req.output is already a string
            output_sr = req.output if hasattr(req, 'output') else "16khz"

            # Generate speech (non-streaming)
            if req.chunking:
                audio_bytes, metrics = await tts_runtime.generate_speech_chunked(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    sample_rate=req.sample_rate,
                    output_sample_rate=output_sr,
                    request_id=request_id,
                )
            else:
                audio_bytes, metrics = await tts_runtime.generate_speech(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    output_sample_rate=output_sr,
                    request_id=request_id,
                )

            if audio_bytes is None:
                return JSONResponse(
                    status_code=500,
                    content={
                        "error": {
                            "code": "GENERATION_FAILED",
                            "message": "Audio generation failed. Check logs for details.",
                            "request_id": request_id,
                        }
                    },
                    headers={"X-Request-ID": request_id},
                )

            # Calculate RTF
            total_time = time.time() - start_time
            audio_duration = metrics.get("audio_duration_seconds", 0)
            rtf = total_time / audio_duration if audio_duration > 0 else 0

            # Build response headers
            runtime = tts_runtime.get_runtime()
            model_version = runtime.model_version if runtime else get_model_version()

            # Use actual output sample rate from metrics
            output_sample_rate = metrics.get("output_sample_rate", req.sample_rate)

            headers = {
                "X-Request-ID": request_id,
                "X-Model-Version": model_version,
                "X-Format": "wav",
                "X-Sample-Rate": str(output_sample_rate),
                "X-Stream": "false",
                "X-Audio-Bytes": str(len(audio_bytes)),
                "X-Audio-Seconds": f"{audio_duration:.2f}",
                "X-TTFB-ms": str(metrics.get("ttfb_ms", 0)),
                "X-RTF": f"{rtf:.3f}",
                "X-Text-Chunked": str(metrics.get("text_chunked", False)).lower(),
                "X-SR-Applied": str(metrics.get("sr_applied", False)).lower(),
            }
            headers.update(_build_perf_headers(metrics))

            if req.seed is not None:
                headers["X-Seed"] = str(req.seed)

            # Fire-and-forget sentence storage (non-blocking, after generation)
            sentence_store = _get_sentence_store()
            sentence_store.store_fire_and_forget(
                request_id=request_id,
                text=text,
                speaker=speaker,
                stream=False,
                format=req.format,
                temperature=req.temperature,
                top_k=req.top_k,
                top_p=req.top_p,
                max_tokens=req.max_tokens,
                repetition_penalty=req.repetition_penalty,
                seed=req.seed,
                text_chunked=metrics.get("text_chunked", False),
                ttfb_ms=metrics.get("ttfb_ms"),
                audio_duration_seconds=audio_duration,
            )

            # Log successful completion (lifecycle event + metrics)
            total_duration_ms = int(total_time * 1000)
            log_request_completed(
                request_id=request_id,
                status_code=200,
                total_duration_ms=total_duration_ms,
                audio_duration_seconds=audio_duration,
                rtf=rtf,
                chunks_sent=1,
                **metrics,
            )
            metrics_request_completed(
                status_code=200,
                duration_seconds=total_time,
                speaker=speaker,
                stream=False,
            )
            record_ttfb(ttfb_seconds=metrics.get("ttfb_ms", 0) / 1000, speaker=speaker, stream=False)
            record_rtf(rtf=rtf, speaker=speaker)
            record_audio_duration(duration_seconds=audio_duration, speaker=speaker)

            clear_request_context()

            # Apply format encoding if not WAV
            media_type = "audio/wav"
            if target_format != AudioFormat.WAV.value:
                try:
                    from veena3modal.audio.encoder import (
                        OpusEncoder, MP3Encoder, MuLawEncoder, FLACEncoder
                    )

                    # Get sample rate from output (after potential SR)
                    sample_rate = metrics.get("output_sample_rate", 16000)

                    # Strip WAV header to get PCM data
                    pcm_data = audio_bytes[44:] if len(audio_bytes) > 44 else audio_bytes

                    if target_format == AudioFormat.OPUS.value:
                        encoder = OpusEncoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/opus"
                    elif target_format == AudioFormat.MP3.value:
                        encoder = MP3Encoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/mpeg"
                    elif target_format == AudioFormat.MULAW.value:
                        encoder = MuLawEncoder()
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/x-wav"
                    elif target_format == AudioFormat.FLAC.value:
                        encoder = FLACEncoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/flac"

                    headers["X-Audio-Bytes"] = str(len(audio_bytes))
                    logger.info(f"Encoded audio to {target_format}: {len(audio_bytes)} bytes")
                except ImportError as e:
                    logger.warning(f"Audio encoder not available: {e}")
                    # Fall back to WAV
                except Exception as e:
                    logger.error(f"Format encoding failed: {e}")
                    # Fall back to WAV

            return Response(
                content=audio_bytes,
                media_type=media_type,
                headers=headers,
            )

        except Exception as e:
            logger.exception(f"TTS generation error: {e}")

            # Log failure (lifecycle event + metrics)
            log_request_failed(
                request_id=request_id,
                status_code=500,
                error_code="INTERNAL_ERROR",
                error_message=str(e),
            )
            metrics_request_failed(
                status_code=500,
                error_code="INTERNAL_ERROR",
                speaker=speaker,
            )
            clear_request_context()

            return JSONResponse(
                status_code=500,
                content={
                    "error": {
                        "code": "INTERNAL_ERROR",
                        "message": str(e),
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

    async def _handle_streaming_request(req, request_id: str, text: str, speaker: str, start_time: float):
        """
        Handle streaming TTS request.

        Returns a StreamingResponse that yields WAV header + PCM chunks.
        True streaming: first bytes sent ASAP before full audio is generated.
        """
        from fastapi.responses import StreamingResponse
        from veena3modal.api.schemas import AudioFormat
        from veena3modal.services import tts_runtime

        # Check if streaming pipeline is available
        runtime = tts_runtime.get_runtime()
        if runtime is None or runtime.streaming_pipeline is None:
            return JSONResponse(
                status_code=503,
                content={
                    "error": {
                        "code": "STREAMING_UNAVAILABLE",
                        "message": "Streaming pipeline not initialized.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Non-WAV formats not implemented yet (M5)
        if req.format != AudioFormat.WAV.value:
            return JSONResponse(
                status_code=501,
                content={
                    "error": {
                        "code": "FORMAT_NOT_IMPLEMENTED",
                        "message": f"Streaming format '{req.format}' not yet implemented. Use format=wav.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Set request context for structured logging
        set_request_context(request_id=request_id)

        # Create async generator for streaming response
        async def audio_stream_generator():
            """Yields audio chunks as they're generated."""
            final_metrics = {}
            first_chunk_sent = False
            chunks_count = 0
            try:
                async for audio_chunk, metrics in tts_runtime.generate_speech_streaming(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    enable_chunking=req.chunking,
                    request_id=request_id,
                ):
                    final_metrics = metrics
                    chunks_count += 1

                    # Fire-and-forget sentence storage after first chunk (don't block TTFB)
                    if not first_chunk_sent:
                        first_chunk_sent = True
                        ttfb_ms = metrics.get("ttfb_ms", 0)

                        # Log first audio emitted (TTFB marker)
                        log_first_audio_emitted(
                            request_id=request_id,
                            ttfb_ms=ttfb_ms,
                            chunk_size_bytes=len(audio_chunk),
                        )
                        record_ttfb(ttfb_seconds=ttfb_ms / 1000, speaker=speaker, stream=True)

                        sentence_store = _get_sentence_store()
                        sentence_store.store_fire_and_forget(
                            request_id=request_id,
                            text=text,
                            speaker=speaker,
                            stream=True,
                            format=req.format,
                            temperature=req.temperature,
                            top_k=req.top_k,
                            top_p=req.top_p,
                            max_tokens=req.max_tokens,
                            repetition_penalty=req.repetition_penalty,
                            seed=req.seed,
                            text_chunked=metrics.get("text_chunked", False),
                            ttfb_ms=ttfb_ms,
                            # audio_duration_seconds not known yet for streaming
                        )

                    yield audio_chunk

                # Stream completed - log metrics
                total_time = time.time() - start_time
                audio_duration = final_metrics.get("audio_duration_seconds", 0)
                rtf = total_time / audio_duration if audio_duration > 0 else 0

                log_request_completed(
                    request_id=request_id,
                    status_code=200,
                    total_duration_ms=int(total_time * 1000),
                    audio_duration_seconds=audio_duration,
                    rtf=rtf,
                    chunks_sent=chunks_count,
                    **final_metrics,
                )
                metrics_request_completed(
                    status_code=200,
                    duration_seconds=total_time,
                    speaker=speaker,
                    stream=True,
                )
                record_rtf(rtf=rtf, speaker=speaker)
                record_audio_duration(duration_seconds=audio_duration, speaker=speaker)
                record_chunks_sent(chunks=chunks_count, speaker=speaker)
                clear_request_context()
            except Exception as e:
                logger.exception(f"Streaming error: {e}")
                # Log streaming failure
                log_request_failed(
                    request_id=request_id,
                    status_code=500,


sed -n '760,980p' veena3modal/api/fastapi_app.py


Output:
Chunk ID: 58417b
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 431
Output:
                    status_code=500,
                    error_code="STREAMING_ERROR",
                    error_message=str(e),
                )
                metrics_request_failed(
                    status_code=500,
                    error_code="STREAMING_ERROR",
                    speaker=speaker,
                )
                clear_request_context()
                # Can't return error once streaming started
                # Best we can do is stop yielding
                return

        # Build response headers
        # Note: Some headers can't be set dynamically during streaming
        # TTFB and RTF will be logged but not in headers for streaming
        model_version = runtime.model_version if runtime else get_model_version()

        headers = {
            "X-Request-ID": request_id,
            "X-Model-Version": model_version,
            "X-Format": "wav",
            "X-Sample-Rate": str(req.sample_rate),
            "X-Stream": "true",
            "X-Chunking-Enabled": str(req.chunking).lower(),
            # Transfer-Encoding: chunked is set automatically by StreamingResponse
        }

        if req.seed is not None:
            headers["X-Seed"] = str(req.seed)

        return StreamingResponse(
            audio_stream_generator(),
            media_type="audio/wav",
            headers=headers,
        )

    # === WebSocket Support (M5b) ===
    try:
        from veena3modal.api.websocket_handler import add_websocket_routes
        add_websocket_routes(app)
        logger.info("WebSocket support enabled at /v1/tts/ws")
    except ImportError as e:
        logger.warning(f"WebSocket support not available: {e}")

    return app


sed -n '1,260p' veena3modal/services/tts_runtime.py


Output:
Chunk ID: 9e596d
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2375
Output:
"""
TTS runtime wrapper (container-scoped singleton).

In Modal, each container should load the model once and then handle many requests.
This wrapper owns:
- SparkTTSModel (vLLM engine)
- IndicPromptBuilder
- BiCodecDecoder
- SparkTTSPipeline + Veena3SlidingWindowPipeline
- optional SuperResolutionService

Imports framework-agnostic inference + processing code from `veena3modal/core`,
`veena3modal/processing`, and `veena3modal/audio`.
"""

from __future__ import annotations

import os
import sys
import time
import asyncio
import logging
from threading import Lock
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

_PROFILE_REQUEST_TIMING = "VEENA3_TIMING_LOGS"
_inflight_counter = 0
_inflight_lock = Lock()

def _timing_enabled() -> bool:
    """Return True when detailed request timing should be emitted."""
    return os.getenv(_PROFILE_REQUEST_TIMING, "").lower() in {"1", "true", "yes", "on"}

def _generation_timeout_seconds() -> float:
    """Optional hard timeout (seconds) for each generation call."""
    raw = os.getenv("VEENA3_REQUEST_TIMEOUT_SECONDS", "").strip()
    if not raw:
        return 0.0
    try:
        value = float(raw)
    except ValueError:
        return 0.0
    return max(0.0, value)

def _bicodec_device(default_device: str = "cuda") -> str:
    """Resolve BiCodec device from env override and return the effective device."""
    raw = os.getenv("BICODEC_DEVICE", "").strip().lower()
    if not raw or raw == "default":
        return default_device
    if raw not in {"cuda", "cpu"}:
        return default_device
    return raw

def _with_inflight_delta(delta: int) -> int:
    """Update in-flight request counter and return the new value."""
    global _inflight_counter
    with _inflight_lock:
        _inflight_counter = max(0, _inflight_counter + delta)
        return _inflight_counter

@dataclass
class TTSRuntime:
    """
    Holds long-lived, per-container inference objects.

    Thread Safety:
        This class is designed for async concurrency within a single container.
        The vLLM engine handles internal batching and scheduling.
        Do NOT share across processes without proper synchronization.
    """
    model: Any = None
    pipeline: Any = None
    streaming_pipeline: Any = None
    prompt_builder: Any = None
    bicodec_decoder: Any = None
    sr_service: Optional[Any] = None

    model_version: str = "not_loaded"
    is_loaded: bool = False
    load_time_ms: float = 0.0

    # Configuration
    model_path: str = ""
    bicodec_path: str = ""
    sr_checkpoint_dir: Optional[str] = None
    device: str = "cuda"

    # OPTIMIZATION: Pre-computed global tokens per speaker (eliminates ~110ms pre-roll in streaming)
    # Maps speaker_name -> list of 32 global token IDs captured at startup
    speaker_global_cache: Dict[str, List[int]] = field(default_factory=dict)

# Module-level singleton (per-container)
_runtime: Optional[TTSRuntime] = None

def get_runtime() -> Optional[TTSRuntime]:
    """Get the current TTS runtime singleton."""
    return _runtime

def is_initialized() -> bool:
    """Check if the runtime is initialized and ready."""
    return _runtime is not None and _runtime.is_loaded

def initialize_runtime(
    model_path: Optional[str] = None,
    bicodec_path: Optional[str] = None,
    sr_checkpoint_dir: Optional[str] = None,
    device: str = "cuda",
    hf_token: Optional[str] = None,
    gpu_memory_utilization: float = 0.25,
    enable_sr: bool = False,
    num_engines: int = 1,
    bicodec_device: Optional[str] = None,
) -> TTSRuntime:
    """
    Initialize the TTS runtime with all components.

    This should be called once per container (e.g., in Modal's @modal.enter).

    Args:
        model_path: Path to Spark TTS model (env: SPARK_TTS_MODEL_PATH)
        bicodec_path: Path to BiCodec model (env: BICODEC_MODEL_PATH, defaults to model_path)
        sr_checkpoint_dir: Path to super-resolution checkpoints (env: AP_BWE_CHECKPOINT_DIR)
        device: Device for inference (cuda/cpu)
        hf_token: HuggingFace token for private models (env: HF_TOKEN)
        gpu_memory_utilization: vLLM GPU memory fraction (default: 0.25)
        enable_sr: Enable super-resolution service
        num_engines: Number of vLLM engine instances (default: 1, set to 2-3 for Tier 3 optimization)

    Returns:
        Initialized TTSRuntime instance

    Raises:
        RuntimeError: If initialization fails
    """
    global _runtime

    start_time = time.time()

    # Resolve paths from env vars if not provided
    model_path = model_path or os.environ.get(
        'SPARK_TTS_MODEL_PATH',
        os.environ.get('MODEL_PATH', '/models/spark_tts_4speaker')
    )
    bicodec_device = bicodec_device or _bicodec_device(default_device=device)
    bicodec_path = bicodec_path or os.environ.get(
        'BICODEC_MODEL_PATH',
        model_path  # BiCodec is usually in the same directory
    )
    sr_checkpoint_dir = sr_checkpoint_dir or os.environ.get('AP_BWE_CHECKPOINT_DIR')
    hf_token = hf_token or os.environ.get('HF_TOKEN') or os.environ.get('HUGGING_FACE_HUB_TOKEN')

    logger.info(f"Initializing TTS runtime...")
    logger.info(f"  Model path: {model_path}")
    logger.info(f"  BiCodec path: {bicodec_path}")
    logger.info(f"  BiCodec device: {bicodec_device}")
    logger.info(f"  SR enabled: {enable_sr}, path: {sr_checkpoint_dir}")

    try:
        # Import inference components (vendored into veena3modal for Modal deployment)
        from veena3modal.core.model_loader import SparkTTSModel
        from veena3modal.core.pipeline import SparkTTSPipeline
        from veena3modal.core.bicodec_decoder import BiCodecDecoder
        from veena3modal.processing.prompt_builder import IndicPromptBuilder

        # Load SparkTTS model with vLLM engine
        if num_engines > 1:
            # TIER 3 OPTIMIZATION: Multiple vLLM engines on same GPU
            # Each engine gets gpu_memory_utilization / num_engines to share GPU fairly.
            # E.g., 3 engines at 0.08 each = 0.24 total, ~8GB per engine, ~24GB total.
            from veena3modal.core.multi_engine import create_multi_engine_model
            per_engine_mem = gpu_memory_utilization / num_engines
            logger.info(f"Loading {num_engines} vLLM engines ({per_engine_mem:.2f} GPU mem each)...")
            model = create_multi_engine_model(
                model_path=model_path,
                num_engines=num_engines,
                hf_token=hf_token,
                gpu_memory_per_engine=per_engine_mem,
            )
        else:
            logger.info("Loading SparkTTS model with vLLM...")
            model = SparkTTSModel(
                model_path=model_path,
                hf_token=hf_token,
                gpu_memory_utilization=gpu_memory_utilization,
            )

        # Initialize prompt builder
        logger.info("Initializing prompt builder...")
        prompt_builder = IndicPromptBuilder(
            tokenizer=model.tokenizer,
            model=model,
        )

        # Initialize BiCodec decoder
        logger.info("Initializing BiCodec decoder...")
        bicodec_decoder = BiCodecDecoder(
            device=bicodec_device,
            model_path=bicodec_path,
        )

        # Initialize pipeline
        logger.info("Initializing TTS pipeline...")
        pipeline = SparkTTSPipeline(
            model=model,
            prompt_builder=prompt_builder,
            bicodec_decoder=bicodec_decoder,
        )

        # Initialize streaming pipeline (for M4)
        streaming_pipeline = None
        try:
            from veena3modal.core.streaming_pipeline import Veena3SlidingWindowPipeline
            logger.info("Initializing streaming pipeline...")
            # NOTE: Parameter is named 'snac_decoder' for legacy reasons, but works with BiCodecDecoder
            streaming_pipeline = Veena3SlidingWindowPipeline(
                model=model,
                prompt_builder=prompt_builder,
                snac_decoder=bicodec_decoder,  # BiCodecDecoder is interface-compatible
            )
        except ImportError as e:
            logger.warning(f"Streaming pipeline not available: {e}")

        # Initialize super-resolution (optional)
        sr_service = None
        if enable_sr and sr_checkpoint_dir:
            try:
                from veena3modal.core.super_resolution import SuperResolutionService
                logger.info(f"Initializing super-resolution from {sr_checkpoint_dir}...")
                sr_service = SuperResolutionService(checkpoint_dir=sr_checkpoint_dir)
                # Load the model explicitly
                if sr_service.load_model(device=device):
                    logger.info("✅ Super-resolution model loaded successfully")
                else:
                    logger.warning("Super-resolution model failed to load")
                    sr_service = None
            except Exception as e:
                logger.warning(f"Super-resolution not available: {e}")

        # Determine model version
        model_version = os.path.basename(model_path.rstrip('/'))
        if not model_version:
            model_version = "spark-tts"

        load_time = (time.time() - start_time) * 1000

        # Create runtime
        _runtime = TTSRuntime(
            model=model,


sed -n '260,620p' veena3modal/services/tts_runtime.py


Output:
Chunk ID: a7b299
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 3454
Output:
            model=model,
            pipeline=pipeline,
            streaming_pipeline=streaming_pipeline,
            prompt_builder=prompt_builder,
            bicodec_decoder=bicodec_decoder,
            sr_service=sr_service,
            model_version=model_version,
            is_loaded=True,
            load_time_ms=load_time,
            model_path=model_path,
            bicodec_path=bicodec_path,
            sr_checkpoint_dir=sr_checkpoint_dir,
            device=device,
        )

        logger.info(f"✅ TTS runtime initialized in {load_time:.0f}ms")
        logger.info(f"   Model version: {model_version}")

        # OPTIMIZATION: Pre-compute global tokens for all speakers
        # Disabled with multi-engine to avoid async event loop conflicts
        if num_engines <= 1:
            _precompute_speaker_globals(_runtime)
        else:
            logger.info("Speaker globals cache skipped (multi-engine mode)")

        # Update FastAPI app with model version
        try:
            from veena3modal.api.fastapi_app import set_model_version
            set_model_version(model_version)
        except ImportError:
            pass

        return _runtime

    except Exception as e:
        logger.error(f"❌ Failed to initialize TTS runtime: {e}")
        import traceback
        traceback.print_exc()
        raise RuntimeError(f"TTS runtime initialization failed: {e}") from e

def _precompute_speaker_globals(runtime: TTSRuntime) -> None:
    """
    Pre-compute and cache global tokens for all 12 speakers at startup.

    OPTIMIZATION: BiCodec streaming requires 32 "global tokens" before any audio can be emitted.
    These encode speaker identity via FSQ quantization. Since we have only 12 fixed speakers,
    we generate one short utterance per speaker at startup, capture the 32 global tokens,
    and cache them. During streaming, we inject cached globals via build_prefix_with_globals(),
    skipping the ~110ms global token pre-roll phase entirely.

    This runs synchronously at startup (adds ~5-10s to cold start, saves ~110ms per streaming TTFB).
    """
    import asyncio
    import re
    from vllm import SamplingParams
    from veena3modal.core.constants import (
        SPEAKER_MAP, TRAINING_STOP_TOKEN_IDS,
        DEFAULT_TEMPERATURE, DEFAULT_TOP_K, DEFAULT_TOP_P,
    )

    logger.info("Pre-computing global tokens for all speakers...")
    t_start = time.time()

    test_text = "Hello, this is a voice test."

    sampling_params = SamplingParams(
        temperature=DEFAULT_TEMPERATURE,
        top_k=DEFAULT_TOP_K,
        top_p=DEFAULT_TOP_P,
        max_tokens=128,  # Only need ~60 tokens (32 global + some semantic)
        stop=TRAINING_STOP_TOKEN_IDS,
        skip_special_tokens=False,
    )

    async def _capture_globals_for_speaker(speaker_name: str) -> Optional[List[int]]:
        """Generate a short utterance and capture the 32 global tokens."""
        prompt = runtime.prompt_builder.build_prefix(speaker_name, test_text)
        request_id = f"warmup-{speaker_name}-{int(time.time() * 1000)}"

        final_output = None
        async for request_output in runtime.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        ):
            final_output = request_output

        if final_output is None:
            return None

        generated_text = final_output.outputs[0].text
        global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", generated_text)
        global_ids = [int(t) for t in global_matches]

        if len(global_ids) >= 32:
            return global_ids[:32]
        return None

    async def _run_all():
        for speaker_name in SPEAKER_MAP.keys():
            try:
                global_ids = await _capture_globals_for_speaker(speaker_name)
                if global_ids and len(global_ids) == 32:
                    runtime.speaker_global_cache[speaker_name] = global_ids
                    logger.info(f"  Cached globals for {speaker_name}")
                else:
                    logger.warning(f"  Failed to capture globals for {speaker_name}")
            except Exception as e:
                logger.warning(f"  Error caching globals for {speaker_name}: {e}")

    # Run async warmup in event loop
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # Already in async context (unlikely at startup, but handle it)
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as pool:
                pool.submit(lambda: asyncio.run(_run_all())).result()
        else:
            loop.run_until_complete(_run_all())
    except RuntimeError:
        # No event loop exists yet
        asyncio.run(_run_all())

    elapsed = (time.time() - t_start) * 1000
    cached = len(runtime.speaker_global_cache)
    total = len(SPEAKER_MAP)
    logger.info(f"Speaker globals cached: {cached}/{total} speakers in {elapsed:.0f}ms")

async def generate_speech(
    text: str,
    speaker: str,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 1.0,
    max_tokens: int = 4096,
    repetition_penalty: float = 1.05,
    seed: Optional[int] = None,
    output_sample_rate: str = "16khz",
    request_id: Optional[str] = None,
) -> Tuple[Optional[bytes], Dict[str, Any]]:
    """
    Generate speech audio (non-streaming).

    Args:
        text: Text to synthesize (already normalized)
        speaker: Internal speaker name (resolved)
        temperature: Sampling temperature
        top_k: Top-k sampling
        top_p: Nucleus sampling
        max_tokens: Maximum tokens to generate
        repetition_penalty: Repetition penalty
        seed: Random seed for reproducibility
        output_sample_rate: "16khz" or "48khz" (triggers super-resolution)

    Returns:
        Tuple of (audio_bytes, metrics_dict)
        audio_bytes: WAV audio bytes (16kHz or 48kHz, 16-bit PCM) or None if failed
        metrics_dict: Dictionary with timing metrics

    Raises:
        RuntimeError: If runtime not initialized
    """
    if not is_initialized():
        raise RuntimeError("TTS runtime not initialized")

    runtime = get_runtime()
    metrics = {
        "request_id": request_id,
        "ttfb_ms": 0,
        "generation_ms": 0,
        "tokens_generated": 0,
        "audio_duration_seconds": 0.0,
        "sr_applied": False,
        "output_sample_rate": 16000,
    }

    if request_id is None:
        import uuid
        request_id = f"runtime-ns-{uuid.uuid4().hex[:10]}"
        metrics["request_id"] = request_id

    active_requests = _with_inflight_delta(1)
    start_time = time.time()

    pipeline_timings = {}

    try:
        logger.info(
            "tts_runtime_nonstream_start",
            extra={
                "speaker": speaker,
                "request_id": request_id,
                "text_length": len(text),
                "request_inflight": active_requests,
                "temperature": temperature,
                "top_k": top_k,
                "top_p": top_p,
                "max_tokens": max_tokens,
            },
        )

        generation_timeout = _generation_timeout_seconds()

        gen_kwargs = dict(
            speaker=speaker,
            text=text,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            repetition_penalty=repetition_penalty,
            seed=seed,
        )

        if _timing_enabled():
            generation_coro = runtime.pipeline.generate_speech_indic(
                **gen_kwargs,
                return_timing=True,
                request_id=request_id,
            )
            if generation_timeout:
                try:
                    audio_bytes, pipeline_timings = await asyncio.wait_for(
                        generation_coro,
                        timeout=generation_timeout,
                    )
                except asyncio.TimeoutError as exc:
                    raise TimeoutError(
                        f"Non-stream generation timed out after {generation_timeout:.2f}s"
                    ) from exc
            else:
                audio_bytes, pipeline_timings = await generation_coro
        else:
            generation_coro = runtime.pipeline.generate_speech_indic(
                **gen_kwargs,
                request_id=request_id,
            )
            if generation_timeout:
                try:
                    audio_bytes = await asyncio.wait_for(
                        generation_coro,
                        timeout=generation_timeout,
                    )
                except asyncio.TimeoutError as exc:
                    raise TimeoutError(
                        f"Non-stream generation timed out after {generation_timeout:.2f}s"
                    ) from exc
            else:
                audio_bytes = await generation_coro

        generation_time = time.time() - start_time
        metrics["generation_ms"] = int(generation_time * 1000)
        metrics["ttfb_ms"] = metrics["generation_ms"]  # Non-streaming: TTFB = total time

        if pipeline_timings:
            metrics.update({
                "llm_token_total": pipeline_timings.get("token_total", 0),
                "llm_batch_count": pipeline_timings.get("batch_count", 0),
                "llm_batch_gpu_ms_total": pipeline_timings.get("llm_gpu_batch_ms_total", 0),
                "llm_batch_gpu_ms_min": pipeline_timings.get("llm_gpu_batch_ms_min", 0),
                "llm_batch_gpu_ms_max": pipeline_timings.get("llm_gpu_batch_ms_max", 0),
                "llm_batch_gpu_ms_p50": pipeline_timings.get("llm_gpu_batch_ms_p50", 0),
                "llm_decode_cpu_ms": pipeline_timings.get("decode_cpu_ms", 0),
                "llm_parse_ms": pipeline_timings.get("parse_ms", 0),
                "llm_wall_ms": pipeline_timings.get("llm_wall_ms", metrics["generation_ms"]),
                "ttfb_ms": int(pipeline_timings.get("ttfb_ms", metrics["ttfb_ms"])),
                "tokens_per_batch_min": pipeline_timings.get("tokens_per_batch_min"),
                "tokens_per_batch_max": pipeline_timings.get("tokens_per_batch_max"),
                "tokens_per_batch_p50": pipeline_timings.get("tokens_per_batch_p50"),
            })

        if audio_bytes:
            # Apply super-resolution if requested and available
            sample_rate = 16000
            logger.info(f"SR check: output_sample_rate={output_sample_rate!r}, "
                       f"sr_service={runtime.sr_service is not None}, "
                       f"is_loaded={runtime.sr_service.is_loaded if runtime.sr_service else False}")

            if output_sample_rate == "48khz" and runtime.sr_service and runtime.sr_service.is_loaded:
                try:
                    sr_start = time.time()
                    logger.info(f"Applying super-resolution to {len(audio_bytes)} bytes...")
                    audio_bytes = _apply_super_resolution(audio_bytes, runtime.sr_service)
                    sr_time = (time.time() - sr_start) * 1000
                    metrics["sr_ms"] = int(sr_time)
                    metrics["sr_applied"] = True
                    sample_rate = 48000
                    logger.info(f"✅ Super-resolution applied in {sr_time:.1f}ms, output={len(audio_bytes)} bytes")
                except Exception as e:
                    logger.warning(f"Super-resolution failed, returning 16kHz: {e}")
                    import traceback
                    traceback.print_exc()

            metrics["output_sample_rate"] = sample_rate

            # Calculate audio duration
            audio_duration = (len(audio_bytes) - 44) / (sample_rate * 2)  # -44 for WAV header
            metrics["audio_duration_seconds"] = max(0.0, audio_duration)

        return audio_bytes, metrics

    except Exception as e:
        logger.error(f"Speech generation failed: {e}")
        logger.exception(
            "tts_runtime_nonstream_error",
            extra={
                "speaker": speaker,
                "request_id": request_id,
                "error": str(e),
            },
        )
        raise
    finally:
        active_requests = _with_inflight_delta(-1)
        if _timing_enabled():
            metrics["request_inflight"] = active_requests
            metrics["active_requests"] = active_requests
            logger.info(
                "tts_runtime_nonstream_complete",
                extra=metrics,
            )
        if _timing_enabled():
            logger.info(
                "tts_runtime_nonstream_inflight",
                extra={
                    "request_id": request_id,
                    "active_requests": active_requests,
                },
            )

def _apply_super_resolution(audio_bytes: bytes, sr_service) -> bytes:
    """
    Apply super-resolution to audio bytes.

    Args:
        audio_bytes: WAV audio at 16kHz
        sr_service: SuperResolutionService instance

    Returns:
        WAV audio at 48kHz
    """
    import numpy as np
    import struct
    import torch

    # Parse WAV header and extract PCM data
    if len(audio_bytes) < 44:
        raise ValueError("Invalid WAV data")

    # Extract PCM data (skip 44-byte WAV header)
    pcm_data = np.frombuffer(audio_bytes[44:], dtype=np.int16)

    # Convert to float32 for SR model
    audio_float = pcm_data.astype(np.float32) / 32768.0

    # Convert numpy to torch tensor for process_chunk
    audio_tensor = torch.from_numpy(audio_float)


rg -n "pipeline_timings|llm_batch|time_per|timing" veena3modal/services/tts_runtime.py


Output:
Chunk ID: a369cc
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 957
Output:
34:def _timing_enabled() -> bool:
35:    """Return True when detailed request timing should be emitted."""
420:        metrics_dict: Dictionary with timing metrics
447:    pipeline_timings = {}
477:        if _timing_enabled():
480:                return_timing=True,
485:                    audio_bytes, pipeline_timings = await asyncio.wait_for(
494:                audio_bytes, pipeline_timings = await generation_coro
517:        if pipeline_timings:
519:                "llm_token_total": pipeline_timings.get("token_total", 0),
520:                "llm_batch_count": pipeline_timings.get("batch_count", 0),
521:                "llm_batch_gpu_ms_total": pipeline_timings.get("llm_gpu_batch_ms_total", 0),
522:                "llm_batch_gpu_ms_min": pipeline_timings.get("llm_gpu_batch_ms_min", 0),
523:                "llm_batch_gpu_ms_max": pipeline_timings.get("llm_gpu_batch_ms_max", 0),
524:                "llm_batch_gpu_ms_p50": pipeline_timings.get("llm_gpu_batch_ms_p50", 0),
525:                "llm_decode_cpu_ms": pipeline_timings.get("decode_cpu_ms", 0),
526:                "llm_parse_ms": pipeline_timings.get("parse_ms", 0),
527:                "llm_wall_ms": pipeline_timings.get("llm_wall_ms", metrics["generation_ms"]),
528:                "ttfb_ms": int(pipeline_timings.get("ttfb_ms", metrics["ttfb_ms"])),
529:                "tokens_per_batch_min": pipeline_timings.get("tokens_per_batch_min"),
530:                "tokens_per_batch_max": pipeline_timings.get("tokens_per_batch_max"),
531:                "tokens_per_batch_p50": pipeline_timings.get("tokens_per_batch_p50"),
577:        if _timing_enabled():
584:        if _timing_enabled():
710:        if _timing_enabled():
730:            if _timing_enabled():
742:                    return_timing=True,
745:                    audio_bytes, chunk_timing = chunk_result
748:                    chunk_timing = {}
762:                chunk_timing = {}
775:            if _timing_enabled():
778:                    return_timing=True,
783:                        audio_bytes, pipeline_timings = await asyncio.wait_for(
792:                    audio_bytes, pipeline_timings = await generation_coro
793:                chunk_timing = pipeline_timings
811:                chunk_timing = {}
817:        if chunk_timing:
819:                "llm_token_total": chunk_timing.get("token_total", 0),
820:                "llm_batch_count": chunk_timing.get("batch_count", 0),
821:                "llm_batch_gpu_ms_total": chunk_timing.get("llm_gpu_batch_ms_total", 0),
822:                "llm_batch_wall_ms_total": chunk_timing.get("llm_wall_ms", metrics["generation_ms"]),
823:                "llm_parse_ms": chunk_timing.get("parse_ms", 0),
824:                "llm_decode_cpu_ms": chunk_timing.get("decode_cpu_ms", 0),
825:                "llm_batch_wall_ms_min": chunk_timing.get("llm_batch_wall_ms_min", 0),
826:                "llm_batch_wall_ms_max": chunk_timing.get("llm_batch_wall_ms_max", 0),
827:                "llm_batch_gpu_ms_min": chunk_timing.get("llm_batch_gpu_ms_min", 0),
828:                "llm_batch_gpu_ms_max": chunk_timing.get("llm_batch_gpu_ms_max", 0),
829:                "llm_batch_gpu_ms_p50": chunk_timing.get("llm_batch_gpu_ms_p50", 0),
830:                "tokens_per_batch_min": chunk_timing.get("tokens_per_batch_min"),
831:                "tokens_per_batch_max": chunk_timing.get("tokens_per_batch_max"),
832:                "tokens_per_batch_p50": chunk_timing.get("tokens_per_batch_p50"),
833:                "llm_batch_tokens": chunk_timing.get("token_batches", []),
836:                metrics["chunks_processed"] = chunk_timing.get("chunks_processed", 0)
860:        if _timing_enabled():
912:        - Final yield: metrics_dict has final timing info
955:    if _timing_enabled():
1058:        if _timing_enabled():


sed -n '660,940p' veena3modal/services/tts_runtime.py


Output:
Chunk ID: d4a3a3
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 2797
Output:
async def generate_speech_chunked(
    text: str,
    speaker: str,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 1.0,
    max_tokens: int = 4096,
    repetition_penalty: float = 1.05,
    seed: Optional[int] = None,
    sample_rate: int = 16000,
    output_sample_rate: str = "16khz",
    request_id: Optional[str] = None,
) -> Tuple[Optional[bytes], Dict[str, Any]]:
    """
    Generate speech with automatic text chunking for long inputs.

    Uses LongTextProcessor to split text and stitch audio.

    Args:
        text: Text to synthesize
        speaker: Internal speaker name
        ... (same as generate_speech)
        sample_rate: Internal generation sample rate (always 16000)
        output_sample_rate: "16khz" or "48khz" (triggers super-resolution)

    Returns:
        Tuple of (audio_bytes, metrics_dict)
    """
    if not is_initialized():
        raise RuntimeError("TTS runtime not initialized")

    runtime = get_runtime()
    if request_id is None:
        import uuid
        request_id = f"runtime-chunk-{uuid.uuid4().hex[:10]}"
    metrics = {
        "request_id": request_id,
        "ttfb_ms": 0,
        "generation_ms": 0,
        "chunks_processed": 0,
        "audio_duration_seconds": 0.0,
        "text_chunked": False,
        "sr_applied": False,
        "output_sample_rate": 16000,
    }

    active_requests = _with_inflight_delta(1)
    start_time = time.time()

    try:
        if _timing_enabled():
            logger.info(
                "tts_runtime_chunking_start",
                extra={
                    "speaker": speaker,
                    "request_id": request_id,
                    "text_length": len(text),
                    "request_inflight": active_requests,
                    "chunking": True,
                },
            )
        # Import long text processor
        from veena3modal.processing.long_text_processor import LongTextProcessor

        long_processor = LongTextProcessor(pipeline=runtime.pipeline)
        generation_timeout = _generation_timeout_seconds()

        # Check if chunking is needed
        if long_processor.should_chunk(text):
            metrics["text_chunked"] = True
            if _timing_enabled():
                chunk_result = await long_processor.generate_with_chunking(
                    text=text,
                    speaker=speaker,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                    sample_rate=sample_rate,
                    request_id=request_id,
                    return_timing=True,
                )
                if isinstance(chunk_result, tuple):
                    audio_bytes, chunk_timing = chunk_result
                else:
                    audio_bytes = chunk_result
                    chunk_timing = {}
            else:
                audio_bytes = await long_processor.generate_with_chunking(
                    text=text,
                    speaker=speaker,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                    sample_rate=sample_rate,
                    request_id=request_id,
                )
                chunk_timing = {}
        else:
            # Short text: direct generation
            gen_kwargs = dict(
                speaker=speaker,
                text=text,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,
            )
            if _timing_enabled():
                generation_coro = runtime.pipeline.generate_speech_indic(
                    **gen_kwargs,
                    return_timing=True,
                    request_id=request_id,
                )
                if generation_timeout:
                    try:
                        audio_bytes, pipeline_timings = await asyncio.wait_for(
                            generation_coro,
                            timeout=generation_timeout,
                        )
                    except asyncio.TimeoutError as exc:
                        raise TimeoutError(
                            f"Chunked generation timed out after {generation_timeout:.2f}s"
                        ) from exc
                else:
                    audio_bytes, pipeline_timings = await generation_coro
                chunk_timing = pipeline_timings
            else:
                generation_coro = runtime.pipeline.generate_speech_indic(
                    **gen_kwargs,
                    request_id=request_id,
                )
                if generation_timeout:
                    try:
                        audio_bytes = await asyncio.wait_for(
                            generation_coro,
                            timeout=generation_timeout,
                        )
                    except asyncio.TimeoutError as exc:
                        raise TimeoutError(
                            f"Chunked generation timed out after {generation_timeout:.2f}s"
                        ) from exc
                else:
                    audio_bytes = await generation_coro
                chunk_timing = {}

        generation_time = time.time() - start_time
        metrics["generation_ms"] = int(generation_time * 1000)
        metrics["ttfb_ms"] = metrics["generation_ms"]

        if chunk_timing:
            metrics.update({
                "llm_token_total": chunk_timing.get("token_total", 0),
                "llm_batch_count": chunk_timing.get("batch_count", 0),
                "llm_batch_gpu_ms_total": chunk_timing.get("llm_gpu_batch_ms_total", 0),
                "llm_batch_wall_ms_total": chunk_timing.get("llm_wall_ms", metrics["generation_ms"]),
                "llm_parse_ms": chunk_timing.get("parse_ms", 0),
                "llm_decode_cpu_ms": chunk_timing.get("decode_cpu_ms", 0),
                "llm_batch_wall_ms_min": chunk_timing.get("llm_batch_wall_ms_min", 0),
                "llm_batch_wall_ms_max": chunk_timing.get("llm_batch_wall_ms_max", 0),
                "llm_batch_gpu_ms_min": chunk_timing.get("llm_batch_gpu_ms_min", 0),
                "llm_batch_gpu_ms_max": chunk_timing.get("llm_batch_gpu_ms_max", 0),
                "llm_batch_gpu_ms_p50": chunk_timing.get("llm_batch_gpu_ms_p50", 0),
                "tokens_per_batch_min": chunk_timing.get("tokens_per_batch_min"),
                "tokens_per_batch_max": chunk_timing.get("tokens_per_batch_max"),
                "tokens_per_batch_p50": chunk_timing.get("tokens_per_batch_p50"),
                "llm_batch_tokens": chunk_timing.get("token_batches", []),
            })
            if metrics["text_chunked"]:
                metrics["chunks_processed"] = chunk_timing.get("chunks_processed", 0)

        if audio_bytes:
            # Apply super-resolution if requested and available
            final_sample_rate = 16000
            if output_sample_rate == "48khz" and runtime.sr_service and runtime.sr_service.is_loaded:
                try:
                    sr_start = time.time()
                    logger.info(f"Applying super-resolution to chunked audio ({len(audio_bytes)} bytes)...")
                    audio_bytes = _apply_super_resolution(audio_bytes, runtime.sr_service)
                    sr_time = (time.time() - sr_start) * 1000
                    metrics["sr_ms"] = int(sr_time)
                    metrics["sr_applied"] = True
                    final_sample_rate = 48000
                    logger.info(f"✅ Super-resolution applied in {sr_time:.1f}ms, output={len(audio_bytes)} bytes")
                except Exception as e:
                    logger.warning(f"Super-resolution failed, returning 16kHz: {e}")
                    import traceback
                    traceback.print_exc()

            metrics["output_sample_rate"] = final_sample_rate
            audio_duration = (len(audio_bytes) - 44) / (final_sample_rate * 2)
            metrics["audio_duration_seconds"] = max(0.0, audio_duration)

        if _timing_enabled():
            metrics["active_requests"] = active_requests - 1
            metrics["request_inflight"] = active_requests - 1
            logger.info("tts_runtime_chunking_complete", extra=metrics)
        return audio_bytes, metrics

    except Exception as e:
        logger.error(f"Chunked speech generation failed: {e}")
        raise
    finally:
        active_requests = _with_inflight_delta(-1)
        metrics["active_requests"] = active_requests
        metrics["request_inflight"] = active_requests

from typing import AsyncGenerator

async def generate_speech_streaming(
    text: str,
    speaker: str,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 1.0,
    max_tokens: int = 4096,
    repetition_penalty: float = 1.05,
    seed: Optional[int] = None,
    enable_chunking: bool = True,
    request_id: Optional[str] = None,
) -> AsyncGenerator[Tuple[bytes, Dict[str, Any]], None]:
    """
    Generate speech audio with true streaming (yields chunks as they're generated).

    This is the core streaming implementation for M4.
    First yield includes WAV header + first PCM chunk.
    Subsequent yields are raw PCM chunks.

    Args:
        text: Text to synthesize (already normalized)
        speaker: Internal speaker name (resolved)
        temperature: Sampling temperature
        top_k: Top-k sampling
        top_p: Nucleus sampling
        max_tokens: Maximum tokens to generate
        repetition_penalty: Repetition penalty
        seed: Random seed for reproducibility
        enable_chunking: Enable text chunking for long inputs with voice consistency

    Yields:
        Tuple of (audio_bytes, metrics_dict)
        - First yield: WAV header (44 bytes) prepended to first PCM chunk
        - Subsequent yields: Raw PCM chunks (int16, 16kHz)
        - Final yield: metrics_dict has final timing info

    Raises:
        RuntimeError: If runtime not initialized or streaming pipeline unavailable
    """
    if not is_initialized():
        raise RuntimeError("TTS runtime not initialized")

    runtime = get_runtime()

    if request_id is None:
        import uuid
        request_id = f"runtime-stream-{uuid.uuid4().hex[:10]}"

    if runtime.streaming_pipeline is None:
        raise RuntimeError("Streaming pipeline not available")

    # Import audio utils for WAV header
    from veena3modal.audio.utils import create_wav_header

    # Metrics tracking
    metrics = {
        "request_id": request_id,
        "speaker": speaker,
        "ttfb_ms": 0,
        "chunks_sent": 0,
        "total_bytes": 0,
        "audio_duration_seconds": 0.0,
        "text_chunked": False,


sed -n '940,1240p' veena3modal/services/tts_runtime.py