feat: 新增 GPU 资源门控模块(探测/需求估算/准入租约)
为"按资源而不是按配置调度"提供判定基础:`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` 调整。 - 测试覆盖互斥、余量边界、无探测退化、真实权重估算与端点判定。
This commit is contained in:
@@ -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
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user