diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index 60ffcf487501..dd2835972fdc 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -105,6 +105,7 @@ struct SchedulerStatsHandles { nixl_num_failed_transfers: U64Counter, nixl_num_failed_notifications: U64Counter, nixl_num_kv_expired_reqs: U64Counter, + nixl_num_notifications_after_expiry: U64Counter, // Non-Prometheus interval accumulator for periodic text-log helpers. log_stats: SchedulerLogStatsAccumulator, @@ -232,6 +233,9 @@ fn resolve_scheduler_stats_handles( .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), + nixl_num_notifications_after_expiry: metrics + .nixl_num_notifications_after_expiry + .get_or_create_owned(&labels), labels, } } @@ -415,13 +419,21 @@ fn record_nixl_stats(handles: &SchedulerStatsHandles, stats: &NixlStats) { 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_transfers.inc_by( + stats + .num_failed_transfers + .iter() + .chain(&stats.num_failed_handshakes) + .chain(&stats.num_failed_notifications) + .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()); + handles + .nixl_num_notifications_after_expiry + .inc_by(stats.num_notifications_after_expiry.iter().sum()); } /// Exports `vllm:lora_requests_info` as a single series covering all LoRA @@ -543,8 +555,10 @@ mod tests { bytes_transferred: vec![4096, 8192], num_descriptors: vec![2, 4], num_failed_transfers: vec![], - num_failed_notifications: vec![], + num_failed_notifications: vec![1], + num_failed_handshakes: vec![1], num_kv_expired_reqs: vec![1], + num_notifications_after_expiry: vec![2], } } @@ -650,6 +664,14 @@ mod tests { "vllm:nixl_num_kv_expired_reqs_total{model_name=\"model\",engine=\"0\"} 1" ) ); + assert!( + rendered.contains( + "vllm:nixl_num_failed_transfers_total{model_name=\"model\",engine=\"0\"} 2" + ) + ); + assert!(rendered.contains( + "vllm:nixl_num_notifications_after_expiry_total{model_name=\"model\",engine=\"0\"} 2" + )); 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 0e362336999a..c56f32eb244e 100644 --- a/rust/src/engine-core-client/src/protocol/stats.rs +++ b/rust/src/engine-core-client/src/protocol/stats.rs @@ -197,7 +197,6 @@ pub mod kv_connector { /// 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, @@ -211,8 +210,14 @@ pub mod kv_connector { pub num_failed_transfers: Vec, /// Notification failure counter increments collected since the previous update. pub num_failed_notifications: Vec, + /// Handshake failure counter increments collected since the previous update. + #[serde(default)] + pub num_failed_handshakes: Vec, /// Expired-request counter increments collected since the previous update. pub num_kv_expired_reqs: Vec, + /// Late notification counter increments collected since the previous update. + #[serde(default)] + pub num_notifications_after_expiry: Vec, } /// Mooncake connector transfer telemetry. @@ -220,7 +225,6 @@ pub mod kv_connector { /// 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, @@ -278,7 +282,6 @@ pub mod kv_connector { /// 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, @@ -315,6 +318,11 @@ pub mod kv_connector { /// /// Variants are tried in order. `Multi` accepts any map, so every /// connector payload with its own variant must come before it. + /// + /// The typed payloads ignore unknown fields: Python adds counters to them + /// over time, and a strict decode would drop every metric of that + /// connector, and of its `MultiConnector` siblings, until Rust caught up. + /// The variants stay distinct through the fields each one requires. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(untagged)] pub enum KvConnectorStats { @@ -365,7 +373,7 @@ pub struct SchedulerStats { #[cfg(test)] mod tests { - use super::{CachedTokensBySource, PrefixCacheStats}; + use super::{CachedTokensBySource, KvConnectorStats, PrefixCacheStats}; #[test] fn prefix_cache_stats_hits_by_source_defaults_and_decodes() { @@ -389,4 +397,39 @@ mod tests { assert_eq!(stats.hits_by_source.counts(), [0, 8, 4, 0, 0]); assert_eq!(CachedTokensBySource::SOURCES[1], "host"); } + + #[test] + fn kv_connector_stats_ignore_unknown_fields() { + let decode = |payload: serde_json::Value| -> KvConnectorStats { + rmp_serde::from_slice(&rmp_serde::to_vec_named(&payload).unwrap()).unwrap() + }; + let nixl = serde_json::json!({ + "transfer_duration": [0.5], "post_duration": [0.1], "bytes_transferred": [4096], + "num_descriptors": [2], "num_failed_transfers": [], "num_failed_notifications": [], + "num_kv_expired_reqs": [], "num_future_counter": [1] + }); + let mooncake_transfer = serde_json::json!({ + "transfer_duration": [0.5], "bytes_transferred": [4096], "num_descriptors": [2], + "num_failed_transfers": [1], "num_failed_recvs": [], "num_kv_expired_reqs": [], + "num_future_counter": [1] + }); + + assert!(matches!( + decode(nixl.clone()), + KvConnectorStats::Nixl(stats) if stats.bytes_transferred == [4096] + )); + assert!(matches!( + decode(mooncake_transfer.clone()), + KvConnectorStats::MooncakeTransfer(stats) if stats.num_failed_transfers == [1] + )); + // A newer child payload must not drop its siblings. + assert!(matches!( + decode(serde_json::json!({ + "NixlConnector": nixl, + "MooncakeConnector": mooncake_transfer, + })), + KvConnectorStats::Multi(stats) + if stats.nixl.is_some() && stats.mooncake_transfer.is_some() + )); + } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 21bf6e119ea7..f5f8678e942d 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -21,6 +21,7 @@ use zeromq::prelude::{Socket, SocketRecv, SocketSend}; use zeromq::util::PeerIdentity; use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, XPubSocket, ZmqMessage}; +use crate::metrics::SchedulerStatsRecorder; use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage}; use crate::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; use crate::protocol::logprobs::MaybeWireLogprobs; @@ -42,6 +43,7 @@ use crate::test_utils::{ IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets, setup_mock_engine_with_init, spawn_mock_engine_task, }; +use crate::transport::ConnectedEngine; use crate::{ CoordinatorMode, ENGINE_CORE_DEAD_SENTINEL, EngineCoreClient, EngineCoreClientConfig, EngineId, Error, TransportMode, @@ -2746,7 +2748,7 @@ async fn reset_prefix_cache_returns_false_when_any_engine_fails() { } #[test] -fn python_msgpack_fixtures_match_rust_encoding() { +fn python_msgpack_fixtures_match_rust_encoding_and_record_nixl_metrics() { init_tracing(); let script = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("src/tests/python_compat.py"); let output = Command::new(&script) @@ -3006,13 +3008,57 @@ fn python_msgpack_fixtures_match_rust_encoding() { 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") + let KvConnectorStats::Multi(multi_connector_stats) = multi_connector_stats else { + panic!("Python MultiConnector stats should decode as Multi"); + }; + assert!(multi_connector_stats.nixl.is_some()); + assert!(multi_connector_stats.mooncake_transfer.is_some()); + assert!(multi_connector_stats.mooncake.is_some()); + assert!(multi_connector_stats.other.contains_key("UnsupportedConnector")); + let nixl_stats = multi_connector_stats.nixl.as_ref().expect("NIXL stats decode"); + assert_eq!(nixl_stats.num_failed_handshakes, [1]); + assert_eq!(nixl_stats.num_notifications_after_expiry, [2]); + + let metrics = vllm_metrics::Metrics::new(); + let recorder = SchedulerStatsRecorder::new( + &metrics.scheduler, + "test-model", + &[ConnectedEngine { + engine_id: EngineId::from_engine_index(0), + ready_response: crate::mock_engine::default_ready_response(), + }], + ); + recorder.record( + 0, + &SchedulerStats { + kv_connector_stats: Some(KvConnectorStats::Multi(multi_connector_stats)), + ..Default::default() + }, + ); + let rendered_metrics = metrics.render().unwrap(); + assert!( + rendered_metrics.contains( + "vllm:nixl_xfer_time_seconds_count{model_name=\"test-model\",engine=\"0\"} 2" + ) + ); + assert!( + rendered_metrics.contains( + "vllm:nixl_bytes_transferred_count{model_name=\"test-model\",engine=\"0\"} 2" + ) + ); + assert!( + rendered_metrics.contains( + "vllm:nixl_bytes_transferred_sum{model_name=\"test-model\",engine=\"0\"} 12288" + ) + ); + assert!(rendered_metrics.contains( + "vllm:nixl_num_failed_transfers_total{model_name=\"test-model\",engine=\"0\"} 2" + )); + assert!(rendered_metrics.contains( + "vllm:nixl_num_failed_notifications_total{model_name=\"test-model\",engine=\"0\"} 1" + )); + assert!(rendered_metrics.contains( + "vllm:nixl_num_notifications_after_expiry_total{model_name=\"test-model\",engine=\"0\"} 2" )); let map_keys = |bytes: &[u8]| -> BTreeSet { 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 89ed5dcd3e5c..7d35ee09eed8 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -566,8 +566,10 @@ class EngineCoreReadyResponse: "bytes_transferred": [4096, 8192], "num_descriptors": [2, 4], "num_failed_transfers": [], - "num_failed_notifications": [], + "num_failed_notifications": [1], + "num_failed_handshakes": [1], "num_kv_expired_reqs": [1], + "num_notifications_after_expiry": [2], } mooncake_stats = { "load_get": [ diff --git a/rust/src/metrics/src/scheduler.rs b/rust/src/metrics/src/scheduler.rs index ed190d88ff11..05392a2cf56e 100644 --- a/rust/src/metrics/src/scheduler.rs +++ b/rust/src/metrics/src/scheduler.rs @@ -268,6 +268,7 @@ pub struct SchedulerMetrics { pub nixl_num_failed_transfers: Family, pub nixl_num_failed_notifications: Family, pub nixl_num_kv_expired_reqs: Family, + pub nixl_num_notifications_after_expiry: Family, /// Non-Prometheus interval accumulators for periodic text-log helpers. pub log_stats: Family, @@ -559,14 +560,17 @@ impl SchedulerMetrics { let nixl_num_failed_transfers = Family::default(); registry.register( "vllm:nixl_num_failed_transfers", - "Number of failed NIXL KV Cache transfers.", + "Number of failed NIXL KV Cache transfers, including handshake and notification \ + failures. NOTE: KV expiry is tracked separately in \ + vllm:nixl_num_kv_expired_reqs.", 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.", + "Number of failed NIXL KV Cache notifications. Retained for compatibility; \ + these failures are also included in vllm:nixl_num_failed_transfers.", nixl_num_failed_notifications.clone(), ); @@ -578,6 +582,15 @@ impl SchedulerMetrics { nixl_num_kv_expired_reqs.clone(), ); + let nixl_num_notifications_after_expiry = Family::default(); + registry.register( + "vllm:nixl_num_notifications_after_expiry", + "Number of completion notifications for requests that were no longer tracked, \ + usually because their KV lease expired. The KV blocks may have been reused \ + before the transfer finished. Counted per notification (one per remote rank).", + nixl_num_notifications_after_expiry.clone(), + ); + Self { scheduler_running, scheduler_waiting, @@ -617,6 +630,7 @@ impl SchedulerMetrics { nixl_num_failed_transfers, nixl_num_failed_notifications, nixl_num_kv_expired_reqs, + nixl_num_notifications_after_expiry, log_stats: Family::default(), } }