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
|
||||
Reference in New Issue
Block a user