cat-shark
df7a91d97c
为"按资源而不是按配置调度"提供判定基础：`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` 调整。 - 测试覆盖互斥、余量边界、无探测退化、真实权重估算与端点判定。
151 lines
5.5 KiB
Python
151 lines
5.5 KiB
Python
"""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
|