src/gpu_nvml.c

   1 #include "daemon/collectd.h"
   2 #include "daemon/common.h"
   3 #include "daemon/plugin.h"
   4
   5 #include <nvml.h>
   6 #include <stdint.h>
   7 #include <stdio.h>
   8
   9 #define MAX_DEVNAME_LEN 256
  10 #define PLUGIN_NAME "gpu_nvml"
  11
  12 static nvmlReturn_t nv_status = NVML_SUCCESS;
  13 static char *nv_errline = "";
  14
  15 #define TRY_CATCH(f, catch)                                                    \
  16   if ((nv_status = f) != NVML_SUCCESS) {                                       \
  17     nv_errline = #f;                                                           \
  18     goto catch;                                                                \
  19   }
  20 #define TRY(f) TRY_CATCH(f, catch)
  21 #define WRAPGAUGE(x) ((value_t){.gauge = (gauge_t)(x)})
  22
  23 static const char *config_keys[] = {
  24     "GPUIndex",
  25     "IgnoreSelected",
  26 };
  27 static const unsigned int n_config_keys = STATIC_ARRAY_SIZE(config_keys);
  28
  29 static uint64_t conf_match_mask = 0;
  30 static bool conf_mask_is_exclude = 0;
  31
  32 static int nvml_config(const char *key, const char *value) {
  33
  34   unsigned long device_ix;
  35   char *eptr;
  36
  37   if (strcasecmp(key, config_keys[0]) == 0) {
  38     device_ix = strtoul(value, &eptr, 10);
  39     if (eptr == value) {
  40       return -1;
  41     }
  42     if (device_ix > 64) {
  43       return -2;
  44     }
  45     conf_match_mask |= (1 << device_ix);
  46   } else if (strcasecmp(key, config_keys[1])) {
  47     if
  48       IS_TRUE(value) { conf_mask_is_exclude = 1; }
  49   } else {
  50     return -10;
  51   }
  52
  53   return 0;
  54 }
  55
  56 static int nvml_init(void) {
  57   TRY(nvmlInit());
  58   return 0;
  59
  60   catch : ERROR("NVML init failed with %d", nv_status);
  61   return -1;
  62 }
  63
  64 static int nvml_shutdown(void) {
  65   TRY(nvmlShutdown())
  66   return 0;
  67
  68   catch : ERROR("NVML shutdown failed with %d", nv_status);
  69   return -1;
  70 }
  71
  72 static void nvml_submit(const char *plugin_instance, const char *type,
  73                         const char *type_instance, value_t nvml) {
  74
  75   value_list_t vl = VALUE_LIST_INIT;
  76
  77   vl.values = &nvml;
  78   vl.values_len = 1;
  79
  80   sstrncpy(vl.plugin, PLUGIN_NAME, sizeof(vl.plugin));
  81   sstrncpy(vl.plugin_instance, plugin_instance, sizeof(vl.plugin_instance));
  82
  83   sstrncpy(vl.type, type, sizeof(vl.type));
  84
  85   if (type_instance != NULL) {
  86     sstrncpy(vl.type_instance, type_instance, sizeof(vl.type_instance));
  87   }
  88
  89   plugin_dispatch_values(&vl);
  90 }
  91
  92 static int nvml_read(void) {
  93
  94   unsigned int device_count;
  95   TRY_CATCH(nvmlDeviceGetCount(&device_count), catch_nocount);
  96
  97   if (device_count > 64) {
  98     device_count = 64;
  99   }
 100
 101   nvmlDevice_t dev;
 102   char dev_name[MAX_DEVNAME_LEN + 1];
 103   unsigned int fan_speed;
 104   nvmlUtilization_t utilization;
 105   nvmlMemory_t meminfo;
 106   unsigned int core_temp;
 107
 108   for (int ix = 0; ix < device_count; ix++) {
 109
 110     int is_match = ((1 << ix) & conf_match_mask) || (conf_match_mask == 0);
 111     if (conf_mask_is_exclude == !!is_match) {
 112       continue;
 113     }
 114
 115     TRY(nvmlDeviceGetHandleByIndex(ix, &dev));
 116
 117     dev_name[0] = '\0';
 118     TRY(nvmlDeviceGetName(dev, dev_name, MAX_DEVNAME_LEN));
 119
 120     TRY(nvmlDeviceGetMemoryInfo(dev, &meminfo))
 121     TRY(nvmlDeviceGetUtilizationRates(dev, &utilization))
 122     TRY(nvmlDeviceGetFanSpeed(dev, &fan_speed))
 123     TRY(nvmlDeviceGetTemperature(dev, NVML_TEMPERATURE_GPU, &core_temp))
 124
 125     double pct_mem_used = 100. * (double)meminfo.used / meminfo.total;
 126
 127     nvml_submit(dev_name, "percent", "GPU", WRAPGAUGE(pct_mem_used));
 128     nvml_submit(dev_name, "percent", "GPU", WRAPGAUGE(utilization.gpu));
 129     nvml_submit(dev_name, "fanspeed", "GPU", WRAPGAUGE(fan_speed));
 130     nvml_submit(dev_name, "temperature", "GPU", WRAPGAUGE(core_temp));
 131     continue;
 132
 133     catch : WARNING("NVML call \"%s\" failed with code %d!", nv_errline,
 134                     nv_status);
 135     continue;
 136   }
 137
 138   return 0;
 139
 140 catch_nocount:
 141   ERROR("Failed to enumerate NVIDIA GPUs (\"%s\" returned %d)", nv_errline,
 142         nv_status);
 143   return -1;
 144 }
 145
 146 void module_register(void) {
 147   plugin_register_init(PLUGIN_NAME, nvml_init);
 148   plugin_register_config(PLUGIN_NAME, nvml_config, config_keys, n_config_keys);
 149   plugin_register_read(PLUGIN_NAME, nvml_read);
 150   plugin_register_shutdown(PLUGIN_NAME, nvml_shutdown);
 151 }