diff --git a/rust/Cargo.lock b/rust/Cargo.lock index b3f9084bcf..5b122ea14c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5903,6 +5903,7 @@ dependencies = [ name = "vllm-metrics" version = "0.1.0" dependencies = [ + "expect-test", "itertools 0.14.0", "prometheus-client", ] diff --git a/rust/clippy.toml b/rust/clippy.toml new file mode 100644 index 0000000000..d3ea524317 --- /dev/null +++ b/rust/clippy.toml @@ -0,0 +1,3 @@ +disallowed-types = [ + { path = "prometheus_client::metrics::histogram::Histogram", reason = "it locks on every observation; use `vllm_metrics::Histogram`" }, +] diff --git a/rust/src/metrics/Cargo.toml b/rust/src/metrics/Cargo.toml index ab1a72098b..74e4d56d5b 100644 --- a/rust/src/metrics/Cargo.toml +++ b/rust/src/metrics/Cargo.toml @@ -8,5 +8,8 @@ license.workspace = true itertools.workspace = true prometheus-client.workspace = true +[dev-dependencies] +expect-test.workspace = true + [lints] workspace = true diff --git a/rust/src/metrics/src/api_server.rs b/rust/src/metrics/src/api_server.rs index 4decaa5573..855e9fbbb9 100644 --- a/rust/src/metrics/src/api_server.rs +++ b/rust/src/metrics/src/api_server.rs @@ -3,10 +3,9 @@ use prometheus_client::encoding::EncodeLabelSet; use prometheus_client::metrics::family::Family; -use prometheus_client::metrics::histogram::Histogram; use prometheus_client::registry::Registry; -use crate::U64Counter; +use crate::{Histogram, U64Counter}; const HTTP_REQUEST_DURATION_BUCKETS: [f64; 3] = [0.1, 0.5, 1.0]; const HTTP_REQUEST_DURATION_HIGHR_BUCKETS: [f64; 21] = [ diff --git a/rust/src/metrics/src/histogram.rs b/rust/src/metrics/src/histogram.rs new file mode 100644 index 0000000000..f6d9afd86c --- /dev/null +++ b/rust/src/metrics/src/histogram.rs @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::fmt; +use std::iter::once; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use prometheus_client::encoding::{EncodeMetric, MetricEncoder, NoLabelSet}; +use prometheus_client::metrics::{MetricType, TypedMetric}; + +/// Prometheus histogram with lock-free observations. +/// +/// `prometheus_client::metrics::histogram::Histogram` takes an exclusive lock for every +/// observation, which serializes all Tokio workers on hot histograms such as inter-token +/// latency. This histogram updates atomic cells instead, so observers never block each other +/// or a concurrent scrape. +/// +/// A scrape may see an observation's bucket count before its sum. `_count` is derived from the +/// bucket counts, so it always equals the `+Inf` bucket. +#[derive(Clone, Debug)] +pub struct Histogram { + inner: Arc, +} + +#[derive(Debug)] +struct Inner { + /// Upper bounds, ending with `f64::MAX`, which is encoded as `+Inf`. + upper_bounds: Box<[f64]>, + /// Non-cumulative observation count of each bucket. + bucket_counts: Box<[AtomicU64]>, + /// Bits of the `f64` sum of all observations. + sum_bits: AtomicU64, +} + +impl Histogram { + /// Create a histogram with the given finite bucket upper bounds, in ascending order. + pub fn new(buckets: impl IntoIterator) -> Self { + let upper_bounds: Box<[f64]> = buckets.into_iter().chain(once(f64::MAX)).collect(); + let bucket_counts = upper_bounds.iter().map(|_| AtomicU64::new(0)).collect(); + Self { + inner: Arc::new(Inner { + upper_bounds, + bucket_counts, + sum_bits: AtomicU64::new(0.0_f64.to_bits()), + }), + } + } + + pub fn observe(&self, value: f64) { + let inner = &*self.inner; + // Unlike `prometheus_client`, which counts `+inf` and NaN in `_count` but in no bucket, + // put them in the `+Inf` bucket as the Go client does. + let index = inner + .upper_bounds + .iter() + .position(|bound| value <= *bound) + .unwrap_or(inner.upper_bounds.len() - 1); + inner.bucket_counts[index].fetch_add(1, Ordering::Relaxed); + // The closure always returns `Some`, so the update cannot fail. + let _ = inner.sum_bits.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |bits| { + Some((f64::from_bits(bits) + value).to_bits()) + }); + } +} + +impl TypedMetric for Histogram { + const TYPE: MetricType = MetricType::Histogram; +} + +impl EncodeMetric for Histogram { + fn encode(&self, mut encoder: MetricEncoder) -> Result<(), fmt::Error> { + let inner = &*self.inner; + let buckets: Vec<(f64, u64)> = inner + .upper_bounds + .iter() + .zip(&inner.bucket_counts) + .map(|(bound, count)| (*bound, count.load(Ordering::Relaxed))) + .collect(); + let count = buckets.iter().map(|(_, count)| count).sum(); + let sum = f64::from_bits(inner.sum_bits.load(Ordering::Relaxed)); + encoder.encode_histogram::(sum, count, &buckets, None) + } + + fn metric_type(&self) -> MetricType { + Self::TYPE + } +} + +#[cfg(test)] +mod tests { + use expect_test::expect; + use prometheus_client::encoding::text::encode; + use prometheus_client::registry::Registry; + + use super::*; + + const BUCKETS: [f64; 3] = [0.1, 1.0, 10.0]; + + fn render(metric: impl EncodeMetric + fmt::Debug + Send + Sync + 'static) -> String { + let mut registry = Registry::default(); + registry.register("latency", "Latency", metric); + let mut output = String::new(); + encode(&mut output, ®istry).unwrap(); + output + } + + #[test] + #[expect( + clippy::disallowed_types, + reason = "compare against the replaced implementation" + )] + fn finite_observations_match_prometheus_client() { + let baseline = prometheus_client::metrics::histogram::Histogram::new(BUCKETS); + let histogram = Histogram::new(BUCKETS); + let mut values = vec![-1.0, 0.0, 100.0, f64::MAX]; + for bound in BUCKETS { + values.extend([bound.next_down(), bound, bound.next_up()]); + } + for value in values { + baseline.observe(value); + histogram.observe(value); + } + assert_eq!(render(histogram), render(baseline)); + } + + #[test] + fn concurrent_observations_are_not_lost() { + let histogram = Histogram::new(BUCKETS); + std::thread::scope(|scope| { + // Dyadic values keep the sum exact in any addition order. + for value in [0.0625, 0.5, 4.0, 64.0] { + let histogram = histogram.clone(); + scope.spawn(move || (0..10_000).for_each(|_| histogram.observe(value))); + } + }); + expect![[r##" + # HELP latency Latency. + # TYPE latency histogram + latency_sum 685625.0 + latency_count 40000 + latency_bucket{le="0.1"} 10000 + latency_bucket{le="1.0"} 20000 + latency_bucket{le="10.0"} 30000 + latency_bucket{le="+Inf"} 40000 + # EOF + "##]] + .assert_eq(&render(histogram)); + } + + #[test] + fn non_finite_observations_count_in_inf_bucket() { + let histogram = Histogram::new(BUCKETS); + histogram.observe(f64::INFINITY); + histogram.observe(f64::NAN); + expect![[r##" + # HELP latency Latency. + # TYPE latency histogram + latency_sum NaN + latency_count 2 + latency_bucket{le="0.1"} 0 + latency_bucket{le="1.0"} 0 + latency_bucket{le="10.0"} 0 + latency_bucket{le="+Inf"} 2 + # EOF + "##]] + .assert_eq(&render(histogram)); + } +} diff --git a/rust/src/metrics/src/lib.rs b/rust/src/metrics/src/lib.rs index 25ab4fc6ba..0324946676 100644 --- a/rust/src/metrics/src/lib.rs +++ b/rust/src/metrics/src/lib.rs @@ -9,14 +9,15 @@ use prometheus_client::encoding::text::encode; use prometheus_client::metrics::counter::Counter; pub use prometheus_client::metrics::family::Family; use prometheus_client::metrics::gauge::Gauge; -use prometheus_client::metrics::histogram::Histogram; use prometheus_client::registry::Registry; mod api_server; +mod histogram; mod request; mod scheduler; pub use api_server::*; +pub use histogram::Histogram; pub use request::*; pub use scheduler::*; diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index 5bfc32e895..b6e7683e76 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -3,10 +3,9 @@ use prometheus_client::encoding::EncodeLabelSet; use prometheus_client::metrics::family::Family; -use prometheus_client::metrics::histogram::Histogram; use prometheus_client::registry::Registry; -use crate::{EngineLabels, HistogramFamily, U64Counter}; +use crate::{EngineLabels, Histogram, HistogramFamily, U64Counter}; const TTFT_BUCKETS: [f64; 22] = [ 0.001, 0.005, 0.01, 0.02, 0.04, 0.06, 0.08, 0.1, 0.25, 0.5, 0.75, 1.0, 2.5, 5.0, 7.5, 10.0, diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index d7f52fd0a6..cdc25a244d 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -7,10 +7,9 @@ use std::sync::{Arc, Mutex}; use itertools::Itertools as _; use prometheus_client::encoding::{EncodeLabelSet, EncodeLabelValue, LabelValueEncoder}; use prometheus_client::metrics::family::Family; -use prometheus_client::metrics::histogram::Histogram; use prometheus_client::registry::Registry; -use crate::{F64Gauge, HistogramFamily, U64Counter, U64Gauge}; +use crate::{F64Gauge, Histogram, HistogramFamily, U64Counter, U64Gauge}; const KV_CACHE_RESIDENCY_BUCKETS: [f64; 21] = [ 0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 60.0,