"""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