cat-shark
8a715a8064
按"测试规则"重写 tests/：一个模块一个目录、用例按数据→过程→验证三段书写、 不保留全局 conftest.py、测试过程只调用真实生产代码。 结构（73 个文件、30 个模块目录、477 用例）： - tests/nodes/ 15 个模块目录（srt/whisper/ass/ffmpeg/frame_extract/vlm/ subtitle_ocr/llm/llm_filter/subtitle_cleanup/subtitle_correction/ proper_nouns/adaptive_pool/vad_profiler/echo）； - tests/app/ 11 个模块目录（db/scheduler/batch/maintenance/registry/seed/ storage/config/logging/main/routers 三组 API）； - tests/sdk/test_models、tests/web/test_crop、tests/shared（公共设施）。 测试数据随模块目录入库（tests/**/data/），删除根级 testdata/；.gitignore 的 data/ 改为 /data/，否则会连带忽略 tests/**/data/ 导致测试数据无法入库。 顺带发现并修复三个真实缺陷： - nodes/srt.py：相邻条目缺少空行时把下一条时间轴吞进正文（静默错位）， 改为正文行遇时间戳行即报错； - src/wov_app/scheduler.py：_file_size 只捕获 OSError，含 \x00 的产物 URI 抛 ValueError 导致任务误判失败，改为同时捕获； - nodes/subtitle_correction.py：生产代码依赖测试包解析 SRT， 改用生产模块 nodes/srt.py。 真实模型/服务集成测试按外部状态跳过：新增 tests/shared/gpu_memory.py （运行时探测显存、CUDA OOM 转跳过）与 tests/shared/llm_service.py （无 Key / 余额 / 限流转跳过）。全量 477 passed。
132 lines
5.2 KiB
Python
132 lines
5.2 KiB
Python
"""GPU 显存探测与不足时跳过（供真实模型集成测试使用）。
|
||
|
||
真实模型（faster-whisper + CUDA）需要足够的显存；显存被其他进程占用或本身
|
||
偏小时，推理会在中途抛 "CUDA failed with error out of memory"。这属于**外部
|
||
运行环境状态**，不是被测代码的行为，按测试规则应跳过而不是判失败。
|
||
|
||
两层防护：
|
||
|
||
1. `require_gpu_memory()`：运行前按模型体量估算所需显存，不足则跳过；
|
||
2. `skip_on_cuda_oom()`：运行中若仍发生显存不足（被其他进程动态抢占），
|
||
也会被识别并转为跳过，而不是让整个套件判红。
|
||
|
||
探测不到 NVIDIA GPU（无 `nvidia-smi`，例如纯 CPU 机器）时**不做显存检查**：
|
||
faster-whisper 的 `device=auto` 会回退 CPU，此时测试应当照常执行。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import shutil
|
||
import subprocess
|
||
from contextlib import contextmanager
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
# 推理时的额外开销系数：CTranslate2 除权重外还需推理工作区、激活与 CUDA 上下文。
|
||
# 实测标定（V2 权重 2.87GB，6GB 卡）：推理前占用约 1120 MiB，峰值约 5217 MiB
|
||
# → 峰值增量约 4097 MiB ≈ 权重 × 1.39。系数取 1.45 覆盖 5% 分配余量，
|
||
# 使 6GB 卡（可用约 4.6GB）刚好放行、而更小的卡会被拦下。
|
||
#
|
||
# 注意：这只给出"是否可能跑完"的预估。实测在临界卡上（余量仅数百 MB）
|
||
# 即使预估通过仍可能因分配碎片化抛 CUDA OOM，所以还需要
|
||
# `skip_on_cuda_oom` / `require_node_result` 作为运行中的兼底。
|
||
_MEMORY_FACTOR = 1.45
|
||
|
||
|
||
def available_gpu_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: list[int] = []
|
||
for line in completed.stdout.splitlines():
|
||
stripped = line.strip()
|
||
if stripped.isdigit():
|
||
values.append(int(stripped))
|
||
return max(values) if values else None
|
||
|
||
|
||
def required_memory_mb(model_dir: Path) -> int:
|
||
"""按模型权重体量估算推理所需显存（MB）。"""
|
||
weights = model_dir / "model.bin"
|
||
size_mb = weights.stat().st_size / (1024 * 1024) if weights.is_file() else 0.0
|
||
return int(size_mb * _MEMORY_FACTOR)
|
||
|
||
|
||
def require_gpu_memory(model_dir: Path) -> None:
|
||
"""显存不足以跑完该模型时跳过测试；无 GPU 时不检查（回退 CPU 执行）。"""
|
||
free = available_gpu_memory_mb()
|
||
if free is None:
|
||
return # 无 NVIDIA GPU：device=auto 会走 CPU，无需显存检查
|
||
needed = required_memory_mb(model_dir)
|
||
if free < needed:
|
||
pytest.skip(
|
||
f"可用显存不足（需约 {needed} MB，当前可用 {free} MB），"
|
||
f"跳过真实模型集成测试以避免 OOM"
|
||
)
|
||
|
||
|
||
def is_cuda_oom(text: str) -> bool:
|
||
"""判断错误文本是否为显存不足（CUDA OOM）。"""
|
||
lowered = text.lower()
|
||
return "out of memory" in lowered or "cuda failed" in lowered
|
||
|
||
|
||
@contextmanager
|
||
def skip_on_cuda_oom():
|
||
"""运行真实推理；中途发生 CUDA OOM 时转为跳过（显存被动态抢占的场景）。"""
|
||
try:
|
||
yield
|
||
except Exception as exc: # noqa: BLE001 - 需要按消息识别 OOM
|
||
if is_cuda_oom(str(exc)):
|
||
pytest.skip(f"GPU 显存不足（{exc}），跳过真实模型集成测试")
|
||
raise
|
||
|
||
|
||
def require_node_result(response, model_dir: Path) -> None:
|
||
"""校验节点响应：显存不足转为跳过，其它失败照常抛出由断言处理。
|
||
|
||
节点（如 whisper）会把推理异常包装成 `status="failed"` 的响应，因此
|
||
OOM 不会以异常形式冒泡；这里统一识别并跳过。
|
||
"""
|
||
if getattr(response, "status", "") != "completed" and is_cuda_oom(
|
||
str(getattr(response, "error", ""))
|
||
):
|
||
pytest.skip(
|
||
f"GPU 显存不足（{getattr(response, 'error', '')}），跳过真实模型集成测试"
|
||
)
|
||
|
||
|
||
def fits_with_margin(model_dir: Path, safety_mb: int = 512) -> bool:
|
||
"""显存是否充裕到可承受分块/多次推理（预留 safety_mb 余量）。
|
||
|
||
用于区分两种用例写法：
|
||
- 显存充裕：按生产默认走分块路径（更接近线上配置）；
|
||
- 显存临界：退化为整段单次推理，或直接跳过——分块会产生更多分配峰值，
|
||
在临界卡上易因碎片化触发 CUDA OOM（实测同一峰值下 chunk 失败而单次成功）。
|
||
|
||
无 NVIDIA GPU 时返回 False（此时应改用 CPU 友好配置，而非依赖显存）。
|
||
"""
|
||
free = available_gpu_memory_mb()
|
||
if free is None:
|
||
return False
|
||
return free >= required_memory_mb(model_dir) + safety_mb
|