[Rust Frontend] Cap request token histogram buckets at max_model_len

rust/src/llm/src/lib.rs

@@ -5,6 +5,7 @@ use std::num::NonZeroU32;
 
 use tracing::Span;
 use vllm_engine_core_client::EngineCoreClient;
+use vllm_metrics::METRICS;
 
 mod error;
 mod inflight;
@@ -45,6 +46,7 @@ impl Llm {
     /// Create a new minimal LLM facade from an already connected engine-core
     /// client.
     pub fn new(client: EngineCoreClient) -> Self {
+        METRICS.request.set_max_model_len(client.max_model_len());
         Self {
             client,
             randomize_request_id: true,

rust/src/llm/src/request_metrics.rs

@@ -406,11 +406,13 @@ pub fn current_unix_timestamp_secs() -> f64 {
 mod tests {
     use vllm_engine_core_client::protocol::output::{EngineCoreEvent, EngineCoreEventType};
     use vllm_engine_core_client::protocol::stats::PrefillStats;
+    use vllm_metrics::METRICS;
 
     use super::{RequestMetricsTracker, RequestTimestamps, diff_or_zero};
 
     #[test]
     fn tracker_updates_timing_state_across_prefill_decode_and_finish() {
+        METRICS.request.set_max_model_len(4096);
         let mut tracker =
             RequestMetricsTracker::new("model".to_string(), 2, 100.0, 64, Some(128), 1, true);
 

rust/src/metrics/src/request.rs

@@ -1,8 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+use std::sync::{Arc, OnceLock};
+
 use prometheus_client::encoding::EncodeLabelSet;
-use prometheus_client::metrics::family::Family;
+use prometheus_client::metrics::family::{Family, MetricConstructor};
 use prometheus_client::registry::Registry;
 
 use crate::{EngineLabels, Histogram, HistogramFamily, U64Counter};
@@ -25,20 +27,19 @@ const ITERATION_TOKENS_BUCKETS: [f64; 13] = [
     1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0,
 ];
 
+/// Build a 1-2-5 series of buckets capped at `max_value`, like Python's
+/// `build_1_2_5_buckets`: `100` gives `[1, 2, 5, 10, 20, 50, 100]`.
 fn build_1_2_5_buckets(max_value: u32) -> Vec<f64> {
+    let max_value = u64::from(max_value);
     let mut buckets = Vec::new();
     let mut exponent = 0;
     loop {
-        for mantissa in [1_u32, 2, 5] {
-            let value = mantissa * 10_u32.pow(exponent);
-            if value <= max_value {
-                buckets.push(value as f64);
-            } else {
-                if buckets.last().copied() != Some(max_value as f64) {
-                    buckets.push(max_value as f64);
-                }
+        for mantissa in [1_u64, 2, 5] {
+            let value = mantissa * 10_u64.pow(exponent);
+            if value > max_value {
                 return buckets;
             }
+            buckets.push(value as f64);
         }
         exponent += 1;
     }
@@ -60,11 +61,6 @@ fn request_latency_histogram() -> Histogram {
     Histogram::new(REQUEST_LATENCY_BUCKETS.iter().copied())
 }
 
-fn request_token_count_histogram() -> Histogram {
-    // TODO: determine max value based on `max_model_len`.
-    Histogram::new(build_1_2_5_buckets(131_072))
-}
-
 fn request_params_n_histogram() -> Histogram {
     Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied())
 }
@@ -93,6 +89,28 @@ pub struct PromptTokenSourceLabels {
 
 pub(crate) type FinishedReasonCounterFamily = Family<FinishedReasonLabels, U64Counter>;
 pub(crate) type PromptTokenSourceCounterFamily = Family<PromptTokenSourceLabels, U64Counter>;
+pub(crate) type RequestTokensHistogramFamily =
+    Family<EngineLabels, Histogram, RequestTokensHistogram>;
+
+/// Constructor of the request token-count histograms, which share Python's
+/// `request_tokens` buckets: a 1-2-5 series capped at `max_model_len`.
+///
+/// The registry is built before any engine is connected, so the buckets are
+/// filled in later by [`RequestMetrics::set_max_model_len`].
+#[derive(Clone, Debug, Default)]
+pub struct RequestTokensHistogram {
+    buckets: Arc<OnceLock<Vec<f64>>>,
+}
+
+impl MetricConstructor<Histogram> for RequestTokensHistogram {
+    fn new_metric(&self) -> Histogram {
+        let buckets = self
+            .buckets
+            .get()
+            .expect("max_model_len must be set before recording request token metrics");
+        Histogram::new(buckets.iter().copied())
+    }
+}
 
 /// Request and output-batch Prometheus families.
 pub struct RequestMetrics {
@@ -110,13 +128,13 @@ pub struct RequestMetrics {
 
     // Request lifecycle counters and histograms.
     pub request_success: FinishedReasonCounterFamily,
-    pub request_prompt_tokens: HistogramFamily,
-    pub request_generation_tokens: HistogramFamily,
-    pub request_max_num_generation_tokens: HistogramFamily,
-    pub request_params_max_tokens: HistogramFamily,
+    pub request_prompt_tokens: RequestTokensHistogramFamily,
+    pub request_generation_tokens: RequestTokensHistogramFamily,
+    pub request_max_num_generation_tokens: RequestTokensHistogramFamily,
+    pub request_params_max_tokens: RequestTokensHistogramFamily,
     pub request_params_n: HistogramFamily,
     pub request_num_preemptions: HistogramFamily,
-    pub request_prefill_kv_computed_tokens: HistogramFamily,
+    pub request_prefill_kv_computed_tokens: RequestTokensHistogramFamily,
     pub time_to_first_token_seconds: HistogramFamily,
     pub inter_token_latency_seconds: HistogramFamily,
     pub e2e_request_latency_seconds: HistogramFamily,
@@ -125,11 +143,16 @@ pub struct RequestMetrics {
     pub request_decode_time_seconds: HistogramFamily,
     pub request_inference_time_seconds: HistogramFamily,
     pub request_time_per_output_token_seconds: HistogramFamily,
+
+    /// Constructor shared by the request token-count histogram families.
+    request_tokens_histogram: RequestTokensHistogram,
 }
 
 impl RequestMetrics {
     /// Register the request-oriented metric families into the shared registry.
     pub(crate) fn register(registry: &mut Registry) -> Self {
+        let request_tokens_histogram = RequestTokensHistogram::default();
+
         let iteration_tokens_total =
             Family::new_with_constructor(iteration_tokens_histogram as fn() -> Histogram);
         registry.register(
@@ -189,32 +212,31 @@ impl RequestMetrics {
             request_success.clone(),
         );
 
-        let request_prompt_tokens =
-            Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
+        let request_prompt_tokens = Family::new_with_constructor(request_tokens_histogram.clone());
         registry.register(
             "vllm:request_prompt_tokens",
             "Number of prefill tokens processed.",
             request_prompt_tokens.clone(),
         );
 
         let request_generation_tokens =
-            Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
+            Family::new_with_constructor(request_tokens_histogram.clone());
         registry.register(
             "vllm:request_generation_tokens",
             "Number of generation tokens processed.",
             request_generation_tokens.clone(),
         );
 
         let request_max_num_generation_tokens =
-            Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
+            Family::new_with_constructor(request_tokens_histogram.clone());
         registry.register(
             "vllm:request_max_num_generation_tokens",
             "Histogram of maximum number of requested generation tokens.",
             request_max_num_generation_tokens.clone(),
         );
 
         let request_params_max_tokens =
-            Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
+            Family::new_with_constructor(request_tokens_histogram.clone());
         registry.register(
             "vllm:request_params_max_tokens",
             "Histogram of the max_tokens request parameter.",
@@ -238,7 +260,7 @@ impl RequestMetrics {
         );
 
         let request_prefill_kv_computed_tokens =
-            Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
+            Family::new_with_constructor(request_tokens_histogram.clone());
         registry.register(
             "vllm:request_prefill_kv_computed_tokens",
             "Histogram of new KV tokens computed during prefill (excluding cached tokens).",
@@ -334,6 +356,50 @@ impl RequestMetrics {
             request_decode_time_seconds,
             request_inference_time_seconds,
             request_time_per_output_token_seconds,
+            request_tokens_histogram,
         }
     }
+
+    /// Cap the request token-count histogram buckets at the engine-reported
+    /// `max_model_len`.
+    ///
+    /// Must be called before the first request is recorded. The registry is
+    /// process-global and serves one model, so only the first call takes
+    /// effect.
+    pub fn set_max_model_len(&self, max_model_len: u32) {
+        let _ = self.request_tokens_histogram.buckets.set(build_1_2_5_buckets(max_model_len));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use expect_test::expect;
+    use itertools::Itertools;
+    use prometheus_client::encoding::text::encode;
+
+    use super::*;
+
+    /// Request token-count histograms use Python's 1-2-5 buckets capped at
+    /// `max_model_len`, with no extra bound at `max_model_len` itself.
+    #[test]
+    fn request_token_buckets_are_capped_at_max_model_len() {
+        let mut registry = Registry::default();
+        let metrics = RequestMetrics::register(&mut registry);
+        metrics.set_max_model_len(4096);
+        metrics.request_prompt_tokens.get_or_create_owned(&EngineLabels {
+            model_name: "model".to_string(),
+            engine: 0,
+        });
+
+        let mut output = String::new();
+        encode(&mut output, &registry).unwrap();
+        let bounds = output
+            .lines()
+            .filter(|line| line.starts_with("vllm:request_prompt_tokens_bucket"))
+            .filter_map(|line| line.split_once("le=\"")?.1.split_once('"'))
+            .map(|(bound, _)| bound)
+            .join(" ");
+        expect!["1.0 2.0 5.0 10.0 20.0 50.0 100.0 200.0 500.0 1000.0 2000.0 +Inf"]
+            .assert_eq(&bounds);
+    }
 }