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