Update C API

triton-inference-server · Aug 13, 2024 · 58eed78 · 58eed78
1 parent c1c4883
commit 58eed78
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git a/src/metric.cc b/src/metric.cc
@@ -280,15 +280,35 @@ Metric::InitializeTritonMetric()
 {
   std::vector<const TRITONSERVER_Parameter*> labels_params;
   ParseLabels(labels_params, labels_);
+  TRITONSERVER_MetricKind kind;
+  THROW_IF_TRITON_ERROR(TRITONSERVER_GetMetricFamilyKind(
+      reinterpret_cast<TRITONSERVER_MetricFamily*>(metric_family_address_),
+      &kind));
+  TRITONSERVER_MetricArgs* args = nullptr;
+  switch (kind) {
+    case TRITONSERVER_METRIC_KIND_COUNTER:
+    case TRITONSERVER_METRIC_KIND_GAUGE:
+      break;
+    case TRITONSERVER_METRIC_KIND_HISTOGRAM: {
+      const std::vector<double>& buckets = buckets_.value();
+      THROW_IF_TRITON_ERROR(TRITONSERVER_MetricArgsNew(&args));
+      THROW_IF_TRITON_ERROR(TRITONSERVER_MetricArgsSetHistogram(
+          args, buckets.data(), buckets.size()));
+      break;
+    }
+    default:
+      break;
+  }
+
   TRITONSERVER_Metric* triton_metric = nullptr;
-  void* buckets_ptr = buckets_.has_value() ? &(buckets_.value()) : nullptr;
-  THROW_IF_TRITON_ERROR(TRITONSERVER_MetricNew(
+  THROW_IF_TRITON_ERROR(TRITONSERVER_MetricNewWithArgs(
       &triton_metric,
       reinterpret_cast<TRITONSERVER_MetricFamily*>(metric_family_address_),
-      labels_params.data(), labels_params.size(), buckets_ptr));
+      labels_params.data(), labels_params.size(), args));
   for (const auto label : labels_params) {
     TRITONSERVER_ParameterDelete(const_cast<TRITONSERVER_Parameter*>(label));
   }
+  THROW_IF_TRITON_ERROR(TRITONSERVER_MetricArgsDelete(args));
   return reinterpret_cast<void*>(triton_metric);
 }
 
@@ -346,7 +366,7 @@ void
 Metric::Observe(const double& value)
 {
   auto triton_metric = reinterpret_cast<TRITONSERVER_Metric*>(metric_address_);
-  THROW_IF_TRITON_ERROR(TRITONSERVER_MetricObserve(triton_metric, value));
+  THROW_IF_TRITON_ERROR(TRITONSERVER_MetricSet(triton_metric, value));
 }
 
 double

diff --git a/src/metric_family.h b/src/metric_family.h
@@ -131,8 +131,8 @@ class MetricFamily {
   std::string name_;
   // The description of the metric family.
   std::string description_;
-  // The metric kind of the metric family. Currently only supports GAUGE and
-  // COUNTER.
+  // The metric kind of the metric family. Currently only supports GAUGE,
+  // COUNTER and HISTOGRAM.
   MetricKind kind_;
   // The address of the TRITONSERVER_MetricFamily object.
   void* metric_family_address_;