async_streaming.py

"""
AsyncLLM-based streaming for Qwen3-ASR — concurrency via vLLM continuous batching.

Ports qwen-asr 0.0.6's streaming-state algorithm (init_streaming_state /
streaming_transcribe / finish_streaming_transcribe) onto vllm.v1.engine.AsyncLLM,
so concurrent generate() calls from many WebSocket streams batch on the GPU
instead of serializing through the offline LLM. Benchmark showed ~3.8x at 8
concurrent and ~28x KV-cache headroom on one L4.

This module owns ONE shared AsyncLLM engine + processor; each utterance is an
independent AsyncUtteranceSession (its own buffer/state), so sessions are
isolated and only the GPU engine is shared.
"""

import os
import itertools
import numpy as np

SAMPLE_RATE = 16000

MODEL_ID                = os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-1.7B")
GPU_MEMORY_UTILIZATION  = float(os.getenv("GPU_MEMORY_UTILIZATION", "0.80"))
STREAMING_MAX_NEW_TOKENS = int(os.getenv("STREAMING_MAX_NEW_TOKENS", "1024"))
MAX_MODEL_LEN           = int(os.getenv("MAX_MODEL_LEN", "0")) or None
CHUNK_SIZE_SEC          = float(os.getenv("CHUNK_SIZE_SEC", "4.0"))
UNFIXED_CHUNK_NUM       = int(os.getenv("UNFIXED_CHUNK_NUM", "5"))
UNFIXED_TOKEN_NUM       = int(os.getenv("UNFIXED_TOKEN_NUM", "15"))

import logging
log = logging.getLogger("qwen3-asr-async")

_engine     = None      # vllm AsyncLLM
_processor  = None       # HF processor (chat template + tokenizer)
_sp         = None       # SamplingParams
_parse      = None       # qwen_asr.inference.utils.parse_asr_output
_req_ids    = itertools.count()


def is_ready() -> bool:
    return _engine is not None


def _pick_dtype() -> str:
    """
    Choose engine precision per GPU. bfloat16 needs compute capability >= 8.0
    (Ampere+, e.g. L4=8.9); Turing (T4=7.5) supports only float16 and vLLM
    hard-errors on bf16 there. Override with the DTYPE env var if needed.
    """
    override = os.getenv("DTYPE", "").strip().lower()
    if override:
        return override
    try:
        import torch
        cap = torch.cuda.get_device_capability()
        return "bfloat16" if cap[0] >= 8 else "float16"
    except Exception:
        return "float16"


def _enforce_eager() -> bool:
    """
    Whether to skip vLLM's torch.compile + CUDA-graph capture at startup.
    Pre-Ampere GPUs (T4 = compute 7.5) crash during that compile step
    ('Engine core initialization failed' / 'Not enough SMs'), so default to
    eager there. Ampere+ (L4/A100) keep the faster compiled path.
    Override with ENFORCE_EAGER=true/false.
    """
    v = os.getenv("ENFORCE_EAGER", "").strip().lower()
    if v in ("1", "true", "yes"):
        return True
    if v in ("0", "false", "no"):
        return False
    try:
        import torch
        return torch.cuda.get_device_capability()[0] < 8
    except Exception:
        return False


async def init_engine():
    """Build the shared AsyncLLM engine + processor (call once, inside the loop)."""
    global _engine, _processor, _sp, _parse
    if _engine is not None:
        return

    from transformers import AutoConfig, AutoModel, AutoProcessor
    from qwen_asr.core.transformers_backend import (
        Qwen3ASRConfig,
        Qwen3ASRForConditionalGeneration as HFModel,
        Qwen3ASRProcessor,
    )
    AutoConfig.register("qwen3_asr", Qwen3ASRConfig)
    AutoModel.register(Qwen3ASRConfig, HFModel)
    AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor)

    from qwen_asr.core.vllm_backend import Qwen3ASRForConditionalGeneration as VLLMModel
    from qwen_asr.inference.utils import parse_asr_output
    from vllm import ModelRegistry, AsyncEngineArgs, SamplingParams
    from vllm.v1.engine.async_llm import AsyncLLM   # V1 async engine (confirmed by spike)

    ModelRegistry.register_model("Qwen3ASRForConditionalGeneration", VLLMModel)

    dtype = _pick_dtype()
    eager = _enforce_eager()
    engine_args = AsyncEngineArgs(
        model=MODEL_ID,
        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
        dtype=dtype,
        max_model_len=MAX_MODEL_LEN,
        enforce_eager=eager,
        limit_mm_per_prompt={"audio": 1},
    )
    log.info(f"Building AsyncLLM for {MODEL_ID} (gpu_mem={GPU_MEMORY_UTILIZATION}, dtype={dtype}, enforce_eager={eager})...")
    _engine = AsyncLLM.from_engine_args(engine_args)
    _processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
    _sp = SamplingParams(temperature=0.0, max_tokens=STREAMING_MAX_NEW_TOKENS)
    _parse = parse_asr_output
    log.info("AsyncLLM engine ready")


