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