scripts/worker.py - master - Max_Cherep/super-resolution

scripts/worker.py
435 строк13 KB
maxim.cherepanov
recover orphan jobs and retry RabbitMQ DNS failures
15 июл 2026, 19:36
15 июл 2026, 19:360abe886
#!/usr/bin/env python3
"""Worker: RabbitMQ → ffmpeg / FlashSR / CUE."""

from __future__ import annotations

import json
import os
import socket

os.environ.setdefault("TQDM_DISABLE", "1")

import shutil
import sys
import time
from pathlib import Path
from typing import Any

import pika

sys.path.insert(0, str(Path(__file__).resolve().parent))

from audio_pipeline import run_pipeline
from cue_split import run_cue_batch, run_cue_split
from db import get_conn, get_job, init_db, try_begin_processing, update_job, utc_now_iso
from ffmpeg_ops import FfmpegError, AudioIntegrityError, validate_input_audio
from input_integrity import record_input_check
from job_cancel import JobCancelled, job_context
from messaging import QUEUE_NAME, publish_job, rabbit_connection_params
from process_options import parse_options
from progress import (
    ProgressReporter,
    build_job_progress,
    clear_job_progress_state,
)

WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", "/app/weights")
MOCK_MODE = os.environ.get("MOCK_MODE", "").lower() in ("1", "true", "yes")
_SKIP_STATUSES = frozenset({"done", "failed", "skipped"})


def _rabbit_params() -> pika.ConnectionParameters:
    return rabbit_connection_params(heartbeat=0)


def _cleanup_work_files(dst: Path, job_id: int) -> None:
    for p in dst.parent.glob(f".work_{dst.stem}_{job_id}*.wav"):
        p.unlink(missing_ok=True)
    legacy = dst.with_name(dst.stem + "_48k.wav")
    if legacy.exists() and not dst.exists():
        legacy.unlink()


def _iso_now() -> str:
    return utc_now_iso()


def _job_options(job: dict) -> dict:
    return parse_options(job.get("options"))


def _cleanup_partial_outputs(paths: list[Path], job_id: int) -> None:
    for path in paths:
        if path.exists():
            path.unlink(missing_ok=True)
        if path.name:
            _cleanup_work_files(path, job_id)


def _finish_job(job_id: int, **fields: Any) -> None:
    fields.setdefault("progress_pct", None)
    fields.setdefault("progress_detail", None)
    update_job(job_id, **fields)
    clear_job_progress_state(job_id)


def _cancel_job(job_id: int, job: dict, *, partial_outputs: list[Path] | None = None) -> None:
    dst = Path(job.get("output_path", ""))
    if dst.name:
        _cleanup_work_files(dst, job_id)
        if dst.exists():
            dst.unlink(missing_ok=True)
    if partial_outputs:
        _cleanup_partial_outputs(partial_outputs, job_id)
    _finish_job(
        job_id,
        status="cancelled",
        finished_at=_iso_now(),
        cancel_requested=1,
        error_message="прервано пользователем",
    )
    print(f"Job {job_id}: cancelled")


def _process_pipeline_job(job_id: int, job: dict, model: Any, device: Any) -> None:
    src = Path(job["input_path"])
    dst = Path(job["output_path"])

    if not src.exists():
        _finish_job(
            job_id,
            status="failed",
            finished_at=_iso_now(),
            error_message=f"input not found: {src}",
        )
        return

    try:
        validate_input_audio(src)
    except AudioIntegrityError as exc:
        record_input_check(src)
        _finish_job(
            job_id,
            status="failed",
            finished_at=_iso_now(),
            error_message=f"битый файл: {exc}"[:2000],
        )
        print(f"Job {job_id}: corrupted input — {exc}")
        return

    _cleanup_work_files(dst, job_id)
    tag = "[MOCK] " if MOCK_MODE else ""
    print(f"Job {job_id}: {tag}processing {src.name}")

    opts = _job_options(job)
    reporter = ProgressReporter(job_id, prefix=f"Job {job_id} ")
    progress = build_job_progress(reporter, opts)

    dur, orig_sr, out_sr = run_pipeline(
        job_id,
        src,
        dst,
        opts,
        model=model,
        device=device,
        progress=progress,
    )
    reporter.finish_line()
    _finish_job(
        job_id,
        status="done",
        finished_at=_iso_now(),
        duration_sec=dur,
        input_sr=orig_sr,
        output_sr=out_sr,
    )
    print(f"Job {job_id}: done ({dur:.1f}s audio)")


