Collector Fixes

sources/agent/src/collectors/nvml_collector.cc

@@ -3,6 +3,7 @@
 #include <algorithm>
 #include <chrono>
 #include <iostream>
+#include <utility>
 
 namespace volta {
 namespace agent {
@@ -31,30 +32,47 @@ bool NvmlCollector::Init() {
     return false;
   }
 
-  result = nvml->DeviceGetHandleByIndex(0, &device_handle_);
-  if (result != NVML_SUCCESS) {
-    std::cerr << "Failed to get device handle: " << nvml->ErrorString(result)
-              << std::endl;
+  unsigned int device_count = 0;
+  result = nvml->DeviceGetCount(&device_count);
+  if (result != NVML_SUCCESS || device_count == 0) {
+    std::cerr << "Failed to get NVML device count: "
+              << nvml->ErrorString(result) << std::endl;
     nvml->Shutdown();
     return false;
   }
 
-  nvmlPciInfo_t pci_info;
-  result = nvml->DeviceGetPciInfo(device_handle_, &pci_info);
-  if (result != NVML_SUCCESS) {
-    std::cerr << "Failed to get PCI info: " << nvml->ErrorString(result)
-              << std::endl;
+  devices_.clear();
+  for (unsigned int i = 0; i < device_count; ++i) {
+    nvmlDevice_t device_handle;
+    result = nvml->DeviceGetHandleByIndex(i, &device_handle);
+    if (result != NVML_SUCCESS) {
+      std::cerr << "Failed to get device handle for GPU " << i << ": "
+                << nvml->ErrorString(result) << std::endl;
+      continue;
+    }
+
+    nvmlPciInfo_t pci_info;
+    result = nvml->DeviceGetPciInfo(device_handle, &pci_info);
+    if (result != NVML_SUCCESS) {
+      std::cerr << "Failed to get PCI info for GPU " << i << ": "
+                << nvml->ErrorString(result) << std::endl;
+      continue;
+    }
+
+    GpuID gpu_id;
+    gpu_id.set_pci_domain(pci_info.domain);
+    gpu_id.set_pci_bus(pci_info.bus);
+    gpu_id.set_pci_device(pci_info.device);
+    gpu_id.set_pci_function(0);
+    devices_.push_back({device_handle, std::move(gpu_id)});
+  }
+
+  if (devices_.empty()) {
+    std::cerr << "Failed to discover any NVML GPUs" << std::endl;
     nvml->Shutdown();
     return false;
   }
 
-  GpuID gpu_id;
-  gpu_id.set_pci_domain(pci_info.domain);
-  gpu_id.set_pci_bus(pci_info.bus);
-  gpu_id.set_pci_device(pci_info.device);
-  gpu_id.set_pci_function(0);
-  gpu_id_ = std::move(gpu_id);
-
   nvml_ = nvml;
   initialized_ = true;
   return true;
@@ -95,57 +113,61 @@ std::vector<Metric> NvmlCollector::Collect() {
   const auto& nvml = *nvml_;
   std::vector<Metric> metrics;
   unsigned int power_mw = 0;
-  auto now = std::chrono::system_clock::now().time_since_epoch().count();
+  auto now = std::chrono::duration_cast<std::chrono::nanoseconds>(
+                 std::chrono::system_clock::now().time_since_epoch())
+                 .count();
 
   auto needs = [&](MetricType type) {
     return std::find(requested_metrics_.begin(), requested_metrics_.end(),
                      type) != requested_metrics_.end();
   };
 
   nvmlReturn_t result;
-  if (needs(MetricType::METRIC_TYPE_GPU_POWER)) {
-    result = nvml.DeviceGetPowerUsage(device_handle_, &power_mw);
-    if (result == NVML_SUCCESS) {
-      metrics.push_back({MetricType::METRIC_TYPE_GPU_POWER, gpu_id_,
-                         static_cast<double>(power_mw) / 1000.0, now});
+  for (const auto& device : devices_) {
+    if (needs(MetricType::METRIC_TYPE_GPU_POWER)) {
+      result = nvml.DeviceGetPowerUsage(device.handle, &power_mw);
+      if (result == NVML_SUCCESS) {
+        metrics.push_back({MetricType::METRIC_TYPE_GPU_POWER, device.id,
+                           static_cast<double>(power_mw) / 1000.0, now});
+      }
     }
-  }
 
-  if (needs(MetricType::METRIC_TYPE_GPU_TEMPERATURE)) {
-    unsigned int temp_c = 0;
-    result = nvml.DeviceGetTemperature(device_handle_, NVML_TEMPERATURE_GPU,
-                                       &temp_c);
-    if (result == NVML_SUCCESS) {
-      metrics.push_back({MetricType::METRIC_TYPE_GPU_TEMPERATURE, gpu_id_,
-                         static_cast<double>(temp_c), now});
+    if (needs(MetricType::METRIC_TYPE_GPU_TEMPERATURE)) {
+      unsigned int temp_c = 0;
+      result = nvml.DeviceGetTemperature(device.handle, NVML_TEMPERATURE_GPU,
+                                         &temp_c);
+      if (result == NVML_SUCCESS) {
+        metrics.push_back({MetricType::METRIC_TYPE_GPU_TEMPERATURE, device.id,
+                           static_cast<double>(temp_c), now});
+      }
     }
-  }
 
-  if (needs(MetricType::METRIC_TYPE_GPU_UTILIZATION) ||
-      needs(MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION)) {
-    nvmlUtilization_t utilization;
-    result = nvml.DeviceGetUtilizationRates(device_handle_, &utilization);
-    if (result == NVML_SUCCESS) {
-      if (needs(MetricType::METRIC_TYPE_GPU_UTILIZATION)) {
-        metrics.push_back({MetricType::METRIC_TYPE_GPU_UTILIZATION, gpu_id_,
-                           static_cast<double>(utilization.gpu), now});
-      }
-      if (needs(MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION)) {
-        metrics.push_back(
-            {MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION, gpu_id_,
-             static_cast<double>(utilization.memory), now});
+    if (needs(MetricType::METRIC_TYPE_GPU_UTILIZATION) ||
+        needs(MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION)) {
+      nvmlUtilization_t utilization;
+      result = nvml.DeviceGetUtilizationRates(device.handle, &utilization);
+      if (result == NVML_SUCCESS) {
+        if (needs(MetricType::METRIC_TYPE_GPU_UTILIZATION)) {
+          metrics.push_back({MetricType::METRIC_TYPE_GPU_UTILIZATION, device.id,
+                             static_cast<double>(utilization.gpu), now});
+        }
+        if (needs(MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION)) {
+          metrics.push_back(
+              {MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION, device.id,
+               static_cast<double>(utilization.memory), now});
+        }
       }
     }
-  }
 
-  if (needs(MetricType::METRIC_TYPE_GPU_VRAM_USED)) {
-    nvmlMemory_t memory;
-    result = nvml.DeviceGetMemoryInfo(device_handle_, &memory);
-    if (result == NVML_SUCCESS) {
-      metrics.push_back(
-          {MetricType::METRIC_TYPE_GPU_VRAM_USED, gpu_id_,
-           static_cast<double>(memory.used) / static_cast<double>(memory.total),
-           now});
+    if (needs(MetricType::METRIC_TYPE_GPU_VRAM_USED)) {
+      nvmlMemory_t memory;
+      result = nvml.DeviceGetMemoryInfo(device.handle, &memory);
+      if (result == NVML_SUCCESS) {
+        metrics.push_back({MetricType::METRIC_TYPE_GPU_VRAM_USED, device.id,
+                           static_cast<double>(memory.used) /
+                               static_cast<double>(memory.total),
+                           now});
+      }
     }
   }
 

sources/agent/src/collectors/nvml_collector.h

@@ -25,11 +25,15 @@ class NvmlCollector final : public RegisteredCollector<NvmlCollector> {
   void SetRequestedMetrics(const std::vector<MetricType>& metrics) override;
 
  private:
+  struct DeviceInfo {
+    nvmlDevice_t handle;
+    GpuID id;
+  };
+
   std::vector<MetricType> requested_metrics_;
-  std::optional<GpuID> gpu_id_;
+  std::vector<DeviceInfo> devices_;
 
   const platform::NvmlApi* nvml_ = nullptr;
-  nvmlDevice_t device_handle_ = nullptr;
   bool initialized_ = false;
 };
 

sources/agent/src/collectors/proc_stat_collector.cc

@@ -45,7 +45,9 @@ std::vector<Metric> ProcStatCollector::Collect() {
   m.type = MetricType::METRIC_TYPE_CPU_UTILIZATION;
   m.devId = std::nullopt;
   m.value = usage_percent;
-  m.timestamp = std::chrono::system_clock::now().time_since_epoch().count();
+  m.timestamp = std::chrono::duration_cast<std::chrono::nanoseconds>(
+                    std::chrono::system_clock::now().time_since_epoch())
+                    .count();
 
   return {m};
 }

sources/agent/src/collectors/ram_collector.cc

@@ -34,7 +34,9 @@ std::vector<Metric> RamCollector::Collect() {
   uint64_t total = 0;
   ReadStats(used, total);
 
-  auto now = std::chrono::system_clock::now().time_since_epoch().count();
+  auto now = std::chrono::duration_cast<std::chrono::nanoseconds>(
+                 std::chrono::system_clock::now().time_since_epoch())
+                 .count();
   std::vector<Metric> metrics;
   if (needs_total) {
     metrics.push_back(