#!/usr/bin/env python3
"""
gb10_clockcheck.py — is this GB10 (DGX Spark / ASUS Ascent GX10 / other GB10) stuck in a low-power state?

Runs a 45-second bf16 matrix-multiply burn on the GPU, measures real throughput (TFLOP/s) and samples the
SM clock, power and throttle reasons while it runs. Read-only: it changes no settings.

PRIVACY: this script makes NO network connections, sends nothing anywhere, collects nothing and writes no
files. It only runs a matmul on your GPU and reads nvidia-smi locally. It's ~160 lines; read it before you
run it. (We verified it with a Python audit hook and strace: zero network connections.)

BEFORE YOU RUN IT: stop anything else using the GPU (model servers, training, ComfyUI...). A busy GPU gives a
falsely low score.

Needs: Python 3 + PyTorch with CUDA (any recent NGC / vLLM / PyTorch container works), and nvidia-smi.
    python3 gb10_clockcheck.py            # 45 s test
    python3 gb10_clockcheck.py --seconds 60

By Anantus. Free to use and share. MIT license.
"""
import argparse, shutil, subprocess, sys, threading, time

# Reference results from our own 8-node GB10 cluster (same test: 8192^3 bf16, 45 s, ~96% utilization).
# Other chassis (Dell, HP, MSI, Gigabyte, Lenovo...) may differ; please send us yours.
REFERENCE = [
    ("ASUS Ascent GX10, healthy (7 units)", "88.8 - 93.7 TFLOP/s", "2275 - 2333 MHz", "~89 - 92 W"),
    ("NVIDIA DGX Spark, healthy (1 unit)",  "~77 TFLOP/s",         "lower power envelope", ""),
    ("ASUS Ascent GX10, STUCK (2 units)",   "42.5 TFLOP/s",        "851 MHz",          "20.7 W"),
]
STUCK_TFLOPS = 60.0   # well below every healthy unit we have seen
STUCK_MHZ = 1400      # median SM clock under full load below this = suspicious


def smi_sample():
    """One nvidia-smi sample: (sm_clock_MHz, power_W, temp_C, throttle_reasons_hex) or None."""
    if not shutil.which("nvidia-smi"):
        return None
    q = "clocks.sm,power.draw,temperature.gpu,clocks_throttle_reasons.active"
    try:
        out = subprocess.run(["nvidia-smi", f"--query-gpu={q}", "--format=csv,noheader,nounits"],
                             capture_output=True, text=True, timeout=10).stdout.strip().splitlines()[0]
        f = [x.strip() for x in out.split(",")]
        num = lambda s: float(s) if s.replace(".", "", 1).isdigit() else None
        return num(f[0]), num(f[1]), num(f[2]), f[3]
    except Exception:
        return None


