diff --git a/.env.example b/.env.example index 39266bb27..525794557 100755 --- a/.env.example +++ b/.env.example @@ -217,6 +217,7 @@ APIZERO_API_KEY= # GPU Worker 长期鉴权 Token，Worker 端 .env 的 GPU_WORKER_TOKEN 必须与此一致 # 留空时 development 环境允许匿名访问（仅本地调试），staging/production 必须配置 GPU_WORKER_TOKEN= -# 单任务超时（秒），超过则回退 pending 或标记 failed -GPU_TASK_TIMEOUT_SECONDS=300 +# 单任务超时（秒），processing 超过此时长无任务心跳才回退 pending 或标记 failed +# #1970：RTX2060 6G 推理 720p 长视频需 5 分钟以上，默认 900 +GPU_TASK_TIMEOUT_SECONDS=900 diff --git a/apps/api/app/api/routes/gpu_lipsync.py b/apps/api/app/api/routes/gpu_lipsync.py index da1b8385f..c00c73ccf 100644 --- a/apps/api/app/api/routes/gpu_lipsync.py +++ b/apps/api/app/api/routes/gpu_lipsync.py @@ -99,6 +99,7 @@ def register_worker( gpu_name=body.gpu_name, free_vram_mb=body.free_vram_mb, capabilities=body.capabilities, + task_id=body.task_id, ) return GpuWorkerRegisterResponse(ok=True, server_time=datetime.now(UTC), message="ok") diff --git a/apps/api/app/schemas/gpu_lipsync.py b/apps/api/app/schemas/gpu_lipsync.py index 592c4227b..fbd26de1c 100644 --- a/apps/api/app/schemas/gpu_lipsync.py +++ b/apps/api/app/schemas/gpu_lipsync.py @@ -22,6 +22,14 @@ class GpuWorkerRegisterRequest(BaseModel): 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'") + task_id: Optional[str] = Field( + None, + max_length=64, + description=( + "当前正在处理的任务 ID。Worker 推理期间定期心跳时携带，" + "服务端同步刷新该任务 last_heartbeat_at，防止长推理被误判超时；空闲时不传" + ), + ) class GpuWorkerRegisterResponse(BaseModel): diff --git a/apps/api/app/services/gpu_lipsync_service.py b/apps/api/app/services/gpu_lipsync_service.py index a3c1cac14..f9491006f 100644 --- a/apps/api/app/services/gpu_lipsync_service.py +++ b/apps/api/app/services/gpu_lipsync_service.py @@ -49,7 +49,15 @@ class GpuLipsyncService: gpu_name: str = "", free_vram_mb: int = 0, capabilities: str = "musetalk", + task_id: Optional[str] = None, ) -> GpuWorkerModel: + """Worker 注册/心跳。 + + task_id 非空时（Worker 推理期间的任务级心跳），同步把对应 processing + 任务的 last_heartbeat_at 续到当前时间，使长推理不会被 + ``_recover_timed_out_tasks`` 误回退。任务已结束 / 不属于该 worker + （如已被超时回收重新派发）时忽略，不报错。 + """ now = datetime.now(UTC) worker = self.db.query(GpuWorkerModel).filter(GpuWorkerModel.worker_id == worker_id).one_or_none() if worker is None: @@ -69,6 +77,8 @@ class GpuLipsyncService: worker.free_vram_mb = free_vram_mb worker.capabilities = capabilities or worker.capabilities worker.last_heartbeat_at = now + if task_id: + self._touch_task_heartbeat(task_id, worker_id, now) self.db.commit() return worker @@ -78,8 +88,10 @@ class GpuLipsyncService: """原子地认领一条最早的 pending 任务，返回给 worker；无任务返回 None. 同时会： - - 把 processing 状态且超时（超过 gpu_task_timeout_seconds 无心跳）的任务 - 回退为 pending（attempt++，超过 MAX_ATTEMPTS 置 failed），让其它 worker 认领。 + - 把 processing 状态且真正超时（任务心跳停滞超过 + gpu_task_timeout_seconds；Worker 推理期会通过 register(task_id=...) + 续心跳，长推理不会误判）的任务回退为 pending（attempt++，超过 + MAX_ATTEMPTS 置 failed），让其它 worker 认领。 - 刷新 worker 心跳。 """ now = datetime.now(UTC) @@ -238,6 +250,28 @@ class GpuLipsyncService: def _result_key(self, task_id: str) -> str: return f"{self.RESULT_PREFIX}{task_id}.mp4" + def _touch_task_heartbeat(self, task_id: str, worker_id: str, now: datetime) -> None: + """Worker 推理期间的任务级心跳：只刷新属于该 worker 且仍在 processing 的任务。 + + 任务不存在 / 已被超时回收重新派发 / 已完成 → 静默忽略（此时旧 worker 的 + 结果上报会被结果接口按最终态处理）。 + """ + task = self.db.get(GpuLipsyncTaskModel, task_id) + if task is None: + return + if task.status != "processing" or task.worker_id != worker_id: + logger.info( + "忽略过期任务心跳 task=%s worker=%s（status=%s owner=%s）", + task_id, + worker_id, + task.status, + task.worker_id, + ) + return + task.last_heartbeat_at = now + task.updated_at = now + self.db.flush() + def _touch_worker(self, worker_id: str, now: datetime) -> None: if not worker_id: return @@ -260,7 +294,13 @@ class GpuLipsyncService: self.db.flush() def _recover_timed_out_tasks(self, now: datetime) -> None: - """扫描 processing 状态且超时（无心跳）的任务，回退 pending 或失败.""" + """扫描 processing 状态且真正超时的任务，回退 pending 或失败。 + + 判定只看任务自身 last_heartbeat_at：claim 时写入，Worker 推理期间通过 + /gpu/register(task_id=...) 每 30s 续期。因此仅在 Worker 崩溃/断网 + （任务心跳停滞超过 gpu_task_timeout_seconds）时才回收， + 不会因 Worker 主循环忙于推理而误回退。 + """ timeout = self.settings.gpu_task_timeout_seconds cutoff = now - timedelta(seconds=timeout) stuck_tasks = ( diff --git a/deploy/configs/.env.production b/deploy/configs/.env.production index f080f85ad..5f3cdaad5 100644 --- a/deploy/configs/.env.production +++ b/deploy/configs/.env.production @@ -258,4 +258,4 @@ APIZERO_API_KEY=${APIZERO_API_KEY} # ==================== GPU MuseTalk Worker（反向轮询） ==================== GPU_WORKER_TOKEN=${GPU_WORKER_TOKEN} -GPU_TASK_TIMEOUT_SECONDS=300 +GPU_TASK_TIMEOUT_SECONDS=900 diff --git a/deploy/configs/.env.staging b/deploy/configs/.env.staging index 93c9d4f79..374f4d91e 100644 --- a/deploy/configs/.env.staging +++ b/deploy/configs/.env.staging @@ -275,4 +275,4 @@ APIZERO_API_KEY=${APIZERO_API_KEY} # ==================== GPU MuseTalk Worker（反向轮询） ==================== GPU_WORKER_TOKEN=${GPU_WORKER_TOKEN} -GPU_TASK_TIMEOUT_SECONDS=300 +GPU_TASK_TIMEOUT_SECONDS=900 diff --git a/deploy/gpu_worker/.env.example b/deploy/gpu_worker/.env.example index 74bd32480..fa55dd842 100644 --- a/deploy/gpu_worker/.env.example +++ b/deploy/gpu_worker/.env.example @@ -19,7 +19,12 @@ MUSE_TALK_URL=http://127.0.0.1:7861 # 轮询/心跳/超时（秒） POLL_INTERVAL=5 HEARTBEAT_INTERVAL=15 -REQUEST_TIMEOUT=300 +# 下载/推理/上传 HTTP 超时，需与服务端 GPU_TASK_TIMEOUT_SECONDS 对齐（默认 900） +REQUEST_TIMEOUT=900 -# 单个任务本地最大重试次数（首次失败后再重试 N 次，默认 2） -TASK_MAX_RETRY=2 +# 单个任务本地最大重试次数（仅网络/MuseTalk 瞬时错误才重试，默认 1） +TASK_MAX_RETRY=1 +# 推理期间任务心跳间隔（秒，独立线程，无需改动） +TASK_HEARTBEAT_INTERVAL=30 +# 输入视频最短时长（秒），小于则直接上报失败，不调用 MuseTalk +MIN_VIDEO_DURATION_SECONDS=3 diff --git a/deploy/gpu_worker/README.md b/deploy/gpu_worker/README.md index 8230a63e8..a5af3420c 100644 --- a/deploy/gpu_worker/README.md +++ b/deploy/gpu_worker/README.md @@ -89,8 +89,9 @@ SaaS 后端部署完成后需配置： | 日志 `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，或限制输入视频时长 | +| 服务端看到任务回退到 pending 重试 | 任务心跳真正超时（默认 900s）：Worker 进程崩溃/断网，或推理彻底卡死；正常长推理期间心跳线程每 30s 续期，不会回退 | +| 日志 `MuseTalk 推理超时或连接失败` | 视频太长或显存不足；可临时调大 REQUEST_TIMEOUT（服务端 GPU_TASK_TIMEOUT_SECONDS 需同步调大），或限制输入视频时长 | +| 日志 `视频过短（x.xxs < 3s）` | 输入视频不足 3s，MuseTalk 对短视频会 division by zero，已在本地直接上报失败；可用 MIN_VIDEO_DURATION_SECONDS 调整阈值 | ## 七、安全注意事项 diff --git a/deploy/gpu_worker/gpu_worker.py b/deploy/gpu_worker/gpu_worker.py index ecdeac35b..8745597b7 100644 --- a/deploy/gpu_worker/gpu_worker.py +++ b/deploy/gpu_worker/gpu_worker.py @@ -9,9 +9,12 @@ 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 + 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 @@ -26,6 +29,7 @@ import platform import socket import sys import tempfile +import threading import time import uuid from pathlib import Path @@ -54,8 +58,17 @@ class Config: 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")) + # #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 @@ -96,8 +109,12 @@ def _check_musetalk_health() -> tuple[bool, dict]: return False, {"error": str(exc)} -def _register() -> bool: - """向服务端注册 / 心跳，附带 GPU 信息.""" +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 "" @@ -111,6 +128,8 @@ def _register() -> bool: "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", @@ -181,11 +200,15 @@ def _download(url: str, path: Path) -> bool: return False -def _call_musetalk(video_path: Path, audio_path: Path, out_path: Path) -> tuple[bool, float, str]: +def _call_musetalk( + video_path: Path, audio_path: Path, out_path: Path +) -> tuple[bool, float, str, bool]: """调用本地 MuseTalk /inference. - 返回 (success, duration_seconds, error_msg). + 返回 (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: @@ -199,17 +222,20 @@ def _call_musetalk(video_path: Path, audio_path: Path, out_path: Path) -> tuple[ timeout=Config.request_timeout, ) if r.status_code != 200: - return False, 0.0, f"MuseTalk HTTP {r.status_code}: {r.text[:500]}" + 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)" + # 确定性失败（推理产物异常），本地重试大概率还是坏的，不重试 + return False, 0.0, f"MuseTalk 返回结果过小 ({out_path.stat().st_size} bytes)", False duration = _probe_duration(out_path) - return True, duration, "" - except requests.exceptions.Timeout: - return False, 0.0, f"MuseTalk 推理超时（>{Config.request_timeout}s）" + 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}" + return False, 0.0, f"MuseTalk 调用异常: {exc}", False def _probe_duration(path: Path) -> float: @@ -276,42 +302,91 @@ def _report_result(task_id: str, success: bool, duration: float = 0.0, error_msg 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 = threading.Event() + + def run(self) -> None: + # 先立即发一次，再按间隔循环（首次心跳失败不影响主流程） + while not self._stop.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.wait(self.interval) + + def stop(self) -> None: + self._stop.set() + + 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" + # 领取任务后立即启动任务级心跳线程，覆盖下载/推理/上报全过程 + 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 + # 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 + # 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. 上报结果（multipart 同时上传文件 → API 代为 PUT 到 OSS，逻辑最稳） - _report_success_with_file(task_id, duration, out_path) + # 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: diff --git a/packages/config/base.py b/packages/config/base.py index e16a4381c..b2c56914e 100755 --- a/packages/config/base.py +++ b/packages/config/base.py @@ -84,8 +84,11 @@ class SharedSettings(BaseSettings): # Worker 用这个长期 Token 鉴权（不是用户 JWT）。多 Worker 共用同一个 Token； # worker_id 用于区分具体机器。生产必须配置；development 留空会跳过校验。 gpu_worker_token: str = "" - # GPU 任务超时（秒）：超过此时长仍未完成则标记为 failed，可重新 poll - gpu_task_timeout_seconds: int = 300 + # GPU 任务超时（秒）：processing 状态超过此时长（以任务心跳为准）才回退 + # pending / failed。#1970：RTX2060 6G 推理 720p 长视频需 5 分钟以上，300→900。 + # Worker 推理期间每 30s 通过 /gpu/register(task_id=...) 续心跳， + # 只有真正超时或 Worker 明确上报 failed 才会回退。 + gpu_task_timeout_seconds: int = 900 # 结果预签名 URL 有效期（秒） gpu_result_url_expires: int = 3600 # 输入预签名 URL 有效期（秒，需留出 Worker 下载时间） diff --git a/tests/unit/test_1970_gpu_worker_heartbeat.py b/tests/unit/test_1970_gpu_worker_heartbeat.py new file mode 100644 index 000000000..082149a45 --- /dev/null +++ b/tests/unit/test_1970_gpu_worker_heartbeat.py @@ -0,0 +1,254 @@ +"""#1970 GPU Worker 修复单测. + +覆盖 deploy/gpu_worker/gpu_worker.py（独立部署脚本，不在 apps/packages 包内， +按文件路径动态加载）： +1. 默认配置：REQUEST_TIMEOUT=900 / TASK_MAX_RETRY=1 / 心跳 30s / 最短 3s； +2. 推理期心跳线程 POST /gpu/register 带 task_id，任务结束能停； +3. <3s 短视频直接上报失败，不调用 MuseTalk； +4. _call_musetalk 仅对 5xx/网络瞬时错误标记 retryable，4xx 不重试； +5. _handle_task 只对 retryable 错误本地重试 1 次。 +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +import time +from pathlib import Path +from unittest import mock + +import pytest + +ROOT = Path(__file__).resolve().parents[2] +WORKER_PATH = ROOT / "deploy" / "gpu_worker" / "gpu_worker.py" + + +def _load_worker_module(): + spec = importlib.util.spec_from_file_location("gpu_worker_standalone_1970", WORKER_PATH) + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +@pytest.fixture +def worker(): + return _load_worker_module() + + +# ── 默认配置 ─────────────────────────────────────────────────────── + + +def test_config_defaults_900_and_retry_one(monkeypatch): + """CI/本机若显式导出过这些 env，说明是运维覆盖，不应拿默认值断言； + 因此只在四个 env 全部缺失时校验脚本内置默认值（#1970：900/1/30/3）。""" + keys = ( + "REQUEST_TIMEOUT", + "TASK_MAX_RETRY", + "TASK_HEARTBEAT_INTERVAL", + "MIN_VIDEO_DURATION_SECONDS", + ) + if any(k in os.environ for k in keys): + pytest.skip("环境显式设置了 worker 超时/重试变量，跳过默认值断言") + for key in keys: + monkeypatch.delenv(key, raising=False) + mod = _load_worker_module() + assert mod.Config.request_timeout == 900.0 + assert mod.Config.task_max_retry == 1 + assert mod.Config.task_heartbeat_interval == 30.0 + assert mod.Config.min_video_duration_seconds == 3.0 + + +# ── register 携带 task_id ────────────────────────────────────────── + + +def test_register_payload_includes_task_id_only_when_provided(worker, monkeypatch): + captured = [] + + class _Resp: + status_code = 200 + text = "" + + def _fake_post(url, json=None, headers=None, timeout=None): + captured.append(json) + return _Resp() + + monkeypatch.setattr(worker.requests, "post", _fake_post) + monkeypatch.setattr(worker, "_check_musetalk_health", lambda: (True, {})) + + assert worker._register("task-abc") is True + assert captured[-1]["task_id"] == "task-abc" + assert captured[-1]["worker_id"] + + worker._register() # 空闲心跳不带 task_id + assert "task_id" not in captured[-1] + + +# ── 推理期心跳线程 ───────────────────────────────────────────────── + + +def test_task_heartbeat_thread_sends_and_stops(worker, monkeypatch): + calls = [] + + def _fake_register(task_id=None): + calls.append(task_id) + return True + + monkeypatch.setattr(worker, "_register", _fake_register) + hb = worker.TaskHeartbeat("task-hb1", interval=5) + hb.start() + time.sleep(0.3) # 启动后立即发一次 + hb.stop() + hb.join(timeout=2) + assert not hb.is_alive() + assert calls and all(c == "task-hb1" for c in calls) + + +# ── 短视频前置拦截 ───────────────────────────────────────────────── + + +def test_handle_task_short_video_reports_failed_without_inference(worker, monkeypatch, tmp_path): + video = tmp_path / "input.mp4" + video.write_bytes(b"fake-mp4-bytes") + audio = tmp_path / "input_audio.bin" + audio.write_bytes(b"fake-audio") + reports = [] + + monkeypatch.setattr(worker, "_register", lambda *a, **k: True) + monkeypatch.setattr(worker, "_download", lambda url, path: True) + # ffprobe 读出 1.2s → 低于 3s 阈值 + monkeypatch.setattr(worker, "_probe_duration", lambda path: 1.2) + + def _boom(*a, **k): + raise AssertionError("短视频不应调用 MuseTalk 推理") + + monkeypatch.setattr(worker, "_call_musetalk", _boom) + monkeypatch.setattr( + worker, + "_report_result", + lambda task_id, success, duration=0.0, error_msg="": reports.append((task_id, success, error_msg)) or True, + ) + + task = { + "task_id": "task-short", + "video_url": "https://example.com/v.mp4", + "audio_url": "https://example.com/a.bin", + } + worker._handle_task(task) + + assert len(reports) == 1 + tid, ok, err = reports[0] + assert tid == "task-short" + assert ok is False + assert "视频过短" in err + assert "3" in err + + +def test_handle_task_probe_failure_does_not_block(worker, monkeypatch): + """ffprobe 不可用（duration=0.0）时不能误杀，应继续推理.""" + reports = [] + monkeypatch.setattr(worker, "_register", lambda *a, **k: True) + monkeypatch.setattr(worker, "_download", lambda url, path: True) + monkeypatch.setattr(worker, "_probe_duration", lambda path: 0.0) + monkeypatch.setattr( + worker, + "_call_musetalk", + lambda v, a, o: (True, 8.0, "", False), + ) + uploaded = [] + monkeypatch.setattr( + worker, + "_report_success_with_file", + lambda task_id, duration, path: uploaded.append((task_id, duration)), + ) + monkeypatch.setattr(worker, "_report_result", lambda *a, **k: True) + + worker._handle_task({"task_id": "task-probe0", "video_url": "u", "audio_url": "u"}) + assert uploaded == [("task-probe0", 8.0)] + assert reports == [] + + +# ── 重试语义：仅瞬时错误重试 ─────────────────────────────────────── + + +def test_call_musetalk_4xx_not_retryable_5xx_retryable(worker, monkeypatch, tmp_path): + video = tmp_path / "v.mp4" + audio = tmp_path / "a.bin" + video.write_bytes(b"v") + audio.write_bytes(b"a") + out = tmp_path / "o.mp4" + + class _Resp: + def __init__(self, code, body=b"x" * 2048): + self.status_code = code + self.content = body + self.text = "err" + + # 4xx：确定性失败，不重试 + monkeypatch.setattr(worker.requests, "post", lambda *a, **k: _Resp(400)) + ok, _, _, retryable = worker._call_musetalk(video, audio, out) + assert ok is False and retryable is False + + monkeypatch.setattr(worker.requests, "post", lambda *a, **k: _Resp(503)) + ok, _, _, retryable = worker._call_musetalk(video, audio, out) + assert ok is False and retryable is True + + # 连接异常：瞬时错误，可重试 + import requests as _requests + + def _conn_err(*a, **k): + raise _requests.exceptions.ConnectionError("reset") + + monkeypatch.setattr(worker.requests, "post", _conn_err) + ok, _, _, retryable = worker._call_musetalk(video, audio, out) + assert ok is False and retryable is True + + +def test_handle_task_retries_once_for_transient_then_succeeds(worker, monkeypatch): + calls = [] + + def _fake_call(v, a, o): + calls.append(1) + if len(calls) == 1: + return False, 0.0, "MuseTalk HTTP 503: busy", True + return True, 6.5, "", False + + monkeypatch.setattr(worker, "_register", lambda *a, **k: True) + monkeypatch.setattr(worker, "_download", lambda url, path: True) + monkeypatch.setattr(worker, "_probe_duration", lambda path: 12.0) + monkeypatch.setattr(worker, "_call_musetalk", _fake_call) + monkeypatch.setattr(worker, "time", mock.MagicMock()) # 重试 sleep 立即返回 + uploaded = [] + monkeypatch.setattr( + worker, + "_report_success_with_file", + lambda task_id, duration, path: uploaded.append((task_id, duration)), + ) + + worker._handle_task({"task_id": "t-retry", "video_url": "u", "audio_url": "u"}) + assert len(calls) == 2 + assert uploaded == [("t-retry", 6.5)] + + +def test_handle_task_no_retry_for_deterministic_failure(worker, monkeypatch): + calls = [] + + def _fake_call(v, a, o): + calls.append(1) + return False, 0.0, "MuseTalk HTTP 400: bad input", False + + reports = [] + monkeypatch.setattr(worker, "_register", lambda *a, **k: True) + monkeypatch.setattr(worker, "_download", lambda url, path: True) + monkeypatch.setattr(worker, "_probe_duration", lambda path: 12.0) + monkeypatch.setattr(worker, "_call_musetalk", _fake_call) + monkeypatch.setattr( + worker, + "_report_result", + lambda task_id, success, duration=0.0, error_msg="": reports.append(error_msg) or True, + ) + + worker._handle_task({"task_id": "t-4xx", "video_url": "u", "audio_url": "u"}) + assert len(calls) == 1 # 4xx 本地不重试，直接交服务端决定 + assert reports and "400" in reports[0] diff --git a/tests/unit/test_gpu_lipsync_service.py b/tests/unit/test_gpu_lipsync_service.py index e3a5303c0..4d9017e23 100644 --- a/tests/unit/test_gpu_lipsync_service.py +++ b/tests/unit/test_gpu_lipsync_service.py @@ -181,6 +181,73 @@ def test_register_worker_creates_then_updates(svc): assert w2.created_at == w.created_at # 没新建 +def test_register_with_task_id_refreshes_task_heartbeat(svc): + """#1970 推理期心跳：register(task_id=...) 只刷新本 worker 的 processing 任务.""" + from packages.adapters.sqlalchemy_impl.models import GpuWorkerModel + + t = svc.create_task(video_url="v", audio_url="a") + svc.poll_task("w-1") + svc.db.refresh(t) + old_hb = t.last_heartbeat_at + assert t.status == "processing" + # 模拟时间流逝后心跳到达 + svc.db.query(GpuWorkerModel).filter_by(worker_id="w-1").update( + {"last_heartbeat_at": old_hb - timedelta(seconds=300)} + ) + svc.db.commit() + svc.register_worker("w-1", task_id=t.id) + svc.db.refresh(t) + assert t.last_heartbeat_at > old_hb + assert t.status == "processing" # 心跳不改变状态 + # worker 表心跳也被刷新 + w = svc.db.query(GpuWorkerModel).filter_by(worker_id="w-1").one() + assert w.last_heartbeat_at > old_hb + + +def test_register_task_heartbeat_ignores_finished_or_foreign_task(svc): + """任务已 done，或已被超时回收重新派发给别的 worker 时，旧心跳必须忽略.""" + from packages.adapters.sqlalchemy_impl.models import GpuLipsyncTaskModel, GpuWorkerModel + + # 场景 1：任务已完成 → register 带 task_id 不得改写任务心跳 + t = svc.create_task(video_url="v", audio_url="a") + svc.poll_task("w-1") + done = svc.report_result(t.id, "w-1", success=True, duration_seconds=10.0) + hb_when_done = done.last_heartbeat_at + svc.register_worker("w-1", task_id=t.id) + svc.db.refresh(t) + assert t.status == "done" + assert t.last_heartbeat_at == hb_when_done # 没被改写 + + # 场景 2：任务超时回收后被 w-2 重新认领，旧 worker w-1 的迟到心跳无效 + t2 = svc.create_task(video_url="v2", audio_url="a2") + svc.poll_task("w-1") + svc.db.refresh(t2) + t2.last_heartbeat_at = datetime.now(UTC) - timedelta(days=1) + svc.db.commit() + claimed = svc.poll_task("w-2") # 触发回收并由 w-2 重新认领 + assert claimed is not None and claimed.id == t2.id + owner_hb = claimed.last_heartbeat_at + # 把 w-2 的 worker 心跳拨早，确认旧心跳不会影响任务归属 + svc.db.query(GpuWorkerModel).filter_by(worker_id="w-2").update( + {"last_heartbeat_at": owner_hb - timedelta(seconds=600)} + ) + svc.db.commit() + svc.register_worker("w-1", task_id=t2.id) # 旧 worker 迟到心跳 + svc.db.refresh(t2) + assert t2.worker_id == "w-2" + assert t2.status == "processing" + assert t2.last_heartbeat_at == owner_hb + + # 场景 3：不存在的 task_id 不报错 + svc.register_worker("w-1", task_id="nonexistent-id") + assert svc.db.get(GpuLipsyncTaskModel, "nonexistent-id") is None + + +def test_default_gpu_task_timeout_is_900(svc): + """#1970 默认超时 300→900，覆盖 RTX2060 长视频推理.""" + assert svc.settings.gpu_task_timeout_seconds == 900 + + # ── get_by_lipsync_job ─────────────────────────────────────────────