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 请求超时秒（下载/推理/上传统一使用），默认 900
                      需与服务端 GPU_TASK_TIMEOUT_SECONDS（默认 900）对齐
  TASK_MAX_RETRY      单任务本地最大重试次数（仅对瞬时错误重试），默认 1
  TASK_HEARTBEAT_INTERVAL  推理期间任务心跳间隔秒，默认 30
  MIN_VIDEO_DURATION_SECONDS  最短输入视频时长秒，小于则直接上报失败，默认 3

用法：
  python gpu_worker.py
"""

from __future__ import annotations

import logging
import os
import platform
import socket
import sys
import tempfile
import threading
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"))
    # #1970：RTX2060 6G 处理 720p 长视频可能 >5min；与服务端
    # GPU_TASK_TIMEOUT_SECONDS 默认值对齐为 900，避免推理被本地/服务端先掐断。
    request_timeout: float = float(_env("REQUEST_TIMEOUT", "900"))
    # 本地只在网络/MuseTalk 瞬时错误时重试 1 次；服务端 MAX_ATTEMPTS=3
    # 负责跨 worker/真正超时后的重派发，总尝试次数不再相乘放大。
    task_max_retry: int = int(_env("TASK_MAX_RETRY", "1"))
    # 推理期间任务心跳间隔（独立线程 POST /gpu/register 带 task_id）
    task_heartbeat_interval: float = float(_env("TASK_HEARTBEAT_INTERVAL", "30"))
    # 输入视频最短时长（秒）：过短（如 1s）MuseTalk 会 division by zero，
    # 本地前置拦截，直接上报 failed，不浪费 GPU 时间
    min_video_duration_seconds: float = float(_env("MIN_VIDEO_DURATION_SECONDS", "3"))
    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(task_id: Optional[str] = None) -> tuple[bool, bool]:
    """向服务端注册 / 心跳，附带 GPU 信息。

    推理期间的心跳线程传 task_id：服务端会同步刷新该 processing 任务的
    last_heartbeat_at，防止长推理被误判超时回收。同时服务端会检查该任务
    是否已被用户取消，若是则返回 cancel_task=True。

    返回 (ok, cancel_task)。
    """
    ok, info = _check_musetalk_health()
    if isinstance(info, dict):
        gpu_info = info.get("gpu", info)
        free_vram = int(gpu_info.get("free_vram_mb", gpu_info.get("memory_free_mb", 0)) or 0)
        gpu_name = gpu_info.get("gpu_name", info.get("gpu_name", ""))
    else:
        free_vram = 0
        gpu_name = ""
    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",
    }
    if task_id:
        payload["task_id"] = task_id
    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:
            resp_body = r.json()
            cancel_task = resp_body.get("cancel_task", False)
            return True, cancel_task
        logger.error("注册/心跳失败: HTTP %d body=%s", r.status_code, r.text[:300])
        return False, False
    except Exception as exc:
        logger.error("注册/心跳异常: %s", exc)
        return False, 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, bool]:
    """调用本地 MuseTalk /inference.

    返回 (success, duration_seconds, error_msg, retryable)。
    duration 用 ffprobe 读结果视频，失败填 0。
    retryable 仅对瞬时错误（连接失败/超时/5xx）为 True；HTTP 4xx、结果过小
    等确定性失败不重试，直接上报服务端（服务端 MAX_ATTEMPTS 再决定是否重派发）。
    """
    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:
            retryable = r.status_code >= 500
            return False, 0.0, f"MuseTalk HTTP {r.status_code}: {r.text[:500]}", retryable
        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)", False
        duration = _probe_duration(out_path)
        return True, duration, "", False
    except (requests.exceptions.Timeout, requests.exceptions.ConnectionError):
        # 瞬时网络/超时错误，允许本地重试 1 次；同时调 /cancel 让服务端终止僵尸推理
        _cancel_musetalk()
        return False, 0.0, f"MuseTalk 推理超时或连接失败（>{Config.request_timeout}s）", True
    except Exception as exc:
        return False, 0.0, f"MuseTalk 调用异常: {exc}", False


def _cancel_musetalk() -> None:
    """调 MuseTalk /cancel 端点终止服务端僵尸推理进程，避免超时后任务还在跑占显存."""
    try:
        r = requests.post(f"{Config.muse_talk_url}/cancel", timeout=10)
        if r.status_code == 200:
            logger.info("已调 MuseTalk /cancel，服务端终止推理")
        else:
            logger.warning("MuseTalk /cancel 返回 %d: %s", r.status_code, r.text[:200])
    except Exception as exc:
        # /cancel 失败不应影响主流程上报
        logger.warning("调 MuseTalk /cancel 异常（忽略）: %s", 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


class TaskHeartbeat(threading.Thread):
    """推理期间的任务心跳线程。

    主循环的空闲心跳在 ``_handle_task`` 同步阻塞（下载/推理/上传最长 900s）
    期间无法发送，服务端会因任务 last_heartbeat_at 停滞而误判超时回退 pending。
    本线程每 task_heartbeat_interval 秒（默认 30s）POST /gpu/register 并
    携带当前 task_id，让服务端持续续期任务心跳；任务处理结束 stop()。

    同时检测服务端返回的 cancel_task 信号：若为 True，说明用户已取消任务，
    立即调用 _cancel_musetalk() 终止本地推理，并设置 cancelled 标志供主流程检查。
    """

    def __init__(self, task_id: str, interval: float):
        super().__init__(daemon=True, name=f"hb-{task_id[:8]}")
        self.task_id = task_id
        self.interval = max(5.0, interval)
        self._stop_event = threading.Event()
        self.cancelled = False  # 外部可读的取消标志

    def run(self) -> None:
        # 先立即发一次，再按间隔循环（首次心跳失败不影响主流程）
        while not self._stop_event.is_set():
            try:
                ok, cancel_task = _register(self.task_id)
                if ok:
                    logger.debug("任务 %s 心跳已发送", self.task_id)
                if cancel_task:
                    logger.warning("任务 %s 已被用户取消，正在终止本地推理...", self.task_id)
                    self.cancelled = True
                    _cancel_musetalk()
                    self._stop_event.set()
                    return
            except Exception as exc:  # noqa: BLE001
                logger.warning("任务 %s 心跳异常（忽略）: %s", self.task_id, exc)
            self._stop_event.wait(self.interval)

    def stop(self) -> None:
        self._stop_event.set()


def _handle_task(task: dict) -> None:
    """处理一条任务（整个串行流程：下载→时长校验→推理→上传→上报）。"""
    task_id = task["task_id"]
    logger.info("开始处理任务 %s", task_id)
    # 领取任务后立即启动任务级心跳线程，覆盖下载/推理/上报全过程
    hb = TaskHeartbeat(task_id, Config.task_heartbeat_interval)
    hb.start()
    try:
        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 hb.cancelled:
                logger.info("任务 %s 在下载阶段被用户取消", task_id)
                _report_result(task_id, False, 0.0, "用户取消任务")
                return
            if not _download(task["audio_url"], audio_path):
                _report_result(task_id, False, 0.0, "下载驱动音频失败")
                return
            if hb.cancelled:
                logger.info("任务 %s 在下载阶段被用户取消", task_id)
                _report_result(task_id, False, 0.0, "用户取消任务")
                return

            # 2. 输入时长前置校验：短视频 MuseTalk 会 division by zero，
            #    直接上报 failed，不浪费 GPU 时间。ffprobe 不可用/读失败（0.0）
            #    时不拦截，交给 MuseTalk 处理，避免误杀。
            video_duration = _probe_duration(video_path)
            if video_duration and video_duration < Config.min_video_duration_seconds:
                msg = (
                    f"视频过短（{video_duration:.2f}s < {Config.min_video_duration_seconds:.0f}s），"
                    "MuseTalk 无法处理"
                )
                logger.error("任务 %s %s", task_id, msg)
                _report_result(task_id, False, 0.0, msg)
                return

            # 3. 推理（本地仅对瞬时错误重试）
            success = False
            duration = 0.0
            err = ""
            retryable = False
            for attempt in range(Config.task_max_retry + 1):
                if hb.cancelled:
                    logger.info("任务 %s 在推理前被用户取消", task_id)
                    _report_result(task_id, False, 0.0, "用户取消任务")
                    return
                if attempt > 0:
                    logger.info("任务 %s 第 %d 次重试（瞬时错误）...", task_id, attempt + 1)
                    time.sleep(2)
                success, duration, err, retryable = _call_musetalk(video_path, audio_path, out_path)
                if success or not retryable:
                    break
            if hb.cancelled:
                logger.info("任务 %s 被用户取消（推理已终止）", task_id)
                _report_result(task_id, False, 0.0, "用户取消任务")
                return
            if not success:
                logger.error("任务 %s 推理失败: %s", task_id, err)
                _report_result(task_id, False, 0.0, err)
                return

            # 4. 上报结果（multipart 同时上传文件 → API 代为 PUT 到 OSS，逻辑最稳）
            _report_success_with_file(task_id, duration, out_path)
    finally:
        hb.stop()


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:
                ok, _ = _register()
                if ok:
                    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())