def _build_prompt(context: str, force_language) -> str:
    """Mirror qwen_asr._build_text_prompt (chat template + optional forced lang)."""
    msgs = [
        {"role": "system", "content": context or ""},
        {"role": "user", "content": [{"type": "audio", "audio": ""}]},
    ]
    base = _processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
    if force_language:
        base = base + f"language {force_language}<asr_text>"
    return base


async def transcribe_audio(audio: np.ndarray, context: str, force_language) -> str:
    """
    One-shot (non-streaming) transcription of a full clip via the shared
    AsyncLLM engine — used by the batch /v1/audio/transcriptions endpoint.
    Returns the parsed transcript text (raw, un-romanized).
    """
    prompt = _build_prompt(context, force_language)
    inp = {"prompt": prompt, "multi_modal_data": {"audio": [np.asarray(audio, dtype=np.float32)]}}
    rid = f"b{next(_req_ids)}"
    out = None
    async for o in _engine.generate(prompt=inp, sampling_params=_sp, request_id=rid):
        out = o
    gen = out.outputs[0].text if (out and out.outputs) else ""
    _lang, txt = _parse(gen, user_language=force_language)
    return txt


class AsyncUtteranceSession:
    """
    One streaming utterance. Faithful async port of qwen-asr's streaming state:
    buffer audio into chunks, re-feed accumulated audio each step, roll back the
    last UNFIXED_TOKEN_NUM tokens for the prefix prompt after UNFIXED_CHUNK_NUM
    chunks, decode via AsyncLLM, parse to (language, text).
    """

    def __init__(self, context: str, force_language):
        self.context        = context
        self.force_language  = force_language          # canonical name or None
        self.prompt_raw     = _build_prompt(context, force_language)
        self.chunk_samples  = max(1, int(round(CHUNK_SIZE_SEC * SAMPLE_RATE)))
        self.buffer         = np.zeros((0,), dtype=np.float32)
        self.audio_accum    = np.zeros((0,), dtype=np.float32)
        self.chunk_id       = 0
        self._raw_decoded   = ""
        self.text           = ""
        self.language       = ""

    async def feed(self, audio: np.ndarray) -> str:
        x = np.asarray(audio, dtype=np.float32)
        if x.ndim != 1:
            x = x.reshape(-1)
        if x.shape[0] > 0:
            self.buffer = np.concatenate([self.buffer, x])
        while self.buffer.shape[0] >= self.chunk_samples:
            chunk = self.buffer[:self.chunk_samples]
            self.buffer = self.buffer[self.chunk_samples:]
            self._accumulate(chunk)
            await self._decode_step(final=False)
        return self.text

    async def finish(self) -> str:
        if self.buffer.shape[0] > 0:
            tail = self.buffer
            self.buffer = np.zeros((0,), dtype=np.float32)
            self._accumulate(tail)
            await self._decode_step(final=True)
        return self.text

    def _accumulate(self, chunk: np.ndarray):
        if self.audio_accum.shape[0] == 0:
            self.audio_accum = chunk
        else:
            self.audio_accum = np.concatenate([self.audio_accum, chunk], axis=0)

    def _compute_prefix(self, final: bool) -> str:
        if self.chunk_id < UNFIXED_CHUNK_NUM:
            return ""
        tok = _processor.tokenizer
        cur_ids = tok.encode(self._raw_decoded)
        k = int(UNFIXED_TOKEN_NUM)
        if final:
            end_idx = max(1, len(cur_ids) - k)
            return tok.decode(cur_ids[:end_idx])
        # per-chunk: grow rollback until the decoded prefix has no broken char
        while True:
            end_idx = max(0, len(cur_ids) - k)
            prefix = tok.decode(cur_ids[:end_idx]) if end_idx > 0 else ""
            if "�" not in prefix:
                return prefix
            if end_idx == 0:
                return ""
            k += 1

    async def _decode_step(self, final: bool):
        prefix = self._compute_prefix(final)
        prompt = self.prompt_raw + prefix
        inp = {"prompt": prompt, "multi_modal_data": {"audio": [self.audio_accum]}}
        rid = f"u{next(_req_ids)}"
        out = None
        async for o in _engine.generate(prompt=inp, sampling_params=_sp, request_id=rid):
            out = o
        gen_text = out.outputs[0].text if (out and out.outputs) else ""
        self._raw_decoded = prefix + gen_text
        self.language, self.text = _parse(self._raw_decoded, user_language=self.force_language)
        self.chunk_id += 1