kvcache-ai-ktransformers/csrc/balance_serve/sched/metrics.cpp

#include "metrics.h"
#include <iostream>

// 构造函数
Metrics::Metrics(const MetricsConfig& config)
    : registry_(std::make_shared<prometheus::Registry>()),
      exposer_(config.endpoint),
      stop_uptime_thread_(false),
      start_time_(std::chrono::steady_clock::now()) {
  // 定义统一的桶大小，最大为 10000 ms (10 s)
  std::vector<double> common_buckets = {0.001, 0.005, 0.01,  0.05,  0.1,    0.5,    1.0,    5.0,
                                        10.0,  50.0,  100.0, 500.0, 1000.0, 5000.0, 10000.0};  // 毫秒

  // 注册 TTFT_ms Histogram
  auto& TTFT_family = prometheus::BuildHistogram()
                          .Name(std::string(METRIC_PREFIX) + "_TTFT_ms")
                          .Help("Time to first token in milliseconds")
                          .Register(*registry_);
  TTFT_ms = &TTFT_family.Add({{"model", config.model_name}}, common_buckets);

  // 注册 TBT_ms Histogram
  auto& TBT_family = prometheus::BuildHistogram()
                         .Name(std::string(METRIC_PREFIX) + "_TBT_ms")
                         .Help("Time between tokens in milliseconds")
                         .Register(*registry_);
  TBT_ms = &TBT_family.Add({{"model", config.model_name}}, common_buckets);

  // 注册 schedule_time Histogram
  auto& schedule_time_family = prometheus::BuildHistogram()
                                   .Name(std::string(METRIC_PREFIX) + "_schedule_time_ms")
                                   .Help("Time to generate schedule in milliseconds")
                                   .Register(*registry_);
  schedule_time = &schedule_time_family.Add({{"model", config.model_name}}, common_buckets);

  // 注册 generated_tokens Counter
  auto& generated_tokens_family = prometheus::BuildCounter()
                                      .Name(std::string(METRIC_PREFIX) + "_generated_tokens_total")
                                      .Help("Total generated tokens")
                                      .Register(*registry_);
  generated_tokens = &generated_tokens_family.Add({{"model", config.model_name}});

  // 注册 throughput_query Gauge
  auto& throughput_query_family = prometheus::BuildGauge()
                                      .Name(std::string(METRIC_PREFIX) + "_throughput_query")
                                      .Help("Throughput per second based on queries")
                                      .Register(*registry_);
  throughput_query = &throughput_query_family.Add({{"model", config.model_name}});

  // 注册 throughput_generated_tokens Gauge
  auto& throughput_generated_tokens_family = prometheus::BuildGauge()
                                                 .Name(std::string(METRIC_PREFIX) + "_throughput_generated_tokens")
                                                 .Help("Throughput per second based on generated tokens")
                                                 .Register(*registry_);
  throughput_generated_tokens = &throughput_generated_tokens_family.Add({{"model", config.model_name}});

  // 注册 event_count Counter family
  event_count_family_ = &prometheus::BuildCounter()
                             .Name(std::string(METRIC_PREFIX) + "_event_count_total")
                             .Help("Count of various events")
                             .Register(*registry_);

  batch_count_family_ = &prometheus::BuildCounter()
                             .Name(std::string(METRIC_PREFIX) + "_batch_count_total")
                             .Help("Count of various batch by status")
                             .Register(*registry_);

  // 注册 query_count Counter family
  query_count_family_ = &prometheus::BuildCounter()
                             .Name(std::string(METRIC_PREFIX) + "_query_count_total")
                             .Help("Count of queries by status")
                             .Register(*registry_);

  // 注册 uptime_ms Gauge
  auto& uptime_family = prometheus::BuildGauge()
                            .Name(std::string(METRIC_PREFIX) + "_uptime_ms")
                            .Help("Uptime of the scheduler in milliseconds")
                            .Register(*registry_);
  uptime_ms = &uptime_family.Add({{"model", config.model_name}});

  // 注册 GPU 利用率 Gauges
  auto& gpu_util_family = prometheus::BuildGauge()
                              .Name(std::string(METRIC_PREFIX) + "_gpu_utilization_ratio")
                              .Help("Current GPU utilization ratio (0 to 1)")
                              .Register(*registry_);
  for (size_t i = 0; i < config.gpu_count; ++i) {
    gpu_utilization_gauges.push_back(
        &gpu_util_family.Add({{"gpu_id", std::to_string(i)}, {"model", config.model_name}}));
  }

  // 将 Registry 注册到 Exposer 中
  exposer_.RegisterCollectable(registry_);

  // 启动 uptime 更新线程
  StartUptimeUpdater();
}

// 析构函数
Metrics::~Metrics() {
  StopUptimeUpdater();
}

// 启动 uptime 更新线程
void Metrics::StartUptimeUpdater() {
  uptime_thread_ = std::thread([this]() {
    while (!stop_uptime_thread_) {
      auto now = std::chrono::steady_clock::now();
      std::chrono::duration<double, std::milli> uptime_duration = now - start_time_;
      uptime_ms->Set(uptime_duration.count());
      // fn_every_sec(this);
      std::this_thread::sleep_for(std::chrono::seconds(1));
    }
  });
}

// 停止 uptime 更新线程
void Metrics::StopUptimeUpdater() {
  stop_uptime_thread_ = true;
  if (uptime_thread_.joinable()) {
    uptime_thread_.join();
  }
}

// 获取 event_count 指标
prometheus::Counter* Metrics::event_count(const std::string& type) {
  return &event_count_family_->Add({{"type", type}});  // 可根据需要添加更多标签
}

// 获取 query_count 指标
prometheus::Counter* Metrics::query_count(const std::string& status) {
  return &query_count_family_->Add({{"status", status}});  // 可根据需要添加更多标签
}

prometheus::Counter* Metrics::batch_count(const std::string& type) {
  return &batch_count_family_->Add({{"type", type}});
}