GPU Power with Context

measure.py

import os
import subprocess
import io
import time

import pandas as pd

class GetPower(object):
	def __enter__(self):
		cmd = "nvidia-smi --query-gpu=index,timestamp,name,power.draw --format=csv --loop-ms=5"
		self.smi = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, text=True)
		return self
	def __exit__(self, type, value, traceback):
		self.smi.kill()
		self.data = io.StringIO(self.smi.stdout.read())
		self.df = pd.read_csv(self.data, parse_dates=[1], skipinitialspace=True)
	def energy(self):
		import numpy as np
		_energy = []
		for df_name, df_group in self.df.groupby('index'):
			_energy.append(np.sum(df_group['power.draw [W]'].str.replace(' W', '').astype(float) * df_group['timestamp'].diff().dt.total_seconds())/3600)
		return _energy

def main():
	with GetPower() as measured_scope:
		print('Sleeping... But actually be a pytorch.train() call')
		time.sleep(0.2)
	print(measured_scope.df.groupby('index').get_group(0))
	print(measured_scope.energy())

if __name__ == '__main__':
	main()