"""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) -> bool: """向服务端注册 / 心跳，附带 GPU 信息。 推理期间的心跳线程传 task_id：服务端会同步刷新该 processing 任务的 last_heartbeat_at，防止长推理被误判超时回收。 """ 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", } 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: 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, 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 次 return False, 0.0, f"MuseTalk 推理超时或连接失败（>{Config.request_timeout}s）", True except Exception as exc: return False, 0.0, f"MuseTalk 调用异常: {exc}", False 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()。 """ 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() def run(self) -> None: # 先立即发一次，再按间隔循环（首次心跳失败不影响主流程） while not self._stop_event.is_set(): try: if _register(self.task_id): logger.debug("任务 %s 心跳已发送", self.task_id) 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 not _download(task["audio_url"], audio_path): _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 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 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: 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())