Custom Diamond Collectors

CPUTempCollector.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Custom collector for

Requires PySensors.

:copyright:
    Lion Krischer, 2017
:license:
    MIT License (https://opensource.org/licenses/MIT)
"""
import diamond.collector
import sensors


class CPUTempCollector(diamond.collector.Collector):
    def collect(self):
        sensors.init()

        for _s in sensors.iter_detected_chips():
            if _s.prefix == "coretemp":
                for _i in _s:
                    if (_i.label == "Physical id 0"):
                        self.publish("cputemp", _i.get_value())
                        break
                break
        sensors.cleanup()

NvidiaGPUCollector.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Custom diamond collector for Nvidia GPUs.

Requires nvidia-smi and nvidia-ml-py.

:copyright:
    Lion Krischer, 2017
:license:
    MIT License (https://opensource.org/licenses/MIT)
"""
import diamond.collector
import pynvml as nv


class NvidiaGPUCollector(diamond.collector.Collector):
    def __init__(self, *args, **kwargs):
        nv.nvmlInit()
        try:
            self._gpu_count = nv.nvmlDeviceGetCount()
            self._handles = [nv.nvmlDeviceGetHandleByIndex(_i)
                             for _i in range(self._gpu_count)]
        finally:
            nv.nvmlShutdown()

        diamond.collector.Collector.__init__(self, *args, **kwargs)

    def collect(self):
        # For some reason this must also be initialized here.
        nv.nvmlInit()
        try:
            self._handles = [nv.nvmlDeviceGetHandleByIndex(_i)
                             for _i in range(self._gpu_count)]
            info = {}
            for _i, handle in enumerate(self._handles):
                mem = nv.nvmlDeviceGetMemoryInfo(handle)
                info["memory.free"] = mem.free
                info["memory.used"] = mem.used
                info["memory.total"] = mem.total

                info["clocks.graphics_current"] = \
                    nv.nvmlDeviceGetClockInfo(handle, 0)
                info["clocks.sm_current"] = \
                    nv.nvmlDeviceGetClockInfo(handle, 1)
                info["clocks.memory_current"] = \
                    nv.nvmlDeviceGetClockInfo(handle, 2)
                info["clocks.video_current"] = \
                    nv.nvmlDeviceGetClockInfo(handle, 3)

                info["clocks.graphics_max"] = \
                    nv.nvmlDeviceGetMaxClockInfo(handle, 0)
                info["clocks.sm_max"] = \
                    nv.nvmlDeviceGetMaxClockInfo(handle, 1)
                info["clocks.memory_max"] = \
                    nv.nvmlDeviceGetMaxClockInfo(handle, 2)
                info["clocks.video_max"] = \
                    nv.nvmlDeviceGetMaxClockInfo(handle, 3)

                info["temperature.current"] = \
                    nv.nvmlDeviceGetTemperature(handle, 0)

                info["power.current"] = nv.nvmlDeviceGetPowerUsage(handle)

                # Publish everything with the chosen prefix.
                prefix = "gpu%i" % _i
                for key, value in info.items():
                    self.publish("%s.%s" % (prefix, key), value)
        finally:
            nv.nvmlShutdown()