xiaoxia-saas/apps/api/app/tasks/lipsync_gpu.py

"""GPU MuseTalk 异步推理任务 — 将 GPU 推理等待从 HTTP 请求移至 Celery 后台执行.

优化目标：将 POST /lipsync/jobs 的 API 响应时间从 >200s 降到 <1s。
任务流程：
1. 加载 LipsyncJob，获取 gpu_task_id
2. 调用 GpuLipsyncService.wait_for_result 轮询等待 GPU 完成
3. 签名结果 URL（7 天），更新 job 为 completed
4. 失败/超时时：尝试 MediaKit 兜底，若仍失败则标记 job 为 failed

使用 @shared_task 确保被 Worker 侧 celery_app 正确注册。
"""

import logging
from datetime import UTC, datetime

from celery import shared_task
from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import LipsyncJobModel
from packages.shared.storage import get_shared_storage_service

logger = logging.getLogger(__name__)

# 与 LipsyncService 保持一致
_MEDIAKIT_URL_TTL_SECONDS = 7 * 24 * 3600


def _get_db_session() -> Session:
    """获取 DB session（兼容 API 和 Worker 两种运行时）."""
    try:
        from worker_app.db import SessionLocal  # type: ignore
    except ImportError:
        from app.db import SessionLocal  # type: ignore
    return SessionLocal()


def _sign_media_url(url: str) -> str:
    """对自家 OSS URL 签 7 天预签名。"""
    if not url:
        return url
    try:
        from urllib.parse import urlparse

        storage = get_shared_storage_service()
        public_base = getattr(storage, "public_url", "")
        if not isinstance(public_base, str) or not public_base:
            return url
        own_host = urlparse(public_base).netloc.lower()
        host = urlparse(url).netloc.lower()
        if not own_host or host != own_host:
            return url
        return storage.get_download_url(url, expires_seconds=_MEDIAKIT_URL_TTL_SECONDS)
    except Exception:
        return url


@shared_task(
    name="lipsync_gpu_process_async",
    bind=True,
    max_retries=0,
    acks_late=True,
)
def lipsync_gpu_process_async(self, job_id: str, user_id: str, gpu_task_id: str) -> None:
    """异步处理 GPU MuseTalk 推理。

    Args:
        job_id: LipsyncJob 的 ID
        user_id: 用户 ID
        gpu_task_id: GpuLipsyncTask 的 ID
    """
    db: Session = _get_db_session()
    try:
        job = db.query(LipsyncJobModel).filter_by(id=job_id, user_id=user_id).first()
        if job is None:
            logger.error("[lipsync_gpu_async] job 不存在: job_id=%s", job_id)
            return

        # 确保状态为 processing
        if job.status not in ("processing", "gpu_processing"):
            logger.warning(
                "[lipsync_gpu_async] job 状态异常，跳过: job_id=%s status=%s",
                job_id,
                job.status,
            )
            return

        from app.services.gpu_lipsync_service import GpuLipsyncService

        gpu_svc = GpuLipsyncService(db)
        final_task = gpu_svc.wait_for_result(gpu_task_id)

        if final_task is None:
            logger.warning(
                "[lipsync_gpu_async] GPU 超时，回退 MediaKit: job_id=%s gpu_task=%s",
                job_id,
                gpu_task_id,
            )
            _fallback_to_mediakit(db, job)
            return

        if final_task.status == "cancelled":
            # 用户已取消任务，不回退 MediaKit，直接标记 job 为 cancelled
            job.status = "cancelled"
            job.updated_at = datetime.now(UTC)
            db.commit()
            logger.info("[lipsync_gpu_async] GPU 任务已被用户取消: job_id=%s", job_id)
            return

        if final_task.status != "done":
            logger.warning(
                "[lipsync_gpu_async] GPU 失败，回退 MediaKit: job_id=%s gpu_task=%s status=%s",
                job_id,
                gpu_task_id,
                final_task.status,
            )
            _fallback_to_mediakit(db, job)
            return

        # 签名结果 URL
        result_url = final_task.result_url or ""
        try:
            storage = get_shared_storage_service()
            signed = storage.get_download_url(result_url, expires_seconds=_MEDIAKIT_URL_TTL_SECONDS)
            if signed:
                result_url = signed
        except Exception as exc:
            logger.warning(
                "[lipsync_gpu_async] 签名失败，用原 URL: job_id=%s err=%s",
                job_id,
                exc,
            )

        job.status = "completed"
        job.output_video_url = result_url
        job.output_duration = final_task.result_duration or 0.0
        job.completed_at = datetime.now(UTC)
        job.updated_at = datetime.now(UTC)
        db.commit()
        logger.info(
            "[lipsync_gpu_async] GPU 完成: job_id=%s duration=%.2f",
            job_id,
            job.output_duration,
        )
    except Exception as exc:
        logger.exception("[lipsync_gpu_async] 异常: job_id=%s err=%s", job_id, exc)
        try:
            job = db.query(LipsyncJobModel).filter_by(id=job_id).first()
            if job:
                job.status = "failed"
                job.error_message = f"GPU 异步处理异常: {exc}"
                job.error_code = "GpuAsyncError"
                job.updated_at = datetime.now(UTC)
                db.commit()
        except Exception:
            pass
    finally:
        db.close()


def _fallback_to_mediakit(db: Session, job: LipsyncJobModel) -> None:
    """GPU 失败时回退到 MediaKit 云端渲染。"""
    try:
        from app.services.mediakit_client import MediaKitError, get_mediakit_client

        client = get_mediakit_client()
        video_url = _sign_media_url(job.video_url)
        audio_url = _sign_media_url(job.audio_url)

        result = client.submit_lipsync(
            video_url=video_url,
            audio_url=audio_url,
            enable_video_loop=job.enable_video_loop,
            client_token=job.id,
        )
        job.mediakit_task_id = result["task_id"]
        job.status = "submitted"
        job.submitted_at = datetime.now(UTC)
        job.updated_at = datetime.now(UTC)
        db.commit()
        logger.info(
            "[lipsync_gpu_async] 已回退 MediaKit: job_id=%s task_id=%s",
            job.id,
            result["task_id"],
        )
    except MediaKitError as exc:
        job.status = "failed"
        job.error_message = str(exc)
        job.error_code = exc.code
        job.updated_at = datetime.now(UTC)
        db.commit()
        logger.error("[lipsync_gpu_async] MediaKit 也失败: job_id=%s err=%s", job.id, exc)
    except Exception as exc:
        job.status = "failed"
        job.error_message = f"GPU+MediaKit 均失败: {exc}"
        job.error_code = "FallbackFailed"
        job.updated_at = datetime.now(UTC)
        db.commit()
        logger.error("[lipsync_gpu_async] 兜底异常: job_id=%s err=%s", job.id, exc)