LLM-infa/api/lib/jobqueue.py

"""Очередь фоновых задач: одно видео или плейлист, элементы обрабатываются
последовательно (чтобы не долбить YouTube и LLM параллельными запросами).

Живой прогресс держится в памяти (_live) и подмешивается в GET /jobs;
устойчивое состояние (счётчики, статус) — в таблице jobs.
"""
from __future__ import annotations

import asyncio
import logging

from api.db.job import get_job, update_job
from api.db.video import get_video_by_url
from api.lib.processing import PipelineError, process_single_video
from api.lib.youtube import expand_url


_ITEM_ATTEMPTS = 3  # попытки на один элемент (YouTube любит разовые 403/429)
_RETRY_DELAY = 30.0


logger = logging.getLogger("jobs")

_queue: asyncio.Queue[str] = asyncio.Queue()
_live: dict[str, str] = {}
_cancel: set[str] = set()
_params: dict[str, dict] = {}


def submit(job_id: str, params: dict) -> None:
    _params[job_id] = params
    _queue.put_nowait(job_id)


def live_status(job_id: str) -> str:
    return _live.get(job_id, "")


def request_cancel(job_id: str) -> None:
    _cancel.add(job_id)


def is_cancelled(job_id: str) -> bool:
    return job_id in _cancel


async def worker() -> None:
    """Единственный воркер; запускается на старте приложения."""
    logger.info("job worker started")
    while True:
        job_id = await _queue.get()
        try:
            await _run_job(job_id)
        except Exception:
            logger.exception("job %s crashed", job_id)
            try:
                await update_job(job_id, status="error", current_item="")
            except Exception:
                pass
        finally:
            _live.pop(job_id, None)
            _cancel.discard(job_id)
            _params.pop(job_id, None)
            _queue.task_done()


async def _run_job(job_id: str) -> None:
    job = await get_job(job_id)
    if job is None or job.status == "cancelled":
        return
    params = _params.get(job_id, {})
    loop = asyncio.get_running_loop()

    await update_job(job_id, status="running", current_item="анализ ссылки…")
    _live[job_id] = "анализ ссылки…"

    plan = await loop.run_in_executor(None, expand_url, job.url)
    entries = plan["entries"]
    if not entries:
        await update_job(job_id, status="error", error="Плейлист пуст", current_item="")
        return
    await update_job(
        job_id, kind=plan["kind"], title=plan["title"], total_items=len(entries)
    )

    done = failed = 0
    errors: list[str] = []
    n = len(entries)

    for i, entry in enumerate(entries, start=1):
        if is_cancelled(job_id):
            await update_job(job_id, status="cancelled", current_item="")
            logger.info("job %s cancelled at item %d/%d", job_id, i, n)
            return

        label = f"{i}/{n} · {entry['title']}"
        await update_job(job_id, current_item=label)
        _live[job_id] = label

        def _notify(msg: str, _label: str = label) -> None:
            _live[job_id] = f"{_label} — {msg}"

        # дедупликация: повторная постановка плейлиста не переделывает готовое
        if await get_video_by_url(entry["url"]) is not None:
            done += 1
            _live[job_id] = f"{label} — уже в библиотеке, пропуск"
            logger.info("job %s item %d/%d already processed, skip", job_id, i, n)
            await update_job(job_id, done_items=done)
            continue

        for attempt in range(1, _ITEM_ATTEMPTS + 1):
            try:
                await process_single_video(
                    entry["url"],
                    section_id=params.get("section_id"),
                    model_path=params.get("model_path"),
                    summary_max_sentences=params.get("summary_max_sentences"),
                    keep_video=params.get("keep_video"),
                    job_id=job_id,
                    progress=_notify,
                )
                done += 1
                break
            except Exception as e:
                retryable = not (isinstance(e, PipelineError) and e.status == 400)
                if retryable and attempt < _ITEM_ATTEMPTS and not is_cancelled(job_id):
                    logger.warning(
                        "job %s item %d/%d attempt %d failed, retrying: %s",
                        job_id, i, n, attempt, e,
                    )
                    _live[job_id] = (
                        f"{label} — ошибка, повтор {attempt}/{_ITEM_ATTEMPTS - 1} через {int(_RETRY_DELAY)} с"
                    )
                    await asyncio.sleep(_RETRY_DELAY)
                    continue
                failed += 1
                errors.append(f"{entry['title']}: {e}")
                logger.warning("job %s item %d/%d failed: %s", job_id, i, n, e)
                break
        await update_job(
            job_id, done_items=done, failed_items=failed, error="\n".join(errors)[:4000]
        )

    status = "error" if done == 0 and failed > 0 else "done"
    await update_job(job_id, status=status, current_item="")
    logger.info("job %s finished: %d ok, %d failed", job_id, done, failed)