main_async.py

"""
Combined ASYNC entry point: one shared vllm AsyncLLM engine, both apps mounted,
served from a single process. Concurrent streams across BOTH endpoints batch on
the GPU (continuous batching) instead of serializing.

URL layout
──────────
  GET /health                              combined health
  POST/WS /en/v1/...                       English / Malay (async)
  POST/WS /chughtai/v1/...                 Chughtai Lab (async, romanized)

Mounted sub-apps do NOT run their own lifespan — this parent lifespan performs
the single engine + VAD init (init_engine is idempotent regardless).
"""

import os
import logging
from contextlib import asynccontextmanager

if not os.environ.get("OMP_NUM_THREADS"):
    os.environ["OMP_NUM_THREADS"] = "4"

from fastapi import FastAPI
import uvicorn

from async_streaming import init_engine, is_ready, MODEL_ID
from shared_model import _ensure_vad_loaded, is_vad_ready

import app_en_async as en
import app_chughtai_async as chughtai

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s [%(levelname)s] %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S")
log = logging.getLogger("qwen3-asr-async-combined")

PORT = int(os.getenv("PORT", "7860"))


@asynccontextmanager
async def lifespan(app: FastAPI):
    log.info("Building shared AsyncLLM engine (used by both /en and /chughtai)...")
    await init_engine()
    try:
        _ensure_vad_loaded()
        log.info("Silero VAD ready")
    except Exception as e:
        log.warning(f"Silero VAD load failed, RMS fallback will be used: {e}")
    log.info(f"Combined async server ready on 0.0.0.0:{PORT}")
    yield


main_app = FastAPI(title="Qwen3-ASR Combined (async)", version="3.0.0-async", lifespan=lifespan)


@main_app.get("/health")
def health():
    return {
        "status": "ok",
        "model": MODEL_ID,
        "model_status": "loaded" if is_ready() else "loading",
        "vad": "silero" if is_vad_ready() else "rms_fallback",
        "engine": "vllm.AsyncLLM (continuous batching)",
        "variants": {
            "english-malay": "/en/v1/realtime",
            "chughtai-lab":  "/chughtai/v1/realtime",
        },
    }


@main_app.get("/")
def root():
    return {
        "service": "Qwen3-ASR Combined async server",
        "model": MODEL_ID,
        "endpoints": {
            "/health": "GET",
            "/en/v1/realtime": "WS — English/Malay (async)",
            "/chughtai/v1/realtime": "WS — Chughtai Lab (async)",
        },
    }


# Mount AFTER root routes so / and /health aren't shadowed.
main_app.mount("/en", en.app)
main_app.mount("/chughtai", chughtai.app)


if __name__ == "__main__":
    log.info("=" * 60)
    log.info("Qwen3-ASR Combined ASYNC Server")
    log.info("=" * 60)
    uvicorn.run(main_app, host="0.0.0.0", port=PORT)