xiaoxia-saas/deploy/gpu_worker/gpu_worker.py

"""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())