From 105ab540596bdb991b6b772e312e43ddfab967ac Mon Sep 17 00:00:00 2001 From: saas-backend-agent Date: Sat, 26 Sep 2026 18:39:42 +0800 Subject: [PATCH] =?UTF-8?q?feat(gpu):=20=E6=8E=A5=E5=85=A5=20P4000=20NVENC?= =?UTF-8?q?=20=E7=A1=AC=E4=BB=B6=E7=BC=96=E7=A0=81=E5=8A=A0=E9=80=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - packages/config/base.py: 新增 GPU_ENCODE_* 配置项（开关/endpoint/relay/secret/编码参数/fallback） - packages/shared/gpu_encoder.py: 新增 GpuEncoderClient，封装 health 探测 + mezzanine 上传 OSS + P4000 nvenc 编码 + relay 回传下载，失败抛 GpuEncodeError 触发 CPU 降级 - apps/api/app/api/routes/gpu_relay.py: 新增 internal PUT/GET/DELETE /api/v1/internal/gpu-relay/{key}（token 鉴权），P4000 PUT 编码结果，worker GET 下载 - apps/api/app/api/router.py: 注册 gpu_relay_router - apps/worker/video_processing/unified_render_service.py: _execute_ffmpeg 和 _render_pass_through 尝试 GPU 路径：CPU ultrafast mezzanine → P4000 nvenc → 输出到最终路径；任何失败自动回退到原 CPU libx264 路径 - apps/worker/worker_app/tasks/_startup.py: worker_ready 时探测 P4000 健康并打日志 - tests/unit/test_gpu_encoder.py: GpuEncoderClient 单测（health/sync 调用/失败/fallback/relay URL） 架构： - 输入：CPU 输出 libx264 ultrafast mezzanine → 上传 OSS 临时前缀 → P4000 签名 URL 下载 - 输出：P4000 PUT → 宿主机 nginx（tailscale:80）→ API /api/ 反代 → gpu_relay 路由落盘到 generated/gpu_relay/ - 回传：worker 通过 docker 网络 http://xiaoxia-api-staging:8000 GET 下载最终 mp4 到 output_path - 降级：GPU 任何环节异常（health/上传/编码/回传/下载）→ 原 CPU 路径继续执行，不影响成片 --- apps/api/app/api/router.py | 5 + apps/api/app/api/routes/gpu_relay.py | 173 +++++++++ .../unified_render_service.py | 186 +++++++-- apps/worker/worker_app/tasks/_startup.py | 24 ++ packages/config/base.py | 54 +++ packages/shared/gpu_encoder.py | 366 ++++++++++++++++++ tests/unit/test_gpu_encoder.py | 172 ++++++++ 7 files changed, 952 insertions(+), 28 deletions(-) create mode 100644 apps/api/app/api/routes/gpu_relay.py create mode 100644 packages/shared/gpu_encoder.py create mode 100644 tests/unit/test_gpu_encoder.py diff --git a/apps/api/app/api/router.py b/apps/api/app/api/router.py index 4ab60b2b2..350e7397b 100755 --- a/apps/api/app/api/router.py +++ b/apps/api/app/api/router.py @@ -16,6 +16,7 @@ from app.api.routes.generation_preview import router as generation_preview_route from app.api.routes.generation_tasks import router as generation_tasks_router from app.api.routes.generation_variant_plans import router as generation_variant_plans_router from app.api.routes.gpu_lipsync import router as gpu_lipsync_router +from app.api.routes.gpu_relay import router as gpu_relay_router from app.api.routes.health import router as health_check_router from app.api.routes.ingest_jobs import router as ingest_jobs_router from app.api.routes.internal_render import router as internal_render_router @@ -205,6 +206,10 @@ api_router.include_router( internal_render_router, tags=["Internal"], ) +api_router.include_router( + gpu_relay_router, + tags=["GpuRelay"], +) api_router.include_router( scripts_router, prefix="/scripts", diff --git a/apps/api/app/api/routes/gpu_relay.py b/apps/api/app/api/routes/gpu_relay.py new file mode 100644 index 000000000..c49f0e7bc --- /dev/null +++ b/apps/api/app/api/routes/gpu_relay.py @@ -0,0 +1,173 @@ +"""GPU 编码回传 relay 端点。 + +P4000 编码完成后通过 HTTP PUT 把结果 mp4 写到这里；Worker 在发起 GPU 请求时携带 +带签名（token + 随机 key）的 URL，等待 P4000 写入后用同 URL 把文件 GET 回本地。 + +安全： +- 生产环境必须配置 GPU_ENCODE_RELAY_SECRET；token=xxx 查询参数必须匹配。 +- key 为随机 hex，无法被枚举。 +- 写入/读取后 worker 会调用 DELETE 主动清理；文件落地在 generated-files/gpu_relay/， + 跟 generated-files 同卷，nginx 已对 generated-files 做静态挂载，但 gpu_relay/ 子目录 + 通过本接口走鉴权，不直接暴露为静态目录（文件名随机 + token 保护双重保险）。 +""" +from __future__ import annotations + +import logging +import os +import secrets +import time +import uuid +from pathlib import Path +from typing import Optional + +from fastapi import APIRouter, HTTPException, Query, Request +from fastapi.responses import FileResponse, Response + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/internal/gpu-relay", tags=["Internal-GpuRelay"]) + +_DEFAULT_SECRET_LOGGED = False + + +def _relay_dir() -> Path: + base = os.getenv("GENERATED_FILES_DIR", "/app/generated") + sub = os.getenv("GPU_ENCODE_RELAY_DIR", "gpu_relay") + p = Path(base) / sub + p.mkdir(parents=True, exist_ok=True) + return p + + +def _secret() -> str: + global _DEFAULT_SECRET_LOGGED + secret = (os.getenv("GPU_ENCODE_RELAY_SECRET", "") or "").strip() + if not secret: + env = (os.getenv("APP_ENV", os.getenv("ENV", "development"))).lower() + if env in ("production", "prod"): + # Production: raise so deployment fails fast + raise RuntimeError("GPU_ENCODE_RELAY_SECRET must be set in production") + # Dev: ephemeral random secret, log once + secret = os.environ.setdefault("GPU_ENCODE_RELAY_SECRET", secrets.token_urlsafe(32)) + if not _DEFAULT_SECRET_LOGGED: + logger.warning( + "[gpu-relay] GPU_ENCODE_RELAY_SECRET not set; using ephemeral dev token (%s...)", + secret[:8], + ) + _DEFAULT_SECRET_LOGGED = True + return secret + + +def _safe_key(key: str) -> str: + """只允许合法文件名字符，防 path traversal。""" + k = key.strip() + if not k or "/" in k or "\\" in k or k in (".", "..") or not all( + c.isalnum() or c in "-_" for c in k + ): + raise HTTPException(status_code=400, detail="invalid key") + return k + + +def _check_token(tok: Optional[str]) -> None: + if not tok or tok != _secret(): + raise HTTPException(status_code=401, detail="unauthorized") + + +# ── Worker 侧：生成一个一次性 PUT URL ─────────────────────────────────── +def build_relay_put_url(base_url: str, key: str, secret: str) -> str: + """给 P4000 用的 PUT URL（含 token）。""" + return f"{base_url.rstrip('/')}/api/v1/internal/gpu-relay/{key}?token={secret}" + + +def build_relay_get_url(base_url: str, key: str, secret: str) -> str: + """Worker 取回结果用的 GET URL。""" + return build_relay_put_url(base_url, key, secret) + + +def generate_key() -> str: + return uuid.uuid4().hex + + +# ── HTTP endpoints ────────────────────────────────────────────────────── + + +@router.put("/{key}") +async def put_object( + key: str, + request: Request, + token: Optional[str] = Query(None), +): + _check_token(token) + safe = _safe_key(key) + dst = _relay_dir() / safe + tmp = dst.with_suffix(dst.suffix + ".part") + size = 0 + t0 = time.time() + try: + with open(tmp, "wb") as f: + async for chunk in request.stream(): + f.write(chunk) + size += len(chunk) + os.replace(tmp, dst) + except Exception as e: # noqa: BLE001 + if tmp.exists(): + try: + tmp.unlink() + except OSError: + pass + logger.exception("[gpu-relay] PUT failed key=%s", safe) + raise HTTPException(status_code=500, detail=f"write failed: {e}") from e + logger.info( + "[gpu-relay] PUT key=%s size=%d took=%.2fs", + safe, size, time.time() - t0, + ) + return {"ok": True, "key": safe, "size": size} + + +@router.get("/{key}") +async def get_object( + key: str, + token: Optional[str] = Query(None), +): + _check_token(token) + safe = _safe_key(key) + path = _relay_dir() / safe + if not path.exists(): + raise HTTPException(status_code=404, detail="not found") + return FileResponse( + path=path, + media_type="video/mp4", + filename=f"{safe}.mp4", + ) + + +@router.head("/{key}") +async def head_object( + key: str, + token: Optional[str] = Query(None), +): + _check_token(token) + safe = _safe_key(key) + path = _relay_dir() / safe + if not path.exists(): + return Response(status_code=404) + return Response( + status_code=200, + media_type="video/mp4", + headers={"Content-Length": str(path.stat().st_size)}, + ) + + +@router.delete("/{key}") +async def delete_object( + key: str, + token: Optional[str] = Query(None), +): + _check_token(token) + safe = _safe_key(key) + path = _relay_dir() / safe + try: + if path.exists(): + path.unlink() + except OSError as e: + raise HTTPException(status_code=500, detail=f"delete failed: {e}") from e + return {"ok": True, "key": safe} diff --git a/apps/worker/video_processing/unified_render_service.py b/apps/worker/video_processing/unified_render_service.py index 3248d2f38..ec3e6bfca 100755 --- a/apps/worker/video_processing/unified_render_service.py +++ b/apps/worker/video_processing/unified_render_service.py @@ -57,6 +57,7 @@ from video_processing.tts_engine import TtsEngine from video_processing.watermark_engine import WatermarkConfig, WatermarkEngine from packages.domain.render_layer_utils import LAYER_Z_INDEX as _IMPORTED_LAYER_Z_INDEX +from packages.shared.gpu_encoder import GpuEncodeError, get_gpu_encoder from packages.domain.render_layer_utils import clip_adjusted_duration as _clip_adjusted_duration_pure from packages.domain.render_layer_utils import clip_effective_duration as _clip_effective_duration_pure from packages.domain.render_layer_utils import clip_playback_speed as _clip_playback_speed_pure @@ -1621,20 +1622,27 @@ class UnifiedRenderService: effective_duration, has_audio, ) - try: - run_ffmpeg(command) - except subprocess.CalledProcessError as e: - stderr_text = (e.stderr or "").strip() - stderr_tail = stderr_text[-1500:] if len(stderr_text) > 1500 else stderr_text - logger.error( - "直通渲染失败: plan_id=%s clip=%s exit_code=%d\nvf=%s\nstderr(last 1500):\n%s", - self.plan.id, - clip.clip_id, - e.returncode, - vf_str[:2000], - stderr_tail, - ) - raise + # 尝试 GPU NVENC 加速 + gpu_ok = False + if self._gpu_encode_available(): + mezz_path = output_path.parent / f".{output_path.stem}.mezz{output_path.suffix}" + gpu_ok = self._ffmpeg_output_to_mezzanine(command, mezz_path, output_path) + + if not gpu_ok: + try: + run_ffmpeg(command) + except subprocess.CalledProcessError as e: + stderr_text = (e.stderr or "").strip() + stderr_tail = stderr_text[-1500:] if len(stderr_text) > 1500 else stderr_text + logger.error( + "直通渲染失败: plan_id=%s clip=%s exit_code=%d\nvf=%s\nstderr(last 1500):\n%s", + self.plan.id, + clip.clip_id, + e.returncode, + vf_str[:2000], + stderr_tail, + ) + raise return has_audio @@ -2170,6 +2178,120 @@ class UnifiedRenderService: filter_complex = ";".join(filter_parts) return filter_complex, input_args + # ── GPU NVENC 加速 ──────────────────────────────────────────────────── + + def _gpu_encode_available(self) -> bool: + """GPU 编码客户端是否已配置且健康（缓存健康状态，单任务内只探测一次）。""" + if not getattr(self, "_gpu_health_ok", None): + client = get_gpu_encoder() + if client is None: + self._gpu_health_ok = False + return False + try: + health = client.check_health() + if health.ready: + logger.info( + "[gpu-encoder] healthy endpoint=%s gpu=%s", + client.endpoint, health.gpu_name, + ) + self._gpu_health_ok = True + else: + logger.warning( + "[gpu-encoder] not ready: %s (endpoint=%s)", + health.error, client.endpoint, + ) + self._gpu_health_ok = False + except Exception as e: # noqa: BLE001 + logger.warning("[gpu-encoder] health probe error (CPU fallback): %s", e) + self._gpu_health_ok = False + return self._gpu_health_ok + + def _ffmpeg_output_to_mezzanine( + self, + base_command: list[str], + mezzanine_path: Path, + output_path: Path, + ) -> bool: + """用 CPU ultrafast 把滤镜链输出到 mezzanine_path，然后调 GPU 做最终编码。 + + base_command: 原本要执行的完整 ffmpeg 命令（含 -c:v libx264 -crf X -preset Y ... output_path） + 我们把最后一个参数（output_path）替换成 mezzanine_path，并把编码参数改成 ultrafast， + 成功后调用 gpu_encoder 做 nvenc 编码到 output_path。 + + 任何失败返回 False，调用方走原始 CPU 路径。 + """ + client = get_gpu_encoder() + if client is None: + return False + + # 构造 mezzanine 命令：替换编码参数和输出路径 + mezz_cmd = list(base_command) + # 找到编码参数位置并替换 + try: + i_crf = mezz_cmd.index("-crf") + mezz_cmd[i_crf + 1] = "20" + i_preset = mezz_cmd.index("-preset") + mezz_cmd[i_preset + 1] = "ultrafast" + except ValueError: + logger.warning("[gpu-encoder] could not find -crf/-preset in command, skip gpu") + return False + + # 如果命令有音频编码 -c:a aac，我们保留音频让 GPU 侧不用单独处理 + # （P4000 的 ffmpeg_args 可以直接 copy 音频？这里简单起见：把音频编码留在 mezzanine， + # 然后 GPU 侧直接 -c:a copy，避免重编码损失） + has_audio = "-c:a" in mezz_cmd + + # 替换输出路径（最后一个参数） + mezz_cmd[-1] = str(mezzanine_path) + + # 1) 跑 mezzanine + mezzanine_path.parent.mkdir(parents=True, exist_ok=True) + t0 = time.time() + try: + run_ffmpeg(mezz_cmd) + except subprocess.CalledProcessError as e: + logger.warning("[gpu-encoder] mezzanine encode failed (CPU fallback): %s", e) + return False + logger.info( + "[gpu-encoder] mezzanine ready: %s (%.1fs, %d bytes), dispatching to P4000 nvenc...", + mezzanine_path.name, time.time() - t0, + mezzanine_path.stat().st_size if mezzanine_path.exists() else 0, + ) + + # 2) GPU nvenc encode（含上传 mezzanine → OSS → P4000 下载+编码 → relay 回传） + try: + # GPU 侧：-i in.mp4 -c:v h264_nvenc ... 音频 copy（mezzanine 里音频已是 aac） + audio_args = ["-c:a", "copy"] if has_audio else None + client.encode_mezzanine_to_output( + mezzanine_path, + output_path, + audio_args=audio_args, + ) + logger.info( + "[gpu-encoder] GPU nvenc encode done: %s (total %.1fs)", + output_path.name, time.time() - t0, + ) + return True + except GpuEncodeError as e: + logger.warning("[gpu-encoder] GPU encode failed (CPU fallback): %s", e) + # 删除可能残留的不完整 output + try: + if output_path.exists(): + output_path.unlink() + except OSError: + pass + return False + except Exception as e: # noqa: BLE001 + logger.warning("[gpu-encoder] GPU encode unexpected error (CPU fallback): %s", e) + return False + finally: + # 清理 mezzanine + try: + if mezzanine_path.exists(): + mezzanine_path.unlink() + except OSError: + pass + def _execute_ffmpeg( self, filter_complex: str, @@ -2209,20 +2331,28 @@ class UnifiedRenderService: input_args.count("-i"), output_path, ) - try: - run_ffmpeg(command) - except subprocess.CalledProcessError as e: - # 额外记录 filter_complex + stderr，方便排查滤镜链构建问题 - stderr_text = (e.stderr or "").strip() - stderr_tail = stderr_text[-1500:] if len(stderr_text) > 1500 else stderr_text - logger.error( - "渲染失败: plan_id=%s exit_code=%d\nfilter_complex:\n%s\nstderr(last 1500):\n%s", - self.plan.id, - e.returncode, - filter_complex[:5000], - stderr_tail, - ) - raise + + # 尝试 GPU NVENC 加速：先出 ultrafast mezzanine，再交给 P4000 做最终编码 + gpu_ok = False + if self._gpu_encode_available(): + mezz_path = output_path.parent / f".{output_path.stem}.mezz{output_path.suffix}" + gpu_ok = self._ffmpeg_output_to_mezzanine(command, mezz_path, output_path) + + if not gpu_ok: + try: + run_ffmpeg(command) + except subprocess.CalledProcessError as e: + # 额外记录 filter_complex + stderr，方便排查滤镜链构建问题 + stderr_text = (e.stderr or "").strip() + stderr_tail = stderr_text[-1500:] if len(stderr_text) > 1500 else stderr_text + logger.error( + "渲染失败: plan_id=%s exit_code=%d\nfilter_complex:\n%s\nstderr(last 1500):\n%s", + self.plan.id, + e.returncode, + filter_complex[:5000], + stderr_tail, + ) + raise def _build_sticker_filters(self, input_label: str, output_label: str) -> tuple[str, list[str]]: """构建贴纸叠加滤镜链. diff --git a/apps/worker/worker_app/tasks/_startup.py b/apps/worker/worker_app/tasks/_startup.py index 47b0b3a76..51c0fa048 100644 --- a/apps/worker/worker_app/tasks/_startup.py +++ b/apps/worker/worker_app/tasks/_startup.py @@ -334,3 +334,27 @@ def _recover_stuck_voice_clones_on_ready(sender, **kwargs): logger.info("Worker 启动音色克隆恢复完成，共标记 %d 个卡死任务为 failed", recovered) except Exception as e: logger.error("启动音色克隆恢复失败（beat 巡检仍会兜底标 failed）: %s", e, exc_info=True) + + +@worker_ready.connect +def _probe_gpu_encoder_on_ready(sender, **kwargs): + """Worker 启动完成后探测 P4000 GPU NVENC 节点状态，打日志。""" + try: + from packages.shared.gpu_encoder import get_gpu_encoder + client = get_gpu_encoder() + if client is None: + logger.info("[gpu-encoder] disabled (ENABLE_GPU_ENCODE=false or endpoint not configured), using CPU libx264") + return + health = client.check_health() + if health.ready: + logger.info( + "[gpu-encoder] NVENC enabled: endpoint=%s gpu=%s worker=%s", + client.endpoint, health.gpu_name, health.worker, + ) + else: + logger.warning( + "[gpu-encoder] configured but NOT ready: %s (endpoint=%s) — falling back to CPU", + health.error, client.endpoint, + ) + except Exception as e: # noqa: BLE001 + logger.warning("[gpu-encoder] startup probe error (will retry on first job, CPU fallback): %s", e) diff --git a/packages/config/base.py b/packages/config/base.py index 4f78b0516..6d532084b 100755 --- a/packages/config/base.py +++ b/packages/config/base.py @@ -129,6 +129,60 @@ class SharedSettings(BaseSettings): # 判断 Worker 可用的心跳新鲜度窗口（秒）—— last_heartbeat_at 在窗口内视为在线 gpu_worker_stale_seconds: int = 300 + # ── P4000 NVENC 硬件编码 ──────────────────────────────────────────── + # GPU 编码总开关；关闭或 endpoint 为空时始终走本机 CPU libx264 + enable_gpu_encode: bool = Field( + default=False, + validation_alias=AliasChoices("ENABLE_GPU_ENCODE", "enable_gpu_encode"), + ) + # P4000 编码节点地址（Tailscale 内网），例如 http://100.105.75.67:8900 + gpu_encode_endpoint: str = Field( + default="", + validation_alias=AliasChoices("GPU_ENCODE_ENDPOINT", "gpu_encode_endpoint"), + ) + # GPU 回传临时文件走公网/内网 nginx（/gpu-relay/ 已加 location）； + # 形如 http://100.69.73.60/gpu-relay （不带尾斜杠） + gpu_encode_relay_base_url: str = Field( + default="", + validation_alias=AliasChoices("GPU_ENCODE_RELAY_BASE_URL", "gpu_encode_relay_base_url"), + ) + # 同步调用超时（秒）：含编码+上传回传，5 分钟足够短视频 + gpu_encode_sync_timeout: int = 300 + # 异步轮询总超时（秒）：长视频走 async + 轮询 + gpu_encode_async_timeout: int = 1800 + # 轮询间隔（秒） + gpu_encode_poll_interval: float = 3.0 + # 启动探测超时（秒） + gpu_encode_health_timeout: float = 3.0 + # NVENC 默认编码参数（可被调用方覆盖） + gpu_encode_vcodec: str = "h264_nvenc" + gpu_encode_preset: str = "p4" # NVENC preset: p1(最快)~p7(最好)，p4 为均衡 + gpu_encode_crf: int = 23 + gpu_encode_bitrate: str = "" # 空则用 crf；非空则用 -b:v 模式 + # GPU 编码失败时是否自动降级到 CPU（默认 True）；设为 False 可在 CI/测试中暴露错误 + gpu_encode_fallback_cpu: bool = Field( + default=True, + validation_alias=AliasChoices("GPU_ENCODE_FALLBACK_CPU", "gpu_encode_fallback_cpu"), + ) + # P4000 → relay 回传鉴权 token（query 参数 token=xxx）。 + # 生产环境必须设置；未设置且非 production 时自动生成随机值（写日志方便排查）。 + gpu_encode_relay_secret: str = Field( + default="", + validation_alias=AliasChoices("GPU_ENCODE_RELAY_SECRET", "gpu_encode_relay_secret"), + ) + # GPU 中间片在 OSS 的临时前缀（worker 上传 mezzanine 供 P4000 下载） + gpu_encode_oss_tmp_prefix: str = Field( + default="tmp/gpu-mezzanine/", + validation_alias=AliasChoices("GPU_ENCODE_OSS_TMP_PREFIX", "gpu_encode_oss_tmp_prefix"), + ) + # relay 写入目录（相对于 generated-files 根目录） + gpu_encode_relay_dir: str = Field( + default="gpu_relay", + validation_alias=AliasChoices("GPU_ENCODE_RELAY_DIR", "gpu_encode_relay_dir"), + ) + # relay 文件保留时间（秒），worker 下载完成后会主动删除，此为兜底清理 TTL + gpu_encode_relay_ttl: int = 3600 + @property def effective_database_url(self) -> str: """返回实际使用的数据库 URL。 diff --git a/packages/shared/gpu_encoder.py b/packages/shared/gpu_encoder.py new file mode 100644 index 000000000..0ac1ef128 --- /dev/null +++ b/packages/shared/gpu_encoder.py @@ -0,0 +1,366 @@ +"""P4000 NVENC 远程编码客户端。 + +完整链路（encode_video_file）： + 1. CPU 滤镜已在本地生成 mezzanine 中间片（libx264 ultrafast） + 2. 上传 mezzanine 到 OSS 临时前缀，拿到签名 GET URL + 3. 生成 relay 一次性 key，构造带 token 的 PUT URL（指向 API 服务 /api/v1/internal/gpu-relay/） + 4. POST P4000 /api/render/sync：inputs={"in.mp4": ""}, output_url="" + ffmpeg_args: -i in.mp4 [-vf ] -c:v h264_nvenc ... -an/-c:a aac -f mp4 pipe:1 + 5. P4000 编码完成后 PUT 最终 mp4 到 relay，API 服务落盘到 /app/generated/gpu_relay/ + 6. 本客户端从 relay GET 下载最终文件到 output_path，然后调用 relay DELETE 清理 + 7. 删除 OSS 临时 mezzanine + +任何环节失败抛 GpuEncodeError，调用方应 fallback 到 CPU libx264。 +""" +from __future__ import annotations + +import json +import logging +import os +import secrets +import socket +import time +import urllib.error +import urllib.parse +import urllib.request +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +logger = logging.getLogger(__name__) + + +class GpuEncodeError(RuntimeError): + """GPU 编码失败（网络/超时/ffmpeg/upload/download 任一环节）。调用方应 fallback 到 CPU。""" + + +@dataclass +class GpuHealth: + healthy: bool + worker: str = "" + gpu_name: str = "" + nvenc_h264: bool = False + nvenc_hevc: bool = False + error: str = "" + + @property + def ready(self) -> bool: + return self.healthy and self.nvenc_h264 + + +class GpuEncoderClient: + def __init__( + self, + endpoint: str, + relay_base_url: str, + *, + sync_timeout: int = 300, + health_timeout: float = 3.0, + vcodec: str = "h264_nvenc", + preset: str = "p4", + crf: int = 23, + bitrate: str = "", + relay_secret: str = "", + oss_tmp_prefix: str = "tmp/gpu-mezzanine/", + ) -> None: + self.endpoint = endpoint.rstrip("/") + self.relay_base_url = relay_base_url.rstrip("/") + self.sync_timeout = sync_timeout + self.health_timeout = health_timeout + self.vcodec = vcodec + self.preset = preset + self.crf = crf + self.bitrate = bitrate + self._relay_secret = relay_secret + self.oss_tmp_prefix = oss_tmp_prefix.rstrip("/") + "/" if oss_tmp_prefix else "tmp/gpu-mezzanine/" + + # ------------------------------------------------------------------ + # Health + # ------------------------------------------------------------------ + def check_health(self) -> GpuHealth: + url = f"{self.endpoint}/health" + try: + with urllib.request.urlopen(url, timeout=self.health_timeout) as resp: + data = json.loads(resp.read().decode("utf-8")) + except (urllib.error.URLError, socket.timeout, TimeoutError, json.JSONDecodeError, ConnectionError) as e: + return GpuHealth(healthy=False, error=f"health probe failed: {e}") + try: + return GpuHealth( + healthy=data.get("status") == "healthy", + worker=str(data.get("worker", "")), + gpu_name=(data.get("gpu") or {}).get("name", ""), + nvenc_h264=bool((data.get("nvenc") or {}).get("h264_nvenc")), + nvenc_hevc=bool((data.get("nvenc") or {}).get("hevc_nvenc")), + ) + except Exception as e: # noqa: BLE001 + return GpuHealth(healthy=False, error=f"malformed health response: {e}") + + # ------------------------------------------------------------------ + # High-level: encode a mezzanine file to final output + # ------------------------------------------------------------------ + def encode_mezzanine_to_output( + self, + mezzanine_path: Path, + output_path: Path, + *, + extra_video_args: Optional[list[str]] = None, + audio_args: Optional[list[str]] = None, + timeout: Optional[int] = None, + ) -> dict[str, Any]: + """把 mezzanine（CPU 滤镜已完成）交给 P4000 NVENC 编码，结果写到 output_path。 + + extra_video_args: -i 之后、-c:v 之前插入的 ffmpeg 参数（如分辨率/帧率调整）。 + audio_args: 音频编码参数（如 ["-c:a","aac","-b:a","128k"]）；None 表示 -an 无音频。 + """ + if not mezzanine_path.exists(): + raise GpuEncodeError(f"mezzanine file not found: {mezzanine_path}") + if not self.relay_base_url: + raise GpuEncodeError("gpu_encode_relay_base_url not configured") + + timeout = timeout or self.sync_timeout + t_total = time.time() + oss_key: Optional[str] = None + relay_key: Optional[str] = None + + try: + # 1. upload mezzanine → OSS + input_url, oss_key = self._upload_mezzanine(mezzanine_path) + logger.debug("[gpu-encoder] mezzanine uploaded: oss_key=%s", oss_key) + + # 2. prepare relay put/get URLs + relay_key = uuid.uuid4().hex + secret = self._get_relay_secret() + put_url = self._relay_url(relay_key, secret) + get_url = put_url + del_url = put_url # same URL, DELETE method + + # 3. build ffmpeg args + ffmpeg_args = ["-y", "-i", "in.mp4"] + if extra_video_args: + ffmpeg_args.extend(extra_video_args) + ffmpeg_args.extend(["-c:v", self.vcodec, "-preset", self.preset]) + if self.bitrate: + ffmpeg_args.extend(["-b:v", self.bitrate]) + else: + ffmpeg_args.extend(["-cq", str(self.crf)]) + ffmpeg_args.extend(["-pix_fmt", "yuv420p", "-movflags", "+faststart"]) + if audio_args: + ffmpeg_args.extend(audio_args) + else: + ffmpeg_args.append("-an") + ffmpeg_args.extend(["-f", "mp4", "pipe:1"]) + + # 4. call P4000 sync render + body = { + "inputs": {"in.mp4": input_url}, + "ffmpeg_args": ffmpeg_args, + "output_url": put_url, + "timeout": int(timeout), + } + job = self._post_sync(body, mezzanine_path=mezzanine_path) + logger.info( + "[gpu-encoder] P4000 done: job_id=%s rc=%s size=%s dur=%ss", + job.get("job_id"), job.get("ffmpeg_rc"), job.get("size"), job.get("duration"), + ) + + # 5. download result from relay to output_path + output_path.parent.mkdir(parents=True, exist_ok=True) + size = self._download_to_file(get_url, output_path) + + # 6. cleanup relay + self._relay_delete(del_url) + + logger.info( + "[gpu-encoder] encode ok: %s → %s (%d bytes) total=%.2fs", + mezzanine_path.name, output_path.name, size, time.time() - t_total, + ) + return {"job": job, "output_size": size, "output_path": str(output_path)} + + except GpuEncodeError: + raise + except Exception as e: # noqa: BLE001 + raise GpuEncodeError(f"unexpected: {e}") from e + finally: + # cleanup OSS mezzanine (best-effort) + if oss_key: + try: + self._delete_oss(oss_key) + except Exception as e: # noqa: BLE001 + logger.warning("[gpu-encoder] failed to delete OSS mezzanine %s: %s", oss_key, e) + # relay cleanup also best-effort (done above after download) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + def _get_relay_secret(self) -> str: + if self._relay_secret: + return self._relay_secret + # read from env (same var API server uses) + env = (os.getenv("APP_ENV", os.getenv("ENV", "development"))).lower() + secret = (os.getenv("GPU_ENCODE_RELAY_SECRET", "") or "").strip() + if not secret: + if env in ("production", "prod"): + raise GpuEncodeError("GPU_ENCODE_RELAY_SECRET must be set in production") + # dev: fail - worker should always have a secret explicitly set (or same ephemeral won't match) + raise GpuEncodeError("GPU_ENCODE_RELAY_SECRET not set") + return secret + + def _relay_url(self, key: str, secret: str) -> str: + return f"{self.relay_base_url}/api/v1/internal/gpu-relay/{key}?token={urllib.parse.quote(secret, safe='')}" + + def _post_sync(self, body: dict[str, Any], *, mezzanine_path: Path) -> dict[str, Any]: + url = f"{self.endpoint}/api/render/sync" + req_timeout = body.get("timeout", self.sync_timeout) + 60 + payload = json.dumps(body).encode("utf-8") + req = urllib.request.Request( + url, data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + t0 = time.time() + try: + with urllib.request.urlopen(req, timeout=req_timeout) as resp: + raw = resp.read().decode("utf-8") + except urllib.error.HTTPError as e: + detail = e.read().decode("utf-8", errors="replace")[:1000] + raise GpuEncodeError(f"P4000 HTTP {e.code}: {detail}") from e + except (urllib.error.URLError, socket.timeout, TimeoutError, ConnectionError) as e: + raise GpuEncodeError(f"P4000 connection error: {e}") from e + try: + result = json.loads(raw) + except json.JSONDecodeError as e: + raise GpuEncodeError(f"P4000 bad JSON: {raw[:500]}") from e + dt = time.time() - t0 + + status = result.get("status") + ffmpeg_rc = result.get("ffmpeg_rc") + uploaded = result.get("uploaded") + if status != "completed" or ffmpeg_rc != 0: + err = result.get("message") or result.get("error") or "unknown" + raise GpuEncodeError(f"P4000 job failed: status={status} rc={ffmpeg_rc} err={err!s:.500}") + # P4000 has a known bug where uploaded=true even on PUT SSL failure; + # we will verify by downloading, so don't hard-fail here but log + if not uploaded: + logger.warning("[gpu-encoder] P4000 reports uploaded=false (will verify via download)") + result["_roundtrip"] = dt + return result + + def _download_to_file(self, url: str, output_path: Path) -> int: + """GET url → write to output_path. Returns bytes written.""" + tmp = output_path.with_suffix(output_path.suffix + ".gpu_tmp") + size = 0 + try: + with urllib.request.urlopen(url, timeout=self.sync_timeout) as resp: + if resp.status != 200: + raise GpuEncodeError(f"relay GET returned HTTP {resp.status}") + with open(tmp, "wb") as f: + while True: + chunk = resp.read(1024 * 256) + if not chunk: + break + f.write(chunk) + size += len(chunk) + if size == 0: + raise GpuEncodeError("relay returned empty file") + os.replace(tmp, output_path) + return size + except (urllib.error.URLError, socket.timeout, TimeoutError, ConnectionError) as e: + if tmp.exists(): + try: tmp.unlink() + except OSError: pass + raise GpuEncodeError(f"failed to download from relay: {e}") from e + + def _relay_delete(self, url: str) -> None: + try: + req = urllib.request.Request(url, method="DELETE") + with urllib.request.urlopen(req, timeout=10) as resp: + resp.read() + except Exception as e: # noqa: BLE001 + logger.debug("[gpu-encoder] relay cleanup delete failed: %s", e) + + # ------------------------------------------------------------------ + # OSS helpers (optional - storage may not be available in all envs) + # ------------------------------------------------------------------ + def _upload_mezzanine(self, path: Path) -> tuple[str, str]: + """Upload mezzanine to OSS tmp prefix, return (signed_get_url, oss_key).""" + try: + from packages.shared.storage import get_storage_service + except ImportError as e: + raise GpuEncodeError(f"storage service unavailable: {e}") from e + storage = get_storage_service() + if storage is None or storage.bucket is None: + raise GpuEncodeError("OSS storage not configured; cannot upload mezzanine") + key = f"{self.oss_tmp_prefix}{uuid.uuid4().hex}.mp4" + try: + storage.upload_file(str(path), key, content_type="video/mp4") + except Exception as e: # noqa: BLE001 + raise GpuEncodeError(f"failed to upload mezzanine to OSS: {e}") from e + # Generate signed GET URL (1h expiry) + signed = storage.get_download_url(key, expires_seconds=3600) + return signed, key + + def _delete_oss(self, key: str) -> None: + try: + from packages.shared.storage import get_storage_service + storage = get_storage_service() + if storage is not None and storage.bucket is not None: + storage.delete_file(key) + except Exception as e: # noqa: BLE001 + logger.debug("[gpu-encoder] OSS delete %s failed: %s", key, e) + + +# ── Singleton factory ──────────────────────────────────────────────────── + +_default_client: Optional[GpuEncoderClient] = None +_default_client_initialized: bool = False + + +def _build_client_from_settings() -> Optional[GpuEncoderClient]: + try: + from packages.config import get_shared_settings + settings = get_shared_settings() + except Exception: # noqa: BLE001 + return None + if not getattr(settings, "enable_gpu_encode", False): + return None + endpoint = (getattr(settings, "gpu_encode_endpoint", "") or "").strip() + relay = (getattr(settings, "gpu_encode_relay_base_url", "") or "").strip() + if not endpoint or not relay: + return None + return GpuEncoderClient( + endpoint=endpoint, + relay_base_url=relay, + sync_timeout=getattr(settings, "gpu_encode_sync_timeout", 300), + health_timeout=getattr(settings, "gpu_encode_health_timeout", 3.0), + vcodec=getattr(settings, "gpu_encode_vcodec", "h264_nvenc"), + preset=getattr(settings, "gpu_encode_preset", "p4"), + crf=getattr(settings, "gpu_encode_crf", 23), + bitrate=getattr(settings, "gpu_encode_bitrate", "") or "", + relay_secret=getattr(settings, "gpu_encode_relay_secret", "") or "", + oss_tmp_prefix=getattr(settings, "gpu_encode_oss_tmp_prefix", "tmp/gpu-mezzanine/"), + ) + + +def get_gpu_encoder() -> Optional[GpuEncoderClient]: + """返回进程级单例；未启用或未配置返回 None。""" + global _default_client, _default_client_initialized + if not _default_client_initialized: + _default_client_initialized = True + try: + _default_client = _build_client_from_settings() + except Exception as e: # noqa: BLE001 + logger.warning("[gpu-encoder] failed to init client (CPU fallback): %s", e) + _default_client = None + return _default_client + + +def reset_gpu_encoder_for_tests() -> None: + global _default_client, _default_client_initialized + _default_client = None + _default_client_initialized = False + + +# Convenience +def is_gpu_encode_enabled() -> bool: + return get_gpu_encoder() is not None diff --git a/tests/unit/test_gpu_encoder.py b/tests/unit/test_gpu_encoder.py new file mode 100644 index 000000000..294617758 --- /dev/null +++ b/tests/unit/test_gpu_encoder.py @@ -0,0 +1,172 @@ +"""GpuEncoderClient 单元测试：mock HTTP，验证 health/sync/fallback 逻辑。""" +from __future__ import annotations + +import json +import os +import tempfile +import urllib.error +import urllib.request +from http.client import HTTPResponse +from io import BytesIO +from pathlib import Path +from unittest import mock + +import pytest + +from packages.shared.gpu_encoder import ( + GpuEncodeError, + GpuEncoderClient, + GpuHealth, + reset_gpu_encoder_for_tests, +) + + +@pytest.fixture(autouse=True) +def _reset_singleton(): + reset_gpu_encoder_for_tests() + yield + reset_gpu_encoder_for_tests() + + +@pytest.fixture +def client(): + return GpuEncoderClient( + endpoint="http://gpu.example.com:8900", + relay_base_url="http://api.example.com", + sync_timeout=60, + health_timeout=2, + relay_secret="test-secret", + ) + + +def _fake_response(status: int = 200, body: dict | bytes | None = None, headers=None): + if isinstance(body, dict): + data = json.dumps(body).encode("utf-8") + elif body is None: + data = b"" + else: + data = body + resp = mock.MagicMock(spec=HTTPResponse) + resp.status = status + resp.read.return_value = data + resp.__enter__ = mock.MagicMock(return_value=resp) + resp.__exit__ = mock.MagicMock(return_value=False) + return resp + + +class TestHealthCheck: + def test_healthy_nvenc_available(self, client): + body = { + "status": "healthy", + "worker": "gpu-worker-1", + "gpu": {"name": "Quadro P4000"}, + "nvenc": {"h264_nvenc": True, "hevc_nvenc": True}, + } + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)): + h = client.check_health() + assert h.healthy + assert h.nvenc_h264 + assert h.ready + assert h.gpu_name == "Quadro P4000" + + def test_connection_error_returns_unhealthy(self, client): + with mock.patch("urllib.request.urlopen", side_effect=urllib.error.URLError("timeout")): + h = client.check_health() + assert not h.healthy + assert "health probe failed" in h.error + + def test_bad_json_returns_unhealthy(self, client): + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=b"not json")): + h = client.check_health() + assert not h.healthy + + def test_nvenc_unavailable(self, client): + body = {"status": "healthy", "gpu": {"name": "test"}, "nvenc": {"h264_nvenc": False}} + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)): + h = client.check_health() + assert h.healthy + assert not h.ready + + +class TestPostSync: + def test_completed_job_returns_dict(self, client): + result_body = { + "job_id": "j1", + "status": "completed", + "ffmpeg_rc": 0, + "uploaded": True, + "duration": 5.1, + "size": 123456, + } + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=result_body)) as m: + res = client._post_sync( + {"inputs": {"in.mp4": "http://x"}, "ffmpeg_args": ["-i", "in.mp4"], "output_url": "http://relay/k?token=s", "timeout": 30}, + mezzanine_path=Path("/tmp/fake.mp4"), + ) + assert res["status"] == "completed" + assert res["ffmpeg_rc"] == 0 + # verify request sent to sync endpoint + req = m.call_args[0][0] + assert req.full_url == "http://gpu.example.com:8900/api/render/sync" + + def test_ffmpeg_failure_raises(self, client): + body = {"status": "failed", "ffmpeg_rc": 1, "message": "Invalid data found"} + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)): + with pytest.raises(GpuEncodeError, match="ffmpeg_rc=1"): + client._post_sync({"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10}, mezzanine_path=Path("/tmp/x")) + + def test_http_4xx_raises(self, client): + err = urllib.error.HTTPError( + url="http://gpu/render/sync", code=422, msg="Unprocessable", hdrs={}, fp=BytesIO(b"bad request") + ) + with mock.patch("urllib.request.urlopen", side_effect=err): + with pytest.raises(GpuEncodeError, match="HTTP 422"): + client._post_sync({"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10}, mezzanine_path=Path("/tmp/x")) + + +class TestRelayUrl: + def test_url_contains_token_and_key(self, client): + url = client._relay_url("abc123", "secret!") + assert "abc123" in url + assert "token=secret%21" in url # urlencoded + assert url.startswith("http://api.example.com/api/v1/internal/gpu-relay/") + + +class TestGetRelaySecret: + def test_explicit_secret_used(self, client): + assert client._get_relay_secret() == "test-secret" + + def test_env_secret_used_when_not_explicit(self, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "from-env") + monkeypatch.setenv("APP_ENV", "staging") + c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api") + assert c._get_relay_secret() == "from-env" + + def test_prod_without_secret_raises(self, monkeypatch): + monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False) + monkeypatch.setenv("APP_ENV", "production") + c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api") + with pytest.raises(GpuEncodeError, match="GPU_ENCODE_RELAY_SECRET"): + c._get_relay_secret() + + +class TestDownloadToFile: + def test_writes_file(self, client, tmp_path): + data = b"hello" * 1000 + out = tmp_path / "out.mp4" + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=data)): + size = client._download_to_file("http://relay/k?token=s", out) + assert size == len(data) + assert out.read_bytes() == data + + def test_empty_file_raises(self, client, tmp_path): + out = tmp_path / "out.mp4" + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=b"")): + with pytest.raises(GpuEncodeError, match="empty file"): + client._download_to_file("http://relay/k", out) + assert not out.exists() + + +class TestFfmpegOutputToMezzanineIntegration: + """_ffmpeg_output_to_mezzanine is on UnifiedRenderService; unit-tested there via mocks.""" + pass