File size: 2,844 Bytes
bc6498b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
#!/usr/bin/env python3
# energy_logger_nvml.py
import argparse, csv, json, os, shlex, subprocess, sys, threading, time
try:
from pynvml import nvmlInit, nvmlShutdown, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex, nvmlDeviceGetPowerUsage
NVML_OK = True
except Exception:
NVML_OK = False
class _Sampler(threading.Thread):
def __init__(self, interval=0.1):
super().__init__(daemon=True)
self.interval = interval
self.samples = []
self.running = False
def run(self):
if not NVML_OK: return
nvmlInit()
try:
n = nvmlDeviceGetCount()
handles = [nvmlDeviceGetHandleByIndex(i) for i in range(n)]
self.running = True
while self.running:
t = time.time()
w = 0.0
for h in handles:
w += nvmlDeviceGetPowerUsage(h)/1000.0
self.samples.append((t, w))
time.sleep(self.interval)
finally:
try: nvmlShutdown()
except: pass
def stop(self): self.running = False
def _trapz(samples):
if len(samples) < 2: return 0.0
E = 0.0
for (t0,p0),(t1,p1) in zip(samples, samples[1:]):
E += 0.5*(p0+p1)*(t1-t0)
return E
class EnergyLogger:
def __init__(self, tag="session", interval=0.1, out_dir="energy_logs"):
self.tag = tag; self.interval=interval; self.out_dir=out_dir
self.sampler=_Sampler(interval=interval); self.t0=None; self.t1=None; self.summary={}
def __enter__(self):
os.makedirs(self.out_dir, exist_ok=True)
self.t0 = time.time(); self.sampler.start(); return self
def __exit__(self, *args):
self.sampler.stop(); self.sampler.join(); self.t1 = time.time()
dur = self.t1 - self.t0; E = _trapz(self.sampler.samples); avg = (E/dur) if dur>0 else 0.0
self.summary = {"duration_s": dur, "energy_J": E, "avg_power_W": avg, "samples": len(self.sampler.samples)}
def samples(self): return list(self.sampler.samples)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--cmd", type=str, required=True)
ap.add_argument("--interval", type=float, default=0.1)
ap.add_argument("--tag", type=str, default="session")
ap.add_argument("--out_dir", type=str, default="energy_logs")
args = ap.parse_args()
if not NVML_OK:
print(json.dumps({"error":"NVML not available; pip install pynvml and ensure NVIDIA driver present."}, indent=2)); sys.exit(2)
el = EnergyLogger(tag=args.tag, interval=args.interval, out_dir=args.out_dir)
with el:
import shlex, subprocess
proc = subprocess.Popen(shlex.split(args.cmd))
ret = proc.wait()
print(json.dumps({"returncode":ret, **el.summary}, indent=2)); sys.exit(ret)
if __name__ == "__main__":
main()
|