src/utils/hardware.py

"""Hardware profiling and max-safe guard helpers for local-only experiments."""

from __future__ import annotations

import ctypes
import contextlib
import json
import os
import platform
import shutil
import subprocess
import tempfile
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Callable, Iterable, Optional


BYTES_PER_GB = 1024**3


@dataclass(frozen=True)
class MemoryInfo:
    total_gb: float
    available_gb: Optional[float]
    used_fraction: Optional[float]


@dataclass(frozen=True)
class GpuInfo:
    name: str
    memory_total_mb: int
    memory_used_mb: Optional[int]
    temperature_c: Optional[int]
    driver_version: Optional[str]


@dataclass(frozen=True)
class TorchInfo:
    installed: bool
    version: Optional[str]
    cuda_available: bool
    cuda_version: Optional[str]
    device_name: Optional[str]
    device_memory_gb: Optional[float]
    error: Optional[str] = None


@dataclass(frozen=True)
class MaxSafeConfig:
    cpu_workers: int
    dataloader_workers: int
    max_dataloader_workers: int
    ram_limit_fraction: float
    ram_target_gb: float
    gpu_vram_fraction: float
    gpu_target_vram_gb: Optional[float]
    gpu_pause_temp_c: int
    gpu_abort_temp_c: int
    mixed_precision: bool
    gradient_accumulation: bool
    batch_candidates: list[int]


@dataclass(frozen=True)
class HardwareProfile:
    profile_name: str
    platform: str
    python: str
    cpu_name: str
    logical_cpu_count: int
    memory: MemoryInfo
    gpus: list[GpuInfo]
    torch: TorchInfo
    max_safe: MaxSafeConfig


class ResourceLimitError(RuntimeError):
    """Raised when a max-safe guard would abort a run."""


def _round_gb(value_bytes: int) -> float:
    return round(value_bytes / BYTES_PER_GB, 2)


def get_memory_info() -> MemoryInfo:
    """Return physical memory using stdlib-only platform APIs."""
    if platform.system().lower() == "windows":
        class MEMORYSTATUSEX(ctypes.Structure):
            _fields_ = [
                ("dwLength", ctypes.c_ulong),
                ("dwMemoryLoad", ctypes.c_ulong),
                ("ullTotalPhys", ctypes.c_ulonglong),
                ("ullAvailPhys", ctypes.c_ulonglong),
                ("ullTotalPageFile", ctypes.c_ulonglong),
                ("ullAvailPageFile", ctypes.c_ulonglong),
                ("ullTotalVirtual", ctypes.c_ulonglong),
                ("ullAvailVirtual", ctypes.c_ulonglong),
                ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
            ]

        status = MEMORYSTATUSEX()
        status.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
        ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status))
        used_fraction = 1.0 - (status.ullAvailPhys / status.ullTotalPhys)
        return MemoryInfo(
            total_gb=_round_gb(status.ullTotalPhys),
            available_gb=_round_gb(status.ullAvailPhys),
            used_fraction=round(used_fraction, 4),
        )

    if hasattr(os, "sysconf") and "SC_PAGE_SIZE" in os.sysconf_names:
        page_size = os.sysconf("SC_PAGE_SIZE")
        pages = os.sysconf("SC_PHYS_PAGES")
        return MemoryInfo(total_gb=_round_gb(page_size * pages), available_gb=None, used_fraction=None)

    return MemoryInfo(total_gb=0.0, available_gb=None, used_fraction=None)


def get_nvidia_gpus() -> list[GpuInfo]:
    """Read NVIDIA GPU status through nvidia-smi when available."""
    if not shutil.which("nvidia-smi"):
        return []
    cmd = [
        "nvidia-smi",
        "--query-gpu=name,memory.total,memory.used,temperature.gpu,driver_version",
        "--format=csv,noheader,nounits",
    ]
    try:
        raw = subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL, timeout=10)
    except Exception:
        return []

    gpus = []
    for line in raw.splitlines():
        parts = [part.strip() for part in line.split(",")]
        if len(parts) < 5:
            continue
        name, total_mb, used_mb, temp_c, driver = parts[:5]
        gpus.append(
            GpuInfo(
                name=name,
                memory_total_mb=int(float(total_mb)),
                memory_used_mb=int(float(used_mb)),
                temperature_c=int(float(temp_c)),
                driver_version=driver,
            )
        )
    return gpus


