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..a92e3bbe8 100755 --- a/apps/worker/video_processing/unified_render_service.py +++ b/apps/worker/video_processing/unified_render_service.py @@ -63,6 +63,7 @@ from packages.domain.render_layer_utils import clip_playback_speed as _clip_play from packages.domain.render_layer_utils import estimate_total_duration as _estimate_total_duration_pure from packages.domain.render_layer_utils import resolve_layer_role as _resolve_layer_role_pure from packages.domain.tts_config import TtsConfig +from packages.shared.gpu_encoder import GpuEncodeError, get_gpu_encoder logger = logging.getLogger(__name__) @@ -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,124 @@ 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 +2335,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..ab232f874 100644 --- a/apps/worker/worker_app/tasks/_startup.py +++ b/apps/worker/worker_app/tasks/_startup.py @@ -334,3 +334,33 @@ 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/infra/nginx/gpu-relay-staging.conf b/infra/nginx/gpu-relay-staging.conf new file mode 100644 index 000000000..a97831867 --- /dev/null +++ b/infra/nginx/gpu-relay-staging.conf @@ -0,0 +1,31 @@ +# Staging GPU relay plain-HTTP vhost (P4000 NVENC 编码回传入口) +# - 监听 8092 端口纯 HTTP(绕开 HTTPS 证书与 P4000 httpx SSL 问题) +# - 代理到本机 staging API 的 /api/ 路径(127.0.0.1:8000 是 docker 映射端口) +# - P4000 通过 Tailscale 直连宿主机 100.69.73.60:8092 PUT 编码结果 +# - Worker 通过 Docker DNS (xiaoxia-api-staging:8000) 直接 GET/DELETE, +# 不经宿主机 nginx,避免 UFW FORWARD DROP 阻断 +# +# 部署:cp infra/nginx/gpu-relay-staging.conf /etc/nginx/conf.d/ && nginx -t && systemctl reload nginx + +server { + listen 8092; + server_name _; + + client_max_body_size 2048m; + + location /api/ { + proxy_pass http://127.0.0.1:8000/api/; + proxy_http_version 1.1; + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_request_buffering off; + proxy_read_timeout 600s; + proxy_send_timeout 600s; + } + + location = /health { + proxy_pass http://127.0.0.1:8000/health; + } +} diff --git a/packages/config/base.py b/packages/config/base.py index 4f78b0516..87f463ef3 100755 --- a/packages/config/base.py +++ b/packages/config/base.py @@ -129,6 +129,67 @@ 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"), + description="P4000 回传结果用的外部 URL(worker 通过该 URL 提供给 P4000 PUT),如 http://100.69.73.60:8092", + ) + # Worker→API 内网直连 URL(Docker DNS),用于 worker 自己下载/清理 relay 文件。 + # 未配置时回退到 relay_base_url(本地开发/单节点)。 + gpu_encode_relay_internal_base_url: str = Field( + default="", + validation_alias=AliasChoices("GPU_ENCODE_RELAY_INTERNAL_BASE_URL", "gpu_encode_relay_internal_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..c93bca60e --- /dev/null +++ b/packages/shared/gpu_encoder.py @@ -0,0 +1,400 @@ +"""P4000 NVENC 远程编码客户端。 + +完整链路(encode_video_file): + 1. CPU 滤镜已在本地生成 mezzanine 中间片(libx264 ultrafast) + 2. 上传 mezzanine 到 OSS 临时前缀,拿到签名 GET URL + 3. 生成 relay 一次性 key,构造两个带 token 的 URL: + - put_url:给 P4000 回传结果,走 relay_base_url(外部可达,通常是 host:port 经 nginx) + - get/del_url:worker 自己下载+清理用,走 relay_internal_base_url(Docker DNS 直连 API) + 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 到 put_url,API 服务落盘到 /app/generated/gpu_relay/ + 6. 本客户端通过 get_url(Docker 内网)下载最终文件到 output_path,然后 DELETE 清理 + 7. 删除 OSS 临时 mezzanine + +任何环节失败抛 GpuEncodeError,调用方应 fallback 到 CPU libx264。 +""" + +from __future__ import annotations + +import json +import logging +import os +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, + *, + relay_internal_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("/") + # Worker→API 内网访问地址(Docker DNS 直连,如 http://xiaoxia-api-staging:8000)。 + # 未配置时回退到 relay_base_url(本地开发/单节点)。 + self.relay_internal_base_url = ( + relay_internal_base_url.rstrip("/") if relay_internal_base_url else self.relay_base_url + ) + 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/" + + RELAY_PATH_PREFIX = "/api/v1/internal/gpu-relay" + + # ------------------------------------------------------------------ + # URL builders + # ------------------------------------------------------------------ + def _relay_url_from_base(self, base_url: str, key: str, secret: str) -> str: + return f"{base_url}{self.RELAY_PATH_PREFIX}/{key}?token={urllib.parse.quote(secret, safe='')}" + + def _relay_put_url(self, key: str, secret: str) -> str: + """给 P4000 回传结果用的 URL(外部可达)。""" + return self._relay_url_from_base(self.relay_base_url, key, secret) + + def _relay_internal_url(self, key: str, secret: str) -> str: + """Worker 自己 GET/DELETE 用的 URL(Docker 内网)。""" + return self._relay_url_from_base(self.relay_internal_base_url, key, secret) + + # ------------------------------------------------------------------ + # 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 URLs (PUT 走外部 URL 给 P4000;GET/DELETE 走内部 Docker 网络) + relay_key = uuid.uuid4().hex + secret = self._get_relay_secret() + put_url = self._relay_put_url(relay_key, secret) + get_url = self._relay_internal_url(relay_key, secret) + del_url = get_url # 内部 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 _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() + relay_internal = (getattr(settings, "gpu_encode_relay_internal_base_url", "") or "").strip() + if not endpoint or not relay: + return None + return GpuEncoderClient( + endpoint=endpoint, + relay_base_url=relay, + relay_internal_base_url=relay_internal, + 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..17d9f4451 --- /dev/null +++ b/tests/unit/test_gpu_encoder.py @@ -0,0 +1,590 @@ +"""GpuEncoderClient 单元测试:mock HTTP,覆盖 health/sync/fallback/singleton 等完整路径。""" + +from __future__ import annotations + +import json +import socket +import sys +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, + _build_client_from_settings, + get_gpu_encoder, + is_gpu_encode_enabled, + 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", + relay_internal_base_url="http://api-internal:8000", + 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 + bio = BytesIO(data) + resp = mock.MagicMock(spec=HTTPResponse) + resp.status = status + resp.read.side_effect = lambda n=-1: bio.read(n) + resp.__enter__ = mock.MagicMock(return_value=resp) + resp.__exit__ = mock.MagicMock(return_value=False) + return resp + + +# ── Health check ──────────────────────────────────────────────────── +class TestHealthCheck: + def test_healthy_nvenc_available(self, client): + body = { + "status": "healthy", + "worker": "w1", + "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 and h.nvenc_h264 and 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": "t"}, "nvenc": {"h264_nvenc": False}} + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)): + h = client.check_health() + assert h.healthy and not h.ready + + def test_malformed_response_inner_exception(self, client): + """data 是合法 JSON 但 gpu 字段类型错(字符串)触发内部 except.""" + body = {"status": "healthy", "gpu": "not-a-dict", "nvenc": {}} + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)): + h = client.check_health() + assert not h.healthy + assert "malformed" in h.error + + +# ── _post_sync ────────────────────────────────────────────────────── +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" and res["ffmpeg_rc"] == 0 + 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"} + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)): + with pytest.raises(GpuEncodeError, match="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") + ) + + def test_connection_error_raises(self, client): + with mock.patch("urllib.request.urlopen", side_effect=urllib.error.URLError("conn refused")): + with pytest.raises(GpuEncodeError, match="connection error"): + client._post_sync( + {"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10}, mezzanine_path=Path("/tmp/x") + ) + + def test_timeout_error_raises(self, client): + with mock.patch("urllib.request.urlopen", side_effect=socket.timeout("timed out")): + with pytest.raises(GpuEncodeError, match="connection error"): + client._post_sync( + {"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10}, mezzanine_path=Path("/tmp/x") + ) + + def test_bad_json_raises(self, client): + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=b"not-json")): + with pytest.raises(GpuEncodeError, match="bad JSON"): + client._post_sync( + {"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10}, mezzanine_path=Path("/tmp/x") + ) + + def test_uploaded_false_logs_warning_but_succeeds(self, client, caplog): + body = {"status": "completed", "ffmpeg_rc": 0, "uploaded": False, "job_id": "j"} + with mock.patch("urllib.request.urlopen", return_value=_fake_response(body=body)), caplog.at_level("WARNING"): + res = client._post_sync( + {"inputs": {}, "ffmpeg_args": [], "output_url": "", "timeout": 10}, mezzanine_path=Path("/tmp/x") + ) + assert res["status"] == "completed" + assert "uploaded=false" in caplog.text + + +# ── Relay URL builders ────────────────────────────────────────────── +class TestRelayUrl: + def test_put_url_uses_external_base(self, client): + url = client._relay_put_url("abc123", "secret!") + assert "abc123" in url + assert "token=secret%21" in url + assert url.startswith("http://api.example.com/api/v1/internal/gpu-relay/") + + def test_internal_url_uses_internal_base(self, client): + url = client._relay_internal_url("abc123", "s") + assert url.startswith("http://api-internal:8000/api/v1/internal/gpu-relay/abc123") + + def test_internal_url_falls_back_to_external_when_not_set(self): + c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api.example.com", relay_secret="s") + put = c._relay_put_url("k", "s") + internal = c._relay_internal_url("k", "s") + assert put.startswith("http://api.example.com/") + assert internal == put + + def test_encode_uses_different_put_and_get_urls(self, client): + put_url = client._relay_put_url("k", "test-secret") + get_url = client._relay_internal_url("k", "test-secret") + assert "api.example.com" in put_url and "api-internal:8000" in get_url and put_url != get_url + + +# ── _get_relay_secret ────────────────────────────────────────────── +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() + + def test_dev_without_secret_raises(self, monkeypatch): + """未设置 secret 且非 production 也 raise(worker 必须显式配置)。""" + monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False) + monkeypatch.setenv("APP_ENV", "development") + c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api") + with pytest.raises(GpuEncodeError, match="GPU_ENCODE_RELAY_SECRET not set"): + c._get_relay_secret() + + +# ── _download_to_file ────────────────────────────────────────────── +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) and 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() + + def test_non_200_status_raises(self, client, tmp_path): + out = tmp_path / "o.mp4" + with mock.patch("urllib.request.urlopen", return_value=_fake_response(status=404, body=b"")): + with pytest.raises(GpuEncodeError, match="HTTP 404"): + client._download_to_file("http://relay/k", out) + + def test_url_error_cleans_up_tmp(self, client, tmp_path): + out = tmp_path / "o.mp4" + tmp_file = out.with_suffix(out.suffix + ".gpu_tmp") + tmp_file.write_bytes(b"partial") + assert tmp_file.exists() + with mock.patch("urllib.request.urlopen", side_effect=urllib.error.URLError("net down")): + with pytest.raises(GpuEncodeError, match="failed to download"): + client._download_to_file("http://relay/k", out) + assert not tmp_file.exists() + + +# ── encode_mezzanine_to_output ───────────────────────────────────── +class TestEncodeMezzanine: + def test_happy_path_with_audio(self, client, tmp_path): + mezz = tmp_path / "mezz.mp4" + mezz.write_bytes(b"M" * 100) + out = tmp_path / "out" / "final.mp4" + with ( + mock.patch.object(client, "_upload_mezzanine", return_value=("http://oss/signed", "osskey1")), + mock.patch.object( + client, + "_post_sync", + return_value={ + "job_id": "j1", + "status": "completed", + "ffmpeg_rc": 0, + "uploaded": True, + "size": 5000, + "duration": 1.2, + }, + ) as m_post, + mock.patch.object(client, "_download_to_file", return_value=5000) as m_dl, + mock.patch.object(client, "_relay_delete") as m_del, + mock.patch.object(client, "_delete_oss") as m_ossdel, + ): + result = client.encode_mezzanine_to_output(mezz, out, audio_args=["-c:a", "aac"]) + assert result["output_size"] == 5000 and str(out) == result["output_path"] + body = m_post.call_args[0][0] + assert "-c:a" in body["ffmpeg_args"] and "aac" in body["ffmpeg_args"] + assert "-an" not in body["ffmpeg_args"] + assert body["output_url"].startswith("http://api.example.com/") + assert "api-internal:8000" in m_dl.call_args[0][0] + m_del.assert_called_once() + m_ossdel.assert_called_once_with("osskey1") + + def test_happy_path_no_audio_uses_an_and_cq(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"M") + out = tmp_path / "o.mp4" + with ( + mock.patch.object(client, "_upload_mezzanine", return_value=("http://oss/u", "k")), + mock.patch.object( + client, + "_post_sync", + return_value={"status": "completed", "ffmpeg_rc": 0, "uploaded": True, "job_id": "j"}, + ) as m_post, + mock.patch.object(client, "_download_to_file", return_value=100), + mock.patch.object(client, "_relay_delete"), + mock.patch.object(client, "_delete_oss"), + ): + client.encode_mezzanine_to_output(mezz, out) + body = m_post.call_args[0][0] + assert "-an" in body["ffmpeg_args"] and "-cq" in body["ffmpeg_args"] + assert str(client.crf) in body["ffmpeg_args"] + + def test_bitrate_set_uses_bv_instead_of_cq(self, tmp_path): + c = GpuEncoderClient( + endpoint="http://gpu", + relay_base_url="http://api", + relay_internal_base_url="http://api-int:8000", + relay_secret="s", + bitrate="2M", + ) + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + out = tmp_path / "o.mp4" + with ( + mock.patch.object(c, "_upload_mezzanine", return_value=("http://oss/u", "k")), + mock.patch.object( + c, "_post_sync", return_value={"status": "completed", "ffmpeg_rc": 0, "uploaded": True, "job_id": "j"} + ) as m_post, + mock.patch.object(c, "_download_to_file", return_value=10), + mock.patch.object(c, "_relay_delete"), + mock.patch.object(c, "_delete_oss"), + ): + c.encode_mezzanine_to_output(mezz, out, extra_video_args=["-vf", "scale=1280:-2"]) + body = m_post.call_args[0][0] + assert "-b:v" in body["ffmpeg_args"] and "2M" in body["ffmpeg_args"] + assert "-cq" not in body["ffmpeg_args"] + assert "-vf" in body["ffmpeg_args"] + + def test_mezzanine_not_found_raises(self, client, tmp_path): + with pytest.raises(GpuEncodeError, match="mezzanine file not found"): + client.encode_mezzanine_to_output(tmp_path / "nope.mp4", tmp_path / "o.mp4") + + def test_relay_base_not_configured_raises(self, tmp_path): + c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="", relay_secret="s") + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + with pytest.raises(GpuEncodeError, match="relay_base_url"): + c.encode_mezzanine_to_output(mezz, tmp_path / "o.mp4") + + def test_unexpected_exception_is_wrapped(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + with ( + mock.patch.object(client, "_upload_mezzanine", return_value=("http://oss/u", "k")), + mock.patch.object(client, "_post_sync", side_effect=RuntimeError("boom")), + mock.patch.object(client, "_delete_oss"), + ): + with pytest.raises(GpuEncodeError, match="unexpected: boom"): + client.encode_mezzanine_to_output(mezz, tmp_path / "o.mp4") + + def test_gpu_encode_error_re_raised_directly(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + with ( + mock.patch.object(client, "_upload_mezzanine", return_value=("http://oss/u", "k")), + mock.patch.object(client, "_post_sync", side_effect=GpuEncodeError("direct fail")), + mock.patch.object(client, "_delete_oss"), + ): + with pytest.raises(GpuEncodeError, match="direct fail"): + client.encode_mezzanine_to_output(mezz, tmp_path / "o.mp4") + + def test_oss_cleanup_runs_on_failure(self, client, tmp_path, caplog): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + with ( + mock.patch.object(client, "_upload_mezzanine", return_value=("http://oss/u", "ossk")), + mock.patch.object(client, "_post_sync", side_effect=GpuEncodeError("enc fail")), + mock.patch.object(client, "_delete_oss", side_effect=Exception("oss down")) as m_ossdel, + caplog.at_level("WARNING"), + ): + with pytest.raises(GpuEncodeError): + client.encode_mezzanine_to_output(mezz, tmp_path / "o.mp4") + m_ossdel.assert_called_once_with("ossk") + + +# ── _relay_delete ────────────────────────────────────────────────── +class TestRelayDelete: + def test_exception_is_swallowed(self, client, caplog): + with mock.patch("urllib.request.urlopen", side_effect=RuntimeError("boom")), caplog.at_level("DEBUG"): + client._relay_delete("http://relay/k?token=s") + assert "cleanup delete failed" in caplog.text + + def test_success_issues_delete(self, client): + with mock.patch("urllib.request.urlopen", return_value=_fake_response(status=204, body=b"")) as m: + client._relay_delete("http://relay/k?token=s") + assert m.call_args[0][0].get_method() == "DELETE" + + +# ── OSS helpers ──────────────────────────────────────────────────── +class TestOssHelpers: + def test_upload_storage_import_error(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + # 删除 sys.modules 中 packages.shared.storage 使导入失败 + saved = sys.modules.pop("packages.shared.storage", None) + try: + real_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__ + + def fake_import(name, *a, **kw): + if name == "packages.shared.storage" or name.startswith("packages.shared.storage."): + raise ImportError("no storage") + return real_import(name, *a, **kw) + + with mock.patch("builtins.__import__", side_effect=fake_import): + with pytest.raises(GpuEncodeError, match="storage service unavailable"): + client._upload_mezzanine(mezz) + finally: + if saved is not None: + sys.modules["packages.shared.storage"] = saved + + def test_upload_storage_none(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + fake_mod = mock.MagicMock() + fake_mod.get_storage_service.return_value = None + with mock.patch.dict("sys.modules", {"packages.shared.storage": fake_mod}): + with pytest.raises(GpuEncodeError, match="OSS storage not configured"): + client._upload_mezzanine(mezz) + + def test_upload_bucket_none(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + svc = mock.MagicMock() + svc.bucket = None + fake_mod = mock.MagicMock() + fake_mod.get_storage_service.return_value = svc + with mock.patch.dict("sys.modules", {"packages.shared.storage": fake_mod}): + with pytest.raises(GpuEncodeError, match="OSS storage not configured"): + client._upload_mezzanine(mezz) + + def test_upload_failure_raises(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + svc = mock.MagicMock() + svc.bucket = object() + svc.upload_file.side_effect = RuntimeError("oss err") + fake_mod = mock.MagicMock() + fake_mod.get_storage_service.return_value = svc + with mock.patch.dict("sys.modules", {"packages.shared.storage": fake_mod}): + with pytest.raises(GpuEncodeError, match="failed to upload mezzanine"): + client._upload_mezzanine(mezz) + + def test_upload_success(self, client, tmp_path): + mezz = tmp_path / "m.mp4" + mezz.write_bytes(b"x") + svc = mock.MagicMock() + svc.bucket = object() + svc.get_download_url.return_value = "https://oss/signed?sig=abc" + fake_mod = mock.MagicMock() + fake_mod.get_storage_service.return_value = svc + with mock.patch.dict("sys.modules", {"packages.shared.storage": fake_mod}): + url, key = client._upload_mezzanine(mezz) + assert url.startswith("https://oss/signed") + assert key.startswith("tmp/gpu-mezzanine/") and key.endswith(".mp4") + svc.upload_file.assert_called_once() + + def test_delete_oss_exception_swallowed(self, client, caplog): + fake_mod = mock.MagicMock() + fake_mod.get_storage_service.side_effect = RuntimeError("svc down") + with mock.patch.dict("sys.modules", {"packages.shared.storage": fake_mod}), caplog.at_level("DEBUG"): + client._delete_oss("somekey") + assert "OSS delete" in caplog.text + + def test_delete_oss_bucket_none_noop(self, client): + svc = mock.MagicMock() + svc.bucket = None + fake_mod = mock.MagicMock() + fake_mod.get_storage_service.return_value = svc + with mock.patch.dict("sys.modules", {"packages.shared.storage": fake_mod}): + client._delete_oss("k") + svc.delete_file.assert_not_called() + + def test_delete_oss_success(self, client): + svc = mock.MagicMock() + svc.bucket = object() + fake_mod = mock.MagicMock() + fake_mod.get_storage_service.return_value = svc + with mock.patch.dict("sys.modules", {"packages.shared.storage": fake_mod}): + client._delete_oss("k") + svc.delete_file.assert_called_once_with("k") + + +# ── Singleton / factory ──────────────────────────────────────────── +class TestSingletonFactory: + def test_build_client_import_error_returns_none(self): + saved = sys.modules.get("packages.config") + sys.modules["packages.config"] = None + try: + with mock.patch("builtins.__import__", side_effect=RuntimeError("no cfg")): + assert _build_client_from_settings() is None + finally: + if saved is not None: + sys.modules["packages.config"] = saved + + def test_build_client_not_enabled_returns_none(self): + s = mock.MagicMock() + s.enable_gpu_encode = False + fake_mod = mock.MagicMock() + fake_mod.get_shared_settings.return_value = s + with mock.patch.dict("sys.modules", {"packages.config": fake_mod}): + assert _build_client_from_settings() is None + + def test_build_client_missing_endpoint(self): + s = mock.MagicMock() + s.enable_gpu_encode = True + s.gpu_encode_endpoint = "" + s.gpu_encode_relay_base_url = "http://api" + s.gpu_encode_relay_internal_base_url = "" + fake_mod = mock.MagicMock() + fake_mod.get_shared_settings.return_value = s + with mock.patch.dict("sys.modules", {"packages.config": fake_mod}): + assert _build_client_from_settings() is None + + def test_build_client_missing_relay(self): + s = mock.MagicMock() + s.enable_gpu_encode = True + s.gpu_encode_endpoint = "http://gpu" + s.gpu_encode_relay_base_url = "" + s.gpu_encode_relay_internal_base_url = "" + fake_mod = mock.MagicMock() + fake_mod.get_shared_settings.return_value = s + with mock.patch.dict("sys.modules", {"packages.config": fake_mod}): + assert _build_client_from_settings() is None + + def test_build_client_success(self): + s = mock.MagicMock() + s.enable_gpu_encode = True + s.gpu_encode_endpoint = "http://gpu" + s.gpu_encode_relay_base_url = "http://api/" + s.gpu_encode_relay_internal_base_url = "http://api-int:8000/" + s.gpu_encode_sync_timeout = 120 + s.gpu_encode_health_timeout = 1.0 + s.gpu_encode_vcodec = "h264_nvenc" + s.gpu_encode_preset = "p7" + s.gpu_encode_crf = 20 + s.gpu_encode_bitrate = "" + s.gpu_encode_relay_secret = "s" + s.gpu_encode_oss_tmp_prefix = "tmp/x/" + fake_mod = mock.MagicMock() + fake_mod.get_shared_settings.return_value = s + with mock.patch.dict("sys.modules", {"packages.config": fake_mod}): + c = _build_client_from_settings() + assert c is not None and c.endpoint == "http://gpu" + assert c.relay_base_url == "http://api" + assert c.relay_internal_base_url == "http://api-int:8000" + assert c.preset == "p7" and c.crf == 20 + + def test_get_gpu_encoder_init_failure_returns_none(self, caplog): + with ( + mock.patch("packages.shared.gpu_encoder._build_client_from_settings", side_effect=RuntimeError("boom")), + caplog.at_level("WARNING"), + ): + assert get_gpu_encoder() is None + assert "failed to init client" in caplog.text + + def test_get_gpu_encoder_returns_singleton_and_enabled(self): + c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api", relay_secret="s") + with mock.patch("packages.shared.gpu_encoder._build_client_from_settings", return_value=c): + assert get_gpu_encoder() is c and get_gpu_encoder() is c + assert is_gpu_encode_enabled() is True + + def test_is_gpu_encode_enabled_when_none(self): + with mock.patch("packages.shared.gpu_encoder._build_client_from_settings", return_value=None): + assert is_gpu_encode_enabled() is False + + +# ── Constructor edge cases ──────────────────────────────────────── +class TestConstructor: + def test_oss_tmp_prefix_empty_uses_default(self): + c = GpuEncoderClient(endpoint="http://gpu", relay_base_url="http://api", relay_secret="s", oss_tmp_prefix="") + assert c.oss_tmp_prefix == "tmp/gpu-mezzanine/" + + def test_oss_tmp_prefix_strips_and_adds_slash(self): + c = GpuEncoderClient( + endpoint="http://gpu", relay_base_url="http://api", relay_secret="s", oss_tmp_prefix="tmp/foo" + ) + assert c.oss_tmp_prefix == "tmp/foo/" diff --git a/tests/unit/test_gpu_relay.py b/tests/unit/test_gpu_relay.py new file mode 100644 index 000000000..b5077b73a --- /dev/null +++ b/tests/unit/test_gpu_relay.py @@ -0,0 +1,240 @@ +"""gpu_relay API 路由单元测试:覆盖 helper 函数 + PUT/GET/HEAD/DELETE handler。""" + +from __future__ import annotations + +import os +from pathlib import Path +from unittest import mock + +import pytest +from fastapi import HTTPException + +from apps.api.app.api.routes import gpu_relay + + +# ── _relay_dir ──────────────────────────────────────────────────────── +class TestRelayDir: + def test_default_dir(self, tmp_path, monkeypatch): + monkeypatch.delenv("GENERATED_FILES_DIR", raising=False) + monkeypatch.delenv("GPU_ENCODE_RELAY_DIR", raising=False) + # 用 tmp_path 作 base + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + p = gpu_relay._relay_dir() + assert p == tmp_path / "gpu_relay" + assert p.exists() + + def test_custom_subdir(self, tmp_path, monkeypatch): + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + monkeypatch.setenv("GPU_ENCODE_RELAY_DIR", "custom_relay") + p = gpu_relay._relay_dir() + assert p == tmp_path / "custom_relay" + assert p.exists() + + +# ── _secret ────────────────────────────────────────────────────────── +class TestSecret: + def test_explicit_secret_returned(self, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "topsecret") + gpu_relay._DEFAULT_SECRET_LOGGED = False + assert gpu_relay._secret() == "topsecret" + + def test_prod_without_secret_raises(self, monkeypatch): + monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False) + monkeypatch.setenv("APP_ENV", "production") + with pytest.raises(RuntimeError, match="GPU_ENCODE_RELAY_SECRET must be set"): + gpu_relay._secret() + + def test_dev_without_secret_generates_ephemeral(self, monkeypatch, caplog): + monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False) + monkeypatch.setenv("APP_ENV", "development") + gpu_relay._DEFAULT_SECRET_LOGGED = False + with caplog.at_level("WARNING"): + secret = gpu_relay._secret() + assert len(secret) > 16 + assert "ephemeral dev token" in caplog.text + # 第二次调用不再 log(_DEFAULT_SECRET_LOGGED=True) + before = len(caplog.records) + secret2 = gpu_relay._secret() + assert secret2 == secret + assert len(caplog.records) == before + # 清理 + monkeypatch.delenv("GPU_ENCODE_RELAY_SECRET", raising=False) + + +# ── _safe_key ──────────────────────────────────────────────────────── +class TestSafeKey: + @pytest.mark.parametrize("bad", ["", "../etc", "a/b", "a\\b", ".", "..", "a b", "a%b"]) + def test_invalid_keys_rejected(self, bad): + with pytest.raises(HTTPException) as ei: + gpu_relay._safe_key(bad) + assert ei.value.status_code == 400 + + @pytest.mark.parametrize("good", ["abc123", "ABC-Def_01", "a" * 32]) + def test_valid_keys_accepted(self, good): + assert gpu_relay._safe_key(good) == good + + def test_strips_whitespace(self): + assert gpu_relay._safe_key(" abc ") == "abc" + + +# ── _check_token ───────────────────────────────────────────────────── +class TestCheckToken: + def test_missing_token_401(self): + with pytest.raises(HTTPException) as ei: + gpu_relay._check_token(None) + assert ei.value.status_code == 401 + + def test_wrong_token_401(self, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "correct") + with pytest.raises(HTTPException) as ei: + gpu_relay._check_token("wrong") + assert ei.value.status_code == 401 + + def test_correct_token_passes(self, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "correct") + assert gpu_relay._check_token("correct") is None + + +# ── build_relay_* helpers ─────────────────────────────────────────── +class TestBuildRelayUrls: + def test_put_url(self): + url = gpu_relay.build_relay_put_url("http://api.example.com/", "k1", "s") + assert url == "http://api.example.com/api/v1/internal/gpu-relay/k1?token=s" + + def test_get_url_same_as_put(self): + assert gpu_relay.build_relay_get_url("http://api", "k", "s") == gpu_relay.build_relay_put_url( + "http://api", "k", "s" + ) + + def test_generate_key_is_hex(self): + k = gpu_relay.generate_key() + assert len(k) == 32 + int(k, 16) # valid hex + + +# ── PUT endpoint ──────────────────────────────────────────────────── +@pytest.mark.asyncio +class TestPutObject: + async def test_put_writes_file(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + + # async request.stream 模拟 + async def _stream(): + yield b"chunk1" + yield b"chunk2" + + req = mock.MagicMock() + req.stream = _stream + resp = await gpu_relay.put_object(key="abc123", request=req, token="s") + assert resp["ok"] is True + assert resp["size"] == len(b"chunk1") + len(b"chunk2") + p = tmp_path / "gpu_relay" / "abc123" + assert p.read_bytes() == b"chunk1chunk2" + # .part 临时文件应已 rename + assert not p.with_suffix(p.suffix + ".part").exists() + + async def test_put_invalid_key_400(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + req = mock.MagicMock() + with pytest.raises(HTTPException) as ei: + await gpu_relay.put_object(key="../bad", request=req, token="s") + assert ei.value.status_code == 400 + + async def test_put_bad_token_401(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "correct") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + req = mock.MagicMock() + with pytest.raises(HTTPException) as ei: + await gpu_relay.put_object(key="abc", request=req, token="wrong") + assert ei.value.status_code == 401 + + async def test_put_write_error_cleans_tmp(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + + async def _bad_stream(): + yield b"x" + raise OSError("disk full") + + req = mock.MagicMock() + req.stream = _bad_stream + with pytest.raises(HTTPException) as ei: + await gpu_relay.put_object(key="abc", request=req, token="s") + assert ei.value.status_code == 500 + # tmp 文件被清理 + part = tmp_path / "gpu_relay" / "abc.part" + assert not part.exists() + + +# ── GET endpoint ──────────────────────────────────────────────────── +@pytest.mark.asyncio +class TestGetObject: + async def test_get_missing_404(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + with pytest.raises(HTTPException) as ei: + await gpu_relay.get_object(key="nope", token="s") + assert ei.value.status_code == 404 + + async def test_get_returns_file(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + p = tmp_path / "gpu_relay" / "exist" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(b"viddata") + resp = await gpu_relay.get_object(key="exist", token="s") + assert resp.media_type == "video/mp4" + + +# ── HEAD endpoint ─────────────────────────────────────────────────── +@pytest.mark.asyncio +class TestHeadObject: + async def test_head_missing_404(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + resp = await gpu_relay.head_object(key="nope", token="s") + assert resp.status_code == 404 + + async def test_head_returns_content_length(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + p = tmp_path / "gpu_relay" / "k" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(b"12345") + resp = await gpu_relay.head_object(key="k", token="s") + assert resp.status_code == 200 + assert resp.headers["Content-Length"] == "5" + + +# ── DELETE endpoint ────────────────────────────────────────────────── +@pytest.mark.asyncio +class TestDeleteObject: + async def test_delete_existing(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + p = tmp_path / "gpu_relay" / "k" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(b"x") + resp = await gpu_relay.delete_object(key="k", token="s") + assert resp["ok"] is True + assert not p.exists() + + async def test_delete_missing_is_noop(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + # 不存在时不应 404,返回 ok + resp = await gpu_relay.delete_object(key="nope", token="s") + assert resp["ok"] is True + + async def test_delete_unlink_error_500(self, tmp_path, monkeypatch): + monkeypatch.setenv("GPU_ENCODE_RELAY_SECRET", "s") + monkeypatch.setenv("GENERATED_FILES_DIR", str(tmp_path)) + p = tmp_path / "gpu_relay" / "k" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes(b"x") + with mock.patch.object(Path, "unlink", side_effect=OSError("perm denied")): + with pytest.raises(HTTPException) as ei: + await gpu_relay.delete_object(key="k", token="s") + assert ei.value.status_code == 500