diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 20b65d3d1cbe..1a759f16515e 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -471,12 +471,6 @@ impl StorageTier { } } - pub fn from_kv_medium_or_default(medium: Option<&str>) -> Self { - medium - .and_then(Self::from_kv_medium) - .unwrap_or(Self::Device) - } - /// Canonical wire-format medium string. `None` for the default GPU tier so /// existing consumers that omit the field continue to round-trip. pub fn to_kv_medium(self) -> Option<&'static str> { diff --git a/lib/kv-router/src/zmq_wire/convert.rs b/lib/kv-router/src/zmq_wire/convert.rs index d988e68d1859..2b4d71f437fe 100644 --- a/lib/kv-router/src/zmq_wire/convert.rs +++ b/lib/kv-router/src/zmq_wire/convert.rs @@ -45,7 +45,24 @@ pub fn convert_event( return None; } - let storage_tier = StorageTier::from_kv_medium_or_default(medium); + // Fail closed on unrecognized media instead of silently indexing them on + // the device (G1) primary tree. vLLM 0.26.0 ships `FS` / `OBJ` (pre-#48123 + // wire); those and any future medium strings are dropped, not misfiled. The + // normalizer's preprocess step classifies these as filtered first (so no + // event id is burned); this guard is a defensive backstop for direct + // convert_event callers that bypass preprocess, mirroring the locality gate. + let storage_tier = match medium { + None => StorageTier::Device, + Some(medium) => match StorageTier::from_kv_medium(medium) { + Some(tier) => tier, + None => { + if warning_count.fetch_add(1, Ordering::Relaxed) < 3 { + tracing::warn!(event_id, medium, "Dropping KV event with unknown medium"); + } + return None; + } + }, + }; let dp_rank = worker.dp_rank; let event = match raw { diff --git a/lib/kv-router/src/zmq_wire/deserialize.rs b/lib/kv-router/src/zmq_wire/deserialize.rs index fbbdd1dd5875..b9540c497dd6 100644 --- a/lib/kv-router/src/zmq_wire/deserialize.rs +++ b/lib/kv-router/src/zmq_wire/deserialize.rs @@ -113,7 +113,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { let (raw_token_ids, is_eagle) = normalize_token_ids(token_ids); let block_size = block_size.ok_or_else(|| de::Error::missing_field("block_size"))?; - let medium = medium.unwrap_or(None); + let medium = normalize_medium(medium.unwrap_or(None)); let lora_name = lora_name.unwrap_or(None); let extra_keys = extra_keys.unwrap_or(None); let cache_namespace = cache_namespace.unwrap_or(None).or_else(|| { @@ -141,7 +141,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { Some("BlockRemoved") => { let block_hashes = block_hashes.ok_or_else(|| de::Error::missing_field("block_hashes"))?; - let medium = medium.unwrap_or(None); + let medium = normalize_medium(medium.unwrap_or(None)); Ok(RawKvEvent::BlockRemoved { block_hashes, medium, @@ -187,7 +187,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { .ok_or_else(|| de::Error::invalid_length(4, &"missing block_size"))?; // Position 5 was lora_id in older formats; consume and discard for compat. let _lora_id: Option = seq.next_element()?.unwrap_or(None); - let medium: Option = seq.next_element()?.unwrap_or(None); + let medium: Option = normalize_medium(seq.next_element()?.unwrap_or(None)); let lora_name: Option = seq.next_element()?.unwrap_or(None); let extra_keys: Option>>> = seq.next_element()?.unwrap_or(None); @@ -236,7 +236,7 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { let block_hashes: Vec = seq .next_element()? .ok_or_else(|| de::Error::invalid_length(1, &"missing block_hashes"))?; - let medium: Option = seq.next_element()?.unwrap_or(None); + let medium: Option = normalize_medium(seq.next_element()?.unwrap_or(None)); let mut metadata = KvCacheEventMetadata::default(); for _ in 0..3 { @@ -274,6 +274,14 @@ impl<'de> Visitor<'de> for RawKvEventVisitor { } } +/// vLLM omits `medium` for device (GPU) events; treat an empty string the same +/// as an absent field so an unset medium stays on the default device tier +/// instead of failing closed as an unknown medium. Mirrors the empty-string +/// normalization applied to `cache_salt`. +fn normalize_medium(medium: Option) -> Option { + medium.filter(|value| !value.is_empty()) +} + fn normalize_token_ids(token_ids: KvTokenIds) -> (Vec, bool) { match token_ids { KvTokenIds::Single(tids) => (tids, false), diff --git a/lib/kv-router/src/zmq_wire/mod.rs b/lib/kv-router/src/zmq_wire/mod.rs index db26c27e622f..1145899b9365 100644 --- a/lib/kv-router/src/zmq_wire/mod.rs +++ b/lib/kv-router/src/zmq_wire/mod.rs @@ -66,6 +66,7 @@ struct KvCacheGroupMetadata { pub enum ZmqEventFilterReason { IgnoredEvent, NonLocalLocality, + UnknownMedium, AmbiguousCacheNamespace, NonMainAttentionKind, UnknownKind, @@ -78,6 +79,7 @@ impl ZmqEventFilterReason { match self { Self::IgnoredEvent => "ignored_event", Self::NonLocalLocality => "non_local_locality", + Self::UnknownMedium => "unknown_medium", Self::AmbiguousCacheNamespace => "ambiguous_cache_namespace", Self::NonMainAttentionKind => "non_main_attention_kind", Self::UnknownKind => "unknown_kind", @@ -139,18 +141,25 @@ impl ZmqEventNormalizer { return Err(ZmqEventFilterReason::NonLocalLocality); } - // Hash-only lower-tier events (STORAGE -> Disk, and External) carry no - // extra_keys/cache_namespace and must not mutate per-group metadata or - // the salted-namespace propagation chain; they are also outside the - // SW/SSM group filter's semantics, so route them straight to conversion. - // Device and host-pinned events (CPU offload, #10368) stay on the - // normalizer path so their salted namespaces still propagate. - if raw - .medium() - .and_then(StorageTier::from_kv_medium) - .is_some_and(|tier| matches!(tier, StorageTier::Disk | StorageTier::External)) - { - return Ok(raw); + // Classify by medium before touching normalizer state: + // - Device / HostPinned (GPU, CPU offload #10368) stay on the normalizer + // path so their salted namespaces still propagate. + // - Disk / External (STORAGE) are hash-only lower-tier events with no + // extra_keys/cache_namespace, so they must not mutate per-group + // metadata or the salted-namespace chain and are outside the SW/SSM + // group filter's semantics; bypass straight to conversion, which keeps + // them (no event id is wasted). + // - Unrecognized media (e.g. vLLM 0.26.0 FS/OBJ) fail closed here so the + // listener records an intentional filter. Bypassing to conversion, + // which drops them, would instead accept the event, burn a + // next_event_id, and leave an id gap the event processor mistakes for + // an engine drop -- the same trap the locality gate above avoids. + if let Some(m) = raw.medium() { + match StorageTier::from_kv_medium(m) { + Some(StorageTier::Device | StorageTier::HostPinned) => {} + Some(_) => return Ok(raw), + None => return Err(ZmqEventFilterReason::UnknownMedium), + } } let metadata = raw.metadata(); diff --git a/lib/kv-router/src/zmq_wire/tests.rs b/lib/kv-router/src/zmq_wire/tests.rs index 833bb1404b61..465c2d7184f5 100644 --- a/lib/kv-router/src/zmq_wire/tests.rs +++ b/lib/kv-router/src/zmq_wire/tests.rs @@ -114,6 +114,89 @@ fn test_deserialize_map_block_stored_cache_salt() { assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); } +/// An empty-string `medium` is normalized to absent (like empty `cache_salt`), +/// so an unset medium stays on the device tier and is indexed rather than +/// failing closed as an unknown medium. Covers store and remove. +#[test] +fn test_deserialize_empty_medium_normalizes_to_absent_device() { + let worker = WorkerWithDpRank::new(3, 0); + + let stored = to_vec_named(&MapBlockStoredFixture { + medium: Some(String::new()), + ..Default::default() + }) + .unwrap(); + let stored: RawKvEvent = from_slice(&stored).unwrap(); + assert_eq!( + stored.medium(), + None, + "empty medium must normalize to absent" + ); + let placement = + convert_placement(stored, worker).expect("empty-medium store must be indexed, not dropped"); + assert_eq!(placement.placement.tier, StorageTier::Device); + + let removed = to_vec_named(&MapBlockRemovedFixtureWithMedium { + event_type: "BlockRemoved", + block_hashes: vec![BlockHashValue::Unsigned(11)], + medium: Some(String::new()), + }) + .unwrap(); + let removed: RawKvEvent = from_slice(&removed).unwrap(); + assert_eq!( + removed.medium(), + None, + "empty medium must normalize to absent" + ); +} + +#[derive(Serialize)] +struct MapBlockRemovedFixtureWithMedium { + #[serde(rename = "type")] + event_type: &'static str, + block_hashes: Vec, + medium: Option, +} + +/// The empty-medium normalization must also cover legacy positional (tuple) +/// events: an empty-string `medium` in the sequence slot decodes as absent and +/// stays on the device tier. Covers store and remove. +#[test] +fn test_deserialize_legacy_sequence_empty_medium_normalizes_to_absent_device() { + let worker = WorkerWithDpRank::new(3, 0); + + let encoded_events = [ + to_vec(&( + "BlockStored", + vec![BlockHashValue::Unsigned(11)], + Option::::None, + vec![10u32, 11], + 2usize, + Option::::None, + Some(String::new()), + )) + .unwrap(), + to_vec(&( + "BlockRemoved", + vec![BlockHashValue::Unsigned(11)], + Some(String::new()), + )) + .unwrap(), + ]; + + for encoded in encoded_events { + let event: RawKvEvent = from_slice(&encoded).unwrap(); + assert_eq!( + event.medium(), + None, + "empty sequence medium must normalize to absent" + ); + let placement = convert_placement(event, worker) + .expect("empty-medium legacy event must be indexed, not dropped"); + assert_eq!(placement.placement.tier, StorageTier::Device); + } +} + #[test] fn test_deserialize_extra_keys_cache_namespace_fallback() { let mm_hash = "0123456789abcdef00112233445566778899aabbccddeefffedcba9876543210"; @@ -1103,6 +1186,44 @@ fn test_preprocess_rejects_non_local_locality_for_every_tier() { } } +/// Unrecognized media must be classified as filtered in `preprocess_with_reason` +/// (reason `UnknownMedium`), so the listener records an intentional filter rather +/// than accepting the event, burning a next_event_id, and dropping it only in +/// conversion — an id gap the event processor misreads as an engine drop. This is +/// the same trap the locality gate avoids. Recognized tiers are never filtered +/// here: Device/HostPinned stay on the normalizer path and Disk/External bypass. +#[test] +fn test_preprocess_rejects_unknown_medium() { + let worker = WorkerWithDpRank::new(3, 0); + + for event_kind in [TestEventKind::BlockStored, TestEventKind::BlockRemoved] { + for medium in ["FS", "OBJ", "XYZ"] { + let mut normalizer = ZmqEventNormalizer::new(2); + let raw = raw_placement_event(event_kind, Some(medium), Some(Locality::Local)); + assert_eq!( + normalizer.preprocess_with_reason(raw, worker).unwrap_err(), + ZmqEventFilterReason::UnknownMedium, + "medium={medium} must filter as unknown_medium" + ); + } + } + + // Recognized media are never rejected as unknown: GPU/CPU stay on the + // normalizer path and STORAGE takes the lower-tier bypass. + for medium in ["GPU", "CPU", "STORAGE"] { + let mut normalizer = ZmqEventNormalizer::new(2); + let raw = raw_placement_event( + TestEventKind::BlockStored, + Some(medium), + Some(Locality::Local), + ); + assert!( + normalizer.preprocess_with_reason(raw, worker).is_ok(), + "recognized medium={medium} must pass preprocess" + ); + } +} + /// Mixed-case locality (`"local"`, `"Remote"`) is not UPPERCASE, so it folds to /// `Unknown` and the event is dropped — it never decodes to a local placement. #[test] @@ -1230,3 +1351,60 @@ fn test_lower_tier_events_do_not_pollute_cache_namespace_state() { }; assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); } + +/// Unrecognized media (e.g. vLLM 0.26.0 `FS` / `OBJ`, or any future string) must +/// be dropped in conversion, never silently indexed on the device (G1) primary +/// tree — for both store and remove, and regardless of locality. +#[test] +fn test_convert_event_rejects_unknown_medium_instead_of_defaulting_to_device() { + let worker = WorkerWithDpRank::new(3, 0); + for event_kind in [TestEventKind::BlockStored, TestEventKind::BlockRemoved] { + for medium in ["FS", "OBJ", "XYZ"] { + for locality in [None, Some(Locality::Local)] { + let raw = raw_placement_event(event_kind, Some(medium), locality); + assert!( + convert_placement(raw, worker).is_none(), + "unrecognized medium {medium} must be dropped (locality {locality:?})" + ); + } + } + } +} + +/// Unrecognized media (vLLM 0.26.0 `FS` / `OBJ`) fail closed in preprocess, so +/// they never reach — let alone mutate — the salted-namespace state: a salted GPU +/// hash stays `Namespaced` and later salted GPU children keep inheriting it. +#[test] +fn test_unrecognized_media_do_not_pollute_cache_namespace_state() { + let worker = WorkerWithDpRank::new(7, 0); + let mut normalizer = ZmqEventNormalizer::new(2); + + let salted_gpu = namespaced_block_stored(1, None, Some("tenant-a"), "GPU"); + assert!(normalizer.preprocess(salted_gpu, worker).is_some()); + + // An FS event reusing the same hash fails closed before any namespace work. + let fs_same_hash = namespaced_block_stored(1, None, None, "FS"); + assert_eq!( + normalizer + .preprocess_with_reason(fs_same_hash, worker) + .unwrap_err(), + ZmqEventFilterReason::UnknownMedium + ); + assert!(matches!( + &normalizer.cache_namespaces[&(worker, 1)], + CacheNamespaceState::Namespaced(ns) if ns.as_ref() == "tenant-a" + )); + + // The salted chain continuation still inherits the parent namespace. + let gpu_child = namespaced_block_stored(2, Some(1), None, "GPU"); + let child = normalizer + .preprocess(gpu_child, worker) + .expect("unrecognized-media event must not make the parent ambiguous"); + let RawKvEvent::BlockStored { + cache_namespace, .. + } = child + else { + panic!("expected BlockStored"); + }; + assert_eq!(cache_namespace.as_deref(), Some("tenant-a")); +} diff --git a/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs b/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs index 04c85b4d50f7..e01da17161fa 100644 --- a/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs +++ b/lib/kvbm-consolidator/src/ingress/zmq_subscriber.rs @@ -92,14 +92,14 @@ pub async fn spawn( fn process_event(tracker: &mut Tracker, event: RawKvEvent, engine_source: EventSource) { // G1-only ingress: this source is the engine's local device (G1) cache. - // Non-local events (REMOTE / unknown locality) and native lower-tier media - // (CPU offload, vLLM STORAGE) belong to other systems — KVBM offload arrives - // via its own source — so they must not be tracked as G1 here. + // Non-local events (REMOTE / unknown locality), native lower-tier media + // (CPU offload, vLLM STORAGE), and unrecognized media (fail-closed) belong + // to other systems — KVBM offload arrives via its own source — so they must + // not be tracked as G1 here. if matches!(event.locality(), Some(Locality::Remote | Locality::Unknown)) || event .medium() - .and_then(StorageTier::from_kv_medium) - .is_some_and(|tier| !tier.is_gpu()) + .is_some_and(|m| StorageTier::from_kv_medium(m) != Some(StorageTier::Device)) { return; } @@ -223,8 +223,8 @@ mod tests { } /// G1-only ingress contract: only local device (G1) events reach the - /// tracker. Native lower-tier media (vLLM STORAGE, CPU offload) and non-local - /// (REMOTE / unknown locality) events are dropped. + /// tracker. Native lower-tier media (vLLM STORAGE, CPU offload), unrecognized + /// media (FS), and non-local (REMOTE / unknown locality) events are dropped. #[test] fn process_event_tracks_only_g1_device_events() { let mut tracker = Tracker::new(None); @@ -239,6 +239,11 @@ mod tests { stored_event(Some("CPU"), None), EventSource::Vllm, ); + process_event( + &mut tracker, + stored_event(Some("FS"), None), + EventSource::Vllm, + ); process_event( &mut tracker, stored_event(Some("GPU"), Some(Locality::Remote)), diff --git a/lib/llm/src/block_manager/kv_consolidator/subscriber.rs b/lib/llm/src/block_manager/kv_consolidator/subscriber.rs index 046530000193..099a266fc960 100644 --- a/lib/llm/src/block_manager/kv_consolidator/subscriber.rs +++ b/lib/llm/src/block_manager/kv_consolidator/subscriber.rs @@ -154,14 +154,14 @@ fn process_event( engine_source: EventSource, ) { // G1-only ingress: this source is the engine's local device (G1) cache. - // Non-local events (REMOTE / unknown locality) and native lower-tier media - // (CPU offload, vLLM STORAGE) belong to other systems — KVBM offload arrives - // via its own source — so they must not be tracked as G1 here. + // Non-local events (REMOTE / unknown locality), native lower-tier media + // (CPU offload, vLLM STORAGE), and unrecognized media (fail-closed) belong + // to other systems — KVBM offload arrives via its own source — so they must + // not be tracked as G1 here. if matches!(event.locality(), Some(Locality::Remote | Locality::Unknown)) - || event - .medium() - .and_then(RouterStorageTier::from_kv_medium) - .is_some_and(|tier| !tier.is_gpu()) + || event.medium().is_some_and(|m| { + RouterStorageTier::from_kv_medium(m) != Some(RouterStorageTier::Device) + }) { return; } @@ -294,8 +294,8 @@ mod tests { } /// G1-only ingress contract: only local device (G1) events reach the - /// tracker. Native lower-tier media (vLLM STORAGE, CPU offload) and non-local - /// (REMOTE / unknown locality) events are dropped. + /// tracker. Native lower-tier media (vLLM STORAGE, CPU offload), unrecognized + /// media (FS), and non-local (REMOTE / unknown locality) events are dropped. #[test] fn process_event_tracks_only_g1_device_events() { let mut tracker = PassthroughCacheStatusTracker::new(); @@ -312,6 +312,12 @@ mod tests { None, EventSource::Vllm, ); + process_event( + &mut tracker, + stored_event(Some("FS"), None), + None, + EventSource::Vllm, + ); process_event( &mut tracker, stored_event(Some("GPU"), Some(Locality::Remote)), diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 3248edc1de3d..d0d1c59fcfca 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -1365,6 +1365,108 @@ mod tests_startup_helpers { let _ = listener_handle.await; } + /// Unknown media (e.g. vLLM 0.26.0 `FS`) must be filtered in `preprocess`, not + /// dropped later in conversion. A conversion-time drop would still burn a + /// next_event_id (the listener increments it before `normalize_preprocessed`), + /// leaving an id gap the event processor misreads as an engine drop. Two GPU + /// stores bracketing an FS store must therefore land on consecutive event ids + /// (0, 1) — direct proof the filtered FS store consumed no id. + #[tokio::test] + async fn test_start_zmq_listener_filters_unknown_medium_without_burning_event_id() { + #[derive(serde::Serialize)] + #[serde(tag = "type")] + enum MediumKvEvent { + BlockStored { + block_hashes: Vec, + parent_block_hash: Option, + token_ids: Vec, + block_size: usize, + medium: Option<&'static str>, + group_idx: Option, + kv_cache_spec_kind: Option<&'static str>, + }, + } + + let (tx, mut rx) = mpsc::unbounded_channel::>(); + let (_ipc_dir, endpoint) = unique_ipc_endpoint(); + let pub_socket = bind_pub_socket(&endpoint).await.unwrap(); + let token = dynamo_runtime::CancellationToken::new(); + let listener_handle = tokio::spawn({ + let token = token.clone(); + start_zmq_listener( + endpoint, + String::new(), + 1, + tx, + token, + 4, + Arc::new(AtomicU64::new(0)), + None, + ) + }); + + let stored = |hash: u64, medium: &'static str| MediumKvEvent::BlockStored { + block_hashes: vec![hash], + parent_block_hash: None, + token_ids: vec![0, 1, 2, 3], + block_size: 4, + medium: Some(medium), + group_idx: None, + kv_cache_spec_kind: None, + }; + + // One native batch: GPU, unknown-medium FS, GPU. + let batch = ( + 0.0, + vec![stored(51, "GPU"), stored(52, "FS"), stored(53, "GPU")], + Some(0_i32), + ); + let payload = rmps::to_vec_named(&batch).unwrap(); + let frames = vec![Vec::new(), 0_u64.to_be_bytes().to_vec(), payload]; + + let event_batch = tokio::time::timeout(tokio::time::Duration::from_secs(5), async { + let mut publish_interval = + tokio::time::interval(tokio::time::Duration::from_millis(50)); + loop { + tokio::select! { + event_batch = rx.recv() => { + return event_batch.expect("listener channel closed"); + } + _ = publish_interval.tick() => { + send_multipart(&pub_socket, frames.clone()) + .await + .expect("failed to send ZMQ test event"); + } + } + } + }) + .await + .expect("timed out waiting for listener event"); + + // Only the two GPU stores survive; the FS store is filtered in preprocess. + assert_eq!( + event_batch.len(), + 2, + "unknown-medium FS store must be filtered, leaving two GPU stores" + ); + // Consecutive ids prove the filtered FS store consumed no next_event_id. + assert_eq!(event_batch[0].event.event_id, 0); + assert_eq!(event_batch[1].event.event_id, 1); + let stored_hash = |placement: &PlacementEvent| { + let KvCacheEventData::Stored(KvCacheStoreData { blocks, .. }) = &placement.event.data + else { + panic!("expected KvCacheStoreData"); + }; + assert_eq!(blocks.len(), 1); + blocks[0].block_hash.0 + }; + assert_eq!(stored_hash(&event_batch[0]), 51); + assert_eq!(stored_hash(&event_batch[1]), 53); + + token.cancel(); + let _ = listener_handle.await; + } + #[tokio::test] async fn test_start_zmq_listener_skips_fully_filtered_native_batch() { #[derive(serde::Serialize)]