From a530ebb1cd7c8defcb003331faa6a09545c4b090 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Fri, 23 Oct 2026 09:12:25 +0000 Subject: [PATCH] [Rust Frontend] Export Mooncake connector transfer metrics like Python vllm-project/vllm#56939 added `MooncakePromMetrics` for the P2P `MooncakeConnector`. Rust had no variant for its stats payload, so it decoded as a MultiConnector map of unknown children and was dropped. Decode it as `MooncakeTransferStats`, also as a MultiConnector child, and export the six `vllm:mooncake_*` transfer families per engine, reusing the NIXL bucket edges Python also uses. Co-authored-by: Claude Opus 5.5 Signed-off-by: Bugen Zhao --- rust/src/engine-core-client/src/metrics.rs | 111 +++++++++++++++++- .../engine-core-client/src/protocol/stats.rs | 27 +++++ .../engine-core-client/src/tests/client.rs | 10 ++ .../src/tests/python_compat.py | 10 ++ rust/src/metrics/src/scheduler.rs | 62 ++++++++++ 5 files changed, 217 insertions(+), 3 deletions(-) diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index e8b0741066f1..60ffcf487501 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -13,8 +13,8 @@ use vllm_metrics::{ }; use crate::protocol::stats::{ - CachedTokensBySource, KvConnectorStats, MooncakeStats, MultiConnectorStats, NixlStats, - SchedulerStats, + CachedTokensBySource, KvConnectorStats, MooncakeStats, MooncakeTransferStats, + MultiConnectorStats, NixlStats, SchedulerStats, }; use crate::transport::ConnectedEngine; @@ -89,6 +89,14 @@ struct SchedulerStatsHandles { mooncake_operation_bytes_total: MooncakeOperationCounterFamily, mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily, + // Mooncake connector transfer telemetry, decoded from `kv_connector_stats`. + mooncake_xfer_time_seconds: HistogramMetric, + mooncake_bytes_transferred: HistogramMetric, + mooncake_num_descriptors: HistogramMetric, + mooncake_num_failed_transfers: U64Counter, + mooncake_num_failed_recvs: U64Counter, + mooncake_num_kv_expired_reqs: U64Counter, + // NIXL connector telemetry, decoded from `kv_connector_stats`. nixl_xfer_time_seconds: HistogramMetric, nixl_post_time_seconds: HistogramMetric, @@ -205,6 +213,16 @@ fn resolve_scheduler_stats_handles( 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(), + mooncake_xfer_time_seconds: metrics.mooncake_xfer_time_seconds.get_or_create_owned(&labels), + mooncake_bytes_transferred: metrics.mooncake_bytes_transferred.get_or_create_owned(&labels), + mooncake_num_descriptors: metrics.mooncake_num_descriptors.get_or_create_owned(&labels), + mooncake_num_failed_transfers: metrics + .mooncake_num_failed_transfers + .get_or_create_owned(&labels), + mooncake_num_failed_recvs: metrics.mooncake_num_failed_recvs.get_or_create_owned(&labels), + mooncake_num_kv_expired_reqs: metrics + .mooncake_num_kv_expired_reqs + .get_or_create_owned(&labels), 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), @@ -313,6 +331,9 @@ fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: & if let Some(kv_connector_stats) = &stats.kv_connector_stats { match kv_connector_stats { KvConnectorStats::Nixl(stats) => record_nixl_stats(handles, stats), + KvConnectorStats::MooncakeTransfer(stats) => { + record_mooncake_transfer_stats(handles, stats) + } KvConnectorStats::Mooncake(stats) => record_mooncake_stats(handles, stats), KvConnectorStats::Multi(stats) => record_multi_connector_stats(handles, stats), KvConnectorStats::Other(_) => {} @@ -324,11 +345,33 @@ fn record_multi_connector_stats(handles: &SchedulerStatsHandles, stats: &MultiCo for nixl in [&stats.nixl, &stats.nixl_pull, &stats.nixl_push].into_iter().flatten() { record_nixl_stats(handles, nixl); } + if let Some(mooncake_transfer) = &stats.mooncake_transfer { + record_mooncake_transfer_stats(handles, mooncake_transfer); + } if let Some(mooncake) = &stats.mooncake { record_mooncake_stats(handles, mooncake); } } +fn record_mooncake_transfer_stats(handles: &SchedulerStatsHandles, stats: &MooncakeTransferStats) { + for value in &stats.transfer_duration { + handles.mooncake_xfer_time_seconds.observe(*value); + } + for value in &stats.bytes_transferred { + handles.mooncake_bytes_transferred.observe(*value as f64); + } + for value in &stats.num_descriptors { + handles.mooncake_num_descriptors.observe(*value as f64); + } + handles + .mooncake_num_failed_transfers + .inc_by(stats.num_failed_transfers.iter().sum()); + handles.mooncake_num_failed_recvs.inc_by(stats.num_failed_recvs.iter().sum()); + handles + .mooncake_num_kv_expired_reqs + .inc_by(stats.num_kv_expired_reqs.iter().sum()); +} + fn record_mooncake_stats(handles: &SchedulerStatsHandles, stats: &MooncakeStats) { for (operation, records) in &stats.0 { for record in records { @@ -437,7 +480,7 @@ mod tests { use crate::metrics::LoraInfoExporter; use crate::protocol::stats::{ KvConnectorStats, MooncakeOperation, MooncakeRecord, MooncakeStats, MooncakeStatus, - MultiConnectorStats, NixlStats, SchedulerStats, + MooncakeTransferStats, MultiConnectorStats, NixlStats, SchedulerStats, }; fn names(values: &[&str]) -> BTreeSet { @@ -518,6 +561,53 @@ mod tests { )])) } + /// Records a bare `MooncakeConnector` payload, in the shape Python's + /// `MooncakeKVConnectorStats` sends, into the `vllm:mooncake_*` transfer + /// metrics. + #[test] + fn mooncake_transfer_stats_are_recorded_into_mooncake_metrics() { + let metrics = Metrics::new(); + let handles = super::resolve_scheduler_stats_handles(&metrics.scheduler, "model", 0); + + let kv_connector_stats: KvConnectorStats = serde_json::from_value(serde_json::json!({ + "transfer_duration": [0.01, 0.02], + "bytes_transferred": [4096, 8192], + "num_descriptors": [2, 4], + "num_failed_transfers": [1, 1], + "num_failed_recvs": [1], + "num_kv_expired_reqs": [], + })) + .unwrap(); + let stats = SchedulerStats { + kv_connector_stats: Some(kv_connector_stats), + ..Default::default() + }; + + super::record_scheduler_stats_with_handles(&handles, &stats); + + let rendered = metrics.render().unwrap(); + let mooncake_transfer_series = rendered + .lines() + .filter(|line| { + line.starts_with("vllm:mooncake_") + && !line.starts_with("vllm:mooncake_store_") + && !line.contains("_bucket{") + }) + .collect::>() + .join("\n"); + expect![[r#" + vllm:mooncake_xfer_time_seconds_sum{model_name="model",engine="0"} 0.03 + vllm:mooncake_xfer_time_seconds_count{model_name="model",engine="0"} 2 + vllm:mooncake_bytes_transferred_sum{model_name="model",engine="0"} 12288.0 + vllm:mooncake_bytes_transferred_count{model_name="model",engine="0"} 2 + vllm:mooncake_num_descriptors_sum{model_name="model",engine="0"} 6.0 + vllm:mooncake_num_descriptors_count{model_name="model",engine="0"} 2 + vllm:mooncake_num_failed_transfers_total{model_name="model",engine="0"} 2 + vllm:mooncake_num_failed_recvs_total{model_name="model",engine="0"} 1 + vllm:mooncake_num_kv_expired_reqs_total{model_name="model",engine="0"} 0"#]] + .assert_eq(&mooncake_transfer_series); + } + /// Records one MultiConnector payload into Mooncake and NIXL metrics. #[test] fn kv_connector_stats_are_recorded_into_mooncake_and_nixl_metrics() { @@ -527,6 +617,13 @@ mod tests { let stats = SchedulerStats { kv_connector_stats: Some(KvConnectorStats::Multi(Box::new(MultiConnectorStats { nixl: Some(nixl_stats()), + mooncake_transfer: Some(MooncakeTransferStats { + transfer_duration: vec![0.03], + bytes_transferred: vec![2048], + num_descriptors: vec![1], + num_failed_recvs: vec![1], + ..Default::default() + }), mooncake: Some(mooncake_stats()), ..Default::default() }))), @@ -557,5 +654,13 @@ mod tests { rendered .contains("vllm:nixl_xfer_time_seconds_count{model_name=\"model\",engine=\"0\"} 2") ); + assert!(rendered.contains( + "vllm:mooncake_xfer_time_seconds_count{model_name=\"model\",engine=\"0\"} 1" + )); + assert!( + rendered.contains( + "vllm:mooncake_num_failed_recvs_total{model_name=\"model\",engine=\"0\"} 1" + ) + ); } } diff --git a/rust/src/engine-core-client/src/protocol/stats.rs b/rust/src/engine-core-client/src/protocol/stats.rs index 7797ecf6e48f..0e362336999a 100644 --- a/rust/src/engine-core-client/src/protocol/stats.rs +++ b/rust/src/engine-core-client/src/protocol/stats.rs @@ -215,6 +215,27 @@ pub mod kv_connector { pub num_kv_expired_reqs: Vec, } + /// Mooncake connector transfer telemetry. + /// + /// Original Python definition: `MooncakeKVConnectorStats` in + /// `vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py`. + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] + #[serde(deny_unknown_fields)] + pub struct MooncakeTransferStats { + /// Time spent transferring each successful request. + pub transfer_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, + /// Receive failure counter increments collected since the previous update. + pub num_failed_recvs: 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")] @@ -281,6 +302,8 @@ pub mod kv_connector { pub nixl_pull: Option, #[serde(rename = "NixlPushConnector")] pub nixl_push: Option, + #[serde(rename = "MooncakeConnector")] + pub mooncake_transfer: Option, #[serde(rename = "MooncakeStoreConnector")] pub mooncake: Option, /// Child connector payloads without a Rust telemetry implementation. @@ -289,10 +312,14 @@ pub mod kv_connector { } /// Connector-specific scheduler telemetry. + /// + /// Variants are tried in order. `Multi` accepts any map, so every + /// connector payload with its own variant must come before it. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(untagged)] pub enum KvConnectorStats { Nixl(NixlStats), + MooncakeTransfer(MooncakeTransferStats), Mooncake(MooncakeStats), Multi(Box), Other(BTreeMap), diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 16e8494ac7bc..24a6456912d6 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2774,6 +2774,8 @@ fn python_msgpack_fixtures_match_rust_encoding() { 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 mooncake_transfer_stats_hex = + lines.next().expect("missing Mooncake transfer 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"); @@ -2991,12 +2993,20 @@ fn python_msgpack_fixtures_match_rust_encoding() { if stats.0.contains_key(&MooncakeOperation::LoadGet) )); + let mooncake_transfer_stats: KvConnectorStats = + rmp_serde::from_slice(&hex::decode(mooncake_transfer_stats_hex).unwrap()).unwrap(); + assert!(matches!( + mooncake_transfer_stats, + KvConnectorStats::MooncakeTransfer(stats) if stats.num_failed_transfers == [1] + )); + 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_transfer.is_some() && stats.mooncake.is_some() && stats.other.contains_key("UnsupportedConnector") )); 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 a784ad67cfb3..c19db0ef7832 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -543,8 +543,17 @@ class EngineCoreReadyResponse: } ] } +mooncake_transfer_stats = { + "transfer_duration": [0.01, 0.02], + "bytes_transferred": [4096, 8192], + "num_descriptors": [2, 4], + "num_failed_transfers": [1], + "num_failed_recvs": [], + "num_kv_expired_reqs": [1], +} multi_connector_stats = { "NixlConnector": nixl_stats, + "MooncakeConnector": mooncake_transfer_stats, "MooncakeStoreConnector": mooncake_stats, "UnsupportedConnector": {"sample_count": 1}, } @@ -570,6 +579,7 @@ class EngineCoreReadyResponse: ) print(msgspec.msgpack.encode(nixl_stats).hex()) print(msgspec.msgpack.encode(mooncake_stats).hex()) +print(msgspec.msgpack.encode(mooncake_transfer_stats).hex()) print(msgspec.msgpack.encode(multi_connector_stats).hex()) print(msgspec.msgpack.encode(ready_response).hex()) print(extended_outputs_bytes.hex()) diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index 6df15a86fcde..ed190d88ff11 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -252,6 +252,14 @@ pub struct SchedulerMetrics { pub mooncake_operation_bytes_total: MooncakeOperationCounterFamily, pub mooncake_operation_failed_keys_total: MooncakeOperationCounterFamily, + // Mooncake connector transfer telemetry. Mirrors `MooncakePromMetrics`. + pub mooncake_xfer_time_seconds: HistogramFamily, + pub mooncake_bytes_transferred: HistogramFamily, + pub mooncake_num_descriptors: HistogramFamily, + pub mooncake_num_failed_transfers: Family, + pub mooncake_num_failed_recvs: Family, + pub mooncake_num_kv_expired_reqs: Family, + // NIXL connector telemetry. Mirrors `NixlPromMetrics`. pub nixl_xfer_time_seconds: HistogramFamily, pub nixl_post_time_seconds: HistogramFamily, @@ -467,6 +475,54 @@ impl SchedulerMetrics { mooncake_operation_failed_keys_total.clone(), ); + // Mooncake connector transfer telemetry. Python's `MooncakePromMetrics` + // uses the same bucket edges as `NixlPromMetrics`. + let mooncake_xfer_time_seconds = + Family::new_with_constructor(nixl_xfer_time_histogram as fn() -> Histogram); + registry.register( + "vllm:mooncake_xfer_time_seconds", + "Histogram of transfer duration for Mooncake KV Cache transfers.", + mooncake_xfer_time_seconds.clone(), + ); + + let mooncake_bytes_transferred = + Family::new_with_constructor(nixl_bytes_transferred_histogram as fn() -> Histogram); + registry.register( + "vllm:mooncake_bytes_transferred", + "Histogram of bytes transferred per Mooncake KV Cache transfer.", + mooncake_bytes_transferred.clone(), + ); + + let mooncake_num_descriptors = + Family::new_with_constructor(nixl_num_descriptors_histogram as fn() -> Histogram); + registry.register( + "vllm:mooncake_num_descriptors", + "Histogram of number of descriptors per Mooncake KV Cache transfer.", + mooncake_num_descriptors.clone(), + ); + + let mooncake_num_failed_transfers = Family::default(); + registry.register( + "vllm:mooncake_num_failed_transfers", + "Number of failed Mooncake KV Cache transfers.", + mooncake_num_failed_transfers.clone(), + ); + + let mooncake_num_failed_recvs = Family::default(); + registry.register( + "vllm:mooncake_num_failed_recvs", + "Number of failed Mooncake KV Cache receives.", + mooncake_num_failed_recvs.clone(), + ); + + let mooncake_num_kv_expired_reqs = Family::default(); + registry.register( + "vllm:mooncake_num_kv_expired_reqs", + "Number of requests that had their KV expire. \ + NOTE: This metric is tracked on the P instance.", + mooncake_num_kv_expired_reqs.clone(), + ); + // NIXL connector telemetry. let nixl_xfer_time_seconds = Family::new_with_constructor(nixl_xfer_time_histogram as fn() -> Histogram); @@ -548,6 +604,12 @@ impl SchedulerMetrics { mooncake_operation_keys_total, mooncake_operation_bytes_total, mooncake_operation_failed_keys_total, + mooncake_xfer_time_seconds, + mooncake_bytes_transferred, + mooncake_num_descriptors, + mooncake_num_failed_transfers, + mooncake_num_failed_recvs, + mooncake_num_kv_expired_reqs, nixl_xfer_time_seconds, nixl_post_time_seconds, nixl_bytes_transferred,