NVIDIA GPU textfile exporter

node_gpu_exporter.py

#!/usr/bin/env python3

import time
import atexit
from collections import OrderedDict

from py3nvml import py3nvml as nv

from prometheus_client import Gauge, CollectorRegistry
from prometheus_client import write_to_textfile


class GPUCollector(object):
    def __init__(self, write_file):
        self.labels = ['gpu', 'name', 'driver']
        self.driver = nv.nvmlSystemGetDriverVersion()
        self.registry = CollectorRegistry()
        self.write_file = write_file

        self.n_gpu = nv.nvmlDeviceGetCount()
        self.hnds = [nv.nvmlDeviceGetHandleByIndex(i) for i in
                     range(self.n_gpu)]
        self.args = []
        for i, hnd in enumerate(self.hnds):
            args = OrderedDict()
            args['gpu'] = 'gpu%d' % i
            args['name'] = nv.nvmlDeviceGetName(hnd)
            args['driver'] = self.driver
            self.args.append(args)

        self.gpu_memused = Gauge('node_gpu_memory_bytes_used',     # bytes
                                 'Used GPU Memory', self.labels,
                                 registry=self.registry)
        self.gpu_memtotal = Gauge('node_gpu_memory_bytes_total',   # bytes
                                  'Total GPU Memory', self.labels,
                                  registry=self.registry)
        self.gpu_powertotal = Gauge('node_gpu_power_watts',   # watts
                                    'GPU Power Utilization', self.labels,
                                    registry=self.registry)
        self.gpu_temp = Gauge('node_gpu_temp_celsius',
                              'GPU Temperature', self.labels,
                              registry=self.registry)
        self.gpu_utilpct = Gauge('node_gpu_usage_ratio',   # percent
                                 'GPU Usage Percentage', self.labels,
                                 registry=self.registry)

    def temperature(self, hnd):
        try:
            return nv.nvmlDeviceGetTemperature(hnd, nv.NVML_TEMPERATURE_GPU)
        except nv.NVMLError_NotSupported:
            return -1

    def usage_ratio(self, hnd):
        return nv.nvmlDeviceGetUtilizationRates(hnd).gpu

    def mem_info(self, hnd):
        return nv.nvmlDeviceGetMemoryInfo(hnd)

    def power_usage(self, hnd):
        return nv.nvmlDeviceGetPowerUsage(hnd) / 1000

    def update(self):
        for hnd, args in zip(self.hnds, self.args):
            mem = self.mem_info(hnd)
            self.gpu_memused.labels(**args).set(mem.used)
            self.gpu_memtotal.labels(**args).set(mem.total)
            self.gpu_powertotal.labels(**args).set(self.power_usage(hnd))
            self.gpu_utilpct.labels(**args).set(self.usage_ratio(hnd))
            self.gpu_temp.labels(**args).set(self.temperature(hnd))
            write_to_textfile(self.write_file, self.registry)


if __name__ == "__main__":
    nv.nvmlInit()
    atexit.register(nv.nvmlShutdown)
    gc = GPUCollector('/var/lib/node_exporter/textfile_collector/gpu.prom')

    while True:
        gc.update()
        time.sleep(15)