xiaoxia-saas/apps/api/app/services/gpu_lipsync_service.py

"""GPU MuseTalk 口型同步服务 — 反向轮询模式.

职责：
1. 创建任务（由 lipsync 业务流程调用），为输入/输出生成预签名 URL，任务入队；
2. Worker 心跳注册（register）：登记/刷新 worker 状态；
3. Worker 轮询拉任务（poll）：原子地 CLAIM 一条 pending 任务，返回预签名 URL；
4. Worker 上报结果（report_result）：标记 done/failed，失败可重试；
5. 业务侧查询状态（get_status）。
"""

from __future__ import annotations

import logging
import uuid
from datetime import UTC, datetime, timedelta
from typing import Optional

from app.core.storage import get_storage_service
from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import GpuLipsyncTaskModel, GpuWorkerModel
from packages.config import get_api_settings

logger = logging.getLogger(__name__)

# 任务在 processing 超过此时长仍未完成 → 超时回退 pending 或置 failed
MAX_ATTEMPTS = 3


class GpuLipsyncService:
    """GPU 口型同步服务（无状态方法，每次调用从 DI 拿 db/storage）."""

    RESULT_PREFIX = "gpu-lipsync/results/"
    INPUT_SIGN_EXPIRES_PAD = 600  # 输入预签名 URL 在任务超时基础上再加 10min 余量

    # ── 公共入口 ────────────────────────────────────────────────────

    def __init__(self, db: Session):
        self.db = db
        self.settings = get_api_settings()
        self.storage = get_storage_service()

    # ── Worker 注册/心跳 ────────────────────────────────────────────

    def register_worker(
        self,
        worker_id: str,
        hostname: str = "",
        gpu_name: str = "",
        free_vram_mb: int = 0,
        capabilities: str = "musetalk",
        task_id: Optional[str] = None,
    ) -> tuple[GpuWorkerModel, bool]:
        """Worker 注册/心跳。

        task_id 非空时（Worker 推理期间的任务级心跳），同步把对应 processing
        任务的 last_heartbeat_at 续到当前时间，使长推理不会被
        ``_recover_timed_out_tasks`` 误回退。任务已结束 / 不属于该 worker
        （如已被超时回收重新派发）时忽略，不报错。

        返回 ``(worker, cancel_task)``：当心跳任务已被用户取消时
        ``cancel_task=True``，Worker 应尽快终止推理并释放 GPU。
        """
        now = datetime.now(UTC)
        worker = self.db.query(GpuWorkerModel).filter(GpuWorkerModel.worker_id == worker_id).one_or_none()
        if worker is None:
            worker = GpuWorkerModel(
                worker_id=worker_id,
                hostname=hostname,
                gpu_name=gpu_name,
                free_vram_mb=free_vram_mb,
                capabilities=capabilities,
                last_heartbeat_at=now,
                created_at=now,
            )
            self.db.add(worker)
        else:
            worker.hostname = hostname or worker.hostname
            worker.gpu_name = gpu_name or worker.gpu_name
            worker.free_vram_mb = free_vram_mb
            worker.capabilities = capabilities or worker.capabilities
            worker.last_heartbeat_at = now
        cancel_task = False
        if task_id:
            cancel_task = self._touch_task_heartbeat(task_id, worker_id, now)
        self.db.commit()
        return worker, cancel_task

    # ── 轮询拉任务（Worker 调用） ──────────────────────────────────

    def poll_task(self, worker_id: str) -> Optional[GpuLipsyncTaskModel]:
        """原子地认领一条最早的 pending 任务，返回给 worker；无任务返回 None.

        同时会：
        - 把 processing 状态且真正超时（任务心跳停滞超过
          gpu_task_timeout_seconds；Worker 推理期会通过 register(task_id=...)
          续心跳，长推理不会误判）的任务回退为 pending（attempt++，超过
          MAX_ATTEMPTS 置 failed），让其它 worker 认领。
        - 刷新 worker 心跳。
        """
        now = datetime.now(UTC)
        self._recover_timed_out_tasks(now)
        # 更新 worker 心跳
        self._touch_worker(worker_id, now)

        # 选一条最早 pending 任务（FOR UPDATE SKIP LOCKED 语义：简单起见先查再锁状态）
        task = (
            self.db.query(GpuLipsyncTaskModel)
            .filter(GpuLipsyncTaskModel.status == "pending")
            .order_by(GpuLipsyncTaskModel.created_at.asc())
            .first()
        )
        if task is None:
            self.db.commit()
            return None

        # 原子 claim：用 UPDATE WHERE status=pending 避免并发
        upd_rows = (
            self.db.query(GpuLipsyncTaskModel)
            .filter(
                GpuLipsyncTaskModel.id == task.id,
                GpuLipsyncTaskModel.status == "pending",
            )
            .update(
                {
                    GpuLipsyncTaskModel.status: "processing",
                    GpuLipsyncTaskModel.worker_id: worker_id,
                    GpuLipsyncTaskModel.started_at: now,
                    GpuLipsyncTaskModel.last_heartbeat_at: now,
                    GpuLipsyncTaskModel.attempt: GpuLipsyncTaskModel.attempt + 1,
                    GpuLipsyncTaskModel.updated_at: now,
                },
                synchronize_session=False,
            )
        )
        self.db.commit()
        if upd_rows == 0:
            # 被其它 worker 抢先了
            return None
        self.db.refresh(task)
        # 生成预签名输入/输出 URL（在 claim 时动态生成，避免长时间过期）
        expires = self.settings.gpu_task_timeout_seconds + self.INPUT_SIGN_EXPIRES_PAD
        task._signed_video_url = self.storage.get_download_url(task.video_url, expires_seconds=expires)
        task._signed_audio_url = self.storage.get_download_url(task.audio_url, expires_seconds=expires)
        task._signed_upload_url = self.storage.get_upload_url(
            self._result_key(task.id),
            expires_seconds=expires,
            content_type="video/mp4",
        )
        task._upload_expires_at = now + timedelta(seconds=expires)
        return task

    # ── 上报结果 ──────────────────────────────────────────────────

    def report_result(
        self,
        task_id: str,
        worker_id: str,
        success: bool,
        duration_seconds: float = 0.0,
        error_msg: str = "",
    ) -> GpuLipsyncTaskModel:
        task = self.db.get(GpuLipsyncTaskModel, task_id)
        if task is None:
            raise KeyError(f"task {task_id} not found")
        now = datetime.now(UTC)
        if success:
            task.status = "done"
            task.result_url = self._result_key(task_id)
            task.result_duration = duration_seconds or 0.0
            task.error_msg = ""
            task.finished_at = now
        elif task.status == "cancelled":
            # 用户已取消的任务，Worker 终止后上报失败，保持 cancelled 状态不回退
            task.finished_at = now
            task.error_msg = (error_msg or "用户取消")[:2000]
            logger.info("GPU 任务 %s 已被用户取消，保持 cancelled 状态", task_id)
        else:
            # 失败：若仍可重试（已尝试次数 < MAX_ATTEMPTS）→ 回退 pending；否则 → failed
            if task.attempt < MAX_ATTEMPTS:
                task.status = "pending"
                task.worker_id = ""
                task.started_at = None
                task.error_msg = error_msg[:2000]
                logger.warning(
                    "GPU 任务 %s 在 worker %s 上失败，回退 pending 等待重试（attempt=%d）: %s",
                    task_id,
                    worker_id,
                    task.attempt,
                    error_msg[:200],
                )
            else:
                task.status = "failed"
                task.error_msg = error_msg[:2000]
                task.finished_at = now
                logger.error(
                    "GPU 任务 %s 失败达到最大重试次数 %d，置为 failed: %s",
                    task_id,
                    MAX_ATTEMPTS,
                    error_msg[:200],
                )
        task.updated_at = now
        task.last_heartbeat_at = now
        self._touch_worker(worker_id, now)
        self.db.commit()
        self.db.refresh(task)
        return task

    # ── 业务侧查询 ────────────────────────────────────────────────

    def get_task(self, task_id: str) -> Optional[GpuLipsyncTaskModel]:
        return self.db.get(GpuLipsyncTaskModel, task_id)

    def get_by_lipsync_job(self, lipsync_job_id: str) -> Optional[GpuLipsyncTaskModel]:
        return (
            self.db.query(GpuLipsyncTaskModel)
            .filter(GpuLipsyncTaskModel.lipsync_job_id == lipsync_job_id)
            .order_by(GpuLipsyncTaskModel.created_at.desc())
            .first()
        )

    # ── 创建任务（业务侧调用） ────────────────────────────────────

    def create_task(
        self,
        video_url: str,
        audio_url: str,
        lipsync_job_id: str = "",
        user_id: str = "",
        project_id: str = "",
    ) -> GpuLipsyncTaskModel:
        task_id = str(uuid.uuid4())
        now = datetime.now(UTC)
        task = GpuLipsyncTaskModel(
            id=task_id,
            lipsync_job_id=lipsync_job_id,
            user_id=user_id,
            project_id=project_id,
            video_url=video_url,
            audio_url=audio_url,
            status="pending",
            attempt=0,
            created_at=now,
            updated_at=now,
        )
        self.db.add(task)
        self.db.commit()
        self.db.refresh(task)
        logger.info(
            "创建 GPU 口型任务 %s (lipsync_job=%s, user=%s)",
            task_id,
            lipsync_job_id,
            user_id,
        )
        return task

    # ── 内部辅助 ──────────────────────────────────────────────────

    def _result_key(self, task_id: str) -> str:
        return f"{self.RESULT_PREFIX}{task_id}.mp4"

    def _touch_task_heartbeat(self, task_id: str, worker_id: str, now: datetime) -> bool:
        """Worker 推理期间的任务级心跳：只刷新属于该 worker 且仍在 processing 的任务。

        任务不存在 / 已被超时回收重新派发 / 已完成 → 静默忽略（此时旧 worker 的
        结果上报会被结果接口按最终态处理）。

        返回 ``cancel_task``：任务已被用户取消时为 True，Worker 应终止推理。
        """
        task = self.db.get(GpuLipsyncTaskModel, task_id)
        if task is None:
            return False
        # 任务已被用户取消 → 通知 Worker 终止推理
        if task.status == "cancelled":
            logger.info("任务心跳检测到已取消 task=%s worker=%s，通知 Worker 终止", task_id, worker_id)
            return True
        if task.status != "processing" or task.worker_id != worker_id:
            logger.info(
                "忽略过期任务心跳 task=%s worker=%s（status=%s owner=%s）",
                task_id,
                worker_id,
                task.status,
                task.worker_id,
            )
            return False
        task.last_heartbeat_at = now
        task.updated_at = now
        self.db.flush()
        return False

    def _touch_worker(self, worker_id: str, now: datetime) -> None:
        if not worker_id:
            return
        worker = self.db.query(GpuWorkerModel).filter(GpuWorkerModel.worker_id == worker_id).one_or_none()
        if worker is not None:
            worker.last_heartbeat_at = now
            self.db.flush()
        else:
            # 自注册（poll 时允许自动建一个空 worker 记录，运维可见）
            worker = GpuWorkerModel(
                worker_id=worker_id,
                hostname="",
                gpu_name="",
                free_vram_mb=0,
                capabilities="musetalk",
                last_heartbeat_at=now,
                created_at=now,
            )
            self.db.add(worker)
            self.db.flush()

    def _recover_timed_out_tasks(self, now: datetime) -> None:
        """扫描 processing 状态且真正超时的任务，回退 pending 或失败。

        判定只看任务自身 last_heartbeat_at：claim 时写入，Worker 推理期间通过
        /gpu/register(task_id=...) 每 30s 续期。因此仅在 Worker 崩溃/断网
        （任务心跳停滞超过 gpu_task_timeout_seconds）时才回收，
        不会因 Worker 主循环忙于推理而误回退。
        """
        timeout = self.settings.gpu_task_timeout_seconds
        cutoff = now - timedelta(seconds=timeout)
        stuck_tasks = (
            self.db.query(GpuLipsyncTaskModel)
            .filter(
                GpuLipsyncTaskModel.status == "processing",
                GpuLipsyncTaskModel.last_heartbeat_at < cutoff,
            )
            .all()
        )
        for t in stuck_tasks:
            if t.attempt >= MAX_ATTEMPTS:
                t.status = "failed"
                t.error_msg = f"worker 心跳超时（{timeout}s），重试次数已耗尽"
                t.finished_at = now
            else:
                t.status = "pending"
                t.worker_id = ""
                t.started_at = None
                t.error_msg = f"worker 心跳超时（{timeout}s），等待重试"
                logger.warning("GPU 任务 %s 心跳超时，回退 pending（attempt=%d）", t.id, t.attempt)
            t.updated_at = now
        if stuck_tasks:
            self.db.flush()

    # ── 业务侧辅助 ──────────────────────────────────────────────────

    def has_available_worker(self) -> bool:
        """判断是否有 Worker 在心跳新鲜窗口内可用."""
        stale_cutoff = datetime.now(UTC) - timedelta(seconds=self.settings.gpu_worker_stale_seconds)
        return (
            self.db.query(GpuWorkerModel).filter(GpuWorkerModel.last_heartbeat_at >= stale_cutoff).first() is not None
        )

    def wait_for_result(
        self,
        task_id: str,
        timeout_seconds: Optional[int] = None,
        poll_interval: Optional[float] = None,
    ) -> Optional[GpuLipsyncTaskModel]:
        """同步轮询等待 GPU 任务完成。

        Args:
            task_id: 任务 ID（由 create_task 返回）
            timeout_seconds: 总超时，默认取 settings.gpu_lipsync_wait_timeout
            poll_interval: 轮询间隔秒，默认取 settings.gpu_lipsync_poll_interval

        Returns:
            终态 task（status=done/failed）；超时返回 None（此时调用方应回退 MediaKit）。
            等待期间会自动调用 _recover_timed_out_tasks 做超时回收。
        """
        import time

        timeout = timeout_seconds if timeout_seconds is not None else self.settings.gpu_lipsync_wait_timeout
        interval = poll_interval if poll_interval is not None else self.settings.gpu_lipsync_poll_interval
        deadline = time.monotonic() + timeout

        while True:
            now = datetime.now(UTC)
            # 顺手回收超时任务
            try:
                self._recover_timed_out_tasks(now)
                self.db.commit()
            except Exception as exc:  # noqa: BLE001 - 回收失败不阻塞主流程
                logger.warning("wait_for_result 回收超时任务异常: %s", exc)
                self.db.rollback()

            task = self.db.get(GpuLipsyncTaskModel, task_id)
            if task is None:
                return None
            if task.status in ("done", "failed", "cancelled"):
                return task
            # pending/processing 继续等
            if time.monotonic() >= deadline:
                logger.warning("GPU 任务 %s 等待超时（%ds），回退 MediaKit", task_id, timeout)
                return None
            time.sleep(interval)