Files
AI-Profile-Router/experiments/ornith15-ab/gpu_monitor.py
T

34 lines
1.2 KiB
Python
Executable File

#!/usr/bin/env python3
"""Sample physical GPU load during a GO-authorized A/B case."""
import argparse
import json
import pathlib
import subprocess
import time
def main():
p = argparse.ArgumentParser()
p.add_argument("--output", type=pathlib.Path, required=True)
p.add_argument("--interval", type=float, default=1.0)
p.add_argument("--go", action="store_true")
args = p.parse_args()
if not args.go:
p.error("No GPU monitoring before explicit --go")
args.output.parent.mkdir(parents=True, exist_ok=True)
cmd = ["nvidia-smi", "--query-gpu=uuid,name,memory.used,utilization.gpu,power.draw",
"--format=csv,noheader,nounits"]
with args.output.open("w") as out:
while True:
sample = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
record = {"time": time.time(), "returncode": sample.returncode,
"rows": [line.strip() for line in sample.stdout.splitlines() if line.strip()],
"error": sample.stderr.strip() if sample.returncode else ""}
out.write(json.dumps(record) + "\n")
out.flush()
time.sleep(args.interval)
if __name__ == "__main__":
main()