ASR-demo/scripts/run_backend.py

#!/usr/bin/env python3
"""Start the local CAM++ model service and FunASR WebSocket backend together."""

from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from pathlib import Path
from urllib.error import URLError
from urllib.request import urlopen

from dotenv import load_dotenv

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from scripts.model_manifest import auxiliary_models, load_manifest, model_directory

load_dotenv(PROJECT_ROOT / ".env")

# FunASR's published short names resolve to these ModelScope local directories.
LOCAL_MODEL_NAMES = {
    "paraformer-zh-streaming": (
        "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
    ),
    "fsmn-vad": (
        "iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
        "damo/speech_fsmn_vad_zh-cn-16k-common-pytorch",
    ),
}


def local_model(requested: str, models_dir: Path) -> Path:
    """Resolve a model ID or path beneath MODEL_DIR without a network fallback."""
    name = requested.strip()
    direct = Path(name)
    candidates = [direct] if direct.is_absolute() else [
        PROJECT_ROOT / direct, models_dir / direct,
        *(models_dir / alias for alias in LOCAL_MODEL_NAMES.get(name, ())),
    ]
    for candidate in candidates:
        if candidate.is_dir() and (candidate / "configuration.json").is_file():
            return candidate.resolve()
    checked = ", ".join(str(path) for path in candidates)
    raise FileNotFoundError(f"Local model '{name}' was not found; checked: {checked}")


def local_cam_model(models_dir: Path) -> Path:
    """Require one complete CAM++ speaker verification asset from the manifest."""
    manifest = load_manifest()
    override = os.getenv("CAM_MODEL_PATH", "").strip()
    if override:
        path = Path(override)
        if not path.is_absolute():
            path = models_dir / path
        path = path.resolve()
        configs = [
            config for config in auxiliary_models(manifest).values()
            if config.get("kind") == "speaker_verification"
        ]
        if path.is_dir() and configs and all(
            (path / relative).is_file() for relative in configs[0].get("required_files", [])
        ):
            return path
        raise FileNotFoundError(f"CAM++ model is missing or incomplete: {path}")
    checked = []
    for model_id, config in auxiliary_models(manifest).items():
        if config.get("kind") != "speaker_verification":
            continue
        path = model_directory(model_id, manifest, models_dir)
        required = [path / relative for relative in config.get("required_files", [])]
        if path.is_dir() and all(item.is_file() for item in required):
            return path.resolve()
        checked.append(str(path))
    raise FileNotFoundError("CAM++ speaker model is missing; checked: " + ", ".join(checked))


def wait_for_health(url: str, process: subprocess.Popen, key: str, seconds: int = 300) -> None:
    """Wait until a child has loaded its models, failing when it exits early."""
    deadline = time.monotonic() + seconds
    while time.monotonic() < deadline:
        code = process.poll()
        if code is not None:
            raise RuntimeError(f"Service exited before it became ready ({url}, exit={code})")
        try:
            with urlopen(url, timeout=2) as response:
                data = json.load(response)
            if data.get(key):
                return
        except (OSError, ValueError, URLError):
            pass
        time.sleep(0.5)
    raise TimeoutError(f"Service did not become ready within {seconds}s: {url}")


def stop_child(process: subprocess.Popen | None) -> None:
    """Stop a supervised model or WebSocket process on shutdown."""
    if process is None or process.poll() is not None:
        return
    process.terminate()
    try:
        process.wait(timeout=10)
    except subprocess.TimeoutExpired:
        process.kill()
        process.wait()


def main() -> None:
    """Require ASR, VAD, and CAM++ before exposing the backend WebSocket."""
    models_dir = Path(os.getenv("MODEL_DIR", "models"))
    if not models_dir.is_absolute():
        models_dir = PROJECT_ROOT / models_dir
    models_dir = models_dir.resolve()
    asr = local_model(os.getenv("FUNASR_ASR_MODEL", "paraformer-zh-streaming"), models_dir)
    vad = local_model(os.getenv("FUNASR_VAD_MODEL", "fsmn-vad"), models_dir)
    cam = local_cam_model(models_dir)
    print(f"Local models: ASR={asr}; VAD={vad}; CAM++={cam}", flush=True)

    env = os.environ.copy()
    env.update({
        "MODEL_DIR": str(models_dir),
        "FUNASR_ASR_MODEL": str(asr),
        "FUNASR_VAD_MODEL": str(vad),
        "CAM_MODEL_PATH": str(cam),
        "AUXILIARY_SERVICE_URL": "http://127.0.0.1:8010",
    })
    auxiliary = None
    websocket = None
    try:
        auxiliary = subprocess.Popen(
            [sys.executable, "-m", "scripts.auxiliary_server"],
            cwd=PROJECT_ROOT, env=env,
        )
        wait_for_health("http://127.0.0.1:8010/health", auxiliary, "speaker_embedding_ready")
        websocket = subprocess.Popen(
            [sys.executable, "-m", "scripts.run_funasr_demo", "--no-browser"],
            cwd=PROJECT_ROOT, env=env,
        )
        web_port = int(env.get("WEB_PORT", "8082"))
        wait_for_health(
            f"http://127.0.0.1:{web_port}/api/config", websocket, "engine"
        )
        print(f"Backend ready: ws://127.0.0.1:{web_port}/ws", flush=True)
        while True:
            for label, process in (("CAM++", auxiliary), ("WebSocket", websocket)):
                code = process.poll()
                if code is not None:
                    raise RuntimeError(f"{label} service exited (exit={code})")
            time.sleep(0.5)
    except KeyboardInterrupt:
        pass
    finally:
        stop_child(websocket)
        stop_child(auxiliary)


if __name__ == "__main__":
    main()