|
|
|
|
|
|
|
|
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() |
|
|
|
