@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8f4c69b222

[Rust Frontend] Cache metric handles for scheduler & request stats (#47444)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jul 6, 2026 at 21:02 UTC 8f4c69b22293424690ea5956d536c20c3d4ec393
9 files changed +343 -237
rust/src/engine-core-client/src/client.rs
+1
@@ -512,6 +512,7 @@ impl EngineCoreClient {
512
513 Ok(EngineCoreOutputStream::new(
514 request_id,
515 + engine_id.engine_index().unwrap_or(0),
516 self.abort_tx.clone(),
517 rx,
518 ))
rust/src/engine-core-client/src/client/imp.rs
+6 -7
@@ -15,7 +15,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti
15 use crate::client::stream::EngineCoreStreamOutput;
16 use crate::client::{AbortCause, AbortRequest};
17 use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
18 -use crate::metrics::{LoraInfoExporter, record_scheduler_stats};
18 +use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
19 use crate::protocol::encode_msgpack;
20 use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
21 use crate::protocol::request::EngineCoreRequestType;
@@ -29,6 +29,7 @@ pub(crate) struct ClientInner {
29 /// The runtime handle used for sending messages to the engine.
30 handle: Handle,
31 model_name: String,
32 + scheduler_stats_recorder: SchedulerStatsRecorder,
33 request_reg: Mutex<RequestRegistry>,
34 utility_reg: Mutex<UtilityRegistry>,
35 health_error: ArcSwapOption<Error>,
@@ -43,10 +44,13 @@ impl ClientInner {
44 model_name: String,
45 engines: &[ConnectedEngine],
46 ) -> Self {
47 + let scheduler_stats_recorder =
48 + SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines);
49 Self {
50 input_send,
51 handle,
52 model_name,
53 + scheduler_stats_recorder,
54 request_reg: Mutex::new(RequestRegistry::new(engines)),
55 utility_reg: Mutex::new(UtilityRegistry::default()),
56 health_error: ArcSwapOption::empty(),
@@ -389,12 +393,7 @@ pub(crate) async fn run_output_dispatcher_loop(
393 "dropping scheduler stats for unknown engine"
394 );
395 }
392 - record_scheduler_stats(
393 - &METRICS.scheduler,
394 - inner.model_name(),
395 - batch.engine_index,
396 - scheduler_stats,
397 - );
396 + inner.scheduler_stats_recorder.record(batch.engine_index, scheduler_stats);
397 }
398
399 // The engine's scheduler stats never carry adapter names;
rust/src/engine-core-client/src/client/stream.rs
+8
@@ -45,6 +45,7 @@ impl Deref for EngineCoreStreamOutput {
45 /// `finish_reason` is non-`None`.
46 pub struct EngineCoreOutputStream {
47 request_id: String,
48 + engine_index: u32,
49 abort_tx: mpsc::UnboundedSender<AbortRequest>,
50 state: State,
51 rx: OutputReceiver,
@@ -53,11 +54,13 @@ pub struct EngineCoreOutputStream {
54 impl EngineCoreOutputStream {
55 pub(crate) fn new(
56 request_id: String,
57 + engine_index: u32,
58 abort_tx: mpsc::UnboundedSender<AbortRequest>,
59 rx: OutputReceiver,
60 ) -> Self {
61 Self {
62 request_id,
63 + engine_index,
64 abort_tx,
65 state: State::Running,
66 rx,
@@ -68,6 +71,11 @@ impl EngineCoreOutputStream {
71 pub fn request_id(&self) -> &str {
72 &self.request_id
73 }
74 +
75 + /// Return the index of the engine that owns this request.
76 + pub fn engine_index(&self) -> u32 {
77 + self.engine_index
78 + }
79 }
80
81 impl Stream for EngineCoreOutputStream {
rust/src/engine-core-client/src/metrics.rs
+160 -74
@@ -1,90 +1,190 @@
1 +use std::collections::BTreeMap;
2 use std::collections::BTreeSet;
3 use std::time::{SystemTime, UNIX_EPOCH};
4
5 use vllm_metrics::{
5 - EngineLabels, EnginePositionLabels, LoraAdapterNames, LoraInfoLabels, SchedulerMetrics,
6 + EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames,
7 + LoraInfoLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
8 WaitingReasonLabels,
9 };
10
11 use crate::protocol::stats::SchedulerStats;
12 +use crate::transport::ConnectedEngine;
13
14 const WAITING_REASON_CAPACITY: &str = "capacity";
15 const WAITING_REASON_DEFERRED: &str = "deferred";
16
14 -/// Record the scheduler-stats-backed metrics for one engine at one point in
15 -/// time.
16 -pub(crate) fn record_scheduler_stats(
17 +/// Cached scheduler-stats metric handles for all engines connected to one
18 +/// frontend client.
19 +pub(crate) struct SchedulerStatsRecorder {
20 + engines: BTreeMap<u32, SchedulerStatsHandles>,
21 +}
22 +
23 +/// Per-engine cached metric handles used while recording `SchedulerStats`.
24 +struct SchedulerStatsHandles {
25 + // Base labels reused for dynamic child labels.
26 + labels: EngineLabels,
27 +
28 + // Scheduler state gauges.
29 + scheduler_running: U64Gauge,
30 + scheduler_waiting: U64Gauge,
31 + scheduler_waiting_capacity: U64Gauge,
32 + scheduler_waiting_deferred: U64Gauge,
33 + kv_cache_usage: F64Gauge,
34 +
35 + // Prefix-cache counters, including the connector-backed external cache path.
36 + prefix_cache_queries: U64Counter,
37 + prefix_cache_hits: U64Counter,
38 + external_prefix_cache_queries: U64Counter,
39 + external_prefix_cache_hits: U64Counter,
40 +
41 + // Speculative decoding counters.
42 + spec_decode_num_drafts: U64Counter,
43 + spec_decode_num_draft_tokens: U64Counter,
44 + spec_decode_num_accepted_tokens: U64Counter,
45 + spec_decode_num_accepted_tokens_per_pos: Family<EnginePositionLabels, U64Counter>,
46 +
47 + // Per-engine performance / MFU counters.
48 + estimated_flops_per_gpu: U64Counter,
49 + estimated_read_bytes_per_gpu: U64Counter,
50 + estimated_write_bytes_per_gpu: U64Counter,
51 +
52 + // Sampled KV-cache residency histograms.
53 + kv_block_lifetime_seconds: HistogramMetric,
54 + kv_block_idle_before_evict_seconds: HistogramMetric,
55 + kv_block_reuse_gap_seconds: HistogramMetric,
56 +
57 + // Non-Prometheus interval accumulator for periodic text-log helpers.
58 + log_stats: SchedulerLogStatsAccumulator,
59 +}
60 +
61 +impl SchedulerStatsRecorder {
62 + /// Resolve the fixed-label metric handles for the connected engines.
63 + pub(crate) fn new(
64 + metrics: &SchedulerMetrics,
65 + model_name: &str,
66 + engines: &[ConnectedEngine],
67 + ) -> Self {
68 + let engines = engines
69 + .iter()
70 + .filter_map(|engine| {
71 + let engine = engine.engine_id.engine_index()?;
72 + Some((
73 + engine,
74 + resolve_scheduler_stats_handles(metrics, model_name, engine),
75 + ))
76 + })
77 + .collect();
78 +
79 + Self { engines }
80 + }
81 +
82 + /// Record one scheduler-stats payload for the given engine index.
83 + pub(crate) fn record(&self, engine_index: u32, stats: &SchedulerStats) {
84 + if let Some(handles) = self.engines.get(&engine_index) {
85 + record_scheduler_stats_with_handles(handles, stats);
86 + }
87 + }
88 +}
89 +
90 +/// Resolve all fixed-label scheduler metrics for one engine.
91 +fn resolve_scheduler_stats_handles(
92 metrics: &SchedulerMetrics,
18 - model_name: impl Into<String>,
93 + model_name: &str,
94 engine: u32,
20 - stats: &SchedulerStats,
21 -) {
22 - let model_name = model_name.into();
95 +) -> SchedulerStatsHandles {
96 let labels = EngineLabels {
24 - model_name: model_name.clone(),
97 + model_name: model_name.to_string(),
98 + engine,
99 + };
100 + let capacity = WaitingReasonLabels {
101 + model_name: model_name.to_string(),
102 + engine,
103 + reason: WAITING_REASON_CAPACITY,
104 + };
105 + let deferred = WaitingReasonLabels {
106 + model_name: model_name.to_string(),
107 engine,
108 + reason: WAITING_REASON_DEFERRED,
109 };
110
111 + SchedulerStatsHandles {
112 + scheduler_running: metrics.scheduler_running.get_or_create_owned(&labels),
113 + scheduler_waiting: metrics.scheduler_waiting.get_or_create_owned(&labels),
114 + scheduler_waiting_capacity: metrics
115 + .scheduler_waiting_by_reason
116 + .get_or_create_owned(&capacity),
117 + scheduler_waiting_deferred: metrics
118 + .scheduler_waiting_by_reason
119 + .get_or_create_owned(&deferred),
120 + kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels),
121 + prefix_cache_queries: metrics.prefix_cache_queries.get_or_create_owned(&labels),
122 + prefix_cache_hits: metrics.prefix_cache_hits.get_or_create_owned(&labels),
123 + external_prefix_cache_queries: metrics
124 + .external_prefix_cache_queries
125 + .get_or_create_owned(&labels),
126 + external_prefix_cache_hits: metrics.external_prefix_cache_hits.get_or_create_owned(&labels),
127 + spec_decode_num_drafts: metrics.spec_decode_num_drafts.get_or_create_owned(&labels),
128 + spec_decode_num_draft_tokens: metrics
129 + .spec_decode_num_draft_tokens
130 + .get_or_create_owned(&labels),
131 + spec_decode_num_accepted_tokens: metrics
132 + .spec_decode_num_accepted_tokens
133 + .get_or_create_owned(&labels),
134 + spec_decode_num_accepted_tokens_per_pos: metrics
135 + .spec_decode_num_accepted_tokens_per_pos
136 + .clone(),
137 + log_stats: metrics.log_stats.get_or_create_owned(&labels),
138 + estimated_flops_per_gpu: metrics.estimated_flops_per_gpu.get_or_create_owned(&labels),
139 + estimated_read_bytes_per_gpu: metrics
140 + .estimated_read_bytes_per_gpu
141 + .get_or_create_owned(&labels),
142 + estimated_write_bytes_per_gpu: metrics
143 + .estimated_write_bytes_per_gpu
144 + .get_or_create_owned(&labels),
145 + kv_block_lifetime_seconds: metrics.kv_block_lifetime_seconds.get_or_create_owned(&labels),
146 + kv_block_idle_before_evict_seconds: metrics
147 + .kv_block_idle_before_evict_seconds
148 + .get_or_create_owned(&labels),
149 + kv_block_reuse_gap_seconds: metrics.kv_block_reuse_gap_seconds.get_or_create_owned(&labels),
150 + labels,
151 + }
152 +}
153 +
154 +/// Record scheduler-stats values through pre-resolved metric handles.
155 +fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &SchedulerStats) {
156 // Scheduler state gauges.
29 - metrics.scheduler_running.get_or_create(&labels).set(stats.num_running_reqs);
30 - metrics
157 + handles.scheduler_running.set(stats.num_running_reqs);
158 + handles
159 .scheduler_waiting
32 - .get_or_create(&labels)
160 .set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs);
34 - metrics
35 - .scheduler_waiting_by_reason
36 - .get_or_create(&WaitingReasonLabels {
37 - model_name: model_name.clone(),
38 - engine,
39 - reason: WAITING_REASON_CAPACITY,
40 - })
41 - .set(stats.num_waiting_reqs);
42 - metrics
43 - .scheduler_waiting_by_reason
44 - .get_or_create(&WaitingReasonLabels {
45 - model_name: model_name.clone(),
46 - engine,
47 - reason: WAITING_REASON_DEFERRED,
48 - })
49 - .set(stats.num_skipped_waiting_reqs);
50 - metrics.kv_cache_usage.get_or_create(&labels).set(stats.kv_cache_usage);
161 + handles.scheduler_waiting_capacity.set(stats.num_waiting_reqs);
162 + handles.scheduler_waiting_deferred.set(stats.num_skipped_waiting_reqs);
163 + handles.kv_cache_usage.set(stats.kv_cache_usage);
164
165 // Prefix-cache counters, including the connector-backed external cache path.
53 - metrics
54 - .prefix_cache_queries
55 - .get_or_create(&labels)
56 - .inc_by(stats.prefix_cache_stats.base.queries);
57 - metrics
58 - .prefix_cache_hits
59 - .get_or_create(&labels)
60 - .inc_by(stats.prefix_cache_stats.base.hits);
166 + handles.prefix_cache_queries.inc_by(stats.prefix_cache_stats.base.queries);
167 + handles.prefix_cache_hits.inc_by(stats.prefix_cache_stats.base.hits);
168
169 if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats {
63 - metrics
170 + handles
171 .external_prefix_cache_queries
65 - .get_or_create(&labels)
172 .inc_by(connector_prefix_cache_stats.base.queries);
67 - metrics
173 + handles
174 .external_prefix_cache_hits
69 - .get_or_create(&labels)
175 .inc_by(connector_prefix_cache_stats.base.hits);
176 }
177
178 // Speculative decoding counters.
179 if let Some(spec_decoding_stats) = &stats.spec_decoding_stats {
75 - metrics
76 - .spec_decode_num_drafts
77 - .get_or_create(&labels)
78 - .inc_by(spec_decoding_stats.num_drafts);
79 - metrics
180 + handles.spec_decode_num_drafts.inc_by(spec_decoding_stats.num_drafts);
181 + handles
182 .spec_decode_num_draft_tokens
81 - .get_or_create(&labels)
183 .inc_by(spec_decoding_stats.num_draft_tokens);
83 - metrics
184 + handles
185 .spec_decode_num_accepted_tokens
85 - .get_or_create(&labels)
186 .inc_by(spec_decoding_stats.num_accepted_tokens);
87 - metrics.log_stats.get_or_create(&labels).observe_spec_decode(
187 + handles.log_stats.observe_spec_decode(
188 spec_decoding_stats.num_drafts,
189 &spec_decoding_stats.num_accepted_tokens_per_pos,
190 );
@@ -92,11 +192,11 @@ pub(crate) fn record_scheduler_stats(
192 for (position, accepted_tokens) in
193 spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate()
194 {
95 - metrics
195 + handles
196 .spec_decode_num_accepted_tokens_per_pos
197 .get_or_create(&EnginePositionLabels {
98 - model_name: model_name.clone(),
99 - engine,
198 + model_name: handles.labels.model_name.clone(),
199 + engine: handles.labels.engine,
200 position: position as u32,
201 })
202 .inc_by(accepted_tokens);
@@ -109,22 +209,13 @@ pub(crate) fn record_scheduler_stats(
209 || perf_stats.num_read_bytes_per_gpu != 0
210 || perf_stats.num_write_bytes_per_gpu != 0)
211 {
112 - metrics
113 - .estimated_flops_per_gpu
114 - .get_or_create(&labels)
115 - .inc_by(perf_stats.num_flops_per_gpu);
116 - metrics
117 - .estimated_read_bytes_per_gpu
118 - .get_or_create(&labels)
119 - .inc_by(perf_stats.num_read_bytes_per_gpu);
120 - metrics
121 - .estimated_write_bytes_per_gpu
122 - .get_or_create(&labels)
123 - .inc_by(perf_stats.num_write_bytes_per_gpu);
212 + handles.estimated_flops_per_gpu.inc_by(perf_stats.num_flops_per_gpu);
213 + handles.estimated_read_bytes_per_gpu.inc_by(perf_stats.num_read_bytes_per_gpu);
214 + handles.estimated_write_bytes_per_gpu.inc_by(perf_stats.num_write_bytes_per_gpu);
215 }
216
217 if let Some(cudagraph_stats) = &stats.cudagraph_stats {
127 - metrics.log_stats.get_or_create(&labels).observe_cudagraph(
218 + handles.log_stats.observe_cudagraph(
219 cudagraph_stats.num_unpadded_tokens,
220 cudagraph_stats.num_padded_tokens,
221 cudagraph_stats.num_paddings,
@@ -134,16 +225,11 @@ pub(crate) fn record_scheduler_stats(
225
226 // Sampled KV-cache residency histograms.
227 if !stats.kv_cache_eviction_events.is_empty() {
137 - let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels);
138 - let kv_block_idle_before_evict_seconds =
139 - metrics.kv_block_idle_before_evict_seconds.get_or_create(&labels);
140 - let kv_block_reuse_gap_seconds = metrics.kv_block_reuse_gap_seconds.get_or_create(&labels);
141 -
228 for event in &stats.kv_cache_eviction_events {
143 - kv_block_lifetime_seconds.observe(event.lifetime_seconds);
144 - kv_block_idle_before_evict_seconds.observe(event.idle_seconds);
229 + handles.kv_block_lifetime_seconds.observe(event.lifetime_seconds);
230 + handles.kv_block_idle_before_evict_seconds.observe(event.idle_seconds);
231 for reuse_gap_seconds in &event.reuse_gaps_seconds {
146 - kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds);
232 + handles.kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds);
233 }
234 }
235 }
rust/src/llm/src/lib.rs
+11 -4
@@ -88,14 +88,21 @@ impl Llm {
88 // Record internal engine-core request ID in the current tracing span.
89 Span::current().record("engine_request_id", &internal_request_id);
90
91 + let arrival_time = prepared.engine_request.arrival_time;
92 + let max_tokens_param =
93 + (prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens);
94 + let prompt_len = prepared.prompt_token_ids().len() as u32;
95 +
96 + let stream = self.client.call(prepared.engine_request).await?;
97 +
98 let request_metrics = RequestMetricsTracker::new(
99 self.client.model_name().to_string(),
93 - prepared.engine_request.arrival_time,
94 - prepared.prompt_token_ids().len() as u32,
95 - (prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens),
100 + stream.engine_index(),
101 + arrival_time,
102 + prompt_len,
103 + max_tokens_param,
104 1,
105 );
98 - let stream = self.client.call(prepared.engine_request).await?;
106 let guard = self.inflight.track(external_request_id, internal_request_id);
107
108 Ok(GenerateOutputStream::new(
rust/src/llm/src/output.rs
+1 -6
@@ -248,12 +248,7 @@ impl Stream for GenerateOutputStream {
248 };
249
250 let received_at = current_unix_timestamp_secs();
251 - self.request_metrics.observe_output(
252 - raw.engine_index,
253 - raw.timestamp,
254 - received_at,
255 - &raw.output,
256 - );
251 + self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output);
252
253 let raw = raw.output;
254
rust/src/llm/src/request_metrics.rs
+150 -142
@@ -5,15 +5,12 @@ use vllm_engine_core_client::protocol::output::{
5 };
6 use vllm_engine_core_client::protocol::stats::PrefillStats;
7 use vllm_metrics::{
8 - EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics,
8 + EngineLabels, Family, FinishedReasonLabels, HistogramMetric, METRICS, PromptTokenSourceLabels,
9 + U64Counter,
10 };
11
12 use crate::FinishReason;
13
13 -fn metrics() -> &'static RequestMetrics {
14 - &METRICS.request
15 -}
16 -
14 const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute";
15 const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit";
16 const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer";
@@ -29,9 +26,11 @@ const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer";
26 ///
27 /// Original Python update flow:
28 /// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/engine/output_processor.py#L600-L677>
32 -#[derive(Debug, Clone)]
29 +#[derive(Clone)]
30 pub(crate) struct RequestMetricsTracker {
34 - model_name: String,
31 + /// Cached request metric handles for this request's model and engine index.
32 + handles: RequestMetricHandles,
33 +
34 arrival_time: f64,
35 prompt_len: u32,
36 max_tokens_param: Option<u32>,
@@ -44,7 +43,38 @@ pub(crate) struct RequestMetricsTracker {
43 first_token_latency: f64,
44 num_generation_tokens: u32,
45 latest_num_cached_tokens: u32,
47 - last_seen_engine_index: u32,
46 +}
47 +
48 +/// Cached request metric handles for one model and engine index.
49 +#[derive(Clone)]
50 +struct RequestMetricHandles {
51 + labels: EngineLabels,
52 +
53 + // Request-derived counters.
54 + num_preemptions: U64Counter,
55 + prompt_tokens: U64Counter,
56 + prompt_tokens_local_compute: U64Counter,
57 + prompt_tokens_local_cache_hit: U64Counter,
58 + prompt_tokens_external_kv_transfer: U64Counter,
59 + prompt_tokens_cached: U64Counter,
60 + generation_tokens: U64Counter,
61 +
62 + // Request lifecycle counters and histograms.
63 + request_success: Family<FinishedReasonLabels, U64Counter>,
64 + request_prompt_tokens: HistogramMetric,
65 + request_generation_tokens: HistogramMetric,
66 + request_max_num_generation_tokens: HistogramMetric,
67 + request_params_max_tokens: HistogramMetric,
68 + request_params_n: HistogramMetric,
69 + request_prefill_kv_computed_tokens: HistogramMetric,
70 + time_to_first_token_seconds: HistogramMetric,
71 + inter_token_latency_seconds: HistogramMetric,
72 + e2e_request_latency_seconds: HistogramMetric,
73 + request_queue_time_seconds: HistogramMetric,
74 + request_prefill_time_seconds: HistogramMetric,
75 + request_decode_time_seconds: HistogramMetric,
76 + request_inference_time_seconds: HistogramMetric,
77 + request_time_per_output_token_seconds: HistogramMetric,
78 }
79
80 impl RequestMetricsTracker {
@@ -52,13 +82,14 @@ impl RequestMetricsTracker {
82 /// context.
83 pub(crate) fn new(
84 model_name: String,
85 + engine_index: u32,
86 arrival_time: f64,
87 prompt_len: u32,
88 max_tokens_param: Option<u32>,
89 n_param: u32,
90 ) -> Self {
91 Self {
61 - model_name,
92 + handles: resolve_request_metric_handles(&model_name, engine_index),
93 arrival_time,
94 prompt_len,
95 max_tokens_param,
@@ -71,7 +102,6 @@ impl RequestMetricsTracker {
102 first_token_latency: 0.0,
103 num_generation_tokens: 0,
104 latest_num_cached_tokens: 0,
74 - last_seen_engine_index: 0,
105 }
106 }
107
@@ -81,23 +111,18 @@ impl RequestMetricsTracker {
111 /// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L331-L384>
112 pub(crate) fn observe_output(
113 &mut self,
84 - engine_index: u32,
114 batch_timestamp: f64,
115 received_at: f64,
116 output: &EngineCoreOutput,
117 ) {
89 - self.last_seen_engine_index = engine_index;
118 if let Some(prefill_stats) = &output.prefill_stats {
119 self.latest_num_cached_tokens = prefill_stats.num_cached_tokens;
120 }
121 self.num_generation_tokens += output.new_token_ids.len() as u32;
94 - metrics()
95 - .generation_tokens
96 - .get_or_create(&engine_labels(&self.model_name, engine_index))
97 - .inc_by(output.new_token_ids.len() as u64);
122 + self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64);
123
124 if let Some(events) = &output.events {
100 - self.observe_events(engine_index, events);
125 + self.observe_events(events);
126 }
127
128 // Only outputs that actually carry tokens drive token-timing metrics.
@@ -107,22 +132,16 @@ impl RequestMetricsTracker {
132 if !output.new_token_ids.is_empty() {
133 if self.is_prefilling {
134 if let Some(prefill_stats) = &output.prefill_stats {
110 - record_prompt_tokens(&self.model_name, engine_index, prefill_stats);
135 + self.record_prompt_tokens(prefill_stats);
136 }
137 self.first_token_latency = received_at - self.arrival_time;
113 - observe_time_to_first_token_seconds(
114 - &self.model_name,
115 - engine_index,
116 - self.first_token_latency,
117 - );
138 + self.handles.time_to_first_token_seconds.observe(self.first_token_latency);
139 self.first_token_ts = batch_timestamp;
140 self.is_prefilling = false;
141 } else if self.last_token_ts > 0.0 {
121 - observe_inter_token_latency_seconds(
122 - &self.model_name,
123 - engine_index,
124 - batch_timestamp - self.last_token_ts,
125 - );
142 + self.handles
143 + .inter_token_latency_seconds
144 + .observe(batch_timestamp - self.last_token_ts);
145 }
146
147 self.last_token_ts = batch_timestamp;
@@ -135,7 +154,6 @@ impl RequestMetricsTracker {
154 /// Original Python finished-request stats:
155 /// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L222-L237>
156 pub(crate) fn record_finished(&self, received_at: f64, finish_reason: FinishReason) {
138 - let labels = engine_labels(&self.model_name, self.last_seen_engine_index);
157 let prefill_kv_computed_tokens =
158 self.prompt_len.saturating_sub(self.latest_num_cached_tokens);
159 let e2e_latency_seconds = received_at - self.arrival_time;
@@ -150,57 +168,47 @@ impl RequestMetricsTracker {
168 0.0
169 };
170
153 - record_request_success(&self.model_name, self.last_seen_engine_index, finish_reason);
154 - metrics()
155 - .request_prompt_tokens
156 - .get_or_create(&labels)
157 - .observe(self.prompt_len as f64);
158 - metrics()
171 + self.record_request_success(finish_reason);
172 +
173 + self.handles.request_prompt_tokens.observe(self.prompt_len as f64);
174 + self.handles
175 .request_generation_tokens
160 - .get_or_create(&labels)
176 .observe(self.num_generation_tokens as f64);
162 - metrics()
177 + self.handles
178 .request_max_num_generation_tokens
164 - .get_or_create(&labels)
179 .observe(self.num_generation_tokens as f64);
180 if let Some(max_tokens_param) = self.max_tokens_param {
167 - metrics()
168 - .request_params_max_tokens
169 - .get_or_create(&labels)
170 - .observe(max_tokens_param as f64);
181 + self.handles.request_params_max_tokens.observe(max_tokens_param as f64);
182 }
172 - metrics().request_params_n.get_or_create(&labels).observe(self.n_param as f64);
173 - metrics()
183 + self.handles.request_params_n.observe(self.n_param as f64);
184 + self.handles
185 .request_prefill_kv_computed_tokens
175 - .get_or_create(&labels)
186 .observe(prefill_kv_computed_tokens as f64);
177 - metrics()
178 - .e2e_request_latency_seconds
179 - .get_or_create(&labels)
180 - .observe(e2e_latency_seconds);
181 - metrics()
182 - .request_queue_time_seconds
183 - .get_or_create(&labels)
184 - .observe(queue_time_seconds);
185 - metrics()
186 - .request_prefill_time_seconds
187 - .get_or_create(&labels)
188 - .observe(prefill_time_seconds);
189 - metrics()
190 - .request_decode_time_seconds
191 - .get_or_create(&labels)
192 - .observe(decode_time_seconds);
193 - metrics()
194 - .request_inference_time_seconds
195 - .get_or_create(&labels)
196 - .observe(inference_time_seconds);
197 - metrics()
187 + self.handles.e2e_request_latency_seconds.observe(e2e_latency_seconds);
188 + self.handles.request_queue_time_seconds.observe(queue_time_seconds);
189 + self.handles.request_prefill_time_seconds.observe(prefill_time_seconds);
190 + self.handles.request_decode_time_seconds.observe(decode_time_seconds);
191 + self.handles.request_inference_time_seconds.observe(inference_time_seconds);
192 + self.handles
193 .request_time_per_output_token_seconds
199 - .get_or_create(&labels)
194 .observe(time_per_output_token_seconds);
195 }
196
203 - fn observe_events(&mut self, engine_index: u32, events: &[EngineCoreEvent]) {
197 + /// Record prompt token counters through cached metric handles.
198 + fn record_prompt_tokens(&self, prefill_stats: &PrefillStats) {
199 + let computed = prefill_stats.num_computed_tokens as u64;
200 + let local_cache_hit = prefill_stats.num_local_cached_tokens as u64;
201 + let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64;
202 +
203 + self.handles.prompt_tokens.inc_by(prefill_stats.num_prompt_tokens as u64);
204 + self.handles.prompt_tokens_local_compute.inc_by(computed);
205 + self.handles.prompt_tokens_local_cache_hit.inc_by(local_cache_hit);
206 + self.handles.prompt_tokens_external_kv_transfer.inc_by(external_kv_transfer);
207 + self.handles.prompt_tokens_cached.inc_by(prefill_stats.num_cached_tokens as u64);
208 + }
209 +
210 + /// Record request event counters through cached metric handles.
211 + fn observe_events(&mut self, events: &[EngineCoreEvent]) {
212 for event in events {
213 match event.r#type {
214 EngineCoreEventType::Queued => {
@@ -212,46 +220,86 @@ impl RequestMetricsTracker {
220 }
221 }
222 EngineCoreEventType::Preempted => {
215 - metrics()
216 - .num_preemptions
217 - .get_or_create(&engine_labels(&self.model_name, engine_index))
218 - .inc();
223 + self.handles.num_preemptions.inc();
224 }
225 }
226 }
227 }
223 -}
228
225 -fn engine_labels(model_name: &str, engine: u32) -> EngineLabels {
226 - EngineLabels {
227 - model_name: model_name.to_string(),
228 - engine,
229 + /// Increment the request-success counter for the terminal finish reason.
230 + fn record_request_success(&self, finish_reason: FinishReason) {
231 + self.handles
232 + .request_success
233 + .get_or_create(&FinishedReasonLabels {
234 + model_name: self.handles.labels.model_name.clone(),
235 + engine: self.handles.labels.engine,
236 + finished_reason: finish_reason.as_str(),
237 + })
238 + .inc();
239 }
240 }
241
232 -fn observe_time_to_first_token_seconds(model_name: &str, engine: u32, seconds: f64) {
233 - metrics()
234 - .time_to_first_token_seconds
235 - .get_or_create(&engine_labels(model_name, engine))
236 - .observe(seconds);
237 -}
238 -
239 -fn observe_inter_token_latency_seconds(model_name: &str, engine: u32, seconds: f64) {
240 - metrics()
241 - .inter_token_latency_seconds
242 - .get_or_create(&engine_labels(model_name, engine))
243 - .observe(seconds);
244 -}
242 +/// Resolve fixed request metric handles for one model and engine index.
243 +fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetricHandles {
244 + let metrics = &METRICS.request;
245 + let labels = EngineLabels {
246 + model_name: model_name.to_string(),
247 + engine,
248 + };
249
246 -fn record_request_success(model_name: &str, engine: u32, finish_reason: FinishReason) {
247 - metrics()
248 - .request_success
249 - .get_or_create(&FinishedReasonLabels {
250 - model_name: model_name.to_string(),
251 - engine,
252 - finished_reason: finish_reason.as_str(),
253 - })
254 - .inc();
250 + RequestMetricHandles {
251 + num_preemptions: metrics.num_preemptions.get_or_create_owned(&labels),
252 + prompt_tokens: metrics.prompt_tokens.get_or_create_owned(&labels),
253 + prompt_tokens_local_compute: metrics.prompt_tokens_by_source.get_or_create_owned(
254 + &prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE),
255 + ),
256 + prompt_tokens_local_cache_hit: metrics.prompt_tokens_by_source.get_or_create_owned(
257 + &prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT),
258 + ),
259 + prompt_tokens_external_kv_transfer: metrics.prompt_tokens_by_source.get_or_create_owned(
260 + &prompt_token_source_labels(
261 + model_name,
262 + engine,
263 + PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER,
264 + ),
265 + ),
266 + prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels),
267 + generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels),
268 + request_success: metrics.request_success.clone(),
269 + request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels),
270 + request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels),
271 + request_max_num_generation_tokens: metrics
272 + .request_max_num_generation_tokens
273 + .get_or_create_owned(&labels),
274 + request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels),
275 + request_params_n: metrics.request_params_n.get_or_create_owned(&labels),
276 + request_prefill_kv_computed_tokens: metrics
277 + .request_prefill_kv_computed_tokens
278 + .get_or_create_owned(&labels),
279 + time_to_first_token_seconds: metrics
280 + .time_to_first_token_seconds
281 + .get_or_create_owned(&labels),
282 + inter_token_latency_seconds: metrics
283 + .inter_token_latency_seconds
284 + .get_or_create_owned(&labels),
285 + e2e_request_latency_seconds: metrics
286 + .e2e_request_latency_seconds
287 + .get_or_create_owned(&labels),
288 + request_queue_time_seconds: metrics.request_queue_time_seconds.get_or_create_owned(&labels),
289 + request_prefill_time_seconds: metrics
290 + .request_prefill_time_seconds
291 + .get_or_create_owned(&labels),
292 + request_decode_time_seconds: metrics
293 + .request_decode_time_seconds
294 + .get_or_create_owned(&labels),
295 + request_inference_time_seconds: metrics
296 + .request_inference_time_seconds
297 + .get_or_create_owned(&labels),
298 + request_time_per_output_token_seconds: metrics
299 + .request_time_per_output_token_seconds
300 + .get_or_create_owned(&labels),
301 + labels,
302 + }
303 }
304
305 fn prompt_token_source_labels(
@@ -266,45 +314,6 @@ fn prompt_token_source_labels(
314 }
315 }
316
269 -fn record_prompt_tokens(model_name: &str, engine: u32, prefill_stats: &PrefillStats) {
270 - let computed = prefill_stats.num_computed_tokens as u64;
271 - let local_cache_hit = prefill_stats.num_local_cached_tokens as u64;
272 - let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64;
273 -
274 - metrics()
275 - .prompt_tokens
276 - .get_or_create(&engine_labels(model_name, engine))
277 - .inc_by(prefill_stats.num_prompt_tokens as u64);
278 - metrics()
279 - .prompt_tokens_by_source
280 - .get_or_create(&prompt_token_source_labels(
281 - model_name,
282 - engine,
283 - PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE,
284 - ))
285 - .inc_by(computed);
286 - metrics()
287 - .prompt_tokens_by_source
288 - .get_or_create(&prompt_token_source_labels(
289 - model_name,
290 - engine,
291 - PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT,
292 - ))
293 - .inc_by(local_cache_hit);
294 - metrics()
295 - .prompt_tokens_by_source
296 - .get_or_create(&prompt_token_source_labels(
297 - model_name,
298 - engine,
299 - PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER,
300 - ))
301 - .inc_by(external_kv_transfer);
302 - metrics()
303 - .prompt_tokens_cached
304 - .get_or_create(&engine_labels(model_name, engine))
305 - .inc_by(prefill_stats.num_cached_tokens as u64);
306 -}
307 -
317 fn diff_or_zero(end: f64, start: f64) -> f64 {
318 if end > 0.0 && start > 0.0 && end >= start {
319 end - start
@@ -337,10 +346,10 @@ mod tests {
346
347 #[test]
348 fn tracker_updates_timing_state_across_prefill_decode_and_finish() {
340 - let mut tracker = RequestMetricsTracker::new("model".to_string(), 100.0, 64, Some(128), 1);
349 + let mut tracker =
350 + RequestMetricsTracker::new("model".to_string(), 2, 100.0, 64, Some(128), 1);
351
352 tracker.observe_output(
343 - 2,
353 10.0,
354 100.2,
355 &vllm_engine_core_client::protocol::output::EngineCoreOutput {
@@ -368,7 +377,6 @@ mod tests {
377 },
378 );
379 tracker.observe_output(
371 - 2,
380 11.5,
381 100.4,
382 &vllm_engine_core_client::protocol::output::EngineCoreOutput {
@@ -384,7 +392,7 @@ mod tests {
392 );
393
394 assert!(!tracker.is_prefilling);
387 - assert_eq!(tracker.last_seen_engine_index, 2);
395 + assert_eq!(tracker.handles.labels.engine, 2);
396 assert_eq!(tracker.num_generation_tokens, 3);
397 assert_eq!(tracker.queued_ts, 8.0);
398 assert_eq!(tracker.scheduled_ts, 9.0);
rust/src/llm/tests/generate.rs
+3 -3
@@ -17,7 +17,7 @@ use vllm_engine_core_client::protocol::request::EngineCoreRequest;
17 use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
18 use vllm_engine_core_client::protocol::stats::PrefillStats;
19 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
20 -use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
20 +use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
21 use vllm_llm::{
22 Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm,
23 };
@@ -699,7 +699,7 @@ async fn abort_by_external_id_aborts_all_internal_requests() {
699 async fn generate_records_request_metrics_in_prometheus_output() {
700 let ipc = IpcNamespace::new().unwrap();
701 let handshake_address = ipc.handshake_endpoint();
702 - let engine_id = b"engine-metrics".to_vec();
702 + let engine_id = EngineId::from_engine_index(4);
703 let model_name = request_metrics_model_name("metrics-model");
704
705 let (shutdown_tx, engine_task) = spawn_mock_engine_task(
@@ -832,7 +832,7 @@ async fn generate_records_request_metrics_in_prometheus_output() {
832 async fn dropping_stream_records_abort_terminal_request_metrics() {
833 let ipc = IpcNamespace::new().unwrap();
834 let handshake_address = ipc.handshake_endpoint();
835 - let engine_id = b"engine-metrics-drop".to_vec();
835 + let engine_id = EngineId::from_engine_index(5);
836 let model_name = request_metrics_model_name("metrics-drop-model");
837
838 let (shutdown_tx, engine_task) = spawn_mock_engine_task(
rust/src/metrics/src/lib.rs
+3 -1
@@ -4,7 +4,7 @@ use std::sync::atomic::AtomicU64;
4
5 use prometheus_client::encoding::text::encode;
6 use prometheus_client::metrics::counter::Counter;
7 -use prometheus_client::metrics::family::Family;
7 +pub use prometheus_client::metrics::family::Family;
8 use prometheus_client::metrics::gauge::Gauge;
9 use prometheus_client::metrics::histogram::Histogram;
10 use prometheus_client::registry::Registry;
@@ -23,6 +23,8 @@ pub use scheduler::*;
23 pub type U64Counter = Counter<u64, AtomicU64>;
24 pub type U64Gauge = Gauge<u64, AtomicU64>;
25 pub type F64Gauge = Gauge<f64, AtomicU64>;
26 +/// Histogram metric handle cloned out of a Prometheus family.
27 +pub type HistogramMetric = Histogram;
28 pub(crate) type HistogramFamily = Family<EngineLabels, Histogram, fn() -> Histogram>;
29
30 /// Shared Prometheus registry for frontend metrics.