#!/usr/bin/env python3
"""Measure how busy your GPUs are, for Compute Check (https://check.ashibaresearch.com).

Samples nvidia-smi every 30 seconds and writes one row per hour to gpu_use.csv:
  usage_hour, gpu_type, gpus, busy_gpu_hours, avg_util_pct, peak_util_pct
Drop the CSV into Compute Check ("Fine-tune the numbers" or "Upload a billing export").

No network calls. No hostnames, process names, or users are recorded. Python 3.8+, standard library only.

  python3 gpu_use.py --hours 168          # one week (stop any time; rows are written each hour)
  python3 gpu_use.py --test               # one sample, printed, nothing written
  nohup python3 gpu_use.py --hours 168 &  # keep running after you log out

Run it on each GPU machine. Drop all the files in at once; Compute Check adds them up.
"""
import argparse, csv, datetime as dt, os, shutil, subprocess, sys, time

QUERY = "name,utilization.gpu"


def sample():
    """One reading per GPU: (name, util %), or (name, None) when the GPU reports no utilization (MIG mode, some drivers)."""
    out = subprocess.run(["nvidia-smi", f"--query-gpu={QUERY}", "--format=csv,noheader,nounits"],
                         capture_output=True, text=True, timeout=20)
    if out.returncode != 0:
        raise RuntimeError(out.stderr.strip() or "nvidia-smi failed")
    rows = []
    for line in out.stdout.strip().splitlines():
        name, util = [x.strip() for x in line.rsplit(",", 1)]
        try:
            u = float(util)
        except ValueError:
            u = None  # "[N/A]" or "[Not Supported]": unknown, not idle
        rows.append((name.replace("NVIDIA ", ""), u))
    return rows


def write_hour(path, hour, samples):
    """samples: list of per-reading lists of (name, util). Writes one summary row from the GPUs that report utilization."""
    samples = [k for k in ([(n, u) for n, u in s if u is not None] for s in samples) if k]
    if not samples:
        return
    gpus = max(len(s) for s in samples)
    names = sorted({n for s in samples for n, _ in s})
    utils = [u for s in samples for _, u in s]
    avg = sum(utils) / len(utils)
    # Busy GPU-hours: the share of GPU capacity in use this hour (utilization 0-100%, averaged).
    busy = sum(sum(u for _, u in s) / 100 for s in samples) / len(samples)
    new = not os.path.exists(path)
    with open(path, "a", newline="") as f:
        w = csv.writer(f)
        if new:
            w.writerow(["usage_hour", "gpu_type", "gpus", "busy_gpu_hours", "avg_util_pct", "peak_util_pct"])
        w.writerow([hour.strftime("%Y-%m-%dT%H:00:00Z"), " / ".join(names), gpus, round(busy, 3), round(avg, 1), round(max(utils), 1)])


def main():
    ap = argparse.ArgumentParser(description="Measure GPU use for Compute Check.")
    ap.add_argument("--hours", type=float, default=168, help="how long to run (default 168 = one week)")
    ap.add_argument("--every", type=int, default=30, help="seconds between samples (default 30)")
    ap.add_argument("--out", default="gpu_use.csv", help="output file (default gpu_use.csv)")
    ap.add_argument("--test", action="store_true", help="take one sample, print it, write nothing")
    a = ap.parse_args()

    if not shutil.which("nvidia-smi"):
        sys.exit("nvidia-smi not found. Run this on a machine with NVIDIA GPUs and drivers.")
    first = sample()
    if a.test:
        for i, (n, u) in enumerate(first):
            print(f"GPU {i}: {n}, " + (f"{u:.0f}% busy" if u is not None else "no utilization reading (MIG mode?)"))
    if first and all(u is None for _, u in first):
        sys.exit("No GPU reports utilization. MIG mode hides it from nvidia-smi: measure with DCGM (DCGM_FI_PROF_GR_ENGINE_ACTIVE) instead.")
    for i, (n, u) in enumerate(first):
        if u is None:
            print(f"GPU {i} ({n}) reports no utilization. It is left out of the counts.", file=sys.stderr, flush=True)
    if a.test:
        return

    end = time.time() + a.hours * 3600
    hour, buf = None, []
    print(f"Sampling every {a.every}s for {a.hours:g} hours. Writing hourly rows to {a.out}. Ctrl-C to stop.", flush=True)
    try:
        while time.time() < end:
            now = dt.datetime.now(dt.timezone.utc).replace(minute=0, second=0, microsecond=0)
            if hour is not None and now != hour:
                write_hour(a.out, hour, buf)
                buf = []
            hour = now
            try:
                buf.append(sample())
            except Exception as e:  # keep going through a driver hiccup
                print(f"sample failed: {e}", file=sys.stderr, flush=True)
            time.sleep(a.every)
    except KeyboardInterrupt:
        pass
    write_hour(a.out, hour, buf)
    print(f"Done. Drop {a.out} into https://check.ashibaresearch.com", flush=True)


if __name__ == "__main__":
    main()