def other_gpu_processes():
    if not shutil.which("nvidia-smi"):
        return []
    try:
        out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,process_name", "--format=csv,noheader"],
                             capture_output=True, text=True, timeout=10).stdout.strip()
        return [l for l in out.splitlines() if l.strip()]
    except Exception:
        return []


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[1])
    ap.add_argument("--seconds", type=float, default=45.0)
    ap.add_argument("--size", type=int, default=8192)
    a = ap.parse_args()

    try:
        import torch
    except ImportError:
        sys.exit("PyTorch not found. Run this inside a PyTorch/NGC/vLLM container, or: pip install torch")
    if not torch.cuda.is_available():
        sys.exit("CUDA is not available to PyTorch here (inside Docker, did you pass --gpus all?)")

    busy = other_gpu_processes()
    if busy:
        print("WARNING: other processes are using the GPU; the score will be low for that reason, not a fault:")
        for l in busy:
            print("   ", l)
        print("Stop them and re-run for a valid result.\n")

    n = a.size
    dev = torch.cuda.get_device_name(0)
    print(f"GPU: {dev}  |  test: {n}x{n} bf16 matmul for {a.seconds:.0f} s  (warming up...)")
    x = torch.randn(n, n, device="cuda", dtype=torch.bfloat16)
    y = torch.randn(n, n, device="cuda", dtype=torch.bfloat16)
    for _ in range(10):
        torch.mm(x, y)
    torch.cuda.synchronize()

    samples, stop = [], threading.Event()

    def sampler():
        time.sleep(3)  # let clocks ramp under load first
        while not stop.is_set():
            s = smi_sample()
            if s:
                samples.append(s)
            stop.wait(2)

    th = threading.Thread(target=sampler, daemon=True)
    th.start()
    t0, iters = time.time(), 0
    while time.time() - t0 < a.seconds:
        for _ in range(20):
            torch.mm(x, y)
        torch.cuda.synchronize()
        iters += 20
    dt = time.time() - t0
    stop.set()
    th.join(timeout=5)

    tflops = 2 * n ** 3 * iters / dt / 1e12
    clocks = sorted(s[0] for s in samples if s[0] is not None)
    powers = sorted(s[1] for s in samples if s[1] is not None)
    temps = [s[2] for s in samples if s[2] is not None]
    reasons = sorted({s[3] for s in samples})
    med = lambda v: v[len(v) // 2] if v else None
    mhz, watts = med(clocks), med(powers)

    print("\n=== RESULT ===")
    print(f"Throughput : {tflops:6.1f} TFLOP/s   ({iters} matmuls in {dt:.1f} s)")
    if mhz is not None:
        print(f"SM clock   : {mhz:6.0f} MHz (median under load)")
    if watts is not None:
        print(f"GPU power  : {watts:6.1f} W   (median under load)")
    if temps:
        print(f"GPU temp   : {max(temps):6.0f} C (max)")
    if reasons:
        print(f"Throttle   : {', '.join(reasons)}  (0x0000000000000000 = none reported)")

    print("\nReference (our cluster, same test):")
    for name, tf, clk, pw in REFERENCE:
        print(f"  {name:40s} {tf:22s} {clk:22s} {pw}")

    no_throttle = all(r.endswith("0000000000000000") or r in ("0x0", "0") for r in reasons) if reasons else True
    low_clock = mhz is not None and mhz < STUCK_MHZ
    print()
    if busy:
        print("VERDICT: NOT VALID - other GPU processes were running. Stop them and re-run.")
    elif tflops < STUCK_TFLOPS and (low_clock or mhz is None) and no_throttle:
        print("VERDICT: LIKELY STUCK in the low-power state (low throughput, low clock, no throttle reason).")
        print("FIX (a reboot will NOT clear it; nvidia-smi -lgc is ignored):")
        print("  1. Shut the unit down completely.")
        print("  2. Unplug the power cord from the wall outlet, unplug it from the power brick,")
        print("     and unplug the brick's cable from the back of the unit. (Also unplug USB devices.)")
        print("  3. Leave everything disconnected for at least 5 minutes.")
        print("  4. Reconnect in this order: cord into the wall outlet -> cord into the power brick ->")
        print("     brick cable into the back of the unit.  Then power on.")
        print("  5. Re-run this check. If it is still stuck, contact your vendor and add your numbers to:")
        print("     https://forums.developer.nvidia.com/t/386060")
    elif tflops < STUCK_TFLOPS:
        print("VERDICT: LOW, but not the classic stuck signature (a throttle reason is active or the clock is high).")
        print("Check cooling/airflow and the throttle reason above, then re-run with nothing else on the GPU.")
    elif tflops < 80:
        print("VERDICT: OK. In the range of lower-power chassis (e.g. NVIDIA DGX Spark ~77). Not stuck.")
    else:
        print("VERDICT: HEALTHY.")
    print("\nPlease share your result + your unit's brand/model so the reference table can grow:")
    print("  https://forums.developer.nvidia.com/t/387029")


if __name__ == "__main__":
    main()