"""显存预算（2026-10-18 本机实测后加）。
用户定的预算：这台 3090 只有 19–20 GB 能给训练用——Windows + 4K 桌面常驻就占掉 4–5 GB（空载 nvidia-smi 4.4 GB）。
  limit()     给本进程设硬帽（VRAM_FRAC，默认 0.78 ≈ 18.7 GiB）。超限时在自己的反向里早失败，
              而不是把桌面挤到没有显存（整机卡顿、别的应用崩），这也是 APG batch=128 原先 OOM 时把显存顶到 23.6 GB 的原因。
  peak_gib()  记进日志的 vram 字段：网页指标卡与 ml/selftest.sh 都按它核对预算，不用再去猜。
实测参考（d=192 同结构、1 步即到位）：预训练 batch=1024 → 2.76；APG batch=48/plans=150 → 13.08、batch=128 → OOM。"""
import os
import torch


def limit():
    if not torch.cuda.is_available(): return False
    f = float(os.environ.get('VRAM_FRAC', '0.78'))
    tot = torch.cuda.get_device_properties(0).total_memory / 2**30
    torch.cuda.set_per_process_memory_fraction(f)
    print('显存上限：本进程 %.1f GiB（整卡 %.1f GiB × %.2f，其余留给桌面）' % (f * tot, tot, f), flush=True)
    return True


def peak_gib():
    return round(torch.cuda.max_memory_allocated() / 2**30, 2) if torch.cuda.is_available() else None