[Rust Frontend][Metrics] Add per-request preemption histogram (#57033)

Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Signed-off-by: reidliu41 <reid201711@gmail.com>

rust/src/llm/src/request_metrics.rs

@@ -45,6 +45,7 @@ pub(crate) struct RequestMetricsTracker {
     last_token_ts: f64,
     first_token_latency: f64,
     num_generation_tokens: u32,
+    num_preemptions: u64,
     latest_num_cached_tokens: u32,
 }
 
@@ -69,6 +70,7 @@ struct RequestMetricHandles {
     request_max_num_generation_tokens: HistogramMetric,
     request_params_max_tokens: HistogramMetric,
     request_params_n: HistogramMetric,
+    request_num_preemptions: HistogramMetric,
     request_prefill_kv_computed_tokens: HistogramMetric,
     time_to_first_token_seconds: HistogramMetric,
     inter_token_latency_seconds: HistogramMetric,
@@ -104,6 +106,7 @@ impl RequestMetricsTracker {
             last_token_ts: 0.0,
             first_token_latency: 0.0,
             num_generation_tokens: 0,
+            num_preemptions: 0,
             latest_num_cached_tokens: 0,
         }
     }
@@ -184,6 +187,7 @@ impl RequestMetricsTracker {
             self.handles.request_params_max_tokens.observe(max_tokens_param as f64);
         }
         self.handles.request_params_n.observe(self.n_param as f64);
+        self.handles.request_num_preemptions.observe(self.num_preemptions as f64);
         self.handles
             .request_prefill_kv_computed_tokens
             .observe(prefill_kv_computed_tokens as f64);
@@ -223,6 +227,7 @@ impl RequestMetricsTracker {
                     }
                 }
                 EngineCoreEventType::Preempted => {
+                    self.num_preemptions += 1;
                     self.handles.num_preemptions.inc();
                 }
             }
@@ -276,6 +281,7 @@ fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetri
             .get_or_create_owned(&labels),
         request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels),
         request_params_n: metrics.request_params_n.get_or_create_owned(&labels),
+        request_num_preemptions: metrics.request_num_preemptions.get_or_create_owned(&labels),
         request_prefill_kv_computed_tokens: metrics
             .request_prefill_kv_computed_tokens
             .get_or_create_owned(&labels),

rust/src/llm/tests/generate.rs

@@ -790,6 +790,12 @@ async fn generate_records_request_metrics_in_prometheus_output() {
     assert!(rendered.contains(&format!(
         "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1"
     )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"4\"}} 1.0"
+    )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
+    )));
     assert!(rendered.contains(&format!(
         "vllm:time_to_first_token_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1"
     )));
@@ -883,6 +889,12 @@ async fn dropping_stream_records_abort_terminal_request_metrics() {
     assert!(rendered.contains(&format!(
         "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"5\"}} 1"
     )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"5\"}} 0.0"
+    )));
+    assert!(rendered.contains(&format!(
+        "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"5\"}} 1"
+    )));
 
     llm.shutdown().await.unwrap();
 }

rust/src/metrics/src/request.rs

@@ -21,6 +21,7 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [
     480.0, 960.0, 1920.0, 7680.0,
 ];
 const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0];
+const REQUEST_NUM_PREEMPTIONS_BUCKETS: [f64; 7] = [1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0];
 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,
 ];
@@ -69,6 +70,10 @@ fn request_params_n_histogram() -> Histogram {
     Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied())
 }
 
+fn request_num_preemptions_histogram() -> Histogram {
+    Histogram::new(REQUEST_NUM_PREEMPTIONS_BUCKETS.iter().copied())
+}
+
 fn iteration_tokens_histogram() -> Histogram {
     Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied())
 }
@@ -110,6 +115,7 @@ pub struct RequestMetrics {
     pub request_max_num_generation_tokens: HistogramFamily,
     pub request_params_max_tokens: HistogramFamily,
     pub request_params_n: HistogramFamily,
+    pub request_num_preemptions: HistogramFamily,
     pub request_prefill_kv_computed_tokens: HistogramFamily,
     pub time_to_first_token_seconds: HistogramFamily,
     pub inter_token_latency_seconds: HistogramFamily,
@@ -216,6 +222,14 @@ impl RequestMetrics {
             request_params_n.clone(),
         );
 
+        let request_num_preemptions =
+            Family::new_with_constructor(request_num_preemptions_histogram as fn() -> Histogram);
+        registry.register(
+            "vllm:request_num_preemptions",
+            "Histogram of the number of times a request was preempted.",
+            request_num_preemptions.clone(),
+        );
+
         let request_prefill_kv_computed_tokens =
             Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram);
         registry.register(
@@ -302,6 +316,7 @@ impl RequestMetrics {
             request_max_num_generation_tokens,
             request_params_max_tokens,
             request_params_n,
+            request_num_preemptions,
             request_prefill_kv_computed_tokens,
             time_to_first_token_seconds,
             inter_token_latency_seconds,