Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
57 lines (46 sloc) 1.6 KB
### Author: Vasilis Vryniotis <bbriniotis@datumbox.com>
class dstat_plugin(dstat):
"""
Total GPU usage for NVIDIA cards. Requires the nvidia-ml-py package.
Usage: dstat --nvidia-gpu
"""
def __init__(self):
self.name = 'gpu usage nvidia'
self.type = 'p'
self.width = 5
self.scale = 34
self.cols = 1
self.samples = 10
def check(self):
try:
import pynvml
pynvml.nvmlInit()
except:
raise Exception('The "pynvml" library is missing from this system.')
def vars(self):
ret = ['total']
if op.full:
import pynvml
deviceCount = pynvml.nvmlDeviceGetCount()
for i in range(0, deviceCount):
ret.append('gpu%d' % i)
return ret
def extract(self):
stats = self._getUsagePerGPU(self.samples)
stats['total'] = self._getTotalUsage(stats)
for name in self.vars:
self.val[name] = stats[name]
def _getUsagePerGPU(self, samples):
import pynvml
usage = {}
deviceCount = pynvml.nvmlDeviceGetCount()
for iter in range(0, samples):
for i in range(0, deviceCount):
name = 'gpu%d' % i
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
if name not in usage:
usage[name] = 0.0
usage[name] += pynvml.nvmlDeviceGetUtilizationRates(handle).gpu / float(samples)
return usage
def _getTotalUsage(self, usage_per_gpu):
return sum(usage_per_gpu.values()) / len(usage_per_gpu)