Skip to content

Instantly share code, notes, and snippets.

@AndiH
Created August 30, 2023 08:44
Show Gist options
  • Save AndiH/867d86c5168ec3f8c70ac4ffc1fd4324 to your computer and use it in GitHub Desktop.
Save AndiH/867d86c5168ec3f8c70ac4ffc1fd4324 to your computer and use it in GitHub Desktop.
GPU Power with Context
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment