tritonserver_rs/
metrics.rs

1use std::{
2    mem::transmute,
3    ptr::{null, null_mut},
4    sync::Arc,
5};
6
7use crate::{parameter::Parameter, sys, to_cstring, Error};
8
9#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
10#[repr(u32)]
11/// Metric format types.
12enum Format {
13    /// Base points to a single multiline
14    /// string that gives a text representation of the metrics in
15    /// prometheus format.
16    Prometheus = sys::tritonserver_metricformat_enum_TRITONSERVER_METRIC_PROMETHEUS,
17}
18
19/// Prometheus metrics object.
20#[derive(Debug, Clone)]
21pub struct PrometheusMetrics(pub(crate) Arc<*mut sys::TRITONSERVER_Metrics>);
22
23unsafe impl Send for PrometheusMetrics {}
24unsafe impl Sync for PrometheusMetrics {}
25
26impl PrometheusMetrics {
27    /// Get a buffer containing the metrics in the specified format.
28    pub fn formatted(&self) -> Result<&[u8], Error> {
29        let format = Format::Prometheus;
30
31        #[cfg(target_arch = "x86_64")]
32        let mut ptr = null::<i8>();
33        #[cfg(target_arch = "aarch64")]
34        let mut ptr = null::<u8>();
35        let mut size: usize = 0;
36
37        triton_call!(sys::TRITONSERVER_MetricsFormatted(
38            *self.0,
39            format as _,
40            &mut ptr as *mut _,
41            &mut size as *mut _,
42        ))?;
43
44        assert!(!ptr.is_null());
45        Ok(unsafe { std::slice::from_raw_parts(ptr as *const u8, size) })
46    }
47}
48
49impl Drop for PrometheusMetrics {
50    fn drop(&mut self) {
51        if !self.0.is_null() && Arc::strong_count(&self.0) == 1 {
52            unsafe { sys::TRITONSERVER_MetricsDelete(*self.0) };
53        }
54    }
55}
56
57/// Types of metrics recognized by TRITONSERVER.
58#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
59#[repr(u32)]
60pub enum MetricKind {
61    Counter = sys::TRITONSERVER_metrickind_enum_TRITONSERVER_METRIC_KIND_COUNTER,
62    Gauge = sys::TRITONSERVER_metrickind_enum_TRITONSERVER_METRIC_KIND_GAUGE,
63    Histogram = sys::TRITONSERVER_metrickind_enum_TRITONSERVER_METRIC_KIND_HISTOGRAM,
64}
65
66/// Family of the metrics.
67///
68/// Author note: the current state of [Metric], [MetricKind], [MetricFamily] is not kinda useful.
69/// Added due to the politic to add every item of [original API](https://github.com/triton-inference-server/core/blob/main/include/triton/core/tritonserver.h).
70/// It's worth considering using dedicated crates like [metrics](https://crates.io/crates/metrics).
71#[derive(Debug, Clone)]
72pub struct MetricFamily(Arc<*mut sys::TRITONSERVER_MetricFamily>);
73
74unsafe impl Send for MetricFamily {}
75unsafe impl Sync for MetricFamily {}
76
77impl MetricFamily {
78    /// Create a new metric family object.
79    /// - `kind`: The type of metric family to create.
80    /// - `name`: The name of the metric family seen when calling the metrics
81    ///     endpoint.
82    /// - `description`: The description of the metric family seen when
83    ///     calling the metrics endpoint.
84    pub fn new<N: AsRef<str>, D: AsRef<str>>(
85        kind: MetricKind,
86        name: N,
87        description: D,
88    ) -> Result<Self, Error> {
89        let c_name = to_cstring(name)?;
90        let descr = to_cstring(description)?;
91        let mut res = null_mut::<sys::TRITONSERVER_MetricFamily>();
92        triton_call!(
93            sys::TRITONSERVER_MetricFamilyNew(
94                &mut res as *mut _,
95                kind as _,
96                c_name.as_ptr(),
97                descr.as_ptr()
98            ),
99            Self(Arc::new(res))
100        )
101    }
102
103    /// Get the kind of the metric family.
104    pub fn kind(&self) -> Result<MetricKind, Error> {
105        let mut res = 0;
106        triton_call!(sys::TRITONSERVER_GetMetricFamilyKind(*self.0, &mut res))?;
107        Ok(unsafe { transmute::<u32, MetricKind>(res) })
108    }
109}
110
111impl Drop for MetricFamily {
112    fn drop(&mut self) {
113        if !self.0.is_null() && Arc::strong_count(&self.0) == 1 {
114            unsafe { sys::TRITONSERVER_MetricFamilyDelete(*self.0) };
115        }
116    }
117}
118
119/// Arguments to pass on creating [Metric]. Currently only [Self::set_histogram] is supported.
120///
121/// Author note: the current state of [Metric], [MetricKind], [MetricFamily] is not kinda useful.
122/// Added due to the politic to add every item of [original API](https://github.com/triton-inference-server/core/blob/main/include/triton/core/tritonserver.h).
123/// It's worth considering using dedicated crates like [metrics](https://crates.io/crates/metrics).
124#[derive(Debug, Clone)]
125pub struct MetricArgs(Arc<*mut sys::TRITONSERVER_MetricArgs>);
126
127unsafe impl Send for MetricArgs {}
128unsafe impl Sync for MetricArgs {}
129
130impl MetricArgs {
131    /// Create a new metric args object.
132    pub fn new() -> Result<Self, Error> {
133        let mut res = null_mut::<sys::TRITONSERVER_MetricArgs>();
134        triton_call!(
135            sys::TRITONSERVER_MetricArgsNew(&mut res),
136            Self(Arc::new(res))
137        )
138    }
139
140    /// Set metric args with histogram metric parameter.
141    /// - `buckets`: The array of bucket boundaries for the expected range of
142    ///     observed values.
143    pub fn set_histogram<B: AsRef<[f64]>>(&mut self, buckets: B) -> Result<&mut Self, Error> {
144        let buckets = buckets.as_ref();
145        triton_call!(
146            sys::TRITONSERVER_MetricArgsSetHistogram(*self.0, buckets.as_ptr(), buckets.len() as _),
147            self
148        )
149    }
150}
151
152impl Drop for MetricArgs {
153    fn drop(&mut self) {
154        if !self.0.is_null() && Arc::strong_count(&self.0) == 1 {
155            unsafe { sys::TRITONSERVER_MetricArgsDelete(*self.0) };
156        }
157    }
158}
159
160#[allow(dead_code)]
161/// Metric: Counter, Gauge or Histogram.
162///
163/// Author note: the current state of [Metric], [MetricKind], [MetricFamily] is not kinda useful.
164/// Added due to the politic to add every item of [original API](https://github.com/triton-inference-server/core/blob/main/include/triton/core/tritonserver.h).
165/// It's worth considering using dedicated crates like [metrics](https://crates.io/crates/metrics).
166#[derive(Debug, Clone)]
167pub struct Metric(Arc<*mut sys::TRITONSERVER_Metric>, MetricFamily);
168
169unsafe impl Send for Metric {}
170unsafe impl Sync for Metric {}
171
172impl Metric {
173    /// Create a new metric object.
174    /// - `family`: The metric family to add this new metric to.
175    /// - `labels`: The array of labels to associate with this new metric.
176    pub fn new<P: AsRef<[Parameter]>>(family: &MetricFamily, labels: P) -> Result<Self, Error> {
177        let mut res = null_mut::<sys::TRITONSERVER_Metric>();
178        let mut labels = labels
179            .as_ref()
180            .iter()
181            .map(|p| *p.ptr as *const _)
182            .collect::<Vec<_>>();
183        triton_call!(
184            sys::TRITONSERVER_MetricNew(
185                &mut res,
186                *family.0,
187                labels.as_mut_ptr(),
188                labels.len() as _
189            ),
190            Self(Arc::new(res), family.clone())
191        )
192    }
193
194    /// Create a new metric object.
195    /// - `family`: The metric family to add this new metric to.
196    /// - `labels`: The array of labels to associate with this new metric.
197    /// - `args`: Metric args that store additional arguments to construct
198    ///     particular metric types, e.g. histogram.
199    pub fn new_with_args<P: AsRef<[Parameter]>>(
200        family: &MetricFamily,
201        labels: P,
202        args: &MetricArgs,
203    ) -> Result<Self, Error> {
204        let mut res = null_mut::<sys::TRITONSERVER_Metric>();
205        let mut labels = labels
206            .as_ref()
207            .iter()
208            .map(|p| *p.ptr as *const _)
209            .collect::<Vec<_>>();
210        triton_call!(
211            sys::TRITONSERVER_MetricNewWithArgs(
212                &mut res,
213                *family.0,
214                labels.as_mut_ptr(),
215                labels.len() as _,
216                *args.0 as *const _
217            ),
218            Self(Arc::new(res), family.clone())
219        )
220    }
221
222    /// Get the kind of metric of its corresponding family.
223    pub fn kind(&self) -> Result<MetricKind, Error> {
224        let mut res = 0;
225        triton_call!(sys::TRITONSERVER_GetMetricKind(*self.0, &mut res))?;
226        Ok(unsafe { transmute::<u32, MetricKind>(res) })
227    }
228
229    /// Get the current value of a metric object.
230    /// Supports metrics of kind [MetricKind::Counter]
231    /// and [MetricKind::Gauge], and returns [Error::Unsupported](crate::ErrorCode::Unsupported)
232    /// for other kinds.
233    pub fn value(&self) -> Result<f64, Error> {
234        let mut res = 0.;
235        triton_call!(sys::TRITONSERVER_MetricValue(*self.0, &mut res), res)
236    }
237
238    /// Increment the current value of metric by value.
239    /// Supports metrics of kind [MetricKind::Gauge] for any value,
240    /// and [MetricKind::Counter] for non-negative values. Returns
241    /// [Error::Unsupported](crate::ErrorCode::Unsupported) for other kinds
242    /// and [Error::InvalidArg](crate::ErrorCode::InvalidArg) for negative values on a
243    /// [MetricKind::Counter] metric.
244    /// - `value`: The amount to increment the metric's value by.
245    pub fn increment_by(&self, value: f64) -> Result<(), Error> {
246        triton_call!(sys::TRITONSERVER_MetricIncrement(*self.0, value))
247    }
248
249    /// Set the current value of metric to value.
250    /// Supports metrics of kind [MetricKind::Gauge] and returns
251    /// [Error::Unsupported](crate::ErrorCode::Unsupported) for other kinds.
252    ///
253    /// - `value`: The amount to set metric's value to.
254    pub fn set(&self, value: f64) -> Result<(), Error> {
255        triton_call!(sys::TRITONSERVER_MetricSet(*self.0, value))
256    }
257
258    /// Sample an observation and count it to the appropriate bucket of a metric.
259    /// Supports metrics of kind [MetricKind::Histogram] and returns
260    /// [Error::Unsupported](crate::ErrorCode::Unsupported) for other kinds.
261    /// - `value`: The amount for metric to sample observation.
262    pub fn observe(&self, value: f64) -> Result<(), Error> {
263        triton_call!(sys::TRITONSERVER_MetricObserve(*self.0, value))
264    }
265}
266
267impl Drop for Metric {
268    fn drop(&mut self) {
269        if !self.0.is_null() && Arc::strong_count(&self.0) == 1 {
270            unsafe { sys::TRITONSERVER_MetricDelete(*self.0) };
271        }
272    }
273}
tritonserver_rs/metrics.rs

tritonserver_rs/
metrics.rs
