[Bugfix][Rust Frontend] Decode current NIXL connector stats

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")

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<f64>,
@@ -211,16 +210,21 @@ pub mod kv_connector {
         pub num_failed_transfers: Vec<u64>,
         /// Notification failure counter increments collected since the previous update.
         pub num_failed_notifications: Vec<u64>,
+        /// Handshake failure counter increments collected since the previous update.
+        #[serde(default)]
+        pub num_failed_handshakes: Vec<u64>,
         /// Expired-request counter increments collected since the previous update.
         pub num_kv_expired_reqs: Vec<u64>,
+        /// Late notification counter increments collected since the previous update.
+        #[serde(default)]
+        pub num_notifications_after_expiry: Vec<u64>,
     }
 
     /// 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<f64>,
@@ -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()
+        ));
+    }
 }

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<String> {

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": [

rust/src/metrics/src/scheduler.rs

@@ -268,6 +268,7 @@ pub struct SchedulerMetrics {
     pub nixl_num_failed_transfers: Family<EngineLabels, U64Counter>,
     pub nixl_num_failed_notifications: Family<EngineLabels, U64Counter>,
     pub nixl_num_kv_expired_reqs: Family<EngineLabels, U64Counter>,
+    pub nixl_num_notifications_after_expiry: Family<EngineLabels, U64Counter>,
 
     /// Non-Prometheus interval accumulators for periodic text-log helpers.
     pub log_stats: Family<EngineLabels, SchedulerLogStatsAccumulator>,
@@ -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(),
         }
     }