ModelTrainEval/utils/device.py

"""设备解析与能力查询：统一处理 cuda / mps / cpu 三种后端。

设计要点：
- `resolve_device("auto")` 按 cuda > mps > cpu 的优先级自动挑最快的后端；
- 显式指定但不可用时直接报错，而不是悄悄退回 CPU——训练慢 30 倍却不吭声是最坏的情况；
- AMP 的默认值随后端而变：CUDA 上开（有 tensor core），MPS 上关
  （实测 tiny_conv @128² 每步 fp32 76 ms、fp16 90 ms，autocast 只剩转换开销反而更慢）；
- 显存统计跨后端统一：CUDA 有内置峰值计数器，MPS 没有，靠采样取最大值。
"""

from __future__ import annotations

import torch

#: 自动选择时的优先级
BACKEND_PRIORITY = ("cuda", "mps", "cpu")


class DeviceNotAvailableError(RuntimeError):
    """请求的设备在当前机器上不可用。"""


def _cuda_ready() -> bool:
    return torch.cuda.is_available() and torch.cuda.device_count() > 0


def _mps_ready() -> bool:
    return torch.backends.mps.is_available() and torch.backends.mps.is_built()


def available_backends() -> list[str]:
    """按优先级返回当前可用的后端，末尾必定有 cpu。"""
    backends = []
    if _cuda_ready():
        backends.append("cuda")
    if _mps_ready():
        backends.append("mps")
    backends.append("cpu")
    return backends


def resolve_device(spec: str = "auto") -> torch.device:
    """把配置里的 device 字符串解析成 `torch.device`。

    Args:
        spec: "auto" 自动挑最优后端；也可写 "cuda" / "cuda:1" / "mps" / "cpu"。

    Raises:
        DeviceNotAvailableError: 显式请求的后端不可用。
    """
    spec = str(spec).strip().lower()
    if spec in ("auto", ""):
        return torch.device(available_backends()[0])

    backend = spec.split(":", 1)[0]
    if backend == "cuda" and not _cuda_ready():
        raise DeviceNotAvailableError(
            f"请求了 {spec}，但当前没有可用的 Nvidia GPU"
            f"（PyTorch {torch.__version__}，CUDA 构建：{torch.version.cuda}）。"
            f"可用后端：{available_backends()}，或用 device=auto 自动选择。"
        )
    if backend == "mps" and not _mps_ready():
        raise DeviceNotAvailableError(
            f"请求了 {spec}，但当前 PyTorch 不支持 MPS"
            f"（available={torch.backends.mps.is_available()}, built={torch.backends.mps.is_built()}）。"
            f"可用后端：{available_backends()}。"
        )
    if backend not in ("cuda", "mps", "cpu"):
        raise DeviceNotAvailableError(f"未知设备 {spec}，支持 cuda / mps / cpu / auto。")
    return torch.device(spec)


def describe_device(device: torch.device | str) -> str:
    """给日志用的人类可读描述。"""
    device = torch.device(device)
    if device.type == "cuda":
        index = device.index or 0
        return f"{torch.cuda.get_device_name(index)} (cuda:{index})"
    if device.type == "mps":
        return "Apple Silicon GPU (Metal / MPS)"
    return "CPU"


def amp_supported(device: torch.device | str) -> bool:
    """该后端是否支持 autocast 混合精度。"""
    return torch.device(device).type in ("cuda", "mps")


def default_amp(device: torch.device | str) -> bool:
    """AMP 的默认开关：只有 CUDA 默认开。

    MPS 支持 autocast 且数值正常，但实测比 fp32 更慢，所以默认关闭；
    需要省显存时可以在配置里显式 `amp: true` 打开。
    """
    return torch.device(device).type == "cuda"


def resolve_amp(setting: bool | str | None, device: torch.device | str) -> bool:
    """解析配置里的 `amp` 字段：`auto`/`None` 随后端，`true`/`false` 显式指定。"""
    device = torch.device(device)
    if setting is None or (isinstance(setting, str) and setting.lower() == "auto"):
        return default_amp(device)
    if not bool(setting):
        return False
    return amp_supported(device)


def pin_memory_supported(device: torch.device | str) -> bool:
    """只有 CUDA 的锁页内存有意义；MPS 上开启会告警且无收益。"""
    return torch.device(device).type == "cuda"


def synchronize(device: torch.device | str) -> None:
    """等待该设备上的异步算子完成——计时前后必须调用，否则测到的是下发时间。"""
    device = torch.device(device)
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    elif device.type == "mps":
        torch.mps.synchronize()


def memory_stats(device: torch.device | str) -> dict[str, float]:
    """统一的显存查询，单位 MB。

    CUDA 返回 allocated / reserved / peak；MPS 只有当前分配量与驱动占用
    （Metal 没有暴露峰值计数器）；CPU 返回空字典。
    """
    device = torch.device(device)
    mb = 1024.0 * 1024.0
    if device.type == "cuda":
        return {
            "allocated_mb": torch.cuda.memory_allocated(device) / mb,
            "reserved_mb": torch.cuda.memory_reserved(device) / mb,
            "peak_allocated_mb": torch.cuda.max_memory_allocated(device) / mb,
            "peak_reserved_mb": torch.cuda.max_memory_reserved(device) / mb,
        }
    if device.type == "mps":
        return {
            "allocated_mb": torch.mps.current_allocated_memory() / mb,
            "driver_allocated_mb": torch.mps.driver_allocated_memory() / mb,
        }
    return {}


class PeakMemory:
    """跨后端的显存峰值统计。

    CUDA 有内置峰值计数器，直接读；MPS 没有，只能在 `update()` 时采样取最大值——
    因此 MPS 上的峰值是**下界**（采样点之间的瞬时尖峰会被漏掉）。
    """

    def __init__(self, device: torch.device | str) -> None:
        self.device = torch.device(device)
        self._peak = 0.0
        self.reset()

    def reset(self) -> None:
        self._peak = 0.0
        if self.device.type == "cuda":
            torch.cuda.reset_peak_memory_stats(self.device)
        self.update()

    def update(self) -> None:
        if self.device.type != "mps":
            return
        self._peak = max(self._peak, torch.mps.current_allocated_memory() / (1024.0 * 1024.0))

    @property
    def peak_mb(self) -> float:
        if self.device.type == "cuda":
            return torch.cuda.max_memory_allocated(self.device) / (1024.0 * 1024.0)
        if self.device.type == "mps":
            self.update()
            return self._peak
        return 0.0

    @property
    def reserved_mb(self) -> float:
        """CUDA 是缓存分配器保留量；MPS 是 Metal 驱动占用；CPU 为 0。"""
        if self.device.type == "cuda":
            return torch.cuda.max_memory_reserved(self.device) / (1024.0 * 1024.0)
        if self.device.type == "mps":
            return torch.mps.driver_allocated_memory() / (1024.0 * 1024.0)
        return 0.0

    @property
    def exact(self) -> bool:
        """峰值是否精确（MPS 靠采样，只能给下界）。"""
        return self.device.type != "mps"