h8rt3rmin8r

dstat_nvidia_gpu.py

Nov 15th, 2018
208
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. ### Author: Vasilis Vryniotis <bbriniotis@datumbox.com>
  2.  
  3. class dstat_plugin(dstat):
  4.     """
  5.    Total GPU usage for NVIDIA cards. Requires the nvidia-ml-py package.
  6.  
  7.    Usage: dstat --nvidia-gpu
  8.    """
  9.  
  10.     def __init__(self):
  11.         self.name = 'gpu usage nvidia'
  12.         self.type = 'p'
  13.         self.width = 5
  14.         self.scale = 34
  15.         self.cols = 1
  16.         self.samples = 10
  17.  
  18.     def check(self):
  19.         try:
  20.             import pynvml
  21.             pynvml.nvmlInit()
  22.         except:
  23.             raise Exception('The "pynvml" library is missing from this system.')
  24.  
  25.     def vars(self):
  26.         ret = ['total']
  27.         if op.full:
  28.             import pynvml
  29.             deviceCount = pynvml.nvmlDeviceGetCount()
  30.             for i in range(0, deviceCount):
  31.                 ret.append('gpu%d' % i)
  32.         return ret
  33.  
  34.     def extract(self):
  35.         stats = self._getUsagePerGPU(self.samples)
  36.         stats['total'] = self._getTotalUsage(stats)
  37.         for name in self.vars:
  38.             self.val[name] = stats[name]
  39.  
  40.     def _getUsagePerGPU(self, samples):
  41.         import pynvml
  42.         usage = {}
  43.         deviceCount = pynvml.nvmlDeviceGetCount()
  44.         for iter in range(0, samples):
  45.             for i in range(0, deviceCount):
  46.                 name = 'gpu%d' % i
  47.                 handle = pynvml.nvmlDeviceGetHandleByIndex(i)
  48.                 if name not in usage:
  49.                     usage[name] = 0.0
  50.                 usage[name] += pynvml.nvmlDeviceGetUtilizationRates(handle).gpu / float(samples)
  51.         return usage
  52.  
  53.     def _getTotalUsage(self, usage_per_gpu):
  54.         return sum(usage_per_gpu.values()) / len(usage_per_gpu)
RAW Paste Data