xiaoxia-saas/apps/api/app/api/routes/gpu_lipsync.py

"""GPU MuseTalk Worker 反向轮询路由 — /api/v1/gpu/lipsync/*.

仅面向部署在用户 RTX2060 本地的 GPU Worker 脚本，不面向前端用户。
鉴权方式：长期 API Token（`Authorization: Bearer <GPU_WORKER_TOKEN>`），不走用户 JWT。

接口：
  POST  /api/v1/gpu/register              Worker 注册/心跳
  GET   /api/v1/gpu/lipsync/poll          Worker 轮询拉任务（无任务返回 204）
  POST  /api/v1/gpu/lipsync/result        Worker multipart 上传结果视频/上报失败
  GET   /api/v1/gpu/lipsync/status/{id}   业务侧查询任务状态（内部接口，暂开放给登录用户）
"""

from __future__ import annotations

import logging
import tempfile
from datetime import UTC, datetime
from pathlib import Path
from typing import Optional

import requests
from app.core.storage import get_storage_service
from app.dependencies import get_db_session
from app.schemas.gpu_lipsync import (
    GpuLipsyncPollResponse,
    GpuLipsyncResultResponse,
    GpuLipsyncStatusResponse,
    GpuLipsyncTaskPayload,
    GpuWorkerRegisterRequest,
    GpuWorkerRegisterResponse,
)
from app.services.gpu_lipsync_service import GpuLipsyncService
from fastapi import (
    APIRouter,
    Depends,
    File,
    Form,
    HTTPException,
    Query,
    Request,
    UploadFile,
    status,
)
from fastapi.responses import Response
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer

from packages.config import get_api_settings

logger = logging.getLogger(__name__)

router = APIRouter()

# 复用 bearer scheme 抽 Token，但不校验用户 JWT
_gpu_bearer = HTTPBearer(auto_error=False)


def _verify_gpu_token(
    credentials: Optional[HTTPAuthorizationCredentials] = Depends(_gpu_bearer),
) -> str:
    """校验 GPU Worker Token，返回 worker 提供的 token 串（仅用于日志，不做身份识别）.

    - development 且未配置 token → 直接放行（方便本地调试）。
    - production/staging 未配置 token → 拒绝（避免裸奔）。
    - token 不匹配 → 401。
    """
    settings = get_api_settings()
    expected = (settings.gpu_worker_token or "").strip()
    is_dev = settings.environment == "development"
    if not expected:
        if is_dev:
            return credentials.credentials if credentials else ""
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="GPU_WORKER_TOKEN not configured on server",
        )
    if credentials is None or credentials.scheme.lower() != "bearer":
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing bearer token")
    if credentials.credentials != expected:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid GPU worker token")
    return credentials.credentials


def _get_svc(db=Depends(get_db_session)) -> GpuLipsyncService:
    return GpuLipsyncService(db)


# ── POST /register — Worker 注册/心跳 ──────────────────────────────


@router.post("/register", response_model=GpuWorkerRegisterResponse)
def register_worker(
    body: GpuWorkerRegisterRequest,
    svc: GpuLipsyncService = Depends(_get_svc),
    _token: str = Depends(_verify_gpu_token),
):
    svc.register_worker(
        worker_id=body.worker_id,
        hostname=body.hostname,
        gpu_name=body.gpu_name,
        free_vram_mb=body.free_vram_mb,
        capabilities=body.capabilities,
        task_id=body.task_id,
    )
    return GpuWorkerRegisterResponse(ok=True, server_time=datetime.now(UTC), message="ok")


# ── GET /lipsync/poll — Worker 轮询拉任务 ─────────────────────────


@router.get("/lipsync/poll")
def poll_task(
    worker_id: str = Query(..., min_length=1, max_length=100, description="Worker 唯一 ID"),
    svc: GpuLipsyncService = Depends(_get_svc),
    _token: str = Depends(_verify_gpu_token),
):
    task = svc.poll_task(worker_id=worker_id)
    if task is None:
        return Response(status_code=status.HTTP_204_NO_CONTENT)
    payload = GpuLipsyncTaskPayload(
        task_id=task.id,
        video_url=getattr(task, "_signed_video_url", task.video_url),
        audio_url=getattr(task, "_signed_audio_url", task.audio_url),
        lipsync_job_id=task.lipsync_job_id or "",
        user_id=task.user_id or "",
        project_id=task.project_id or "",
        created_at=task.created_at,
        upload_url=getattr(task, "_signed_upload_url", ""),
        upload_method="PUT",
        expires_at=getattr(task, "_upload_expires_at", datetime.now(UTC)),
    )
    return GpuLipsyncPollResponse(task=payload)