def _process_cue_split_job(job_id: int, job: dict) -> None:
    cue_path = Path(job["input_path"])
    opts = _job_options(job)
    split_format = opts.get("split_format", "wav")
    audio_path = opts.get("audio_path")
    audio = Path(audio_path) if audio_path else None

    if not cue_path.is_file():
        raise FileNotFoundError(f"cue not found: {cue_path}")

    reporter = ProgressReporter(job_id, prefix=f"Job {job_id} ")
    reporter.report(5.0, "CUE · нарезка")

    count, dur, out_dir = run_cue_split(
        cue_path,
        audio_path=audio,
        split_format=split_format,
    )
    reporter.finish_line()
    _finish_job(
        job_id,
        status="done",
        finished_at=_iso_now(),
        duration_sec=dur,
        output_path=str(out_dir),
        error_message=None,
    )
    print(f"Job {job_id}: cue split → {count} tracks in {out_dir}")


def _process_cue_batch_job(job_id: int, job: dict, model: Any, device: Any) -> None:
    cue_path = Path(job["input_path"])
    opts = _job_options(job)
    pipeline_opts = opts.get("pipeline", {})
    output_files: list[str] = []
    batch_outputs: list[Path] = []

    from cue_sheet import parse_cue

    sheet = parse_cue(cue_path, input_dir=cue_path.parent)
    total = len(sheet.files)
    reporter = ProgressReporter(job_id, prefix=f"Job {job_id} ")
    progress = build_job_progress(reporter, pipeline_opts)

    def _run_one(audio: Path, track_index: int) -> None:
        out_fmt = pipeline_opts.get("output_format", "wav")
        out_path = Path(job["output_path"]).parent / f"{audio.stem}.{out_fmt}"
        progress.set_batch(track_index, total)
        run_pipeline(
            job_id,
            audio,
            out_path,
            pipeline_opts,
            model=model,
            device=device,
            progress=progress,
        )
        batch_outputs.append(out_path)
        output_files.append(str(out_path))

    try:
        ok, total, errors = run_cue_batch(cue_path, _run_one)
    except JobCancelled:
        _cancel_job(job_id, job, partial_outputs=batch_outputs)
        return
    reporter.finish_line()
    opts["output_files"] = output_files
    msg = f"batch {ok}/{total}"
    if errors:
        msg += ": " + "; ".join(errors[:5])
    status = "done" if ok == total else "failed"
    _finish_job(
        job_id,
        status=status,
        finished_at=_iso_now(),
        error_message=msg if errors else None,
        options=json.dumps(opts, ensure_ascii=False),
    )
    print(f"Job {job_id}: {msg}")


def process_job(job_id: int, model: Any, device: Any) -> None:
    job = get_job(job_id)
    if job is None:
        print(f"Job {job_id}: not found in DB, skip")
        return

    if job.get("status") == "cancelled":
        print(f"Job {job_id}: already cancelled, skip")
        return

    status = job.get("status")
    if status in _SKIP_STATUSES:
        if status == "done" and (
            job.get("progress_pct") is not None or job.get("progress_detail")
        ):
            update_job(job_id, progress_pct=None, progress_detail=None)
        print(f"Job {job_id}: already {status}, skip")
        return

    job_type = job.get("job_type") or "process"

    with job_context(job_id):
        if job.get("status") == "queued":
            if not try_begin_processing(job_id, _iso_now()):
                print(f"Job {job_id}: not started (cancelled or taken)")
                return
            job = get_job(job_id) or job

        try:
            from job_cancel import is_cancel_requested
            if is_cancel_requested(job_id):
                _cancel_job(job_id, job)
                return
            if job_type == "cue_split":
                _process_cue_split_job(job_id, job)
            elif job_type == "cue_batch":
                _process_cue_batch_job(job_id, job, model, device)
            else:
                _process_pipeline_job(job_id, job, model, device)
        except JobCancelled:
            fresh = get_job(job_id) or job
            _cancel_job(job_id, fresh)
        except FfmpegError as exc:
            dst = Path(job.get("output_path", ""))
            if dst.name:
                _cleanup_work_files(dst, job_id)
            _finish_job(
                job_id,
                status="failed",
                finished_at=_iso_now(),
                error_message=str(exc)[:2000],
            )
            print(f"Job {job_id}: ffmpeg failed — {exc}")
        except Exception as exc:
            dst = Path(job.get("output_path", ""))
            if dst.name:
                _cleanup_work_files(dst, job_id)
            _finish_job(
                job_id,
                status="failed",
                finished_at=_iso_now(),
                error_message=str(exc)[:2000],
            )
            print(f"Job {job_id}: failed — {exc}")


