Agent Metric Config

.github/workflows/agent_ci.yml

@@ -4,12 +4,12 @@ on:
   push:
     branches: [ "main", "devel" ]
     paths:
-      - 'sources/agent/'
+      - 'sources/agent/**'
       - '.github/workflows/agent_ci.yml'
   pull_request:
     branches: [ "main", "devel" ]
     paths:
-      - 'sources/agent/'
+      - 'sources/agent/**'
       - '.github/workflows/agent_ci.yml'
 
 jobs:
@@ -62,4 +62,5 @@ jobs:
         --inline-suppr \
         --force \
         -i sources/agent/build \
-        sources/agent
+        -i sources/agent/gen \
+        sources/agent

sources/agent/.gitignore

@@ -0,0 +1,2 @@
+build/*
+gen/*

sources/agent/CMakeLists.txt

@@ -20,19 +20,21 @@ find_package(Threads REQUIRED)
 # --- Proto / gRPC code generation ---
 set(PROTO_DIR ${CMAKE_SOURCE_DIR}/../proto)
 set(PROTO_FILE ${PROTO_DIR}/volta.proto)
+set(GENERATED_DIR ${CMAKE_SOURCE_DIR}/gen)
 
 get_target_property(GRPC_CPP_PLUGIN gRPC::grpc_cpp_plugin LOCATION)
 
-set(PROTO_SRC "${CMAKE_CURRENT_BINARY_DIR}/volta.pb.cc")
-set(PROTO_HDR "${CMAKE_CURRENT_BINARY_DIR}/volta.pb.h")
-set(GRPC_SRC  "${CMAKE_CURRENT_BINARY_DIR}/volta.grpc.pb.cc")
-set(GRPC_HDR  "${CMAKE_CURRENT_BINARY_DIR}/volta.grpc.pb.h")
+set(PROTO_SRC "${GENERATED_DIR}/volta.pb.cc")
+set(PROTO_HDR "${GENERATED_DIR}/volta.pb.h")
+set(GRPC_SRC  "${GENERATED_DIR}/volta.grpc.pb.cc")
+set(GRPC_HDR  "${GENERATED_DIR}/volta.grpc.pb.h")
 
 add_custom_command(
     OUTPUT  "${PROTO_SRC}" "${PROTO_HDR}" "${GRPC_SRC}" "${GRPC_HDR}"
+    COMMAND ${CMAKE_COMMAND} -E make_directory ${GENERATED_DIR}
     COMMAND protobuf::protoc
-    ARGS    --cpp_out="${CMAKE_CURRENT_BINARY_DIR}"
-            --grpc_out="${CMAKE_CURRENT_BINARY_DIR}"
+    ARGS    --cpp_out="${GENERATED_DIR}"
+            --grpc_out="${GENERATED_DIR}"
             --plugin=protoc-gen-grpc="${GRPC_CPP_PLUGIN}"
             -I "${PROTO_DIR}"
             "${PROTO_FILE}"
@@ -42,7 +44,7 @@ add_custom_command(
 
 add_library(volta_proto STATIC ${PROTO_SRC} ${PROTO_HDR} ${GRPC_SRC} ${GRPC_HDR})
 
-target_include_directories(volta_proto PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
+target_include_directories(volta_proto PUBLIC ${GENERATED_DIR})
 
 target_link_libraries(volta_proto PUBLIC
     protobuf::libprotobuf
@@ -164,4 +166,4 @@ if(GIT_EXE)
              PERMISSIONS OWNER_READ OWNER_WRITE OWNER_EXECUTE)
         message(STATUS "Git pre-commit hook installed")
     endif()
-endif()
+endif()

sources/agent/agent.example.conf

@@ -1 +1,7 @@
 core_affinity = [2, "3-10"]
+
+# Collect CPU and GPU power metrics
+metrics = [
+  "METRIC_TYPE_CPU_POWER_PACKAGE",
+  "METRIC_TYPE_GPU_POWER",
+]

sources/agent/src/collectors/collector.h

@@ -4,6 +4,7 @@
 #include <vector>
 
 #include "metric.h"
+#include "platform/hardware_info.h"
 
 namespace volta {
 namespace agent {
@@ -15,9 +16,66 @@ class Collector {
 
   virtual std::vector<Metric> Collect() = 0;
 
+  virtual bool IsSupported() = 0;
+
+  virtual std::vector<v1::MetricType> Satisfiable() = 0;
+
+  virtual void SetRequestedMetrics(
+      const std::vector<v1::MetricType>& metrics) = 0;
+
   virtual bool Init() { return true; }
 };
 
+class CollectorRegistry {
+ public:
+  static CollectorRegistry& Instance() {
+    static CollectorRegistry instance;
+    return instance;
+  }
+
+  void Register(std::unique_ptr<Collector> collector) {
+    entries_.push_back(std::move(collector));
+  }
+
+  std::vector<Collector*> Resolve(const std::vector<v1::MetricType>& desired) {
+    std::vector<Collector*> result;
+    for (auto& collector : entries_) {
+      if (!collector->IsSupported()) continue;
+
+      std::vector<v1::MetricType> active;
+      for (const auto& type : collector->Satisfiable()) {
+        if (std::find(desired.begin(), desired.end(), type) != desired.end()) {
+          active.push_back(type);
+        }
+      }
+
+      if (!active.empty()) {
+        collector->SetRequestedMetrics(active);
+        result.push_back(collector.get());
+      }
+    }
+    return result;
+  }
+
+ private:
+  CollectorRegistry() = default;
+  std::vector<std::unique_ptr<Collector>> entries_;
+};
+
+template <typename Derived>
+class RegisteredCollector : public Collector {
+  static bool Register() {
+    std::unique_ptr<Collector> c = std::make_unique<Derived>();
+    CollectorRegistry::Instance().Register(std::move(c));
+    return true;
+  }
+  static inline bool registered_ = Register();
+
+  static constexpr std::integral_constant<const bool*, &registered_> force_{};
+
+ public:
+};
+
 }  // namespace collectors
 }  // namespace agent
 }  // namespace volta

sources/agent/src/collectors/nvml_collector.cc

@@ -1,15 +1,14 @@
 #include "collectors/nvml_collector.h"
 
+#include <algorithm>
 #include <chrono>
 #include <iostream>
 
 namespace volta {
 namespace agent {
 namespace collectors {
 
-NvmlCollector::NvmlCollector() {
-  // Initialization moved to Init()
-}
+NvmlCollector::NvmlCollector() = default;
 
 NvmlCollector::~NvmlCollector() {
   if (initialized_) {
@@ -29,61 +28,117 @@ bool NvmlCollector::Init() {
   if (result != NVML_SUCCESS) {
     std::cerr << "Failed to get device handle: " << nvmlErrorString(result)
               << std::endl;
+    nvmlShutdown();
     return false;
   }
 
   initialized_ = true;
   return true;
 }
 
+bool NvmlCollector::IsSupported() {
+  nvmlReturn_t result = nvmlInit();
+  if (result != NVML_SUCCESS) {
+    return false;
+  }
+
+  unsigned int device_count = 0;
+  result = nvmlDeviceGetCount(&device_count);
+  if (result != NVML_SUCCESS || device_count == 0) {
+    nvmlShutdown();
+    return false;
+  }
+
+  nvmlDevice_t device;
+  result = nvmlDeviceGetHandleByIndex(0, &device);
+  nvmlShutdown();
+  return result == NVML_SUCCESS;
+}
+
+void NvmlCollector::SetRequestedMetrics(
+    const std::vector<v1::MetricType>& metrics) {
+  requested_metrics_ = metrics;
+}
+
 std::vector<Metric> NvmlCollector::Collect() {
-  if (!initialized_) return {};
+  if (!initialized_ || requested_metrics_.empty()) return {};
 
   std::vector<Metric> metrics;
   unsigned int power_mw = 0;
   auto now = std::chrono::system_clock::now().time_since_epoch().count();
 
-  nvmlReturn_t result = nvmlDeviceGetPowerUsage(device_handle_, &power_mw);
-
-  if (result == NVML_SUCCESS) {
-    metrics.push_back({"gpu_0_power_watts",
-                       static_cast<double>(power_mw) / 1000.0,  // mW -> W
-                       now});
+  auto needs = [&](v1::MetricType type) {
+    return std::find(requested_metrics_.begin(), requested_metrics_.end(),
+                     type) != requested_metrics_.end();
+  };
+
+  nvmlReturn_t result;
+  if (needs(v1::MetricType::METRIC_TYPE_GPU_POWER)) {
+    result = nvmlDeviceGetPowerUsage(device_handle_, &power_mw);
+    if (result == NVML_SUCCESS) {
+      metrics.push_back({v1::MetricType::METRIC_TYPE_GPU_POWER,
+                         {.index = 0},
+                         static_cast<double>(power_mw) / 1000.0,
+                         now});
+    }
   }
 
-  unsigned int temp_c = 0;
-  result =
-      nvmlDeviceGetTemperature(device_handle_, NVML_TEMPERATURE_GPU, &temp_c);
-  if (result == NVML_SUCCESS) {
-    metrics.push_back({"gpu_0_temp_celsius", static_cast<double>(temp_c), now});
+  if (needs(v1::MetricType::METRIC_TYPE_GPU_TEMPERATURE)) {
+    unsigned int temp_c = 0;
+    result =
+        nvmlDeviceGetTemperature(device_handle_, NVML_TEMPERATURE_GPU, &temp_c);
+    if (result == NVML_SUCCESS) {
+      metrics.push_back({v1::MetricType::METRIC_TYPE_GPU_TEMPERATURE,
+                         {.index = 0},
+                         static_cast<double>(temp_c),
+                         now});
+    }
   }
 
-  nvmlUtilization_t utilization;
-  result = nvmlDeviceGetUtilizationRates(device_handle_, &utilization);
-
-  if (result == NVML_SUCCESS) {
-    metrics.push_back({"gpu_0_utilization_percent",
-                       static_cast<double>(utilization.gpu), now});
-
-    metrics.push_back({"gpu_0_memory_activity_percent",
-                       static_cast<double>(utilization.memory), now});
+  if (needs(v1::MetricType::METRIC_TYPE_GPU_UTILIZATION) ||
+      needs(v1::MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION)) {
+    nvmlUtilization_t utilization;
+    result = nvmlDeviceGetUtilizationRates(device_handle_, &utilization);
+    if (result == NVML_SUCCESS) {
+      if (needs(v1::MetricType::METRIC_TYPE_GPU_UTILIZATION)) {
+        metrics.push_back({v1::MetricType::METRIC_TYPE_GPU_UTILIZATION,
+                           {.index = 0},
+                           static_cast<double>(utilization.gpu),
+                           now});
+      }
+      if (needs(v1::MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION)) {
+        metrics.push_back(
+            {v1::MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION,
+             {.index = 0},
+             static_cast<double>(utilization.memory),
+             now});
+      }
+    }
   }
 
-  nvmlMemory_t memory;
-  result = nvmlDeviceGetMemoryInfo(device_handle_, &memory);
-
-  if (result == NVML_SUCCESS) {
-    metrics.push_back(
-        {"gpu_0_memory_total_bytes", static_cast<double>(memory.total), now});
-    metrics.push_back(
-        {"gpu_0_memory_used_bytes", static_cast<double>(memory.used), now});
-    double used_percent = (double)memory.used / memory.total * 100.0;
-    metrics.push_back({"gpu_0_memory_used_percent", used_percent, now});
+  if (needs(v1::MetricType::METRIC_TYPE_GPU_VRAM_USED)) {
+    nvmlMemory_t memory;
+    result = nvmlDeviceGetMemoryInfo(device_handle_, &memory);
+    if (result == NVML_SUCCESS) {
+      metrics.push_back(
+          {v1::MetricType::METRIC_TYPE_GPU_VRAM_USED,
+           {.index = 0},
+           static_cast<double>(memory.used) / static_cast<double>(memory.total),
+           now});
+    }
   }
 
   return metrics;
 }
 
+std::vector<v1::MetricType> NvmlCollector::Satisfiable() {
+  return {v1::MetricType::METRIC_TYPE_GPU_VRAM_USED,
+          v1::MetricType::METRIC_TYPE_GPU_UTILIZATION,
+          v1::MetricType::METRIC_TYPE_GPU_SHARED_MEMORY_UTILIZATION,
+          v1::MetricType::METRIC_TYPE_GPU_TEMPERATURE,
+          v1::MetricType::METRIC_TYPE_GPU_POWER};
+}
+
 }  // namespace collectors
 }  // namespace agent
 }  // namespace volta

sources/agent/src/collectors/nvml_collector.h

@@ -3,21 +3,31 @@
 
 #include <nvml.h>
 
+#include <vector>
+
 #include "collectors/collector.h"
 
 namespace volta {
 namespace agent {
 namespace collectors {
 
-class NvmlCollector : public Collector {
+class NvmlCollector : public RegisteredCollector<NvmlCollector> {
  public:
   NvmlCollector();
   ~NvmlCollector() override;
 
+  NvmlCollector(NvmlCollector&&) = default;
+  NvmlCollector& operator=(NvmlCollector&&) = default;
+
   bool Init() override;
   std::vector<Metric> Collect() override;
+  bool IsSupported() override;
+  std::vector<v1::MetricType> Satisfiable() override;
+  void SetRequestedMetrics(const std::vector<v1::MetricType>& metrics) override;
 
  private:
+  std::vector<v1::MetricType> requested_metrics_;
+
   nvmlDevice_t device_handle_;
   bool initialized_ = false;
 };