Skip to content

Commit 10d5641

Browse files
yinggehmc-nv
authored andcommitted
feat: Add new histogram metric type (#374)
1 parent 1393d6e commit 10d5641

File tree

9 files changed

+169
-29
lines changed

9 files changed

+169
-29
lines changed

README.md

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1656,12 +1656,12 @@ import triton_python_backend_utils as pb_utils
16561656
class TritonPythonModel:
16571657
def initialize(self, args):
16581658
# Create a MetricFamily object to report the latency of the model
1659-
# execution. The 'kind' parameter must be either 'COUNTER' or
1660-
# 'GAUGE'.
1659+
# execution. The 'kind' parameter must be either 'COUNTER',
1660+
# 'GAUGE' or 'HISTOGRAM'.
16611661
self.metric_family = pb_utils.MetricFamily(
16621662
name="preprocess_latency_ns",
16631663
description="Cumulative time spent pre-processing requests",
1664-
kind=pb_utils.MetricFamily.COUNTER # or pb_utils.MetricFamily.GAUGE
1664+
kind=pb_utils.MetricFamily.COUNTER
16651665
)
16661666

16671667
# Create a Metric object under the MetricFamily object. The 'labels'

src/ipc_message.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -63,6 +63,7 @@ typedef enum PYTHONSTUB_commandtype_enum {
6363
PYTHONSTUB_MetricRequestValue,
6464
PYTHONSTUB_MetricRequestIncrement,
6565
PYTHONSTUB_MetricRequestSet,
66+
PYTHONSTUB_MetricRequestObserve,
6667
PYTHONSTUB_LoadModelRequest,
6768
PYTHONSTUB_UnloadModelRequest,
6869
PYTHONSTUB_ModelReadinessRequest,

src/metric.cc

Lines changed: 97 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -32,9 +32,12 @@
3232

3333
namespace triton { namespace backend { namespace python {
3434

35-
Metric::Metric(const std::string& labels, void* metric_family_address)
36-
: labels_(labels), operation_value_(0), metric_address_(nullptr),
37-
metric_family_address_(metric_family_address), is_cleared_(false)
35+
Metric::Metric(
36+
const std::string& labels, std::optional<const std::vector<double>> buckets,
37+
void* metric_family_address)
38+
: labels_(labels), buckets_(buckets), operation_value_(0),
39+
metric_address_(nullptr), metric_family_address_(metric_family_address),
40+
is_cleared_(false)
3841
{
3942
#ifdef TRITON_PB_STUB
4043
SendCreateMetricRequest();
@@ -62,6 +65,20 @@ Metric::SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool)
6265
custom_metric_shm_ptr_->metric_family_address = metric_family_address_;
6366
custom_metric_shm_ptr_->metric_address = metric_address_;
6467

68+
// Histogram specific case
69+
if (buckets_.has_value()) {
70+
auto buckets_size = buckets_.value().size() * sizeof(double);
71+
std::unique_ptr<PbMemory> buckets_shm = PbMemory::Create(
72+
shm_pool, TRITONSERVER_MemoryType::TRITONSERVER_MEMORY_CPU, 0,
73+
buckets_size, reinterpret_cast<char*>(buckets_.value().data()),
74+
false /* copy_gpu */);
75+
custom_metric_shm_ptr_->buckets_shm_handle = buckets_shm->ShmHandle();
76+
buckets_shm_ = std::move(buckets_shm);
77+
} else {
78+
custom_metric_shm_ptr_->buckets_shm_handle = 0;
79+
buckets_shm_ = nullptr;
80+
}
81+
6582
// Save the references to shared memory.
6683
custom_metric_shm_ = std::move(custom_metric_shm);
6784
labels_shm_ = std::move(labels_shm);
@@ -80,17 +97,40 @@ Metric::LoadFromSharedMemory(
8097
std::unique_ptr<PbString> labels_shm = PbString::LoadFromSharedMemory(
8198
shm_pool, custom_metric_shm_ptr->labels_shm_handle);
8299

83-
return std::unique_ptr<Metric>(new Metric(custom_metric_shm, labels_shm));
100+
std::unique_ptr<PbMemory> buckets_shm = nullptr;
101+
if (custom_metric_shm_ptr->buckets_shm_handle != 0) {
102+
buckets_shm = PbMemory::LoadFromSharedMemory(
103+
shm_pool, custom_metric_shm_ptr->buckets_shm_handle,
104+
false /* open_cuda_handle */);
105+
}
106+
107+
return std::unique_ptr<Metric>(
108+
new Metric(custom_metric_shm, labels_shm, buckets_shm));
84109
}
85110

86111
Metric::Metric(
87112
AllocatedSharedMemory<MetricShm>& custom_metric_shm,
88-
std::unique_ptr<PbString>& labels_shm)
113+
std::unique_ptr<PbString>& labels_shm,
114+
std::unique_ptr<PbMemory>& buckets_shm)
89115
: custom_metric_shm_(std::move(custom_metric_shm)),
90-
labels_shm_(std::move(labels_shm))
116+
labels_shm_(std::move(labels_shm)), buckets_shm_(std::move(buckets_shm))
91117
{
92118
custom_metric_shm_ptr_ = custom_metric_shm_.data_.get();
119+
120+
// FIXME: This constructor is called during each
121+
// set/increment/observe/get_value call. It only needs the pointers.
93122
labels_ = labels_shm_->String();
123+
if (buckets_shm_ != nullptr) { // Histogram
124+
size_t bucket_size = buckets_shm_->ByteSize() / sizeof(double);
125+
std::vector<double> buckets;
126+
buckets.reserve(bucket_size);
127+
for (size_t i = 0; i < bucket_size; ++i) {
128+
buckets.emplace_back(
129+
reinterpret_cast<double*>(buckets_shm_->DataPtr())[i]);
130+
}
131+
buckets_ = std::move(buckets);
132+
}
133+
94134
operation_value_ = custom_metric_shm_ptr_->operation_value;
95135
metric_family_address_ = custom_metric_shm_ptr_->metric_family_address;
96136
metric_address_ = custom_metric_shm_ptr_->metric_address;
@@ -161,6 +201,24 @@ Metric::SendSetValueRequest(const double& value)
161201
}
162202
}
163203

204+
void
205+
Metric::SendObserveRequest(const double& value)
206+
{
207+
try {
208+
CheckIfCleared();
209+
std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
210+
operation_value_ = value;
211+
SaveToSharedMemory(stub->ShmPool());
212+
AllocatedSharedMemory<CustomMetricsMessage> custom_metrics_shm;
213+
stub->SendMessage<CustomMetricsMessage>(
214+
custom_metrics_shm, PYTHONSTUB_MetricRequestObserve, shm_handle_);
215+
}
216+
catch (const PythonBackendException& pb_exception) {
217+
throw PythonBackendException(
218+
"Failed to observe metric value: " + std::string(pb_exception.what()));
219+
}
220+
}
221+
164222
double
165223
Metric::SendGetValueRequest()
166224
{
@@ -222,14 +280,35 @@ Metric::InitializeTritonMetric()
222280
{
223281
std::vector<const TRITONSERVER_Parameter*> labels_params;
224282
ParseLabels(labels_params, labels_);
283+
TRITONSERVER_MetricKind kind;
284+
THROW_IF_TRITON_ERROR(TRITONSERVER_GetMetricFamilyKind(
285+
reinterpret_cast<TRITONSERVER_MetricFamily*>(metric_family_address_),
286+
&kind));
287+
TRITONSERVER_MetricArgs* args = nullptr;
288+
switch (kind) {
289+
case TRITONSERVER_METRIC_KIND_COUNTER:
290+
case TRITONSERVER_METRIC_KIND_GAUGE:
291+
break;
292+
case TRITONSERVER_METRIC_KIND_HISTOGRAM: {
293+
const std::vector<double>& buckets = buckets_.value();
294+
THROW_IF_TRITON_ERROR(TRITONSERVER_MetricArgsNew(&args));
295+
THROW_IF_TRITON_ERROR(TRITONSERVER_MetricArgsSetHistogram(
296+
args, buckets.data(), buckets.size()));
297+
break;
298+
}
299+
default:
300+
break;
301+
}
302+
225303
TRITONSERVER_Metric* triton_metric = nullptr;
226-
THROW_IF_TRITON_ERROR(TRITONSERVER_MetricNew(
304+
THROW_IF_TRITON_ERROR(TRITONSERVER_MetricNewWithArgs(
227305
&triton_metric,
228306
reinterpret_cast<TRITONSERVER_MetricFamily*>(metric_family_address_),
229-
labels_params.data(), labels_params.size()));
307+
labels_params.data(), labels_params.size(), args));
230308
for (const auto label : labels_params) {
231309
TRITONSERVER_ParameterDelete(const_cast<TRITONSERVER_Parameter*>(label));
232310
}
311+
THROW_IF_TRITON_ERROR(TRITONSERVER_MetricArgsDelete(args));
233312
return reinterpret_cast<void*>(triton_metric);
234313
}
235314

@@ -262,6 +341,8 @@ Metric::HandleMetricOperation(
262341
Increment(operation_value_);
263342
} else if (command_type == PYTHONSTUB_MetricRequestSet) {
264343
SetValue(operation_value_);
344+
} else if (command_type == PYTHONSTUB_MetricRequestObserve) {
345+
Observe(operation_value_);
265346
} else {
266347
throw PythonBackendException("Unknown metric operation");
267348
}
@@ -281,6 +362,13 @@ Metric::SetValue(const double& value)
281362
THROW_IF_TRITON_ERROR(TRITONSERVER_MetricSet(triton_metric, value));
282363
}
283364

365+
void
366+
Metric::Observe(const double& value)
367+
{
368+
auto triton_metric = reinterpret_cast<TRITONSERVER_Metric*>(metric_address_);
369+
THROW_IF_TRITON_ERROR(TRITONSERVER_MetricObserve(triton_metric, value));
370+
}
371+
284372
double
285373
Metric::GetValue()
286374
{

src/metric.h

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -26,9 +26,11 @@
2626

2727
#pragma once
2828

29+
#include <optional>
2930
#include <string>
3031

3132
#include "ipc_message.h"
33+
#include "pb_memory.h"
3234
#include "pb_string.h"
3335
#include "pb_utils.h"
3436

@@ -47,6 +49,8 @@ namespace triton { namespace backend { namespace python {
4749
struct MetricShm {
4850
// The shared memory handle of the labels in PbString format.
4951
bi::managed_external_buffer::handle_t labels_shm_handle;
52+
// The shared memory handle of the buckets in PbMemory format.
53+
bi::managed_external_buffer::handle_t buckets_shm_handle;
5054
// The value used for incrementing or setting the metric.
5155
double operation_value;
5256
// The address of the TRITONSERVER_Metric object.
@@ -58,7 +62,10 @@ struct MetricShm {
5862

5963
class Metric {
6064
public:
61-
Metric(const std::string& labels, void* metric_family_address);
65+
Metric(
66+
const std::string& labels,
67+
std::optional<const std::vector<double>> buckets,
68+
void* metric_family_address);
6269

6370
~Metric();
6471

@@ -97,6 +104,10 @@ class Metric {
97104
/// \param value The value to set the metric to.
98105
void SendSetValueRequest(const double& value);
99106

107+
/// Send the request to the parent process to observe the value to the metric.
108+
/// \param value The value to set the metric to.
109+
void SendObserveRequest(const double& value);
110+
100111
/// Send the request to the parent process to get the value of the metric.
101112
/// \return Returns the value of the metric.
102113
double SendGetValueRequest();
@@ -132,6 +143,10 @@ class Metric {
132143
/// \param value The value to set the metric to.
133144
void SetValue(const double& value);
134145

146+
/// Use Triton C API to sample the observation to the metric.
147+
/// \param value The value to sample observation to the metric.
148+
void Observe(const double& value);
149+
135150
/// Use Triton C API to get the value of the metric.
136151
double GetValue();
137152

@@ -146,10 +161,14 @@ class Metric {
146161
// The private constructor for creating a Metric object from shared memory.
147162
Metric(
148163
AllocatedSharedMemory<MetricShm>& custom_metric_shm,
149-
std::unique_ptr<PbString>& labels_shm);
164+
std::unique_ptr<PbString>& labels_shm,
165+
std::unique_ptr<PbMemory>& buckets);
150166

151167
// The labels of the metric, which is the identifier of the metric.
152168
std::string labels_;
169+
// Monotonically increasing values representing bucket boundaries for creating
170+
// histogram metric.
171+
std::optional<std::vector<double>> buckets_;
153172
// The value used for incrementing or setting the metric.
154173
double operation_value_;
155174
// The address of the TRITONSERVER_Metric object.
@@ -168,6 +187,7 @@ class Metric {
168187
MetricShm* custom_metric_shm_ptr_;
169188
bi::managed_external_buffer::handle_t shm_handle_;
170189
std::unique_ptr<PbString> labels_shm_;
190+
std::unique_ptr<PbMemory> buckets_shm_;
171191
};
172192

173193
}}}; // namespace triton::backend::python

src/metric_family.cc

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -166,19 +166,39 @@ MetricFamily::SendCreateMetricFamilyRequest()
166166
}
167167

168168
std::shared_ptr<Metric>
169-
MetricFamily::CreateMetric(const py::object& labels)
169+
MetricFamily::CreateMetric(const py::object& labels, const py::object& buckets)
170170
{
171171
if (!labels.is_none()) {
172172
if (!py::isinstance<py::dict>(labels)) {
173173
throw PythonBackendException(
174-
"Failed to create metric. Labels must be a "
175-
"dictionary.");
174+
"Failed to create metric. Labels must be a dictionary.");
176175
}
177176
}
178177

179178
py::module json = py::module_::import("json");
180179
std::string labels_str = std::string(py::str(json.attr("dumps")(labels)));
181-
auto metric = std::make_shared<Metric>(labels_str, metric_family_address_);
180+
181+
std::optional<std::vector<double>> buckets_vec;
182+
if (!buckets.is_none()) {
183+
if (!py::isinstance<py::list>(buckets)) {
184+
throw PythonBackendException(
185+
"Failed to create metric. Buckets must be a list.");
186+
}
187+
if (kind_ == kCounter || kind_ == kGauge) {
188+
throw PythonBackendException(
189+
"Failed to create metric. Unexpected buckets found.");
190+
}
191+
buckets_vec = buckets.cast<std::vector<double>>();
192+
} else {
193+
if (kind_ == kHistogram) {
194+
throw PythonBackendException(
195+
"Failed to create metric. Missing required buckets.");
196+
}
197+
buckets_vec = std::nullopt;
198+
}
199+
200+
auto metric =
201+
std::make_shared<Metric>(labels_str, buckets_vec, metric_family_address_);
182202
{
183203
std::lock_guard<std::mutex> lock(metric_map_mu_);
184204
metric_map_.insert({metric->MetricAddress(), metric});
@@ -205,6 +225,8 @@ MetricFamily::ToTritonServerMetricKind(const MetricKind& kind)
205225
return TRITONSERVER_METRIC_KIND_COUNTER;
206226
case kGauge:
207227
return TRITONSERVER_METRIC_KIND_GAUGE;
228+
case kHistogram:
229+
return TRITONSERVER_METRIC_KIND_HISTOGRAM;
208230
default:
209231
throw PythonBackendException("Unknown metric kind");
210232
}

src/metric_family.h

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -97,8 +97,11 @@ class MetricFamily {
9797

9898
/// Create a metric from the metric family and store it in the metric map.
9999
/// \param labels The labels of the metric.
100+
/// \param buckets Monotonically increasing values representing bucket
101+
/// boundaries for creating histogram metric.
100102
/// \return Returns the shared pointer to the created metric.
101-
std::shared_ptr<Metric> CreateMetric(const py::object& labels);
103+
std::shared_ptr<Metric> CreateMetric(
104+
const py::object& labels, const py::object& buckets);
102105
#else
103106
/// Initialize the TRITONSERVER_MetricFamily object.
104107
/// \return Returns the address of the TRITONSERVER_MetricFamily object.
@@ -128,8 +131,8 @@ class MetricFamily {
128131
std::string name_;
129132
// The description of the metric family.
130133
std::string description_;
131-
// The metric kind of the metric family. Currently only supports GAUGE and
132-
// COUNTER.
134+
// The metric kind of the metric family. Currently only supports GAUGE,
135+
// COUNTER and HISTOGRAM.
133136
MetricKind kind_;
134137
// The address of the TRITONSERVER_MetricFamily object.
135138
void* metric_family_address_;

0 commit comments

Comments
 (0)