def get_torch_info() -> TorchInfo:
    """Return PyTorch/CUDA availability without making Torch mandatory."""
    try:
        with open(os.devnull, "w", encoding="utf-8") as devnull, contextlib.redirect_stderr(devnull):
            import torch

        cuda_available = bool(torch.cuda.is_available())
        device_name = None
        device_memory_gb = None
        if cuda_available:
            props = torch.cuda.get_device_properties(0)
            device_name = torch.cuda.get_device_name(0)
            device_memory_gb = round(props.total_memory / BYTES_PER_GB, 2)
        return TorchInfo(
            installed=True,
            version=str(torch.__version__),
            cuda_available=cuda_available,
            cuda_version=getattr(torch.version, "cuda", None),
            device_name=device_name,
            device_memory_gb=device_memory_gb,
        )
    except Exception as exc:
        return TorchInfo(
            installed=False,
            version=None,
            cuda_available=False,
            cuda_version=None,
            device_name=None,
            device_memory_gb=None,
            error=repr(exc),
        )


def build_max_safe_config(memory: MemoryInfo, gpus: list[GpuInfo]) -> MaxSafeConfig:
    """Choose conservative max-safe defaults for this laptop class."""
    logical = os.cpu_count() or 1
    cpu_workers = max(1, min(12, logical - 4 if logical > 4 else logical))
    gpu_target_vram_gb = None
    if gpus:
        gpu_target_vram_gb = round((gpus[0].memory_total_mb / 1024) * 0.85, 2)
    return MaxSafeConfig(
        cpu_workers=cpu_workers,
        dataloader_workers=2,
        max_dataloader_workers=4,
        ram_limit_fraction=0.90,
        ram_target_gb=min(13.0, round(memory.total_gb * 0.82, 2)) if memory.total_gb else 13.0,
        gpu_vram_fraction=0.85,
        gpu_target_vram_gb=gpu_target_vram_gb,
        gpu_pause_temp_c=82,
        gpu_abort_temp_c=87,
        mixed_precision=True,
        gradient_accumulation=True,
        batch_candidates=[512, 256, 128, 64, 32, 16, 8],
    )


def build_hardware_profile(profile_name: str = "max_safe") -> HardwareProfile:
    memory = get_memory_info()
    gpus = get_nvidia_gpus()
    return HardwareProfile(
        profile_name=profile_name,
        platform=platform.platform(),
        python=platform.python_version(),
        cpu_name=platform.processor() or platform.machine(),
        logical_cpu_count=os.cpu_count() or 1,
        memory=memory,
        gpus=gpus,
        torch=get_torch_info(),
        max_safe=build_max_safe_config(memory, gpus),
    )


def assert_within_max_safe(profile: HardwareProfile) -> None:
    """Abort before a run if current memory/temperature already violates guardrails."""
    if profile.memory.used_fraction is not None and profile.memory.used_fraction >= profile.max_safe.ram_limit_fraction:
        raise ResourceLimitError(
            f"System RAM usage {profile.memory.used_fraction:.1%} exceeds "
            f"{profile.max_safe.ram_limit_fraction:.0%} guard"
        )
    if profile.gpus:
        gpu = profile.gpus[0]
        if gpu.temperature_c is not None and gpu.temperature_c >= profile.max_safe.gpu_abort_temp_c:
            raise ResourceLimitError(
                f"GPU temperature {gpu.temperature_c}C exceeds "
                f"{profile.max_safe.gpu_abort_temp_c}C abort guard"
            )


def calibrate_batch_size(
    trial_fn: Callable[[int], None],
    candidates: Iterable[int],
    *,
    repeats: int = 2,
) -> Optional[int]:
    """
    Choose the largest candidate batch that survives repeated trial calls.

    `trial_fn` should raise RuntimeError for OOM-like failures. This helper is
    intentionally framework-agnostic so tests can simulate OOM without Torch.
    """
    for batch_size in candidates:
        try:
            for _ in range(repeats):
                trial_fn(int(batch_size))
            return int(batch_size)
        except RuntimeError:
            continue
    return None


def atomic_write_json(path: str | Path, payload: object) -> None:
    """Write JSON atomically so interrupted runs do not corrupt manifests."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if hasattr(payload, "__dataclass_fields__"):
        payload = asdict(payload)
    with tempfile.NamedTemporaryFile("w", delete=False, dir=str(path.parent), encoding="utf-8") as tmp:
        json.dump(payload, tmp, indent=2, sort_keys=True)
        tmp.write("\n")
        tmp_path = Path(tmp.name)
    tmp_path.replace(path)


def dataclass_to_dict(value: object) -> dict:
    if hasattr(value, "__dataclass_fields__"):
        return asdict(value)
    if isinstance(value, dict):
        return value
    raise TypeError(f"Unsupported JSON payload: {type(value)}")