From ae2d1ca96fa2a98f92c775acfe8a154748495a27 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Fri, 4 Sep 2026 08:21:54 +0200 Subject: [PATCH] [Rust Frontend] Record Mooncake/NIXL KV-connector metrics (#52755) Co-authored-by: Bugen Zhao Signed-off-by: Markov Ilya Signed-off-by: Bugen Zhao --- rust/src/engine-core-client/src/metrics.rs | 184 ++++++++++++++++- .../engine-core-client/src/protocol/stats.rs | 120 ++++++++++- .../engine-core-client/src/tests/client.rs | 28 ++- .../src/tests/python_compat.py | 29 +++ rust/src/metrics/src/scheduler.rs | 195 ++++++++++++++++++ .../kv_connector/unit/test_multi_connector.py | 81 ++++---- .../kv_connector/unit/test_nixl_connector.py | 3 +- .../kv_transfer/kv_connector/v1/metrics.py | 4 + .../kv_connector/v1/multi_connector.py | 16 +- vllm/v1/core/sched/scheduler.py | 2 +- 10 files changed, 600 insertions(+), 62 deletions(-) diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index 8b2eb9e588..e15258d608 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -7,11 +7,14 @@ use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, - LoraInfoLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, + LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, + MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, WaitingReasonLabels, }; -use crate::protocol::stats::SchedulerStats; +use crate::protocol::stats::{ + KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats, SchedulerStats, +}; use crate::transport::ConnectedEngine; const WAITING_REASON_CAPACITY: &str = "capacity"; @@ -57,6 +60,24 @@ struct SchedulerStatsHandles { kv_block_idle_before_evict_seconds: HistogramMetric, kv_block_reuse_gap_seconds: HistogramMetric, + // Mooncake store connector telemetry, decoded from `kv_connector_stats`. + // Kept as `Family` (not pre-resolved) because `operation`/`status` are + // dynamic per-record labels. + mooncake_operation_time_seconds: MooncakeOperationHistogramFamily, + mooncake_operation_total: MooncakeOperationCounterFamily, + mooncake_operation_keys_total: MooncakeOperationCounterFamily, + mooncake_operation_bytes_total: MooncakeOperationCounterFamily, + mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily, + + // NIXL connector telemetry, decoded from `kv_connector_stats`. + nixl_xfer_time_seconds: HistogramMetric, + nixl_post_time_seconds: HistogramMetric, + nixl_bytes_transferred: HistogramMetric, + nixl_num_descriptors: HistogramMetric, + nixl_num_failed_transfers: U64Counter, + nixl_num_failed_notifications: U64Counter, + nixl_num_kv_expired_reqs: U64Counter, + // Non-Prometheus interval accumulator for periodic text-log helpers. log_stats: SchedulerLogStatsAccumulator, } @@ -150,6 +171,20 @@ fn resolve_scheduler_stats_handles( .kv_block_idle_before_evict_seconds .get_or_create_owned(&labels), kv_block_reuse_gap_seconds: metrics.kv_block_reuse_gap_seconds.get_or_create_owned(&labels), + mooncake_operation_time_seconds: metrics.mooncake_operation_time_seconds.clone(), + mooncake_operation_total: metrics.mooncake_operation_total.clone(), + mooncake_operation_keys_total: metrics.mooncake_operation_keys_total.clone(), + mooncake_operation_bytes_total: metrics.mooncake_operation_bytes_total.clone(), + mooncake_operation_failed_keys_total: metrics.mooncake_operation_failed_keys_total.clone(), + nixl_xfer_time_seconds: metrics.nixl_xfer_time_seconds.get_or_create_owned(&labels), + nixl_post_time_seconds: metrics.nixl_post_time_seconds.get_or_create_owned(&labels), + nixl_bytes_transferred: metrics.nixl_bytes_transferred.get_or_create_owned(&labels), + nixl_num_descriptors: metrics.nixl_num_descriptors.get_or_create_owned(&labels), + nixl_num_failed_transfers: metrics.nixl_num_failed_transfers.get_or_create_owned(&labels), + nixl_num_failed_notifications: metrics + .nixl_num_failed_notifications + .get_or_create_owned(&labels), + nixl_num_kv_expired_reqs: metrics.nixl_num_kv_expired_reqs.get_or_create_owned(&labels), labels, } } @@ -236,6 +271,79 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: & } } } + + // Connector-specific KV transfer stats. A bare connector reports its own + // flat payload; MultiConnector reports connector class name -> flat child + // payload. + if let Some(kv_connector_stats) = &stats.kv_connector_stats { + match kv_connector_stats { + KvConnectorStats::Nixl(stats) => record_nixl_stats(handles, stats), + KvConnectorStats::Mooncake(stats) => record_mooncake_stats(handles, stats), + KvConnectorStats::Multi(stats) => record_multi_connector_stats(handles, stats), + KvConnectorStats::Other(_) => {} + } + } +} + +fn record_multi_connector_stats(handles: &SchedulerStatsHandles, stats: &MultiConnectorStats) { + for nixl in [&stats.nixl, &stats.nixl_pull, &stats.nixl_push].into_iter().flatten() { + record_nixl_stats(handles, nixl); + } + if let Some(mooncake) = &stats.mooncake { + record_mooncake_stats(handles, mooncake); + } +} + +fn record_mooncake_stats(handles: &SchedulerStatsHandles, stats: &MooncakeStats) { + for (operation, records) in &stats.0 { + for record in records { + let labels = MooncakeOperationLabels { + model_name: handles.labels.model_name.clone(), + engine: handles.labels.engine, + operation: operation.as_str().to_string(), + status: record.status.as_str().to_string(), + }; + handles + .mooncake_operation_time_seconds + .get_or_create(&labels) + .observe(record.duration_seconds); + handles.mooncake_operation_total.get_or_create(&labels).inc(); + handles + .mooncake_operation_keys_total + .get_or_create(&labels) + .inc_by(record.num_keys); + handles + .mooncake_operation_bytes_total + .get_or_create(&labels) + .inc_by(record.num_bytes); + handles + .mooncake_operation_failed_keys_total + .get_or_create(&labels) + .inc_by(record.num_failed_keys); + } + } +} + +fn record_nixl_stats(handles: &SchedulerStatsHandles, stats: &NixlStats) { + for value in &stats.transfer_duration { + handles.nixl_xfer_time_seconds.observe(*value); + } + for value in &stats.post_duration { + handles.nixl_post_time_seconds.observe(*value); + } + for value in &stats.bytes_transferred { + handles.nixl_bytes_transferred.observe(*value as f64); + } + for value in &stats.num_descriptors { + handles.nixl_num_descriptors.observe(*value as f64); + } + handles + .nixl_num_failed_transfers + .inc_by(stats.num_failed_transfers.iter().sum()); + handles + .nixl_num_failed_notifications + .inc_by(stats.num_failed_notifications.iter().sum()); + handles.nixl_num_kv_expired_reqs.inc_by(stats.num_kv_expired_reqs.iter().sum()); } /// Exports `vllm:lora_requests_info` as a single series covering all LoRA @@ -286,12 +394,16 @@ fn now_unix_secs() -> f64 { #[cfg(test)] mod tests { - use std::collections::BTreeSet; + use std::collections::{BTreeMap, BTreeSet}; use expect_test::expect; use vllm_metrics::Metrics; use crate::metrics::LoraInfoExporter; + use crate::protocol::stats::{ + KvConnectorStats, MooncakeOperation, MooncakeRecord, MooncakeStats, MooncakeStatus, + MultiConnectorStats, NixlStats, SchedulerStats, + }; fn names(values: &[&str]) -> BTreeSet { values.iter().map(|name| (*name).to_string()).collect() @@ -345,4 +457,70 @@ mod tests { exporter.update(&metrics.scheduler, names(&[]), names(&[])); expect![[""]].assert_eq(&lora_series(&metrics.render().unwrap())); } + + fn nixl_stats() -> NixlStats { + NixlStats { + transfer_duration: vec![0.01, 0.02], + post_duration: vec![0.001, 0.002], + bytes_transferred: vec![4096, 8192], + num_descriptors: vec![2, 4], + num_failed_transfers: vec![], + num_failed_notifications: vec![], + num_kv_expired_reqs: vec![1], + } + } + + fn mooncake_stats() -> MooncakeStats { + MooncakeStats(BTreeMap::from([( + MooncakeOperation::LoadGet, + vec![MooncakeRecord { + duration_seconds: 0.05, + num_keys: 3, + num_bytes: 1024, + status: MooncakeStatus::Ok, + num_failed_keys: 0, + }], + )])) + } + + /// Records one MultiConnector payload into Mooncake and NIXL metrics. + #[test] + fn kv_connector_stats_are_recorded_into_mooncake_and_nixl_metrics() { + let metrics = Metrics::new(); + let handles = super::resolve_scheduler_stats_handles(&metrics.scheduler, "model", 0); + + let stats = SchedulerStats { + kv_connector_stats: Some(KvConnectorStats::Multi(Box::new(MultiConnectorStats { + nixl: Some(nixl_stats()), + mooncake: Some(mooncake_stats()), + ..Default::default() + }))), + ..Default::default() + }; + + super::record_scheduler_stats_with_handles(&handles, &stats); + + let rendered = metrics.render().unwrap(); + assert!(rendered.contains( + "vllm:mooncake_store_operation_total{model_name=\"model\",engine=\"0\",\ + operation=\"load_get\",status=\"ok\"} 1" + )); + assert!(rendered.contains( + "vllm:mooncake_store_operation_keys_total{model_name=\"model\",engine=\"0\",\ + operation=\"load_get\",status=\"ok\"} 3" + )); + assert!(rendered.contains( + "vllm:mooncake_store_operation_bytes_total{model_name=\"model\",engine=\"0\",\ + operation=\"load_get\",status=\"ok\"} 1024" + )); + assert!( + rendered.contains( + "vllm:nixl_num_kv_expired_reqs_total{model_name=\"model\",engine=\"0\"} 1" + ) + ); + assert!( + rendered + .contains("vllm:nixl_xfer_time_seconds_count{model_name=\"model\",engine=\"0\"} 2") + ); + } } diff --git a/rust/src/engine-core-client/src/protocol/stats.rs b/rust/src/engine-core-client/src/protocol/stats.rs index f02da161ef..021fd7c26a 100644 --- a/rust/src/engine-core-client/src/protocol/stats.rs +++ b/rust/src/engine-core-client/src/protocol/stats.rs @@ -5,8 +5,6 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; -use crate::protocol::OpaqueValue; - /// Stores cache hit statistics. /// /// Original Python definition: @@ -158,6 +156,120 @@ pub struct CudagraphStats { pub runtime_mode: String, } +/// KV connector telemetry DTOs carried by [`SchedulerStats`]. +pub mod kv_connector { + use std::collections::BTreeMap; + + use serde::{Deserialize, Serialize}; + + use crate::protocol::OpaqueValue; + + /// NIXL connector transfer telemetry. + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] + #[serde(deny_unknown_fields)] + pub struct NixlStats { + /// Time spent transferring each successful request. + pub transfer_duration: Vec, + /// Time spent posting each successful transfer. + pub post_duration: Vec, + /// Bytes transferred by each successful request. + pub bytes_transferred: Vec, + /// Descriptor count for each successful request. + pub num_descriptors: Vec, + /// Failure counter increments collected since the previous update. + pub num_failed_transfers: Vec, + /// Notification failure counter increments collected since the previous update. + pub num_failed_notifications: Vec, + /// Expired-request counter increments collected since the previous update. + pub num_kv_expired_reqs: Vec, + } + + /// Mooncake store operation name. + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] + #[serde(rename_all = "snake_case")] + pub enum MooncakeOperation { + SaveExists, + SavePut, + LoadGet, + LookupExists, + } + + impl MooncakeOperation { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::SaveExists => "save_exists", + Self::SavePut => "save_put", + Self::LoadGet => "load_get", + Self::LookupExists => "lookup_exists", + } + } + } + + /// Mooncake store operation status. + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] + #[serde(rename_all = "snake_case")] + pub enum MooncakeStatus { + Ok, + Error, + PartialFailure, + } + + impl MooncakeStatus { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Ok => "ok", + Self::Error => "error", + Self::PartialFailure => "partial_failure", + } + } + } + + /// One Mooncake store operation sample. + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + #[serde(deny_unknown_fields)] + pub struct MooncakeRecord { + pub duration_seconds: f64, + pub num_keys: u64, + pub num_bytes: u64, + pub status: MooncakeStatus, + pub num_failed_keys: u64, + } + + /// Mooncake store telemetry grouped by operation. + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] + #[serde(transparent)] + pub struct MooncakeStats(pub BTreeMap>); + + /// Telemetry emitted by a `MultiConnector` with a flat child payload map. + #[serde_with::skip_serializing_none] + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] + pub struct MultiConnectorStats { + #[serde(rename = "NixlConnector")] + pub nixl: Option, + #[serde(rename = "NixlPullConnector")] + pub nixl_pull: Option, + #[serde(rename = "NixlPushConnector")] + pub nixl_push: Option, + #[serde(rename = "MooncakeStoreConnector")] + pub mooncake: Option, + /// Child connector payloads without a Rust telemetry implementation. + #[serde(flatten)] + pub other: BTreeMap, + } + + /// Connector-specific scheduler telemetry. + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] + #[serde(untagged)] + pub enum KvConnectorStats { + Nixl(NixlStats), + Mooncake(MooncakeStats), + Multi(Box), + Other(BTreeMap), + } +} + +pub use kv_connector::*; + /// Stats associated with the scheduler. /// /// Original Python definition: @@ -185,8 +297,8 @@ pub struct SchedulerStats { pub kv_cache_eviction_events: Vec, /// Speculative decoding scheduler stats, when enabled. pub spec_decoding_stats: Option, - /// Connector-specific KV transfer stats, kept opaque for now. - pub kv_connector_stats: Option>, + /// Connector-specific KV transfer stats. + pub kv_connector_stats: Option, /// CUDA graph runtime stats when graph metrics are enabled. pub cudagraph_stats: Option, /// Estimated MFU/performance stats, when enabled. diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 6f4e53ff1b..d2d52c7294 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -31,7 +31,7 @@ use crate::protocol::output::{ }; use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; use crate::protocol::sampling::EngineCoreSamplingParams; -use crate::protocol::stats::SchedulerStats; +use crate::protocol::stats::{KvConnectorStats, MooncakeOperation, SchedulerStats}; use crate::protocol::tensor::{WireArrayData, WireTensor}; use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; use crate::test_utils::{ @@ -2635,6 +2635,10 @@ fn python_msgpack_fixtures_match_rust_encoding() { let inline_prompt_frames = lines.next().expect("missing inline prompt logprobs fixture line"); let multipart_prompt_frames = lines.next().expect("missing multipart prompt logprobs fixture line"); + let nixl_stats_hex = lines.next().expect("missing NIXL stats fixture line"); + let mooncake_stats_hex = lines.next().expect("missing Mooncake stats fixture line"); + let multi_connector_stats_hex = + lines.next().expect("missing MultiConnector stats fixture line"); let ready_response_hex = lines.next().expect("missing ready response fixture line"); let request_bytes = hex::decode(request_hex).unwrap(); @@ -2799,6 +2803,28 @@ fn python_msgpack_fixtures_match_rust_encoding() { .expect("multipart prompt logprobs decoded"), ); + let nixl_stats: KvConnectorStats = + rmp_serde::from_slice(&hex::decode(nixl_stats_hex).unwrap()).unwrap(); + assert!(matches!(nixl_stats, KvConnectorStats::Nixl(_))); + + let mooncake_stats: KvConnectorStats = + rmp_serde::from_slice(&hex::decode(mooncake_stats_hex).unwrap()).unwrap(); + assert!(matches!( + mooncake_stats, + KvConnectorStats::Mooncake(stats) + if stats.0.contains_key(&MooncakeOperation::LoadGet) + )); + + let multi_connector_stats: KvConnectorStats = + rmp_serde::from_slice(&hex::decode(multi_connector_stats_hex).unwrap()).unwrap(); + assert!(matches!( + multi_connector_stats, + KvConnectorStats::Multi(stats) + if stats.nixl.is_some() + && stats.mooncake.is_some() + && stats.other.contains_key("UnsupportedConnector") + )); + let map_keys = |bytes: &[u8]| -> BTreeSet { match decode_value(bytes) { Value::Map(entries) => entries diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index d3705cc391..fd6cb0c86b 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -456,6 +456,32 @@ ready_response = EngineCoreReadyResponse( ), ) +nixl_stats = { + "transfer_duration": [0.01, 0.02], + "post_duration": [0.001, 0.002], + "bytes_transferred": [4096, 8192], + "num_descriptors": [2, 4], + "num_failed_transfers": [], + "num_failed_notifications": [], + "num_kv_expired_reqs": [1], +} +mooncake_stats = { + "load_get": [ + { + "duration_seconds": 0.05, + "num_keys": 3, + "num_bytes": 1024, + "status": "ok", + "num_failed_keys": 0, + } + ] +} +multi_connector_stats = { + "NixlConnector": nixl_stats, + "MooncakeStoreConnector": mooncake_stats, + "UnsupportedConnector": {"sample_count": 1}, +} + print(msgspec.msgpack.encode(request).hex()) print(msgspec.msgpack.encode(defaults_request).hex()) print(msgpack.packb(multimodal_request_wire, use_bin_type=True).hex()) @@ -475,4 +501,7 @@ print( for frame in encode_output_frames(multipart_prompt_logprobs, size_threshold=1) ) ) +print(msgspec.msgpack.encode(nixl_stats).hex()) +print(msgspec.msgpack.encode(mooncake_stats).hex()) +print(msgspec.msgpack.encode(multi_connector_stats).hex()) print(msgspec.msgpack.encode(ready_response).hex()) diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index af43961b9a..d7f52fd0a6 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -29,6 +29,58 @@ fn kv_block_reuse_gap_histogram() -> Histogram { Histogram::new(KV_CACHE_RESIDENCY_BUCKETS.iter().copied()) } +// Buckets copied verbatim from the Python `MooncakeStorePromMetrics` / +// `NixlPromMetrics` classes so dashboards built against either frontend are +// interchangeable. +const MOONCAKE_OPERATION_TIME_BUCKETS: [f64; 15] = [ + 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 2e-1, 3e-1, 4e-1, 5e-1, 7.5e-1, 1.0, 1.5, 2.0, 3.0, 4.0, +]; +const NIXL_XFER_TIME_BUCKETS: [f64; 12] = [ + 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0, 5.0, +]; +const NIXL_POST_TIME_BUCKETS: [f64; 13] = [ + 0.001, 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0, 5.0, +]; +// Uniform 2KB to 8GB range: 2**(10+i) for i in range(1, 25, 2). +const NIXL_BYTES_TRANSFERRED_BUCKETS: [f64; 12] = [ + 2048.0, + 8192.0, + 32768.0, + 131072.0, + 524288.0, + 2097152.0, + 8388608.0, + 33554432.0, + 134217728.0, + 536870912.0, + 2147483648.0, + 8589934592.0, +]; +const NIXL_NUM_DESCRIPTORS_BUCKETS: [f64; 14] = [ + 10.0, 20.0, 30.0, 50.0, 75.0, 100.0, 200.0, 400.0, 1000.0, 2000.0, 4000.0, 10000.0, 20000.0, + 50000.0, +]; + +fn mooncake_operation_time_histogram() -> Histogram { + Histogram::new(MOONCAKE_OPERATION_TIME_BUCKETS.iter().copied()) +} + +fn nixl_xfer_time_histogram() -> Histogram { + Histogram::new(NIXL_XFER_TIME_BUCKETS.iter().copied()) +} + +fn nixl_post_time_histogram() -> Histogram { + Histogram::new(NIXL_POST_TIME_BUCKETS.iter().copied()) +} + +fn nixl_bytes_transferred_histogram() -> Histogram { + Histogram::new(NIXL_BYTES_TRANSFERRED_BUCKETS.iter().copied()) +} + +fn nixl_num_descriptors_histogram() -> Histogram { + Histogram::new(NIXL_NUM_DESCRIPTORS_BUCKETS.iter().copied()) +} + #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] pub struct EngineLabels { pub model_name: String, @@ -49,6 +101,24 @@ pub struct WaitingReasonLabels { pub reason: &'static str, } +/// Labels for per-operation Mooncake store connector telemetry. `operation` +/// and `status` are dynamic (not known ahead of time), unlike the other +/// fixed-label families in this module. +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] +pub struct MooncakeOperationLabels { + pub model_name: String, + pub engine: u32, + pub operation: String, + pub status: String, +} + +/// Family type aliases for the Mooncake connector telemetry, exposed so +/// `engine-core-client` can hold onto the `Family` directly (rather than a +/// pre-resolved handle) since `operation`/`status` vary per record. +pub type MooncakeOperationHistogramFamily = + Family Histogram>; +pub type MooncakeOperationCounterFamily = Family; + /// Adapter names encoded as a deterministic comma-joined Prometheus label value. #[derive(Clone, Debug, Hash, PartialEq, Eq)] pub struct LoraAdapterNames(pub BTreeSet); @@ -184,6 +254,22 @@ pub struct SchedulerMetrics { pub kv_block_idle_before_evict_seconds: HistogramFamily, pub kv_block_reuse_gap_seconds: HistogramFamily, + // Mooncake store connector telemetry. Mirrors `MooncakeStorePromMetrics`. + pub mooncake_operation_time_seconds: MooncakeOperationHistogramFamily, + pub mooncake_operation_total: MooncakeOperationCounterFamily, + pub mooncake_operation_keys_total: MooncakeOperationCounterFamily, + pub mooncake_operation_bytes_total: MooncakeOperationCounterFamily, + pub mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily, + + // NIXL connector telemetry. Mirrors `NixlPromMetrics`. + pub nixl_xfer_time_seconds: HistogramFamily, + pub nixl_post_time_seconds: HistogramFamily, + pub nixl_bytes_transferred: HistogramFamily, + pub nixl_num_descriptors: HistogramFamily, + pub nixl_num_failed_transfers: Family, + pub nixl_num_failed_notifications: Family, + pub nixl_num_kv_expired_reqs: Family, + /// Non-Prometheus interval accumulators for periodic text-log helpers. pub log_stats: Family, } @@ -336,6 +422,103 @@ impl SchedulerMetrics { kv_block_reuse_gap_seconds.clone(), ); + // Mooncake store connector telemetry. + let mooncake_operation_time_seconds = + Family::new_with_constructor(mooncake_operation_time_histogram as fn() -> Histogram); + registry.register( + "vllm:mooncake_store_operation_time_seconds", + "Histogram of Mooncake store communication time.", + mooncake_operation_time_seconds.clone(), + ); + + // NOTE: registration names below omit the `_total` suffix that the + // Python metric names carry (e.g. `vllm:mooncake_store_operation_total`) + // because this crate's `prometheus-client` (unlike Python's, which + // strips a user-supplied `_total` before re-adding it) appends + // `_total` unconditionally for counters; see the note in `lib.rs`. + let mooncake_operation_total = Family::default(); + registry.register( + "vllm:mooncake_store_operation", + "Number of Mooncake store communication operations.", + mooncake_operation_total.clone(), + ); + + let mooncake_operation_keys_total = Family::default(); + registry.register( + "vllm:mooncake_store_operation_keys", + "Number of Mooncake store keys touched by operations.", + mooncake_operation_keys_total.clone(), + ); + + let mooncake_operation_bytes_total = Family::default(); + registry.register( + "vllm:mooncake_store_operation_bytes", + "Number of bytes transferred by Mooncake store operations.", + mooncake_operation_bytes_total.clone(), + ); + + let mooncake_operation_failed_keys_total = Family::default(); + registry.register( + "vllm:mooncake_store_operation_failed_keys", + "Number of Mooncake store keys that failed in operations.", + mooncake_operation_failed_keys_total.clone(), + ); + + // NIXL connector telemetry. + let nixl_xfer_time_seconds = + Family::new_with_constructor(nixl_xfer_time_histogram as fn() -> Histogram); + registry.register( + "vllm:nixl_xfer_time_seconds", + "Histogram of transfer duration for NIXL KV Cache transfers.", + nixl_xfer_time_seconds.clone(), + ); + + let nixl_post_time_seconds = + Family::new_with_constructor(nixl_post_time_histogram as fn() -> Histogram); + registry.register( + "vllm:nixl_post_time_seconds", + "Histogram of transfer post time for NIXL KV Cache transfers.", + nixl_post_time_seconds.clone(), + ); + + let nixl_bytes_transferred = + Family::new_with_constructor(nixl_bytes_transferred_histogram as fn() -> Histogram); + registry.register( + "vllm:nixl_bytes_transferred", + "Histogram of bytes transferred per NIXL KV Cache transfers.", + nixl_bytes_transferred.clone(), + ); + + let nixl_num_descriptors = + Family::new_with_constructor(nixl_num_descriptors_histogram as fn() -> Histogram); + registry.register( + "vllm:nixl_num_descriptors", + "Histogram of number of descriptors per NIXL KV Cache transfers.", + nixl_num_descriptors.clone(), + ); + + let nixl_num_failed_transfers = Family::default(); + registry.register( + "vllm:nixl_num_failed_transfers", + "Number of failed NIXL KV Cache transfers.", + nixl_num_failed_transfers.clone(), + ); + + let nixl_num_failed_notifications = Family::default(); + registry.register( + "vllm:nixl_num_failed_notifications", + "Number of failed NIXL KV Cache notifications.", + nixl_num_failed_notifications.clone(), + ); + + let nixl_num_kv_expired_reqs = Family::default(); + registry.register( + "vllm:nixl_num_kv_expired_reqs", + "Number of requests that had their KV expire. \ + NOTE: This metric is tracked on the P instance.", + nixl_num_kv_expired_reqs.clone(), + ); + Self { scheduler_running, scheduler_waiting, @@ -356,6 +539,18 @@ impl SchedulerMetrics { kv_block_lifetime_seconds, kv_block_idle_before_evict_seconds, kv_block_reuse_gap_seconds, + mooncake_operation_time_seconds, + mooncake_operation_total, + mooncake_operation_keys_total, + mooncake_operation_bytes_total, + mooncake_operation_failed_keys_total, + nixl_xfer_time_seconds, + nixl_post_time_seconds, + nixl_bytes_transferred, + nixl_num_descriptors, + nixl_num_failed_transfers, + nixl_num_failed_notifications, + nixl_num_kv_expired_reqs, log_stats: Family::default(), } } diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 0c13b1f578..bac0b03738 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( MultiConnector, + MultiKVConnectorPromMetrics, MultiKVConnectorStats, MultiKVConnectorWorkerMetadata, ) @@ -559,14 +560,12 @@ class TestMultiConnectorStats: correct data.""" serialized_data = { "NixlConnector": { - "data": { - "transfer_duration": [1.5, 2.3], - "post_duration": [0.1, 0.2], - "bytes_transferred": [1024, 2048], - "num_descriptors": [10, 20], - "num_failed_transfers": [], - "num_failed_notifications": [], - } + "transfer_duration": [1.5, 2.3], + "post_duration": [0.1, 0.2], + "bytes_transferred": [1024, 2048], + "num_descriptors": [10, 20], + "num_failed_transfers": [], + "num_failed_notifications": [], } } @@ -584,16 +583,14 @@ class TestMultiConnectorStats: """Test reconstruction with multiple connector types that have custom stats.""" serialized_data = { "NixlConnector": { - "data": { - "transfer_duration": [1.5], - "post_duration": [0.1], - "bytes_transferred": [1024], - "num_descriptors": [10], - "num_failed_transfers": [], - "num_failed_notifications": [], - } + "transfer_duration": [1.5], + "post_duration": [0.1], + "bytes_transferred": [1024], + "num_descriptors": [10], + "num_failed_transfers": [], + "num_failed_notifications": [], }, - "MockConnector": {"data": {"mock_field": [1, 2, 3]}}, + "MockConnector": {"mock_field": [1, 2, 3]}, } stats = MultiConnector.build_kv_connector_stats(data=serialized_data) @@ -608,20 +605,19 @@ class TestMultiConnectorStats: assert isinstance(stats.data["MockConnector"], MockConnectorStats) # Verify data is preserved assert stats.data["MockConnector"].data == {"mock_field": [1, 2, 3]} + assert stats.to_dict() == serialized_data def test_build_kv_connector_stats_raises_error_for_unknown_connector(self): """Test that unknown connectors raise an error.""" serialized_data = { - "UnknownConnector": {"data": {"some_field": [1, 2, 3]}}, + "UnknownConnector": {"some_field": [1, 2, 3]}, "NixlConnector": { - "data": { - "transfer_duration": [1.5], - "post_duration": [0.1], - "bytes_transferred": [1024], - "num_descriptors": [10], - "num_failed_transfers": [], - "num_failed_notifications": [], - } + "transfer_duration": [1.5], + "post_duration": [0.1], + "bytes_transferred": [1024], + "num_descriptors": [10], + "num_failed_transfers": [], + "num_failed_notifications": [], }, } @@ -675,7 +671,7 @@ class TestMultiConnectorStats: mixed_data = { "NixlConnector": nixl_stats, # Already instantiated - "MockConnector": {"data": {"mock_field": [1, 2, 3]}}, # Serialized + "MockConnector": {"mock_field": [1, 2, 3]}, # Serialized } stats = MultiConnector.build_kv_connector_stats(data=mixed_data) @@ -695,16 +691,14 @@ class TestMultiConnectorStats: # so it returns None and should be skipped serialized_data = { "NixlConnector": { - "data": { - "transfer_duration": [1.5], - "post_duration": [0.1], - "bytes_transferred": [1024], - "num_descriptors": [10], - "num_failed_transfers": [], - "num_failed_notifications": [], - } + "transfer_duration": [1.5], + "post_duration": [0.1], + "bytes_transferred": [1024], + "num_descriptors": [10], + "num_failed_transfers": [], + "num_failed_notifications": [], }, - "ExampleConnector": {"data": {"some_field": [1, 2, 3]}}, + "ExampleConnector": {"some_field": [1, 2, 3]}, } stats = MultiConnector.build_kv_connector_stats(data=serialized_data) @@ -718,14 +712,15 @@ class TestMultiConnectorStats: # ExampleConnector should be skipped (returns None) assert "ExampleConnector" not in stats.data - def test_build_kv_connector_stats_handles_malformed_data(self): - """Test that malformed data raises appropriate errors.""" - serialized_data = { - "NixlConnector": {"wrong_field": {"transfer_duration": [1.5]}} - } + def test_prom_metrics_passes_flat_data_to_children(self): + child_metrics = MagicMock() + metrics = object.__new__(MultiKVConnectorPromMetrics) + metrics._prom_metrics = {"NixlConnector": child_metrics} + payload = {"NixlConnector": {"transfer_duration": [1.5]}} - with pytest.raises(AssertionError, match="Expected a dict with a 'data' field"): - MultiConnector.build_kv_connector_stats(data=serialized_data) + metrics.observe(payload, engine_idx=2) + + child_metrics.observe.assert_called_once_with({"transfer_duration": [1.5]}, 2) def test_aggregate_same_connector(self): """Test aggregating stats from the same connector type.""" diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 1f9cfc7b25..839d36a9a0 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1709,8 +1709,9 @@ def test_scheduler_kv_connector_stats_aggregation(): final_stats = next( iter(engine_core_outputs.values()) ).scheduler_stats.kv_connector_stats + # The scheduler stats payload carries serialized per-connector dicts. nixl_stats = final_stats["NixlConnector"] - assert nixl_stats.num_successful_transfers == 2 + assert len(nixl_stats["transfer_duration"]) == 2 @pytest.mark.parametrize("distributed_executor_backend", ["ray", None]) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py index faaffd72ec..e102272bc7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/metrics.py @@ -26,6 +26,10 @@ class KVConnectorStats: data: dict[str, Any] = field(default_factory=dict) + def to_dict(self) -> dict[str, Any]: + """Return the serializable connector stats payload.""" + return self.data + def reset(self): """Reset the stats, clear the state.""" raise NotImplementedError diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 348f26acfe..591ded1c67 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -75,6 +75,11 @@ class MultiKVConnectorStats(KVConnectorStats): This is used to aggregate the stats from all connectors separately. """ + def to_dict(self) -> dict[str, Any]: + return { + connector_id: stats.to_dict() for connector_id, stats in self.data.items() + } + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: for connector_id, stats in other.data.items(): if connector_id not in self.data: @@ -122,7 +127,7 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics): f"{connector_id} is not contained in the list of registered connectors " f"with Prometheus metrics support: {self._prom_metrics.keys()}" ) - self._prom_metrics[connector_id].observe(stats_data["data"], engine_idx) + self._prom_metrics[connector_id].observe(stats_data, engine_idx) class MultiConnector(KVConnectorBase_V1, SupportsHMA): @@ -617,16 +622,9 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): connector_name ) - # stats_value is the serialized dataclass which contains {'data': {...}} - # We need to extract the inner 'data' field to avoid double-nesting - assert isinstance(stats_value, dict) and "data" in stats_value, ( - f"Expected a dict with a 'data' field, got {stats_value}" - ) - inner_data = stats_value["data"] - # Use the connector's build_kv_connector_stats to reconstruct if reconstructed_stats := connector_cls.build_kv_connector_stats( - data=inner_data + data=stats_value ): reconstructed_data[connector_name] = reconstructed_stats diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index d9d86668a6..dbc697a798 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2704,7 +2704,7 @@ class Scheduler(SchedulerInterface): ) spec_stats = spec_decoding_stats connector_stats_payload = ( - kv_connector_stats.data if kv_connector_stats else None + kv_connector_stats.to_dict() if kv_connector_stats else None ) return SchedulerStats( num_running_reqs=len(self.running),