diff --git a/.env.example b/.env.example index 3f3eac461..39266bb27 100755 --- a/.env.example +++ b/.env.example @@ -213,3 +213,10 @@ TIKHUB_API_KEY= # apizero.cn API Key (https://v1.apizero.cn) — 国内抖音解析服务 APIZERO_API_KEY= +# ==================== GPU MuseTalk Worker（反向轮询口型同步）==================== +# GPU Worker 长期鉴权 Token，Worker 端 .env 的 GPU_WORKER_TOKEN 必须与此一致 +# 留空时 development 环境允许匿名访问（仅本地调试），staging/production 必须配置 +GPU_WORKER_TOKEN= +# 单任务超时（秒），超过则回退 pending 或标记 failed +GPU_TASK_TIMEOUT_SECONDS=300 + diff --git a/.gitea/workflows/ci-pipeline.yml b/.gitea/workflows/ci-pipeline.yml index 95564ae3b..a0d5d9a17 100755 --- a/.gitea/workflows/ci-pipeline.yml +++ b/.gitea/workflows/ci-pipeline.yml @@ -1190,6 +1190,7 @@ jobs: WECHAT_APP_SECRET: "${{ secrets.WECHAT_APP_SECRET }}" TIKHUB_API_KEY: "${{ secrets.TIKHUB_API_KEY }}" APIZERO_API_KEY: "${{ secrets.APIZERO_API_KEY }}" + GPU_WORKER_TOKEN: "${{ secrets.GPU_WORKER_TOKEN }}" run: | set -eu echo "Rendering .env from template + secrets..." @@ -1644,6 +1645,7 @@ jobs: WECHAT_APP_SECRET: "${{ secrets.WECHAT_APP_SECRET }}" TIKHUB_API_KEY: "${{ secrets.TIKHUB_API_KEY }}" APIZERO_API_KEY: "${{ secrets.APIZERO_API_KEY }}" + GPU_WORKER_TOKEN: "${{ secrets.GPU_WORKER_TOKEN }}" run: | set -eu echo "Rendering .env from template + secrets..." diff --git a/alembic/versions/081_add_gpu_lipsync_tasks.py b/alembic/versions/081_add_gpu_lipsync_tasks.py new file mode 100644 index 000000000..f15e0234d --- /dev/null +++ b/alembic/versions/081_add_gpu_lipsync_tasks.py @@ -0,0 +1,58 @@ +"""add gpu_lipsync_tasks and gpu_workers tables for MuseTalk reverse-poll worker + +Revision ID: 081_add_gpu_lipsync +Revises: 080_edit_plan_clips_atom_clip_id +Create Date: 2026-09-18 +""" + +import sqlalchemy as sa + +from alembic import op + +revision = "081_add_gpu_lipsync" +down_revision = "080_edit_plan_clips_atom_clip_id" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # GPU Worker 注册表 + op.create_table( + "gpu_workers", + sa.Column("worker_id", sa.String(100), primary_key=True), + sa.Column("hostname", sa.String(200), nullable=False, server_default=""), + sa.Column("gpu_name", sa.String(200), nullable=False, server_default=""), + sa.Column("free_vram_mb", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column("capabilities", sa.String(500), nullable=False, server_default=""), + sa.Column("last_heartbeat_at", sa.DateTime(), nullable=True, index=True), + sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + ) + + # GPU 口型同步任务表 + op.create_table( + "gpu_lipsync_tasks", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("lipsync_job_id", sa.String(36), nullable=False, server_default="", index=True), + sa.Column("user_id", sa.String(36), nullable=False, server_default="", index=True), + sa.Column("project_id", sa.String(36), nullable=False, server_default="", index=True), + sa.Column("video_url", sa.Text(), nullable=False), + sa.Column("audio_url", sa.Text(), nullable=False), + sa.Column("result_url", sa.Text(), nullable=False, server_default=""), + sa.Column("result_duration", sa.Float(), nullable=False, server_default=sa.text("0.0")), + sa.Column("status", sa.String(20), nullable=False, server_default="pending", index=True), + sa.Column("worker_id", sa.String(100), nullable=False, server_default="", index=True), + sa.Column("attempt", sa.Integer(), nullable=False, server_default=sa.text("0")), + sa.Column("error_msg", sa.Text(), nullable=False, server_default=""), + sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.Column("started_at", sa.DateTime(), nullable=True), + sa.Column("finished_at", sa.DateTime(), nullable=True), + sa.Column("updated_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.Column("last_heartbeat_at", sa.DateTime(), nullable=True), + ) + op.create_index("ix_gpu_lipsync_status_created", "gpu_lipsync_tasks", ["status", "created_at"]) + + +def downgrade() -> None: + op.drop_index("ix_gpu_lipsync_status_created", table_name="gpu_lipsync_tasks") + op.drop_table("gpu_lipsync_tasks") + op.drop_table("gpu_workers") diff --git a/apps/api/app/api/router.py b/apps/api/app/api/router.py index ab8df9da4..0a9d7e857 100755 --- a/apps/api/app/api/router.py +++ b/apps/api/app/api/router.py @@ -14,6 +14,7 @@ from app.api.routes.generation_cover import router as generation_cover_router from app.api.routes.generation_preview import router as generation_preview_router from app.api.routes.generation_tasks import router as generation_tasks_router from app.api.routes.generation_variant_plans import router as generation_variant_plans_router +from app.api.routes.gpu_lipsync import router as gpu_lipsync_router from app.api.routes.health import router as health_check_router from app.api.routes.ingest_jobs import router as ingest_jobs_router from app.api.routes.internal_render import router as internal_render_router @@ -211,3 +212,8 @@ api_router.include_router( prefix="/usage", tags=["Usage"], ) +api_router.include_router( + gpu_lipsync_router, + prefix="/gpu", + tags=["GPU Worker"], +) diff --git a/apps/api/app/api/routes/gpu_lipsync.py b/apps/api/app/api/routes/gpu_lipsync.py new file mode 100644 index 000000000..db3a83821 --- /dev/null +++ b/apps/api/app/api/routes/gpu_lipsync.py @@ -0,0 +1,227 @@ +"""GPU MuseTalk Worker 反向轮询路由 — /api/v1/gpu/lipsync/*. + +仅面向部署在用户 RTX2060 本地的 GPU Worker 脚本，不面向前端用户。 +鉴权方式：长期 API Token（`Authorization: Bearer `），不走用户 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, + ) + 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，直接标记完成 + + task = svc.report_result( + task_id=task_id, + worker_id=worker_id, + success=success, + duration_seconds=duration_seconds, + error_msg=error_msg, + ) + 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, + ) diff --git a/apps/api/app/schemas/gpu_lipsync.py b/apps/api/app/schemas/gpu_lipsync.py new file mode 100644 index 000000000..592c4227b --- /dev/null +++ b/apps/api/app/schemas/gpu_lipsync.py @@ -0,0 +1,103 @@ +"""GPU MuseTalk 反向轮询 API Schema 定义. + +面向部署在用户 RTX2060 本地的 GPU Worker 脚本，不面向前端用户。 +Worker 用长期 GPU_WORKER_TOKEN 鉴权（不是用户 JWT）。 +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Optional + +from pydantic import BaseModel, Field + +# ── Worker 注册/心跳 ────────────────────────────────────────────── + + +class GpuWorkerRegisterRequest(BaseModel): + """Worker 启动/心跳时上报自身信息.""" + + worker_id: str = Field(..., min_length=1, max_length=100, description="Worker 唯一 ID（机器名+UUID 等）") + hostname: str = Field("", max_length=200, description="主机名，用于运维排查") + gpu_name: str = Field("", max_length=200, description="GPU 型号，如 'NVIDIA GeForce RTX 2060'") + free_vram_mb: int = Field(0, ge=0, description="当前空闲显存（MB）") + capabilities: str = Field("musetalk", max_length=500, description="能力列表，逗号分隔，如 'musetalk'") + + +class GpuWorkerRegisterResponse(BaseModel): + ok: bool = True + server_time: datetime + message: str = "ok" + + +# ── 轮询任务 ──────────────────────────────────────────────────── + + +class GpuLipsyncTaskPayload(BaseModel): + """下发给 Worker 的任务载荷（含预签名下载 URL）.""" + + task_id: str + video_url: str = Field(..., description="人物视频预签名下载 URL（GET）") + audio_url: str = Field(..., description="驱动音频预签名下载 URL（GET）") + lipsync_job_id: str = "" + user_id: str = "" + project_id: str = "" + created_at: datetime + upload_url: str = Field(..., description="结果视频预签名上传 URL（PUT, video/mp4）") + upload_method: str = Field("PUT", description="上传方式，目前只支持 PUT") + expires_at: datetime + + +class GpuLipsyncPollResponse(BaseModel): + """Worker poll 的返回：200 带任务，204 无任务.""" + + task: Optional[GpuLipsyncTaskPayload] = None + + +# ── Worker 上报结果 ────────────────────────────────────────────── + + +class GpuLipsyncResultRequest(BaseModel): + """Worker 通过 multipart 上传结果时携带的字段（非文件字段）.""" + + task_id: str = Field(..., min_length=1, max_length=64) + worker_id: str = Field(..., min_length=1, max_length=100) + success: bool = Field(True, description="true=成功（此时必须上传 result 视频文件）；false=失败") + duration_seconds: float = Field(0.0, ge=0, description="合成后视频时长（秒），成功时应填入") + error_msg: str = Field("", max_length=2000, description="失败原因，success=false 时必填") + + +class GpuLipsyncResultResponse(BaseModel): + ok: bool = True + task_id: str + status: str # done / failed + message: str = "ok" + + +# ── 业务侧查询任务状态 ──────────────────────────────────────────── + + +class GpuLipsyncStatusResponse(BaseModel): + task_id: str + status: str + result_url: str = "" + result_duration: float = 0.0 + error_msg: str = "" + worker_id: str = "" + attempt: int = 0 + created_at: datetime + started_at: Optional[datetime] = None + finished_at: Optional[datetime] = None + + +# ── 创建任务（内部服务调用） ────────────────────────────────────── + + +class GpuLipsyncCreateRequest(BaseModel): + """服务层内部创建 GPU 任务用（不通过 HTTP 暴露给 Worker/前端）.""" + + video_url: str # 已可访问的 OSS key 或公网 URL（API 侧会转预签名） + audio_url: str + lipsync_job_id: str = "" + user_id: str = "" + project_id: str = "" diff --git a/apps/api/app/services/gpu_lipsync_service.py b/apps/api/app/services/gpu_lipsync_service.py new file mode 100644 index 000000000..a3c1cac14 --- /dev/null +++ b/apps/api/app/services/gpu_lipsync_service.py @@ -0,0 +1,287 @@ +"""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", + ) -> GpuWorkerModel: + 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 + self.db.commit() + return worker + + # ── 轮询拉任务（Worker 调用） ────────────────────────────────── + + def poll_task(self, worker_id: str) -> Optional[GpuLipsyncTaskModel]: + """原子地认领一条最早的 pending 任务，返回给 worker；无任务返回 None. + + 同时会： + - 把 processing 状态且超时（超过 gpu_task_timeout_seconds 无心跳）的任务 + 回退为 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 + 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_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 或失败.""" + 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() diff --git a/deploy/configs/.env.production b/deploy/configs/.env.production index 8a212f919..2048200e8 100644 --- a/deploy/configs/.env.production +++ b/deploy/configs/.env.production @@ -252,3 +252,7 @@ DOUYIN_DEBUG_ERRORS=false TIKHUB_API_KEY=${TIKHUB_API_KEY} # P2: apizero.cn（国内付费，https://apizero.cn） APIZERO_API_KEY=${APIZERO_API_KEY} + +# ==================== GPU MuseTalk Worker（反向轮询） ==================== +GPU_WORKER_TOKEN=${GPU_WORKER_TOKEN} +GPU_TASK_TIMEOUT_SECONDS=300 diff --git a/deploy/configs/.env.staging b/deploy/configs/.env.staging index e8b5f2412..e25c04ec4 100644 --- a/deploy/configs/.env.staging +++ b/deploy/configs/.env.staging @@ -269,3 +269,7 @@ DOUYIN_DEBUG_ERRORS=false TIKHUB_API_KEY=${TIKHUB_API_KEY} # P2: apizero.cn（国内付费，https://apizero.cn） APIZERO_API_KEY=${APIZERO_API_KEY} + +# ==================== GPU MuseTalk Worker（反向轮询） ==================== +GPU_WORKER_TOKEN=${GPU_WORKER_TOKEN} +GPU_TASK_TIMEOUT_SECONDS=300 diff --git a/deploy/gpu_worker/.env.example b/deploy/gpu_worker/.env.example new file mode 100644 index 000000000..74bd32480 --- /dev/null +++ b/deploy/gpu_worker/.env.example @@ -0,0 +1,25 @@ +# ============================================================ +# MuseTalk GPU Worker 环境变量 +# 部署到 RTX2060 电脑后，复制为 .env 并修改值 +# ============================================================ + +# SaaS API 基础 URL（staging / production） +API_BASE_URL=https://staging-api.xiaoxiajianji.com +# API_BASE_URL=https://api.xiaoxiajianji.com # 生产 + +# 长期 API Token，必须与服务端 GPU_WORKER_TOKEN 一致（找后端拿） +GPU_WORKER_TOKEN=replace-with-real-token + +# 本机 Worker 唯一 ID（默认自动生成 hostname+MAC 后4位，可手动指定） +# WORKER_ID=rtx2060-0193 + +# 本地 MuseTalk 地址（默认 http://127.0.0.1:7861） +MUSE_TALK_URL=http://127.0.0.1:7861 + +# 轮询/心跳/超时（秒） +POLL_INTERVAL=5 +HEARTBEAT_INTERVAL=15 +REQUEST_TIMEOUT=300 + +# 单个任务本地最大重试次数（首次失败后再重试 N 次，默认 2） +TASK_MAX_RETRY=2 diff --git a/deploy/gpu_worker/README.md b/deploy/gpu_worker/README.md new file mode 100644 index 000000000..8230a63e8 --- /dev/null +++ b/deploy/gpu_worker/README.md @@ -0,0 +1,99 @@ +# MuseTalk GPU Worker — 部署指南 + +本目录包含 RTX2060 本地电脑上运行的 GPU Worker 脚本。 +Worker 采用 **反向轮询模式**：主动向 SaaS API 拉取待处理的口型同步任务 → 调用本地 MuseTalk 推理 → 把结果视频回传到 SaaS。不需要内网穿透。 + +## 目录文件 + +| 文件 | 作用 | +|---|---| +| `gpu_worker.py` | Worker 主程序（单文件，零项目代码依赖，仅依赖 `requests`） | +| `requirements.txt` | Python 依赖（只有 `requests`） | +| `xiaoxia-gpu-worker.service` | systemd 服务单元（开机自启、异常自动重启） | +| `.env.example` | 环境变量样例，复制为 `.env` 后填入真实值 | + +## 一、环境准备 + +1. **Python 3.10+**（Windows 建议从 python.org 安装；Linux 自带） +2. **本地 MuseTalk 服务** 已启动在 `http://127.0.0.1:7861`，health 接口返回 `{"status":"ok","free_vram_mb":...}` +3. **ffmpeg**（可选，用于读取输出视频时长；未装则 duration 报 0，不影响功能） +4. 网络能访问 staging / 生产 API（`curl https://staging-api.xiaoxiajianji.com/health` 应返回 `{"status":"healthy"}`） + +## 二、部署步骤（Linux，推荐 systemd） + +```bash +# 1. 创建部署目录 +sudo mkdir -p /opt/xiaoxia-gpu-worker +sudo chown $USER:$USER /opt/xiaoxia-gpu-worker +cd /opt/xiaoxia-gpu-worker + +# 2. 拷贝脚本和依赖 +cp /path/to/deploy/gpu_worker/{gpu_worker.py,requirements.txt,xiaoxia-gpu-worker.service,.env.example} . +cp .env.example .env +# 编辑 .env，填入 API_BASE_URL 和 GPU_WORKER_TOKEN + +# 3. 创建虚拟环境并安装依赖 +python3 -m venv venv +./venv/bin/pip install -r requirements.txt + +# 4. 前台先跑一次，确认日志正常 +./venv/bin/python gpu_worker.py +# 看到 "MuseTalk 健康检查通过" 和 "注册/心跳" 成功即可 Ctrl+C 退出 + +# 5. 安装 systemd 服务 +sudo cp xiaoxia-gpu-worker.service /etc/systemd/system/ +sudo systemctl daemon-reload +sudo systemctl enable --now xiaoxia-gpu-worker + +# 6. 查看日志 +sudo journalctl -u xiaoxia-gpu-worker -f +``` + +## 三、部署步骤（Windows，快速测试） + +```bat +:: 创建虚拟环境 +python -m venv venv +venv\Scripts\pip install -r requirements.txt + +:: 复制并编辑 .env +copy .env.example .env +notepad .env + +:: 运行 +venv\Scripts\python gpu_worker.py +``` + +可在任务计划程序中添加开机启动项：程序选 `venv\Scripts\python.exe`，参数填 `gpu_worker.py`，起始目录填脚本所在目录。 + +## 四、SaaS 侧配套配置 + +SaaS 后端部署完成后需配置： + +1. 服务端环境变量 `GPU_WORKER_TOKEN` 设为一个随机强 Token（和 Worker `.env` 中一致） +2. 数据库已跑迁移 `081_add_gpu_lipsync_tasks`（自动随 API 启动的 alembic upgrade head 完成） +3. OSS bucket 中 `gpu-lipsync/results/` 路径可写（默认 bucket 已配） + +## 五、验证联调 + +1. Worker 启动后日志看到 `注册/心跳` 成功 +2. 后端调用 `GpuLipsyncService.create_task(video_url=..., audio_url=...)` 放入一条测试任务 +3. Worker 在 5 秒内拉到任务，下载 → 推理 → 上传 → 上报 +4. 后端 `GET /api/v1/gpu/lipsync/status/{task_id}` 返回 `status=done`，`result_url` 非空 + +## 六、故障排查 + +| 现象 | 可能原因 / 排查 | +|---|---| +| 日志 401 `Invalid GPU worker token` | `.env` 的 `GPU_WORKER_TOKEN` 与服务端不一致 | +| 日志 `MuseTalk 健康检查未通过` | 本地 MuseTalk 没启动，或端口不是 7861；`curl http://127.0.0.1:7861/health` 验证 | +| 任务长时间不被拉取 | Worker 和服务端连不上；检查 API_BASE_URL 是否可达、Token 是否正确 | +| 推理后上传 OSS 失败 | 本地出口网络被防火墙拦截 OSS 域名（oss-cn-hangzhou.aliyuncs.com） | +| 服务端看到任务回退到 pending 重试 | Worker 心跳超时（默认 5 分钟）；Worker 进程崩溃或推理卡死超过 5 分钟 | +| 日志 `MuseTalk 推理超时` | 视频太长或显存不足；可临时调大 REQUEST_TIMEOUT，或限制输入视频时长 | + +## 七、安全注意事项 + +- `.env` 包含长期 Token，文件权限设为 600（`chmod 600 .env`） +- Token 泄露要立即在服务端更换 `GPU_WORKER_TOKEN` 并重启 Worker +- Worker 只需要出站访问 SaaS API 和 OSS，不需要开放任何入站端口 diff --git a/deploy/gpu_worker/gpu_worker.py b/deploy/gpu_worker/gpu_worker.py new file mode 100644 index 000000000..ecdeac35b --- /dev/null +++ b/deploy/gpu_worker/gpu_worker.py @@ -0,0 +1,395 @@ +"""MuseTalk GPU Worker — 反向轮询模式. + +部署在有 RTX2060 的本地电脑上（192.168.0.193）， +主动轮询 SaaS API 拉取口型任务、调用本地 MuseTalk 推理、上传结果回 SaaS。 + +环境变量： + API_BASE_URL SaaS API 基础 URL（不含 /api/v1），如 https://staging-api.xiaoxiajianji.com + GPU_WORKER_TOKEN 长期 API Token（服务端 GPU_WORKER_TOKEN 需一致） + WORKER_ID 本机唯一 ID（默认 hostname+网卡MAC 后4位） + MUSE_TALK_URL 本地 MuseTalk 地址，默认 http://127.0.0.1:7861 + POLL_INTERVAL 轮询间隔秒，默认 5 + HEARTBEAT_INTERVAL 心跳间隔秒，默认 15 + REQUEST_TIMEOUT HTTP 请求超时秒，默认 60 + TASK_MAX_RETRY 单个任务最大重试次数（在 Worker 本地的重试），默认 2 + +用法： + python gpu_worker.py +""" + +from __future__ import annotations + +import json +import logging +import os +import platform +import socket +import sys +import tempfile +import time +import uuid +from pathlib import Path +from typing import Optional + +import requests + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger("musetalk-worker") + +# ── 配置 ──────────────────────────────────────────────────────────── + + +def _env(name: str, default: str = "") -> str: + v = os.environ.get(name, default) + return v.strip() if isinstance(v, str) else default + + +class Config: + api_base_url: str = _env("API_BASE_URL", "https://staging-api.xiaoxiajianji.com").rstrip("/") + gpu_worker_token: str = _env("GPU_WORKER_TOKEN") + muse_talk_url: str = _env("MUSE_TALK_URL", "http://127.0.0.1:7861").rstrip("/") + poll_interval: float = float(_env("POLL_INTERVAL", "5")) + heartbeat_interval: float = float(_env("HEARTBEAT_INTERVAL", "15")) + request_timeout: float = float(_env("REQUEST_TIMEOUT", "300")) + task_max_retry: int = int(_env("TASK_MAX_RETRY", "2")) + worker_id: str = _env("WORKER_ID", "") + + @classmethod + def derived_worker_id(cls) -> str: + if cls.worker_id: + return cls.worker_id + # hostname + MAC 后4位 → 稳定唯一 ID + try: + mac = uuid.getnode() + mac_suffix = f"{mac:012x}"[-4:] + except Exception: + mac_suffix = "0000" + host = platform.node() or socket.gethostname() or "rtx2060" + return f"{host}-{mac_suffix}" + + +# ── 辅助 ───────────────────────────────────────────────────────────── + + +def _api_headers() -> dict[str, str]: + token = Config.gpu_worker_token + if not token: + logger.warning("GPU_WORKER_TOKEN 未配置，开发模式下会被服务端拒绝（生产环境必须配置）") + return {"Authorization": f"Bearer {token}"} if token else {} + + +def _check_musetalk_health() -> tuple[bool, dict]: + """检查本地 MuseTalk 健康状态，返回 (ok, info).""" + try: + r = requests.get(f"{Config.muse_talk_url}/health", timeout=5) + if r.status_code == 200: + try: + return True, r.json() + except Exception: + return True, {} + return False, {"status_code": r.status_code, "body": r.text[:200]} + except Exception as exc: + return False, {"error": str(exc)} + + +def _register() -> bool: + """向服务端注册 / 心跳，附带 GPU 信息.""" + ok, info = _check_musetalk_health() + free_vram = int(info.get("free_vram_mb", 0) or 0) if isinstance(info, dict) else 0 + gpu_name = info.get("gpu_name", "") if isinstance(info, dict) else "" + if not gpu_name: + # 尝试在 Windows 上读 nvidia-smi + gpu_name = _probe_gpu_name() + payload = { + "worker_id": Config.derived_worker_id(), + "hostname": platform.node(), + "gpu_name": gpu_name, + "free_vram_mb": free_vram, + "capabilities": "musetalk", + } + try: + r = requests.post( + f"{Config.api_base_url}/api/v1/gpu/register", + json=payload, + headers=_api_headers(), + timeout=15, + ) + if r.status_code == 200: + return True + logger.error("注册/心跳失败: HTTP %d body=%s", r.status_code, r.text[:300]) + return False + except Exception as exc: + logger.error("注册/心跳异常: %s", exc) + return False + + +def _probe_gpu_name() -> str: + """尽力探测 GPU 型号（不强制依赖 pynvml）.""" + try: + import subprocess + + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + stderr=subprocess.DEVNULL, + timeout=5, + ) + return out.decode("utf-8", errors="ignore").strip().splitlines()[0].strip() + except Exception: + return "" + + +def _poll_task() -> Optional[dict]: + """轮询拉取一条待处理任务；无任务返回 None.""" + try: + r = requests.get( + f"{Config.api_base_url}/api/v1/gpu/lipsync/poll", + params={"worker_id": Config.derived_worker_id()}, + headers=_api_headers(), + timeout=30, + ) + if r.status_code == 204: + return None + if r.status_code == 200: + data = r.json() + return data.get("task") + logger.error("poll 返回 %d: %s", r.status_code, r.text[:300]) + return None + except Exception as exc: + logger.error("poll 异常: %s", exc) + return None + + +def _download(url: str, path: Path) -> bool: + """下载文件到本地，支持预签名 URL.""" + try: + with requests.get(url, stream=True, timeout=Config.request_timeout) as r: + if r.status_code >= 400: + logger.error("下载失败 HTTP %d: %s", r.status_code, url[:120]) + return False + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "wb") as f: + for chunk in r.iter_content(chunk_size=1024 * 256): + if chunk: + f.write(chunk) + return path.stat().st_size > 0 + except Exception as exc: + logger.error("下载异常 %s: %s", url[:120], exc) + return False + + +def _call_musetalk(video_path: Path, audio_path: Path, out_path: Path) -> tuple[bool, float, str]: + """调用本地 MuseTalk /inference. + + 返回 (success, duration_seconds, error_msg). + duration 用 ffprobe 读结果视频，失败填 0。 + """ + try: + with open(video_path, "rb") as vf, open(audio_path, "rb") as af: + files = { + "video": (video_path.name, vf, "video/mp4"), + "audio": (audio_path.name, af, "application/octet-stream"), + } + r = requests.post( + f"{Config.muse_talk_url}/inference", + files=files, + timeout=Config.request_timeout, + ) + if r.status_code != 200: + return False, 0.0, f"MuseTalk HTTP {r.status_code}: {r.text[:500]}" + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_bytes(r.content) + if out_path.stat().st_size < 1024: + return False, 0.0, f"MuseTalk 返回结果过小 ({out_path.stat().st_size} bytes)" + duration = _probe_duration(out_path) + return True, duration, "" + except requests.exceptions.Timeout: + return False, 0.0, f"MuseTalk 推理超时（>{Config.request_timeout}s）" + except Exception as exc: + return False, 0.0, f"MuseTalk 调用异常: {exc}" + + +def _probe_duration(path: Path) -> float: + """用 ffprobe 读视频时长（若系统装了 ffmpeg）；否则返回 0.""" + try: + import subprocess + + out = subprocess.check_output( + [ + "ffprobe", "-v", "error", + "-show_entries", "format=duration", + "-of", "default=noprint_wrappers=1:nokey=1", + str(path), + ], + stderr=subprocess.DEVNULL, + timeout=10, + ) + return float(out.decode().strip() or 0) + except Exception: + return 0.0 + + +def _upload_result(upload_url: str, file_path: Path) -> bool: + """PUT 上传结果视频到预签名 URL.""" + try: + with open(file_path, "rb") as f: + r = requests.put( + upload_url, + data=f, + headers={"Content-Type": "video/mp4"}, + timeout=Config.request_timeout, + ) + if r.status_code >= 400: + logger.error("上传结果失败 HTTP %d: %s", r.status_code, r.text[:500]) + return False + return True + except Exception as exc: + logger.error("上传结果异常: %s", exc) + return False + + +def _report_result(task_id: str, success: bool, duration: float = 0.0, error_msg: str = "") -> bool: + """通知服务端结果。失败时也尝试上报错误（不含视频文件）.""" + try: + data = { + "task_id": task_id, + "worker_id": Config.derived_worker_id(), + "success": "true" if success else "false", + "duration_seconds": str(duration), + "error_msg": error_msg, + } + r = requests.post( + f"{Config.api_base_url}/api/v1/gpu/lipsync/result", + data=data, + headers=_api_headers(), + timeout=30, + ) + if r.status_code != 200: + logger.error("上报结果失败 HTTP %d: %s", r.status_code, r.text[:300]) + return False + return True + except Exception as exc: + logger.error("上报结果异常: %s", exc) + return False + + +def _handle_task(task: dict) -> None: + """处理一条任务（整个串行流程：下载→推理→上传→上报）.""" + task_id = task["task_id"] + logger.info("开始处理任务 %s", task_id) + with tempfile.TemporaryDirectory(prefix="musetalk_") as tmpdir: + tmp = Path(tmpdir) + video_path = tmp / "input.mp4" + audio_path = tmp / "input_audio.bin" + out_path = tmp / "output.mp4" + + # 1. 下载 + if not _download(task["video_url"], video_path): + _report_result(task_id, False, 0.0, "下载人物视频失败") + return + if not _download(task["audio_url"], audio_path): + _report_result(task_id, False, 0.0, "下载驱动音频失败") + return + + # 2. 推理（本地重试） + success = False + duration = 0.0 + err = "" + for attempt in range(Config.task_max_retry + 1): + if attempt > 0: + logger.info("任务 %s 第 %d 次重试...", task_id, attempt + 1) + time.sleep(2) + success, duration, err = _call_musetalk(video_path, audio_path, out_path) + if success: + break + if not success: + logger.error("任务 %s 推理失败: %s", task_id, err) + _report_result(task_id, False, 0.0, err) + return + + # 3. 上报结果（multipart 同时上传文件 → API 代为 PUT 到 OSS，逻辑最稳） + _report_success_with_file(task_id, duration, out_path) + + +def _report_success_with_file(task_id: str, duration: float, file_path: Path) -> None: + """上报成功并 multipart 附带结果视频.""" + try: + data = { + "task_id": task_id, + "worker_id": Config.derived_worker_id(), + "success": "true", + "duration_seconds": str(duration), + "error_msg": "", + } + with open(file_path, "rb") as f: + files = {"result": (f"{task_id}.mp4", f, "video/mp4")} + r = requests.post( + f"{Config.api_base_url}/api/v1/gpu/lipsync/result", + data=data, + files=files, + headers=_api_headers(), + timeout=Config.request_timeout, + ) + if r.status_code != 200: + logger.error("上报成功结果失败 HTTP %d: %s", r.status_code, r.text[:300]) + return + logger.info("任务 %s 完成，duration=%.1fs", task_id, duration) + except Exception as exc: + logger.error("上报成功结果异常: %s", exc) + + +# ── 主循环 ────────────────────────────────────────────────────────── + + +def main() -> int: + logger.info("=" * 60) + logger.info("MuseTalk GPU Worker 启动") + logger.info(" worker_id = %s", Config.derived_worker_id()) + logger.info(" api_base = %s", Config.api_base_url) + logger.info(" muse_talk = %s", Config.muse_talk_url) + logger.info(" poll = %.1fs / heartbeat = %.1fs", Config.poll_interval, Config.heartbeat_interval) + logger.info("=" * 60) + + if not Config.gpu_worker_token: + logger.warning("GPU_WORKER_TOKEN 未配置（开发模式），生产环境必须设置") + + # 先检查一次 MuseTalk + ok, info = _check_musetalk_health() + if ok: + logger.info("MuseTalk 健康检查通过: %s", info) + else: + logger.warning("MuseTalk 健康检查未通过: %s（继续运行，等待服务可用）", info) + + # 启动时立即注册 + _register() + last_heartbeat = time.time() + + while True: + try: + # 心跳 + now = time.time() + if now - last_heartbeat >= Config.heartbeat_interval: + if _register(): + last_heartbeat = now + + # 轮询任务 + task = _poll_task() + if task is not None: + _handle_task(task) + # 处理完立即再 poll（不 sleep），尽可能拉满 GPU + continue + + time.sleep(Config.poll_interval) + except KeyboardInterrupt: + logger.info("收到中断信号，退出") + return 0 + except Exception as exc: + logger.exception("主循环异常: %s", exc) + time.sleep(Config.poll_interval) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/deploy/gpu_worker/requirements.txt b/deploy/gpu_worker/requirements.txt new file mode 100644 index 000000000..0eb8cae7f --- /dev/null +++ b/deploy/gpu_worker/requirements.txt @@ -0,0 +1 @@ +requests>=2.31.0 diff --git a/deploy/gpu_worker/xiaoxia-gpu-worker.service b/deploy/gpu_worker/xiaoxia-gpu-worker.service new file mode 100644 index 000000000..41cb5ac1f --- /dev/null +++ b/deploy/gpu_worker/xiaoxia-gpu-worker.service @@ -0,0 +1,21 @@ +[Unit] +Description=MuseTalk GPU Worker (xiaoxia-saas 反向轮询) +After=network.target musetalk.service +# 本地 MuseTalk 服务启动后再启动本 Worker；若 MuseTalk 没有 systemd 服务则删除 musetalk.service + +[Service] +Type=simple +User=%i +WorkingDirectory=/opt/xiaoxia-gpu-worker +# 读取环境变量（API 地址、Token、轮询间隔等） +EnvironmentFile=/opt/xiaoxia-gpu-worker/.env +ExecStart=/opt/xiaoxia-gpu-worker/venv/bin/python /opt/xiaoxia-gpu-worker/gpu_worker.py +Restart=always +RestartSec=10 +# 日志走 journal，用 journalctl -u xiaoxia-gpu-worker -f 查看 +StandardOutput=journal +StandardError=journal +SyslogIdentifier=xiaoxia-gpu-worker + +[Install] +WantedBy=multi-user.target diff --git a/packages/adapters/sqlalchemy_impl/models.py b/packages/adapters/sqlalchemy_impl/models.py index 5285a5d82..5c11de106 100755 --- a/packages/adapters/sqlalchemy_impl/models.py +++ b/packages/adapters/sqlalchemy_impl/models.py @@ -854,3 +854,61 @@ class DailyUsageRecordModel(Base): usage_type = Column(String(50), nullable=False, default="free_clip") count = Column(Integer, nullable=False, default=0) updated_at = Column(DateTime, nullable=False, default=lambda: datetime.now(UTC)) + + +class GpuLipsyncTaskModel(Base): + """GPU 口型同步任务 ORM 模型 — MuseTalk 反向轮询模式. + + 业务侧（AI 数字人生成/lipsync 流程）提交任务后，GPU Worker 主动 poll 拉取、 + 调用本地 MuseTalk 推理、再通过 result 接口回传结果视频。 + """ + + __tablename__ = "gpu_lipsync_tasks" + + id = Column(String(36), primary_key=True) + # 业务关联（原 lipsync_job_id，方便双向查询） + lipsync_job_id = Column(String(36), nullable=False, default="", index=True) + user_id = Column(String(36), nullable=False, default="", index=True) + project_id = Column(String(36), nullable=False, default="", index=True) + + # 输入（预签名下载 URL，由 API 侧生成） + video_url = Column(Text, nullable=False) + audio_url = Column(Text, nullable=False) + + # 结果 + result_url = Column(Text, nullable=False, default="") + result_duration = Column(Float, nullable=False, default=0.0) + + # 任务状态 + status = Column( + String(20), + nullable=False, + default="pending", + index=True, + ) # pending → processing → done / failed / timeout + worker_id = Column(String(100), nullable=False, default="", index=True) + attempt = Column(Integer, nullable=False, default=0) + error_msg = Column(Text, nullable=False, default="") + + # 时间戳 + created_at = Column(DateTime, nullable=False, default=lambda: datetime.now(UTC)) + started_at = Column(DateTime, nullable=True) + finished_at = Column(DateTime, nullable=True) + updated_at = Column(DateTime, nullable=False, default=lambda: datetime.now(UTC)) + + # 心跳：worker 最近一次 poll/result 的时间，用于判定 worker 失联 + last_heartbeat_at = Column(DateTime, nullable=True) + + +class GpuWorkerModel(Base): + """GPU Worker 注册表 — 反向轮询模式下用于心跳与监控.""" + + __tablename__ = "gpu_workers" + + worker_id = Column(String(100), primary_key=True) + hostname = Column(String(200), nullable=False, default="") + gpu_name = Column(String(200), nullable=False, default="") + free_vram_mb = Column(Integer, nullable=False, default=0) + capabilities = Column(String(500), nullable=False, default="") # 逗号分隔，如 "musetalk" + last_heartbeat_at = Column(DateTime, nullable=True, index=True) + created_at = Column(DateTime, nullable=False, default=lambda: datetime.now(UTC)) diff --git a/packages/config/base.py b/packages/config/base.py index f56b5e2ed..a6af3e538 100755 --- a/packages/config/base.py +++ b/packages/config/base.py @@ -79,6 +79,17 @@ class SharedSettings(BaseSettings): # `if settings.points_enabled:` 包裹，防止未完善的扣点逻辑影响现有用户。 points_enabled: bool = False + # ── GPU MuseTalk 反向轮询 Worker ──────────────────────────────────── + # Worker 用这个长期 Token 鉴权（不是用户 JWT）。多 Worker 共用同一个 Token； + # worker_id 用于区分具体机器。生产必须配置；development 留空会跳过校验。 + gpu_worker_token: str = "" + # GPU 任务超时（秒）：超过此时长仍未完成则标记为 failed，可重新 poll + gpu_task_timeout_seconds: int = 300 + # 结果预签名 URL 有效期（秒） + gpu_result_url_expires: int = 3600 + # 输入预签名 URL 有效期（秒，需留出 Worker 下载时间） + gpu_input_url_expires: int = 3600 + @property def effective_database_url(self) -> str: """返回实际使用的数据库 URL。 diff --git a/packages/shared/storage.py b/packages/shared/storage.py index 11d6ad57e..9c5c4b2f6 100755 --- a/packages/shared/storage.py +++ b/packages/shared/storage.py @@ -339,6 +339,44 @@ class SharedStorageService(StoragePort): # ── 浏览器直传 POST ──────────────────────────────────────────────── + def get_upload_url( + self, + storage_key_or_url: str, + expires_seconds: int = 3600, + content_type: str = "video/mp4", + ) -> str: + """获取预签名 PUT 上传 URL（供外部 Worker 上传结果文件）。 + + bucket未配置时降级为 public_url（本地/开发环境）； + 本地产物 key 原样返回。 + """ + if self.bucket is None: + if self._is_local_generated_url(storage_key_or_url): + return storage_key_or_url + logger.warning( + "get_upload_url: OSS bucket not configured, returning raw URL. key=%s", + storage_key_or_url[:200], + ) + return self.get_url(self.normalize_storage_key(storage_key_or_url)) + + storage_key = self.normalize_storage_key(storage_key_or_url) + try: + # oss2 sign_url 支持 'PUT'，需指定 headers 才能限定 Content-Type + headers = {"Content-Type": content_type} if content_type else None + signed = self.bucket.sign_url("PUT", storage_key, expires_seconds, headers=headers) + logger.info( + "get_upload_url: signed PUT URL generated. key=%s url_prefix=%s", + storage_key[:80], + signed[:60], + ) + return signed + except Exception: + logger.exception( + "get_upload_url: sign_url failed, falling back to raw URL. key=%s", + storage_key[:200], + ) + return self.get_url(storage_key) + def create_direct_upload_post( self, storage_key: str, diff --git a/scripts/render_env.sh b/scripts/render_env.sh index 15bdd76b4..b2e68c3ea 100644 --- a/scripts/render_env.sh +++ b/scripts/render_env.sh @@ -57,7 +57,7 @@ if [ "$TARGET_ENV" = "staging" ]; then fi # 共用 secrets 直接导出（如果存在） -SHARED_SECRETS="OSS_ACCESS_KEY_ID OSS_ACCESS_KEY_SECRET COSYVOICE_API_KEY DASHSCOPE_API_KEY MEDIAKIT_API_KEY DOUBAO_API_KEY DOUBAO_MODEL DOUBAO_BASE_URL WECHAT_APP_ID WECHAT_APP_SECRET TIKHUB_API_KEY APIZERO_API_KEY" +SHARED_SECRETS="OSS_ACCESS_KEY_ID OSS_ACCESS_KEY_SECRET COSYVOICE_API_KEY DASHSCOPE_API_KEY MEDIAKIT_API_KEY DOUBAO_API_KEY DOUBAO_MODEL DOUBAO_BASE_URL WECHAT_APP_ID WECHAT_APP_SECRET TIKHUB_API_KEY APIZERO_API_KEY GPU_WORKER_TOKEN" for var in $SHARED_SECRETS; do value="${!var:-}" # 已经在环境中了，无需额外操作 diff --git a/tests/unit/test_gpu_lipsync_service.py b/tests/unit/test_gpu_lipsync_service.py new file mode 100644 index 000000000..e3a5303c0 --- /dev/null +++ b/tests/unit/test_gpu_lipsync_service.py @@ -0,0 +1,191 @@ +"""GpuLipsyncService 单元测试 — 覆盖任务创建、轮询认领、结果上报、超时回退等核心逻辑. + +使用 SQLite 内存数据库，mock 掉存储层（不真实调用 OSS）。 +""" + +from __future__ import annotations + +import os +import sys +from datetime import UTC, datetime, timedelta +from unittest import mock + +import pytest + +# 确保 packages / apps/api 可导入 +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +for p in (ROOT, os.path.join(ROOT, "apps", "api"), os.path.join(ROOT, "packages")): + if p not in sys.path: + sys.path.insert(0, p) + +# 强制使用内存 SQLite（避免依赖 PG） +os.environ["APP_ENV"] = "development" +os.environ["JWT_SECRET_KEY"] = "dev-secret-key-for-testing-00000000" +os.environ["DATABASE_URL"] = "sqlite:///:memory:" +os.environ["USE_IN_MEMORY_DB"] = "1" +os.environ["GPU_WORKER_TOKEN"] = "" # development 空 token 放行 + + +def _build_session(): + from sqlalchemy import create_engine + from sqlalchemy.orm import sessionmaker + + # 使用 packages 的 Base + from packages.adapters.sqlalchemy_impl import models as _ # noqa: F401 # 触发 ORM 注册 + from packages.adapters.sqlalchemy_impl.models import Base + + engine = create_engine("sqlite:///:memory:", future=True) + Base.metadata.create_all(engine) + Session = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True) + return Session() + + +@pytest.fixture +def svc(): + from app.services.gpu_lipsync_service import GpuLipsyncService + + db = _build_session() + service = GpuLipsyncService(db) + # mock 存储签名（SQLite 测试无 OSS） + service.storage = mock.MagicMock() + service.storage.get_download_url.side_effect = ( + lambda k, expires_seconds=3600: f"https://signed.example.com/download/{k}?e={expires_seconds}" + ) + service.storage.get_upload_url.side_effect = ( + lambda k, expires_seconds=3600, content_type="video/mp4": f"https://signed.example.com/upload/{k}?e={expires_seconds}" + ) + return service + + +# ── 创建任务 ────────────────────────────────────────────────────── + + +def test_create_task(svc): + task = svc.create_task( + video_url="uploads/v.mp4", + audio_url="uploads/a.mp3", + lipsync_job_id="lip-1", + user_id="u-1", + project_id="p-1", + ) + assert task.id + assert task.status == "pending" + assert task.lipsync_job_id == "lip-1" + assert task.attempt == 0 + assert task.video_url == "uploads/v.mp4" + + +# ── 轮询认领 ────────────────────────────────────────────────────── + + +def test_poll_returns_none_when_empty(svc): + assert svc.poll_task("w-1") is None + + +def test_poll_claims_pending_task(svc): + svc.create_task(video_url="uploads/v.mp4", audio_url="uploads/a.mp3") + claimed = svc.poll_task("w-1") + assert claimed is not None + assert claimed.status == "processing" + assert claimed.worker_id == "w-1" + assert claimed.attempt == 1 + # 带签名 URL + assert claimed._signed_video_url.startswith("https://signed.example.com/download/") + assert claimed._signed_upload_url.startswith("https://signed.example.com/upload/") + # 再 poll 无任务 + assert svc.poll_task("w-1") is None + + +def test_poll_concurrent_claim_only_one_wins(svc): + """并发场景：两个 worker 同时 poll 只有一个能拿到任务（借助 update where status=pending）。""" + svc.create_task(video_url="v", audio_url="a") + t1 = svc.poll_task("w-1") + t2 = svc.poll_task("w-2") + assert t1 is not None + assert t2 is None + + +# ── 结果上报 ────────────────────────────────────────────────────── + + +def test_report_result_success(svc): + t = svc.create_task(video_url="v", audio_url="a") + svc.poll_task("w-1") # claim + done = svc.report_result(t.id, "w-1", success=True, duration_seconds=12.5) + assert done.status == "done" + assert done.result_duration == 12.5 + assert done.result_url.startswith("gpu-lipsync/results/") + assert done.finished_at is not None + + +def test_report_result_failure_requeues(svc): + t = svc.create_task(video_url="v", audio_url="a") + svc.poll_task("w-1") + failed = svc.report_result(t.id, "w-1", success=False, error_msg="MuseTalk crash") + assert failed.status == "pending" # 仍在重试次数内 → 回队 + assert failed.worker_id == "" + assert failed.started_at is None + assert "MuseTalk crash" in failed.error_msg + + +def test_report_failure_exhausted_goes_failed(svc): + """失败达到 MAX_ATTEMPTS 后标记 failed，不再回队. + + poll 成功会将 attempt 从 0 开始自增； + 第 1/2 次失败回队，第 3 次失败（attempt==MAX_ATTEMPTS）置 failed。 + """ + from app.services import gpu_lipsync_service as mod + + t = svc.create_task(video_url="v", audio_url="a") + # 模拟失败到上限：poll + fail 重复 MAX_ATTEMPTS 次 + for i in range(mod.MAX_ATTEMPTS): + claimed = svc.poll_task(f"w-{i}") + assert claimed is not None, f"第 {i} 次 poll 应能拿到任务" + svc.report_result(t.id, claimed.worker_id, success=False, error_msg=f"fail {i}") + svc.db.refresh(t) + if i == mod.MAX_ATTEMPTS - 1: + assert t.status == "failed" + else: + assert t.status == "pending" + + +# ── 心跳/超时回退 ───────────────────────────────────────────────── + + +def test_timed_out_task_is_redispatched(svc): + """processing 超过 gpu_task_timeout_seconds 无心跳 → 回退 pending.""" + t = svc.create_task(video_url="v", audio_url="a") + svc.poll_task("w-1") + svc.db.refresh(t) + assert t.status == "processing" + # 手动把 last_heartbeat_at 设到很久以前 + t.last_heartbeat_at = datetime.now(UTC) - timedelta(seconds=svc.settings.gpu_task_timeout_seconds + 10) + svc.db.commit() + # 再次 poll 会触发 _recover_timed_out_tasks 把它回队 + claimed = svc.poll_task("w-2") + assert claimed is not None + assert claimed.id == t.id + assert claimed.worker_id == "w-2" + assert claimed.attempt == 2 # 又认领了一次 + + +# ── Worker 注册 ─────────────────────────────────────────────────── + + +def test_register_worker_creates_then_updates(svc): + w = svc.register_worker("w-1", hostname="pc1", gpu_name="RTX2060", free_vram_mb=3500) + assert w.worker_id == "w-1" + assert w.gpu_name == "RTX2060" + w2 = svc.register_worker("w-1", hostname="pc1", gpu_name="RTX2060", free_vram_mb=2000) + assert w2.free_vram_mb == 2000 # 更新 + assert w2.created_at == w.created_at # 没新建 + + +# ── get_by_lipsync_job ───────────────────────────────────────────── + + +def test_get_by_lipsync_job_returns_latest(svc): + svc.create_task(video_url="v", audio_url="a", lipsync_job_id="lip-1") + svc.create_task(video_url="v", audio_url="a", lipsync_job_id="lip-1") + latest = svc.get_by_lipsync_job("lip-1") + assert latest is not None