def _load_model() -> tuple[Any, Any]:
    import torch
    from super_resolve import build_model

    dev = torch.device(
        os.environ.get("DEVICE", "cuda") if torch.cuda.is_available() else "cpu"
    )
    print(f"Worker device: {dev}")
    print("Loading model...")
    t0 = time.monotonic()
    model = build_model(WEIGHTS_DIR, dev)
    print(f"Model loaded in {time.monotonic() - t0:.1f}s")
    return model, dev


def _safe_ack(ch: Any, delivery_tag: int) -> None:
    try:
        if ch.is_open:
            ch.basic_ack(delivery_tag=delivery_tag)
    except pika.exceptions.ChannelWrongStateError:
        print("Ack skipped: channel closed (reconnecting)")
    except pika.exceptions.ChannelClosedByBroker as exc:
        print(f"Ack skipped: broker closed channel ({exc})")
    except Exception as exc:
        print(f"Ack error: {exc}")


def _recover_orphaned_jobs(channel: Any) -> None:
    """
    После рестарта worker/Rabbit:
    - processing без живого процесса → cancelled (если cancel) или снова queued;
    - queued в SQLite при пустой/неполной очереди Rabbit → republish.
    """
    with get_conn() as conn:
        orphans = [
            dict(row)
            for row in conn.execute("SELECT * FROM jobs WHERE status = 'processing'")
        ]

    for job in orphans:
        job_id = int(job["id"])
        if int(job.get("cancel_requested") or 0):
            _cancel_job(job_id, job)
            continue
        dst = Path(job.get("output_path", ""))
        if dst.name:
            _cleanup_work_files(dst, job_id)
            if dst.exists():
                dst.unlink(missing_ok=True)
        update_job(
            job_id,
            status="queued",
            started_at=None,
            finished_at=None,
            progress_pct=None,
            progress_detail=None,
            error_message=None,
        )
        print(f"Job {job_id}: orphaned processing → queued")

    declared = channel.queue_declare(queue=QUEUE_NAME, durable=True)
    msg_count = int(declared.method.message_count)

    with get_conn() as conn:
        queued_ids = [
            int(row["id"])
            for row in conn.execute(
                "SELECT id FROM jobs WHERE status = 'queued' ORDER BY id ASC"
            )
        ]

    if not queued_ids:
        return
    if msg_count >= len(queued_ids):
        print(
            f"Queue sync OK: rabbit={msg_count}, queued_jobs={len(queued_ids)}"
        )
        return

    print(
        f"Queue sync: rabbit={msg_count} < queued_jobs={len(queued_ids)}, republishing"
    )
    for job_id in queued_ids:
        publish_job(job_id)
        print(f"Job {job_id}: published")


def main() -> None:
    init_db()

    model, dev = None, None
    if MOCK_MODE:
        if not shutil.which("ffmpeg"):
            sys.exit("MOCK_MODE: ffmpeg не найден в PATH")
        print("*** Режим «Только обработка»: FlashSR отключён ***")
    else:
        model, dev = _load_model()

    while True:
        connection = None
        try:
            connection = pika.BlockingConnection(_rabbit_params())
            channel = connection.channel()
            channel.queue_declare(queue=QUEUE_NAME, durable=True)
            channel.basic_qos(prefetch_count=1)
            _recover_orphaned_jobs(channel)

            def on_message(ch, method, _props, body: bytes) -> None:
                try:
                    payload = json.loads(body)
                    job_id = int(payload["job_id"])
                    process_job(job_id, model, dev)
                except Exception as exc:
                    print(f"Message error: {exc}")
                finally:
                    _safe_ack(ch, method.delivery_tag)

            channel.basic_consume(queue=QUEUE_NAME, on_message_callback=on_message)
            print(f"Waiting for messages on queue '{QUEUE_NAME}'...")
            channel.start_consuming()
        except (
            pika.exceptions.AMQPConnectionError,
            pika.exceptions.ChannelWrongStateError,
            pika.exceptions.ChannelClosedByBroker,
            pika.exceptions.StreamLostError,
            socket.gaierror,
            ConnectionError,
        ) as exc:
            print(f"RabbitMQ not ready ({exc}), retry in 5s...")
            time.sleep(5)
        finally:
            if connection is not None and not connection.is_closed:
                try:
                    connection.close()
                except Exception:
                    pass


if __name__ == "__main__":
    main()