vram_guard.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""VRAM helpers for heavy jobs (simulate) that load their own HF weights."""
 
from __future__ import annotations
 
import gc
from typing import Callable
 
from aquin.compute.device import (
    empty_device_cache,
    is_oom_error,
    probe_backend,
    resolve_compute_device,
)
 
 
def accelerator_vram_gib() -> tuple[float, float] | None:
    """Free/total GiB for the active accelerator (CUDA/ROCm only; MPS has no public API)."""
    import torch
 
    if torch.cuda.is_available():
        free_b, total_b = torch.cuda.mem_get_info()
        return free_b / (1024**3), total_b / (1024**3)
    return None
 
 
def cuda_vram_gib() -> tuple[float, float] | None:
    """Backward-compatible alias."""
    return accelerator_vram_gib()
 
 
def log_accelerator_vram(log: Callable[[str], None] | None = None) -> None:
    info = accelerator_vram_gib()
    dev = resolve_compute_device()
    if info is None:
        if dev == "mps":
            line = "VRAM: MPS (Metal) — no per-process memory query"
        else:
            return
    else:
        free_g, total_g = info
        line = f"VRAM: {free_g:.1f} GiB free / {total_g:.1f} GiB total ({dev})"
    if log:
        log(line)
    else:
        print(f"[vram] {line}", flush=True)
 
 
def log_cuda_vram(log: Callable[[str], None] | None = None) -> None:
    log_accelerator_vram(log)
 
 
def release_inspection_models() -> None:
    """Unload HookedTransformer / embedding caches so simulate can own the GPU."""
    from aquin.compute.model_runtime import unload_weights
 
    unload_weights(clear_active=False)
    gc.collect()
    empty_device_cache()
 
 
def cleanup_heavy_job_vram() -> None:
    """Drop any models loaded during a heavy job (simulate TL cache, HF weights)."""
    release_inspection_models()
 
 
def restore_session_model(model_id: str | None) -> None:
    if not (model_id or "").strip():
        return
    try:
        from aquin.compute.model_loader import reload_vram_for_model
 
        reload_vram_for_model(model_id.strip())
        print(f"[vram] reloaded session model: {model_id}", flush=True)
    except Exception as exc:
        print(
            f"[vram] could not reload session model '{model_id}': {exc}\n"
            f"       Run: aquin load --model {model_id}",
            flush=True,
        )
 
 
def oom_message(*, job: str, model_id: str) -> str:
    info = accelerator_vram_gib()
    dev = resolve_compute_device()
    vram_bit = ""
    if info:
        free_g, total_g = info
        vram_bit = f" {dev} has {free_g:.1f} GiB free of {total_g:.1f} GiB."
    elif dev == "mps":
        vram_bit = " MPS (Metal) unified memory."
    return (
        f"Out of memory during {job} for '{model_id}'.{vram_bit} "
        "Simulate loads its own copy of the model — use a smaller model "
        "(e.g. gpt2-small, llama-3.2-1b), free memory from other processes, "
        "or run: aquin load --model <smaller-id>."
    )
 
 
def cuda_oom_message(*, job: str, model_id: str) -> str:
    return oom_message(job=job, model_id=model_id)
 
 
def raise_if_cuda_oom(exc: BaseException, *, job: str, model_id: str) -> None:
    if is_oom_error(exc):
        raise RuntimeError(oom_message(job=job, model_id=model_id)) from exc
    raise exc
 
 
def backend_summary() -> str:
    return probe_backend()["summary"]