vrsub/tests/shared/gpu_memory.py

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