[Bugfix][Rust Frontend] Export corrupted request metric (#59937)

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

rust/src/llm/src/request_metrics.rs

@@ -47,6 +47,7 @@ pub(crate) struct RequestMetricsTracker {
     num_generation_tokens: u32,
     num_preemptions: u64,
     latest_num_cached_tokens: u32,
+    is_corrupted: bool,
 }
 
 /// Cached request metric handles for one model and engine index.
@@ -61,6 +62,7 @@ struct RequestMetricHandles {
     prompt_tokens_external_kv_transfer: U64Counter,
     prompt_tokens_cached: U64Counter,
     generation_tokens: U64Counter,
+    corrupted_requests: U64Counter,
 
     // Request lifecycle counters and histograms.
     request_success: Family<FinishedReasonLabels, U64Counter>,
@@ -123,6 +125,7 @@ impl RequestMetricsTracker {
             num_generation_tokens: 0,
             num_preemptions: 0,
             latest_num_cached_tokens: 0,
+            is_corrupted: false,
         }
     }
 
@@ -141,6 +144,7 @@ impl RequestMetricsTracker {
         }
         self.num_generation_tokens += output.new_token_ids.len() as u32;
         self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64);
+        self.is_corrupted |= output.num_nans_in_logits > 0;
 
         if let Some(events) = &output.events {
             self.observe_events(events);
@@ -187,6 +191,9 @@ impl RequestMetricsTracker {
         };
 
         self.record_request_success(finish_reason);
+        if self.is_corrupted {
+            self.handles.corrupted_requests.inc();
+        }
 
         self.handles.request_prompt_tokens.observe(self.prompt_len as f64);
         self.handles
@@ -297,6 +304,7 @@ fn resolve_request_metric_handles(
         ),
         prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels),
         generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels),
+        corrupted_requests: metrics.corrupted_requests.get_or_create_owned(&labels),
         request_success: metrics.request_success.clone(),
         request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels),
         request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels),

rust/src/llm/tests/generate.rs

@@ -898,6 +898,7 @@ async fn generate_records_request_metrics_in_prometheus_output() {
                         engine_index: 4,
                         timestamp: 10.0,
                         outputs: vec![EngineCoreOutput {
+                            num_nans_in_logits: 2,
                             prefill_stats: Some(Box::new(PrefillStats {
                                 num_prompt_tokens: 2,
                                 num_computed_tokens: 2,
@@ -982,6 +983,9 @@ async fn generate_records_request_metrics_in_prometheus_output() {
     assert!(rendered.contains(&format!(
         "vllm:generation_tokens_total{{model_name=\"{model_name}\",engine=\"4\"}} 3"
     )));
+    assert!(rendered.contains(&format!(
+        "vllm:corrupted_requests_total{{model_name=\"{model_name}\",engine=\"4\"}} 1"
+    )));
     assert!(rendered.contains(&format!(
         "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1"
     )));

rust/src/metrics/src/request.rs

@@ -106,6 +106,7 @@ pub struct RequestMetrics {
     pub prompt_tokens_by_source: PromptTokenSourceCounterFamily,
     pub prompt_tokens_cached: Family<EngineLabels, U64Counter>,
     pub generation_tokens: Family<EngineLabels, U64Counter>,
+    pub corrupted_requests: Family<EngineLabels, U64Counter>,
 
     // Request lifecycle counters and histograms.
     pub request_success: FinishedReasonCounterFamily,
@@ -173,6 +174,13 @@ impl RequestMetrics {
             generation_tokens.clone(),
         );
 
+        let corrupted_requests = Family::default();
+        registry.register(
+            "vllm:corrupted_requests",
+            "Corrupted requests, in terms of total number of requests with NaNs in logits.",
+            corrupted_requests.clone(),
+        );
+
         // Request lifecycle counters and histograms.
         let request_success = Family::default();
         registry.register(
@@ -309,6 +317,7 @@ impl RequestMetrics {
             prompt_tokens_by_source,
             prompt_tokens_cached,
             generation_tokens,
+            corrupted_requests,
             request_success,
             request_prompt_tokens,
             request_generation_tokens,