[Rust Frontend] Record Mooncake/NIXL KV-connector metrics (#52755)
Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Markov Ilya <markovilya197@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Ilya Markov
co-authored by
Bugen Zhao
Bugen Zhao
parent
8a728663c1
commit
ae2d1ca96f
@@ -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<String> {
|
||||
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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<f64>,
|
||||
/// Time spent posting each successful transfer.
|
||||
pub post_duration: Vec<f64>,
|
||||
/// Bytes transferred by each successful request.
|
||||
pub bytes_transferred: Vec<u64>,
|
||||
/// Descriptor count for each successful request.
|
||||
pub num_descriptors: Vec<u64>,
|
||||
/// Failure counter increments collected since the previous update.
|
||||
pub num_failed_transfers: Vec<u64>,
|
||||
/// Notification failure counter increments collected since the previous update.
|
||||
pub num_failed_notifications: Vec<u64>,
|
||||
/// Expired-request counter increments collected since the previous update.
|
||||
pub num_kv_expired_reqs: Vec<u64>,
|
||||
}
|
||||
|
||||
/// 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<MooncakeOperation, Vec<MooncakeRecord>>);
|
||||
|
||||
/// 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<NixlStats>,
|
||||
#[serde(rename = "NixlPullConnector")]
|
||||
pub nixl_pull: Option<NixlStats>,
|
||||
#[serde(rename = "NixlPushConnector")]
|
||||
pub nixl_push: Option<NixlStats>,
|
||||
#[serde(rename = "MooncakeStoreConnector")]
|
||||
pub mooncake: Option<MooncakeStats>,
|
||||
/// Child connector payloads without a Rust telemetry implementation.
|
||||
#[serde(flatten)]
|
||||
pub other: BTreeMap<String, OpaqueValue>,
|
||||
}
|
||||
|
||||
/// Connector-specific scheduler telemetry.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum KvConnectorStats {
|
||||
Nixl(NixlStats),
|
||||
Mooncake(MooncakeStats),
|
||||
Multi(Box<MultiConnectorStats>),
|
||||
Other(BTreeMap<String, OpaqueValue>),
|
||||
}
|
||||
}
|
||||
|
||||
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<KvCacheEvictionEvent>,
|
||||
/// Speculative decoding scheduler stats, when enabled.
|
||||
pub spec_decoding_stats: Option<SpecDecodingStats>,
|
||||
/// Connector-specific KV transfer stats, kept opaque for now.
|
||||
pub kv_connector_stats: Option<BTreeMap<String, OpaqueValue>>,
|
||||
/// Connector-specific KV transfer stats.
|
||||
pub kv_connector_stats: Option<KvConnectorStats>,
|
||||
/// CUDA graph runtime stats when graph metrics are enabled.
|
||||
pub cudagraph_stats: Option<CudagraphStats>,
|
||||
/// Estimated MFU/performance stats, when enabled.
|
||||
|
||||
@@ -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<String> {
|
||||
match decode_value(bytes) {
|
||||
Value::Map(entries) => entries
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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<MooncakeOperationLabels, Histogram, fn() -> Histogram>;
|
||||
pub type MooncakeOperationCounterFamily = Family<MooncakeOperationLabels, U64Counter>;
|
||||
|
||||
/// Adapter names encoded as a deterministic comma-joined Prometheus label value.
|
||||
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
|
||||
pub struct LoraAdapterNames(pub BTreeSet<String>);
|
||||
@@ -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<EngineLabels, U64Counter>,
|
||||
pub nixl_num_failed_notifications: Family<EngineLabels, U64Counter>,
|
||||
pub nixl_num_kv_expired_reqs: Family<EngineLabels, U64Counter>,
|
||||
|
||||
/// Non-Prometheus interval accumulators for periodic text-log helpers.
|
||||
pub log_stats: Family<EngineLabels, SchedulerLogStatsAccumulator>,
|
||||
}
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user