# ── POST /lipsync/result — Worker 上报结果（multipart） ─────────────


@router.post("/lipsync/result", response_model=GpuLipsyncResultResponse)
async def report_result(
    request: Request,
    task_id: str = Form(...),
    worker_id: str = Form(...),
    success: bool = Form(True),
    duration_seconds: float = Form(0.0),
    error_msg: str = Form(""),
    result: Optional[UploadFile] = File(None),
    svc: GpuLipsyncService = Depends(_get_svc),
    _token: str = Depends(_verify_gpu_token),
):
    # 参数校验：
    # - success=true + result 文件 → API 代为上传到 OSS（方便 Worker 端实现）
    # - success=true + 无文件 → Worker 已经自己 PUT 到预签名 upload_url，直接确认
    # - success=false → 不上传文件，错误信息通过 error_msg 传递
    if success and result is not None:
        # 把文件落盘到临时目录，然后 PUT 到预签名 URL
        storage = get_storage_service()
        result_key = svc._result_key(task_id)
        upload_url = storage.get_upload_url(result_key, expires_seconds=3600, content_type="video/mp4")
        try:
            with tempfile.TemporaryDirectory(prefix="gpu_result_") as tmpdir:
                tmp_path = Path(tmpdir) / "result.mp4"
                content = await result.read()
                if not content:
                    raise HTTPException(status_code=400, detail="上传的 result 文件为空")
                tmp_path.write_bytes(content)
                headers = {"Content-Type": "video/mp4"}
                with open(tmp_path, "rb") as f:
                    resp = requests.put(upload_url, data=f, headers=headers, timeout=300)
                if resp.status_code >= 400:
                    logger.error(
                        "上传 GPU 结果到 OSS 失败: status=%d body=%s",
                        resp.status_code,
                        resp.text[:500],
                    )
                    raise HTTPException(
                        status_code=502,
                        detail=f"上传结果视频到 OSS 失败 (HTTP {resp.status_code})",
                    )
        except HTTPException:
            raise
        except Exception as exc:
            logger.exception("上传 GPU 结果视频异常: %s", exc)
            raise HTTPException(status_code=500, detail=f"上传结果视频异常: {exc}") from exc
    elif not success:
        # 失败时忽略 result 文件（即便传了也没用）
        pass
    # 其他情况：success=true 且无文件 → Worker 已自行 PUT 到预签名 URL，直接标记完成

    try:
        task = svc.report_result(
            task_id=task_id,
            worker_id=worker_id,
            success=success,
            duration_seconds=duration_seconds,
            error_msg=error_msg,
        )
    except KeyError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    return GpuLipsyncResultResponse(
        ok=True,
        task_id=task.id,
        status=task.status,
        message="ok",
    )


# ── GET /lipsync/status/{task_id} — 业务侧查询状态 ─────────────────
# 说明：此接口会被 lipsync_service 内部在业务流程里直接读 DB，不通过 HTTP。
# 但仍暴露一个简单查询接口，方便调试和前端轮询（如后续需要）。暂不做用户权限校验，
# task_id 本身是 UUID，不可枚举。


@router.get("/lipsync/status/{task_id}", response_model=GpuLipsyncStatusResponse)
def get_task_status(
    task_id: str,
    svc: GpuLipsyncService = Depends(_get_svc),
):
    task = svc.get_task(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail="task not found")
    return GpuLipsyncStatusResponse(
        task_id=task.id,
        status=task.status,
        result_url=task.result_url,
        result_duration=task.result_duration,
        error_msg=task.error_msg,
        worker_id=task.worker_id,
        attempt=task.attempt,
        created_at=task.created_at,
        started_at=task.started_at,
        finished_at=task.finished_at,
    )