From df7a91d97c5bdaf02ed231b5e0c03d58e41d912a Mon Sep 17 00:00:00 2001 From: cat <1716967236@qq.com> Date: Fri, 18 Sep 2026 22:40:37 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=20GPU=20=E8=B5=84?= =?UTF-8?q?=E6=BA=90=E9=97=A8=E6=8E=A7=E6=A8=A1=E5=9D=97=EF=BC=88=E6=8E=A2?= =?UTF-8?q?=E6=B5=8B/=E9=9C=80=E6=B1=82=E4=BC=B0=E7=AE=97/=E5=87=86?= =?UTF-8?q?=E5=85=A5=E7=A7=9F=E7=BA=A6=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为"按资源而不是按配置调度"提供判定基础：`stage_gpu_need_mb` 给出节点的显存 需求（whisper 按权重 ×1.45；本机 LLM/VLM 端点按预留；线上端点为 0）， `GpuGate` 负责准入（GPU 阶段互斥 + 显存余量检查；探测不到 nvidia-smi 时退化为 互斥，与串行等价）。 - 判定只看资源现状，调用方不需要判断"是否云端"。 - 局域网地址上的本机 Ollama 用 `WOV_LOCAL_MODEL_HOSTS` 声明；预留值可用 `WOV_LOCAL_LLM_RESERVE_MB` / `WOV_LOCAL_VLM_RESERVE_MB` 调整。 - 测试覆盖互斥、余量边界、无探测退化、真实权重估算与端点判定。 --- src/wov_app/resources.py | 217 ++++++++++++++++++++++++++ tests/app/test_resources/__init__.py | 0 tests/app/test_resources/test_gate.py | 150 ++++++++++++++++++ 3 files changed, 367 insertions(+) create mode 100644 src/wov_app/resources.py create mode 100644 tests/app/test_resources/__init__.py create mode 100644 tests/app/test_resources/test_gate.py diff --git a/src/wov_app/resources.py b/src/wov_app/resources.py new file mode 100644 index 0000000..1057ca6 --- /dev/null +++ b/src/wov_app/resources.py @@ -0,0 +1,217 @@ +"""GPU 资源门控：按"当前可用显存 + 在途申请"决定阶段能否启动。 + +批量流水线里各阶段对 GPU 的需求不同：`faster-whisper` 必须独占显存；走本机 +Ollama 端点的 `llm-*`/OCR 阶段（模型常驻显存）同样必须独占；而**线上端点**的 +`llm-translate`、提音与 ASS 合成都不占显存。用资源现状判定，就不需要在代码里 +按"是否走云端"分叉： + +- 显存不够或已有 GPU 阶段在跑 → 后续 GPU 阶段排队等待（与串行等价，不抢显存）； +- 不需要 GPU → 立刻放行，于是转写能与线上翻译并行，GPU 不再空转。 + +探测不到 `nvidia-smi`（纯 CPU 机器、无 GPU 容器）时无法核对显存，退化为 +"GPU 阶段互斥"这一保守规则：仍不并发，但与串行结果一致。 + +本模块只做准入判定与需求估算，不加载/卸载任何模型：显存的实际释放由节点与 +引擎负责（本地模型常驻的让位见 docs/operations.md#文件夹批量处理）。 +""" + +from __future__ import annotations + +import os +import shutil +import socket +import subprocess +import threading +import urllib.parse +from pathlib import Path + +# 安全余量：显存刚好"够"时仍可能因分配碎片化抛 CUDA OOM，留 800MB 缓冲。 +DEFAULT_SAFETY_MB = 800 + +# 权重之外的推理开销系数：CTranslate2 除权重外还需推理工作区、激活与 CUDA 上下文。 +# 实测（large-v2 权重 2.87GB）峰值增量约 4097MB ≈ 权重 × 1.39，取 1.45 覆盖余量。 +WEIGHTS_MEMORY_FACTOR = 1.45 + +# 本机 LLM/VLM 端点（Ollama）的显存预留：模型常驻期间 whisper 必须让位。 +# 默认按整卡预留——本机大模型加载后实测占用 20–21GB，宁可保守也不要 OOM。 +DEFAULT_LOCAL_LLM_RESERVE_MB = 22528 +DEFAULT_LOCAL_VLM_RESERVE_MB = 8192 + +# 判定"端点在本机"的主机名集合：默认只看回环；局域网地址的本机 Ollama 需要 +# 通过 WOV_LOCAL_MODEL_HOSTS 显式声明（见 docs/configuration.md）。 +_LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1", "0.0.0.0"} + + +def _env_int(name: str, default: int) -> int: + """读取整数环境变量；空值或非法值回退默认值。""" + raw = os.getenv(name) + if raw is None or not str(raw).strip(): + return default + try: + return int(float(raw)) + except ValueError: + return default + + +def local_model_hosts() -> set[str]: + """返回被视为"本机模型"的主机名集合（回环 + WOV_LOCAL_MODEL_HOSTS）。""" + hosts = set(_LOOPBACK_HOSTS) + for item in (os.getenv("WOV_LOCAL_MODEL_HOSTS") or "").split(","): + item = item.strip().lower() + if item: + hosts.add(item) + return hosts + + +def _endpoint_host(value: str) -> str: + """从 URL 或 host:port 字符串里取出主机名（小写）。""" + text = (value or "").strip() + if not text: + return "" + if "://" not in text: + text = "http://" + text + host = urllib.parse.urlsplit(text).hostname or "" + return host.lower() + + +def is_local_endpoint(value: str) -> bool: + """判断 LLM/VLM 端点是否指向本机（模型占本机显存）。 + + 回环地址一律算本机；局域网地址（如另一张网卡上的 Ollama）需要显式列入 + WOV_LOCAL_MODEL_HOSTS——本机多网卡时无法可靠推断，宁可要求配置。 + """ + host = _endpoint_host(value) + if not host: + return False + if host in local_model_hosts(): + return True + # 主机名解析成本低，顺带覆盖 localhost 之外的别名（如容器内的 host.docker.internal）。 + try: + if socket.gethostbyname(host) in {ip for name in _LOOPBACK_HOSTS + for ip in [socket.gethostbyname(name)] if ip}: + return True + except OSError: + return False + return False + + +def probe_free_memory_mb() -> int | None: + """返回当前最大可用显存（MB）；无 NVIDIA GPU 或探测失败时返回 None。 + + 多卡时取可用显存最大的那张卡（推理默认只用一张卡）。 + """ + if shutil.which("nvidia-smi") is None: + return None + try: + completed = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=memory.free", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=10, + ) + except (OSError, subprocess.SubprocessError): + return None + if completed.returncode != 0: + return None + values = [int(line.strip()) for line in completed.stdout.splitlines() if line.strip().isdigit()] + return max(values) if values else None + + +def whisper_need_mb(model_dir: Path) -> int: + """按 CTranslate2 权重体量估算转写所需显存（MB）。""" + weights = Path(model_dir) / "model.bin" + size_mb = weights.stat().st_size / (1024 * 1024) if weights.is_file() else 0.0 + return int(size_mb * WEIGHTS_MEMORY_FACTOR) + + +def _llm_need_mb() -> int: + """LLM 阶段的显存需求：本机端点按预留整卡，线上端点为 0。""" + base = os.getenv("LLM_API_BASE", "") + if base and is_local_endpoint(base): + return _env_int("WOV_LOCAL_LLM_RESERVE_MB", DEFAULT_LOCAL_LLM_RESERVE_MB) + return 0 + + +def _vlm_need_mb() -> int: + """帧 OCR 阶段的显存需求：本机 Ollama 按预留，远端端点为 0。""" + base = os.getenv("OLLAMA_HOST") or os.getenv("VLM_API_BASE") or "" + if base and is_local_endpoint(base): + return _env_int("WOV_LOCAL_VLM_RESERVE_MB", DEFAULT_LOCAL_VLM_RESERVE_MB) + return 0 + + +def stage_gpu_need_mb(node_type: str, params: dict) -> int: + """估算某个节点在当前配置下的显存需求（MB）；0 表示不需要 GPU。 + + 只认节点类型本身的能力，不看"这是批量还是单任务"——主调度器与批量引擎 + 可以共用同一套判定。 + """ + if node_type == "faster-whisper": + # 延迟导入：本地模型解析会读环境变量与盘上权重，不在模块导入期执行。 + from nodes.whisper import resolve_model_path + + return whisper_need_mb(Path(resolve_model_path(params))) + if node_type.startswith("llm"): + return _llm_need_mb() + if node_type in ("vlm-ocr", "subtitle-ocr"): + return _vlm_need_mb() + # 提音/抽帧/ASS/echo 等走 CPU 或本地磁盘，不占显存。 + return 0 + + +class GpuGate: + """进程内 GPU 准入：GPU 阶段互斥 + 显存余量检查，非 GPU 阶段不受限。 + + 互斥（而不是按显存叠加并发）是刻意的：GPU 算力才是瓶颈，两个转写同时跑 + 只会互相抢算力而整体更慢；需要并发时再按实测调整（见 operations.md）。 + """ + + def __init__(self, probe=None, safety_mb: int = DEFAULT_SAFETY_MB) -> None: + self._probe = probe if probe is not None else probe_free_memory_mb + self._safety_mb = safety_mb + self._lock = threading.Lock() + self._holder: str | None = None + self._holder_need_mb = 0 + + def try_acquire(self, key: str, memory_mb: int) -> bool: + """尝试为 key 取得准入；False 表示当前资源不满足，调用方稍后重试。 + + 同一 key 重复申请视为成功（重入），避免调用方因重复申请把自己挡住。 + """ + if memory_mb <= 0: + # 不需要 GPU：不占车道，直接放行（提音/线上翻译/ASS 可并行）。 + return True + with self._lock: + if self._holder == key: + return True + if self._holder is not None: + return False + free = self._probe() + if free is not None and free < memory_mb + self._safety_mb: + return False + self._holder = key + self._holder_need_mb = memory_mb + return True + + def release(self, key: str) -> None: + """释放 key 持有的 GPU 车道（非持有者调用无副作用）。""" + with self._lock: + if self._holder == key: + self._holder = None + self._holder_need_mb = 0 + + @property + def holder(self) -> str | None: + """当前持有 GPU 车道的阶段标识（无人持有时为 None）。""" + with self._lock: + return self._holder + + @property + def holder_need_mb(self) -> int: + """当前持有者声明的显存需求（MB）。""" + with self._lock: + return self._holder_need_mb diff --git a/tests/app/test_resources/__init__.py b/tests/app/test_resources/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/app/test_resources/test_gate.py b/tests/app/test_resources/test_gate.py new file mode 100644 index 0000000..257b318 --- /dev/null +++ b/tests/app/test_resources/test_gate.py @@ -0,0 +1,150 @@ +"""src/wov_app/resources.py 的模块级测试（数据 → 测试过程 → 验证结果）。 + +被测模块：GPU 资源门控（显存探测 + 需求估算 + 准入租约）。批量流水线用它决定 +"某个阶段现在能不能起"——用资源现状而不是"是否走云端"来分支。 + +`nvidia-smi` 探测属 I/O 边界：用例注入确定的探测结果；需求估算用**真实权重文件** +验证（写入指定大小的文件，不是伪造结构）。 +""" + +from __future__ import annotations + +from pathlib import Path + +from wov_app.resources import ( + GpuGate, + is_local_endpoint, + stage_gpu_need_mb, + whisper_need_mb, +) + + +def _probe(free_mb: int | None): + """构造固定返回值的显存探测函数（None 表示探测不到 GPU）。""" + return lambda: free_mb + + +def test_stage_without_gpu_never_waits() -> None: + """数据：一个已持有 GPU 租约的阶段 + 一个不需要 GPU 的阶段。 + + 过程：不需要 GPU 的阶段申请准入。 + + 验证：直接放行——线上模型阶段与 ffmpeg/ASS 阶段不该被 GPU 占用挡住。 + """ + gate = GpuGate(probe=_probe(24000)) + assert gate.try_acquire("video-a:asr", 4200) is True + + assert gate.try_acquire("video-b:translate", 0) is True + + +def test_gpu_stages_are_exclusive() -> None: + """数据：显存充足（探测 24000MB），但已有 GPU 阶段在跑。 + + 过程：另一个 GPU 阶段申请准入。 + + 验证：拒绝——显存够也不并发跑两个 GPU 阶段（GPU 算力才是瓶颈，两个 whisper + 同时跑只会互相争抢而整体更慢）。 + """ + gate = GpuGate(probe=_probe(24000)) + assert gate.try_acquire("video-a:asr", 4200) is True + + assert gate.try_acquire("video-b:asr", 4200) is False + + +def test_gpu_stage_needs_free_memory_above_need_plus_margin() -> None: + """数据：显存探测结果分别低于/高于"需求 + 安全余量"。 + + 过程：用同一需求申请准入。 + + 验证：不足时拒绝，充足时放行——避免临界卡上因碎片化直接 CUDA OOM。 + """ + need, margin = 4200, 800 + tight = GpuGate(probe=_probe(need + margin - 1), safety_mb=margin) + roomy = GpuGate(probe=_probe(need + margin), safety_mb=margin) + + assert tight.try_acquire("video-a:asr", need) is False + assert roomy.try_acquire("video-a:asr", need) is True + + +def test_probe_missing_falls_back_to_exclusivity() -> None: + """数据：探测不到 GPU（无 nvidia-smi，例如纯 CPU 机器或容器）。 + + 过程：连续申请两个 GPU 阶段。 + + 验证：无法核对显存时退化为"GPU 阶段互斥"——仍然不会并发，行为与串行等价。 + """ + gate = GpuGate(probe=_probe(None)) + + assert gate.try_acquire("video-a:asr", 4200) is True + assert gate.try_acquire("video-b:asr", 4200) is False + gate.release("video-a:asr") + assert gate.try_acquire("video-b:asr", 4200) is True + + +def test_release_frees_the_gpu_lane() -> None: + """数据：一个已释放租约的阶段。 + + 过程：释放后再次申请。 + + 验证：GPU 车道重新可用（视频间不会互相永久阻塞）。 + """ + gate = GpuGate(probe=_probe(24000)) + gate.try_acquire("video-a:asr", 4200) + gate.release("video-a:asr") + + assert gate.try_acquire("video-b:asr", 4200) is True + + +def test_whisper_need_scales_with_real_weights(tmp_path: Path) -> None: + """数据：真实写入的 100MB 权重文件（CTranslate2 的 model.bin）。 + + 过程：按权重体量估算推理显存需求。 + + 验证：需求 = 权重 × 系数（推理工作区与 CUDA 上下文另占约 45%）。 + """ + model_dir = tmp_path / "faster-whisper-large-v2" + model_dir.mkdir() + (model_dir / "model.bin").write_bytes(b"\x00" * (100 * 1024 * 1024)) + + assert whisper_need_mb(model_dir) == int(100 * 1.45) + + +def test_stage_need_zero_for_non_gpu_nodes(tmp_path: Path) -> None: + """数据：ffmpeg 提音、ASS 合成、echo 三类不需要 GPU 的节点。 + + 过程：估算其显存需求。 + + 验证：都是 0，因此可以与 GPU 阶段并行（阶段内提音/合成不再阻塞转写）。 + """ + for node_type in ("ffmpeg-extract", "srt-to-dual-eye-ass", "echo"): + assert stage_gpu_need_mb(node_type, {}) == 0 + + +def test_stage_need_for_local_and_remote_llm(monkeypatch) -> None: + """数据：LLM 端点分别是本机（占显存常驻）与远端（不占显存）。 + + 过程：估算 llm-translate 阶段的显存需求。 + + 验证：本机端点按预留显存计入（挡住 whisper），远端端点为 0（放行转写）。 + """ + monkeypatch.setenv("WOV_LOCAL_MODEL_HOSTS", "") + monkeypatch.setenv("LLM_API_BASE", "http://localhost:11434/v1/chat/completions") + local_need = stage_gpu_need_mb("llm-translate", {}) + monkeypatch.setenv("LLM_API_BASE", "https://api.siliconflow.cn/v1/chat/completions") + remote_need = stage_gpu_need_mb("llm-translate", {}) + + assert local_need > 0 + assert remote_need == 0 + + +def test_local_endpoint_accepts_configured_lan_host(monkeypatch) -> None: + """数据：Ollama 跑在本机另一张网卡的局域网地址上（如 192.168.123.70）。 + + 过程：判断端点是否本机模型。 + + 验证：通过 WOV_LOCAL_MODEL_HOSTS 显式声明后视为本机，按其显存需求预留。 + """ + monkeypatch.setenv("WOV_LOCAL_MODEL_HOSTS", "192.168.123.70") + + assert is_local_endpoint("http://192.168.123.70:11434/v1/chat/completions") is True + assert is_local_endpoint("https://api.siliconflow.cn/v1/chat/completions") is False