jobqueue.py

"""内部任务队列:submit 立刻返回 job_id,后台按并发上限逐个跑,调用方轮询取结果。

为什么必须异步:MCP 是请求/响应模型,客户端(weft)侧超时上限只有 ~120s,而一条视频转写要几分钟到
几十分钟。所以 submit 秒回、结果另行轮询 —— 这是 MCP 上跑长任务的标准做法。
"""
from __future__ import annotations

import asyncio
import contextlib
import logging
import shutil
import time
import uuid
from dataclasses import dataclass, field, asdict

import engines
from config import CFG
from media import MediaError, extract_audio, fetch_media, probe_duration, workspace

logger = logging.getLogger("vt.queue")


@dataclass
class Job:
    id: str
    media_url: str
    title: str = ""
    engine: str = ""
    model: str = ""                 # 本任务用的 ASR 模型;空=服务端默认(CFG.model)
    hotwords: str = ""              # 热词/上下文偏置,仅对支持的模型生效(见 models.supports_hotwords)
    status: str = "queued"          # queued | running | done | failed | canceled
    text: str = ""
    error: str = ""
    audio_seconds: float = 0.0      # 音频时长
    asr_seconds: float = 0.0        # 转写耗时 → 与上者比即"实时率",用来评估算力
    created_at: float = field(default_factory=time.time)
    started_at: float = 0.0
    finished_at: float = 0.0

    def public(self) -> dict:
        d = asdict(self)
        for k in ("audio_seconds", "asr_seconds"):
            d[k] = round(d[k], 1)
        return d


class JobQueue:
    def __init__(self) -> None:
        self._jobs: dict[str, Job] = {}
        self._q: asyncio.Queue[str] = asyncio.Queue()
        self._workers: list[asyncio.Task] = []
        self._started = False

    def start(self) -> None:
        if self._started:
            return
        self._started = True
        for i in range(max(1, CFG.concurrency)):
            self._workers.append(asyncio.create_task(self._worker(i)))
        self._workers.append(asyncio.create_task(self._gc()))
        logger.info("任务队列启动(并发 %d)", max(1, CFG.concurrency))

    def submit(self, media_url: str, title: str = "", engine: str = "",
               model: str = "", hotwords: str = "") -> Job:
        job = Job(id=uuid.uuid4().hex[:16], media_url=media_url, title=title,
                  engine=engine or CFG.engine, model=(model or CFG.model),
                  hotwords=hotwords or "")
        self._jobs[job.id] = job
        self._q.put_nowait(job.id)
        return job

    def get(self, job_id: str) -> Job | None:
        return self._jobs.get(job_id)

    def cancel(self, job_id: str) -> bool:
        """只能取消尚未开跑的任务(已在跑的让它跑完,避免中断留下半截文件)。"""
        j = self._jobs.get(job_id)
        if j and j.status == "queued":
            j.status = "canceled"
            j.finished_at = time.time()
            return True
        return False

    def stats(self) -> dict:
        by: dict[str, int] = {}
        for j in self._jobs.values():
            by[j.status] = by.get(j.status, 0) + 1
        return {"total": len(self._jobs), "by_status": by, "pending": self._q.qsize()}

    async def _worker(self, idx: int) -> None:
        while True:
            job_id = await self._q.get()
            job = self._jobs.get(job_id)
            try:
                if job and job.status == "queued":       # 可能已被取消
                    await self._run(job)
            except Exception:                            # noqa: BLE001  worker 必须长活
                logger.exception("任务处理异常 job=%s", job_id)
                if job:
                    job.status, job.error = "failed", "内部错误"
                    job.finished_at = time.time()
            finally:
                self._q.task_done()

    async def _run(self, job: Job) -> None:
        job.status, job.started_at = "running", time.time()
        work = workspace()
        try:
            src = await asyncio.wait_for(fetch_media(job.media_url, work), timeout=CFG.job_timeout_s)
            wav = await asyncio.wait_for(extract_audio(src, work), timeout=CFG.job_timeout_s)
            job.audio_seconds = await probe_duration(wav)
            eng = await engines.get(job.engine, job.model)
            t0 = time.time()
            text = await asyncio.wait_for(eng.transcribe(wav, job.hotwords),
                                          timeout=CFG.job_timeout_s)
            job.asr_seconds = time.time() - t0
            job.text = (text or "").strip()
            job.status = "done" if job.text else "failed"
            if not job.text:
                job.error = "转写结果为空(可能音频无人声)"
            rt = (job.asr_seconds / job.audio_seconds) if job.audio_seconds else 0
            logger.info("完成 job=%s 音频%.0fs 转写%.0fs 实时率%.2fx 字数%d",
                        job.id, job.audio_seconds, job.asr_seconds, rt, len(job.text))
        except (MediaError, ValueError) as e:
            job.status, job.error = "failed", str(e)[:500]
            logger.warning("任务失败 job=%s: %s", job.id, job.error)
        except asyncio.TimeoutError:
            job.status, job.error = "failed", f"处理超时(>{CFG.job_timeout_s}s)"
        except Exception as e:                           # noqa: BLE001
            job.status, job.error = "failed", f"{type(e).__name__}: {str(e)[:400]}"
            logger.exception("任务失败 job=%s", job.id)
        finally:
            job.finished_at = time.time()
            with contextlib.suppress(Exception):         # 临时文件用完即删(视频很占地方)
                shutil.rmtree(work, ignore_errors=True)

    async def _gc(self) -> None:
        """定期清掉过期的已完成任务,防止长期运行内存无限涨。"""
        while True:
            await asyncio.sleep(600)
            now = time.time()
            stale = [k for k, j in self._jobs.items()
                     if j.finished_at and now - j.finished_at > CFG.result_ttl_s]
            for k in stale:
                self._jobs.pop(k, None)
            if stale:
                logger.info("清理过期任务 %d 条", len(stale))


QUEUE = JobQueue()