From 79cc148f0e4877362298cece1d1edb3a7b311c0d Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Tue, 18 Aug 2026 17:19:33 +0800 Subject: [PATCH 01/17] feat(epp): recover embedded KV index from peers Signed-off-by: Peter Pan --- .../ext-proc/examples/onramp/README.md | 8 +- .../ext-proc/examples/onramp/agg.yaml | 16 +- .../ext-proc/src/epp_router.rs | 28 +- deploy/inference-gateway/ext-proc/src/lib.rs | 1 + .../ext-proc/src/peer_discovery.rs | 893 ++++++++++++++++-- .../ext-proc/src/peer_http.rs | 142 +++ .../inference-gateway/ext-proc/src/runner.rs | 17 +- .../ext-proc/src/selector.rs | 178 +++- .../kv-aware-routing/vanilla-vllm-onramp.mdx | 31 +- 9 files changed, 1157 insertions(+), 157 deletions(-) create mode 100644 deploy/inference-gateway/ext-proc/src/peer_http.rs diff --git a/deploy/inference-gateway/ext-proc/examples/onramp/README.md b/deploy/inference-gateway/ext-proc/examples/onramp/README.md index 1bd650bf8ebd..d05263faaa79 100644 --- a/deploy/inference-gateway/ext-proc/examples/onramp/README.md +++ b/deploy/inference-gateway/ext-proc/examples/onramp/README.md @@ -23,7 +23,8 @@ KV-aware selection is provided by the runtime-free [selection service](../../../../../docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-selection.md), which the EPP runs **in-process**: the EPP and the selection service are compiled into one binary, so there is no separate selector Deployment and no HTTP hop. The EPP can run single-replica, or -**replicated** with cross-replica active-load sync between EPP pods (see +**replicated** with cross-replica active-load sync and startup KV-index recovery between EPP pods +(see [Replicated mode](../../../../../docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx#epp-replication)). Whether the EPP uses the Dynamo runtime or not is controlled with the `DYN_EPP_MODE` environment @@ -66,8 +67,9 @@ flowchart LR - Operator-managed lifecycle for Workers, Services, `InferencePool`, and EPP resources. - Request migration, rejection, cancellation - overall admission control - Data parallelism (The standalone mode which targets DP=1.) -- Cross-replica KV-index warm-up when new replica re-warms from live traffic + replay. -- Initial worker cache-state synchronization instead of rebuilding the index only from live traffic. +- Active-reservation or load-state snapshots. Replica lifecycle events converge only after a replica + joins. +- An atomic KV-index snapshot plus live-event handoff or event-cursor protocol. - Per-tenant KV cache isolation with x-tenant-id / cache_salt. This requires per-engine support and as such is not supported in the Standalone mode. - Management of Transient disconnects. In the Dynamo mode the KV-cache updates the worker sent during the gap are recovered from the worker's **replay** socket when `DYN_EPP_KV_EVENT_REPLAY_PORT` is set (and the vLLM worker exposes one); otherwise the index refreshes from new traffic. - Dropped events / gaps management. The `SelectionCore` indexer does seq-watermark gap detection and replays missed events from the worker's replay socket when `DYN_EPP_KV_EVENT_REPLAY_PORT` is configured. Without a replay socket, gaps are dropped and the index re-warms from new traffic. diff --git a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml index 31a6d8d7ed15..2244a4f798f6 100644 --- a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml +++ b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml @@ -213,6 +213,13 @@ metadata: app: dynamo-epp spec: replicas: 2 + # During the first upgrade from an image without the peer dump endpoint, two + # new EPP Pods must start together so they can recover from each other. + strategy: + type: RollingUpdate + rollingUpdate: + maxUnavailable: 0 + maxSurge: 100% selector: matchLabels: app: dynamo-epp @@ -282,6 +289,9 @@ spec: # ZMQ replica-sync between EPP replicas (embedded replication). - name: replica-agg containerPort: 9092 + # Internal HTTP dump endpoint for joining EPP KV-index recovery. + - name: selection-http + containerPort: 9093 readinessProbe: grpc: port: 9003 @@ -295,7 +305,8 @@ spec: drop: - ALL --- -# GAIE calls the gRPC port. EPP replicas use `replica-agg` for load-state sync. +# GAIE calls the gRPC port. EPP replicas use `replica-agg` for lifecycle sync +# and `selection-http` for startup KV-index recovery. apiVersion: v1 kind: Service metadata: @@ -311,6 +322,9 @@ spec: - name: replica-agg port: 9092 targetPort: replica-agg + - name: selection-http + port: 9093 + targetPort: selection-http --- # InferencePool selects the raw vLLM pods and delegates endpoint selection to # the EPP. The EPP reads this object (read-only) to learn the pod selector and diff --git a/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index dceceaeb991a..4130bd0453f6 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -47,14 +47,8 @@ pub struct EppRouter { // Kept alive for the lifetime of the router; the reconcile loop runs on it. _adapter: TopologyAdapter, reflector_ready: Arc, - /// Peer-discovery readiness (replicated mode only): `None` when replication - /// is off, else a flag that latches `true` after the initial peer-set sync - /// (EndpointSlice LIST + reconcile). ANDed with `reflector_ready` to form the - /// health signal, so a replica does not serve before its peers are discovered - /// and its replica-sync sockets are connected. Note: this proves only that - /// future load deltas will flow — it does NOT bootstrap the load already in - /// flight on peers; that converges from live deltas as pre-existing requests - /// drain (same warm-up shape as the KV index). + /// Replication bootstrap readiness (replicated mode only): initial peer + /// discovery plus KV-index recovery, or authoritative no-peer bootstrap. peer_ready: Option>, model_name: String, /// Bounds total concurrent in-flight `pick()`s. HTTP/2 stream multiplexing @@ -69,7 +63,15 @@ pub struct EppRouter { impl EppRouter { /// Assemble the standalone runtime from the validated selector config. pub async fn from_selector(cfg: EppStandaloneConfig) -> Result { - let selector = Arc::new(Selector::new(&cfg).await?); + let selector = Selector::new(&cfg).await?; + Self::from_built_selector(cfg, selector).await + } + + pub(crate) async fn from_built_selector( + cfg: EppStandaloneConfig, + selector: Selector, + ) -> Result { + let selector = Arc::new(selector); let (renderer, reflector, reflector_ready) = Self::dependencies(&cfg).await?; Ok(Self::from_selector_parts( cfg, @@ -139,9 +141,8 @@ impl EppRouter { } } - /// Overall EPP readiness for the gRPC health signal: the pod reflector is - /// ready (workers synced + pool resolved) AND, in replicated mode, the peer - /// set has finished its initial sync. Polled by the health mirror in `main`. + /// Overall EPP readiness: worker discovery is ready and replicated mode has + /// completed peer discovery plus KV-index recovery/bootstrap. pub fn is_ready(&self) -> bool { compute_ready( self.reflector_ready.load(Ordering::Acquire), @@ -195,8 +196,7 @@ impl EppRouter { } } -/// Overall EPP health: pod readiness AND, when replicated (`peer_ready = Some`), -/// the initial peer sync. `None` means no replication → pod readiness alone. +/// Overall EPP health: pod readiness AND replication bootstrap readiness. fn compute_ready(pod_ready: bool, peer_ready: Option) -> bool { pod_ready && peer_ready.unwrap_or(true) } diff --git a/deploy/inference-gateway/ext-proc/src/lib.rs b/deploy/inference-gateway/ext-proc/src/lib.rs index bfd1d8d137ff..609921a9f011 100644 --- a/deploy/inference-gateway/ext-proc/src/lib.rs +++ b/deploy/inference-gateway/ext-proc/src/lib.rs @@ -19,6 +19,7 @@ pub mod epp_standalone_config; pub mod inference_pool; pub mod metrics; pub mod peer_discovery; +mod peer_http; pub mod picker; pub mod pod_discovery; pub mod proto; diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index 95458e9d6ebc..ee8c177a3ef8 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -7,8 +7,10 @@ //! in-process [`SelectionService`] as sibling EPP replicas join or leave. use std::collections::BTreeSet; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::AtomicBool; use anyhow::{Context, Result}; use k8s_openapi::api::discovery::v1::EndpointSlice; @@ -23,12 +25,23 @@ const SERVICE_NAME_LABEL: &str = "kubernetes.io/service-name"; /// Named Service/EndpointSlice port used for aggregated replica synchronization. pub const REPLICA_AGG_PORT_NAME: &str = "replica-agg"; +/// Named Service/EndpointSlice port used for startup KV-index recovery. +pub const SELECTION_HTTP_PORT_NAME: &str = "selection-http"; + +const INITIAL_RECOVERY_BACKOFF: std::time::Duration = std::time::Duration::from_secs(1); +const MAX_RECOVERY_BACKOFF: std::time::Duration = std::time::Duration::from_secs(30); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct PeerPorts { + pub(crate) replica_sync: u16, + pub(crate) selection_http: u16, +} + type Store = kube::runtime::reflector::Store; +type RecoveryAttempt<'a> = Pin> + Send + 'a>>; -/// Resolve the required aggregated replica-sync port from the peer Service's -/// EndpointSlices. Every slice must expose the same named `replica-agg` port; -/// missing or inconsistent ports fail EPP startup before replica sync is built. -pub async fn resolve_replica_sync_port(namespace: &str, service_name: &str) -> Result { +/// Resolve both peer-plane ports from one authoritative EndpointSlice snapshot. +pub(crate) async fn resolve_peer_ports(namespace: &str, service_name: &str) -> Result { use kube::{Api, Client, api::ListParams}; let client = Client::try_default() @@ -42,81 +55,83 @@ pub async fn resolve_replica_sync_port(namespace: &str, service_name: &str) -> R format!("listing EndpointSlices for EPP peer Service {namespace}/{service_name}") })?; - replica_sync_port(list.items.iter()).with_context(|| { - format!( - "resolving named port {REPLICA_AGG_PORT_NAME:?} for EPP peer Service \ - {namespace}/{service_name}" - ) + peer_ports(list.items.iter()) + .with_context(|| format!("resolving peer ports for EPP Service {namespace}/{service_name}")) +} + +fn peer_ports<'a>(slices: impl Iterator) -> Result { + let slices: Vec<_> = slices.collect(); + Ok(PeerPorts { + replica_sync: named_tcp_port(&slices, REPLICA_AGG_PORT_NAME)?, + selection_http: named_tcp_port(&slices, SELECTION_HTTP_PORT_NAME)?, }) } -fn replica_sync_port<'a>(slices: impl Iterator) -> Result { +fn named_tcp_port(slices: &[&EndpointSlice], port_name: &str) -> Result { let mut resolved = BTreeSet::new(); - let mut slice_count = 0usize; for slice in slices { - slice_count += 1; let slice_name = slice.metadata.name.as_deref().unwrap_or(""); let mut matches = slice .ports .as_deref() .unwrap_or_default() .iter() - // Only a TCP `replica-agg` port satisfies the contract: the replica - // plane binds and dials `tcp://`. Kubernetes defaults `protocol` to - // TCP when absent, so treat `None` as TCP and reject explicit - // UDP/SCTP rather than let a mismatched port through. - .filter(|port| { - port.name.as_deref() == Some(REPLICA_AGG_PORT_NAME) - && port - .protocol - .as_deref() - .is_none_or(|protocol| protocol.eq_ignore_ascii_case("TCP")) - }); + .filter(|port| port.name.as_deref() == Some(port_name)); let endpoint_port = matches.next().with_context(|| { format!( "EndpointSlice {slice_name} does not expose named port \ - {REPLICA_AGG_PORT_NAME:?}" + {port_name:?}" ) })?; anyhow::ensure!( matches.next().is_none(), - "EndpointSlice {slice_name} exposes named port {REPLICA_AGG_PORT_NAME:?} more than once" + "EndpointSlice {slice_name} exposes named port {port_name:?} more than once" + ); + anyhow::ensure!( + endpoint_port + .protocol + .as_deref() + .is_none_or(|protocol| protocol.eq_ignore_ascii_case("TCP")), + "EndpointSlice {slice_name} named port {port_name:?} must use TCP" ); let raw_port = endpoint_port.port.with_context(|| { - format!( - "EndpointSlice {slice_name} named port {REPLICA_AGG_PORT_NAME:?} has no port number" - ) + format!("EndpointSlice {slice_name} named port {port_name:?} has no port number") })?; let port = u16::try_from(raw_port).with_context(|| { format!( - "EndpointSlice {slice_name} named port {REPLICA_AGG_PORT_NAME:?} has invalid port {raw_port}" + "EndpointSlice {slice_name} named port {port_name:?} has invalid port {raw_port}" ) })?; anyhow::ensure!( port > 0, - "named port {REPLICA_AGG_PORT_NAME:?} must be greater than zero" + "named port {port_name:?} must be greater than zero" ); resolved.insert(port); } - anyhow::ensure!(slice_count > 0, "peer Service has no EndpointSlices"); + anyhow::ensure!(!slices.is_empty(), "peer Service has no EndpointSlices"); anyhow::ensure!( resolved.len() == 1, - "named port {REPLICA_AGG_PORT_NAME:?} resolves to inconsistent ports {resolved:?}" + "named port {port_name:?} resolves to inconsistent ports {resolved:?}" ); - Ok(*resolved.first().expect("validated one resolved port")) + resolved + .first() + .copied() + .ok_or_else(|| anyhow::anyhow!("named port {port_name:?} did not resolve")) } /// Starts peer discovery for the EPP's own Kubernetes Service, keeping /// replica-sync peers registered on `service` and excluding `self_ip`. /// -/// Returns a readiness flag that becomes `true` after the initial reconciliation. +/// This call does not return until initial KV-index recovery succeeds or the +/// authoritative sibling set is empty. The dump server must already be bound. pub async fn spawn( service: Arc, namespace: &str, service_name: &str, sync_port: u16, + selection_http_port: u16, self_ip: String, cancel: CancellationToken, ) -> Result> { @@ -133,12 +148,13 @@ pub async fn spawn( let writer = reflector::store::Writer::default(); let store = writer.as_reader(); let reflect = reflector::reflector(writer, watcher(slices, cfg_watch).default_backoff()); - let (changes_tx, changes_rx) = watch::channel(0u64); + let (changes_tx, mut changes_rx) = watch::channel(0u64); tracing::info!( %namespace, service = %service_name, sync_port, + selection_http_port, %self_ip, "Starting EPP peer EndpointSlice watch (embedded replication)" ); @@ -174,66 +190,150 @@ pub async fn spawn( } }); - let peer_ready = Arc::new(AtomicBool::new(false)); + // Block on the first authoritative LIST before the initial reconcile so we + // never latch readiness on an empty snapshot. The reflector retries watch + // errors with backoff, so this resolves once the LIST lands; a writer drop + // (watch task gone) means we can't sync, so bail without latching. + tokio::select! { + _ = cancel.cancelled() => anyhow::bail!("EPP peer discovery cancelled before initial LIST"), + result = store.wait_until_ready() => { + result.context("EPP peer EndpointSlice writer dropped before initial LIST")?; + } + } + + let mut known: BTreeSet = BTreeSet::new(); + // InitDone generated the snapshot we just consumed; do not mistake it for a + // peer change after the first failed recovery attempt. + changes_rx.borrow_and_update(); + recover_initial_index( + &service, + &store, + sync_port, + selection_http_port, + &self_ip, + &mut known, + &mut changes_rx, + &cancel, + INITIAL_RECOVERY_BACKOFF, + MAX_RECOVERY_BACKOFF, + ) + .await?; + + let peer_ready = Arc::new(AtomicBool::new(true)); + tracing::info!("EPP peer discovery and KV-index bootstrap complete"); + + tokio::spawn(async move { + loop { + tokio::select! { + _ = cancel.cancelled() => break, + changed = changes_rx.changed() => { + if changed.is_err() { + break; + } + } + } + reconcile_once(&service, &store, sync_port, &self_ip, &mut known).await; + } + }); + Ok(peer_ready) +} - tokio::spawn(reconcile_loop( +#[allow(clippy::too_many_arguments)] +async fn recover_initial_index( + service: &SelectionService, + store: &Store, + sync_port: u16, + selection_http_port: u16, + self_ip: &str, + known: &mut BTreeSet, + changes_rx: &mut watch::Receiver, + cancel: &CancellationToken, + initial_backoff: std::time::Duration, + max_backoff: std::time::Duration, +) -> Result<()> { + recover_initial_index_with_attempt( service, store, sync_port, + selection_http_port, self_ip, + known, changes_rx, cancel, - peer_ready.clone(), - )); - Ok(peer_ready) + initial_backoff, + max_backoff, + |service, peers| Box::pin(service.recover_indexer_from_peers(peers)), + ) + .await } -/// React to EndpointSlice changes: diff the live sibling set against the peers -/// currently registered and apply the delta. Exits when `cancel` fires or the -/// change channel closes. -async fn reconcile_loop( - service: Arc, - store: Store, +#[allow(clippy::too_many_arguments)] +async fn recover_initial_index_with_attempt( + service: &SelectionService, + store: &Store, sync_port: u16, - self_ip: String, - mut changes_rx: watch::Receiver, - cancel: CancellationToken, - peer_ready: Arc, -) { - // Block on the first authoritative LIST before the initial reconcile so we - // never latch readiness on an empty snapshot. The reflector retries watch - // errors with backoff, so this resolves once the LIST lands; a writer drop - // (watch task gone) means we can't sync, so bail without latching. - tokio::select! { - _ = cancel.cancelled() => return, - result = store.wait_until_ready() => { - if result.is_err() { - tracing::warn!( - "EPP peer EndpointSlice writer dropped before initial LIST; \ - peer discovery never became ready" - ); - return; - } + selection_http_port: u16, + self_ip: &str, + known: &mut BTreeSet, + changes_rx: &mut watch::Receiver, + cancel: &CancellationToken, + initial_backoff: std::time::Duration, + max_backoff: std::time::Duration, + mut recover: F, +) -> Result<()> +where + F: for<'a> FnMut(&'a SelectionService, &'a [String]) -> RecoveryAttempt<'a>, +{ + let mut backoff = initial_backoff; + + loop { + reconcile_once(service, store, sync_port, self_ip, known).await; + let peers = recovery_peer_urls(store, self_ip, selection_http_port); + if peers.is_empty() { + tracing::info!("No sibling EPP peers found; bootstrapping an empty KV index"); + return Ok(()); } - } - let mut known: BTreeSet = BTreeSet::new(); - reconcile_once(&service, &store, sync_port, &self_ip, &mut known).await; - // Set readiness to true after the initial reconciliation. - // Subsequent transient watch failures keep the last-known peers and must not clear it. - peer_ready.store(true, Ordering::Release); - tracing::info!("EPP peer discovery initial sync complete"); + let attempt = recover(service, &peers); + tokio::pin!(attempt); + let result = tokio::select! { + biased; + _ = cancel.cancelled() => { + anyhow::bail!("EPP peer discovery cancelled during KV-index recovery") + } + changed = changes_rx.changed() => { + changed.context("EPP peer EndpointSlice watch ended during KV-index recovery")?; + backoff = initial_backoff; + continue; + } + result = &mut attempt => result, + }; + + match result { + Ok(true) => return Ok(()), + Ok(false) => tracing::warn!( + retry_ms = backoff.as_millis(), + "No reachable EPP peer dump; retrying KV-index recovery" + ), + Err(error) => tracing::warn!( + %error, + retry_ms = backoff.as_millis(), + "EPP peer KV-index recovery failed; retrying" + ), + } - loop { tokio::select! { - _ = cancel.cancelled() => break, + _ = cancel.cancelled() => { + anyhow::bail!("EPP peer discovery cancelled during KV-index recovery") + } changed = changes_rx.changed() => { - if changed.is_err() { - break; - } + changed.context("EPP peer EndpointSlice watch ended during KV-index recovery")?; + backoff = initial_backoff; + } + _ = tokio::time::sleep(backoff) => { + backoff = backoff.saturating_mul(2).min(max_backoff); } } - reconcile_once(&service, &store, sync_port, &self_ip, &mut known).await; } } @@ -275,6 +375,44 @@ fn live_peer_ips(store: &Store, self_ip: &str) -> BTreeSet { ips } +/// Recovery prefers an already-serving sibling, but falls back to not-ready +/// siblings so same-generation replicas can bootstrap during a full surge. +fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { + let want_ipv6 = is_ipv6(self_ip); + let mut preferred = BTreeSet::new(); + let mut fallback = BTreeSet::new(); + + for slice in store.state() { + if !matches_address_family(&slice.address_type, want_ipv6) { + continue; + } + for endpoint in &slice.endpoints { + let is_preferred = endpoint.conditions.as_ref().is_some_and(|conditions| { + conditions.ready == Some(true) || conditions.serving == Some(true) + }); + for address in &endpoint.addresses { + if address.is_empty() || address == self_ip { + continue; + } + if is_preferred { + preferred.insert(address.clone()); + } else { + fallback.insert(address.clone()); + } + } + } + } + for address in &preferred { + fallback.remove(address); + } + + preferred + .into_iter() + .chain(fallback) + .map(|ip| format!("http://{}", authority(&ip, port))) + .collect() +} + /// Format `host:port`, bracketing IPv6 literals (`fd00::1` -> `[fd00::1]`) so the /// resulting `tcp://` endpoint stays valid on dual-stack clusters. fn authority(ip: &str, port: u16) -> String { @@ -289,6 +427,14 @@ fn is_ipv6(ip: &str) -> bool { ip.contains(':') } +fn matches_address_family(address_type: &str, want_ipv6: bool) -> bool { + match address_type { + address_type if address_type.eq_ignore_ascii_case("IPv4") => !want_ipv6, + address_type if address_type.eq_ignore_ascii_case("IPv6") => want_ipv6, + _ => false, + } +} + /// Collects peer IPs for the requested address family. Includes not-ready peers /// and terminating peers that are still serving to preserve synchronization /// while they start or drain. @@ -298,7 +444,7 @@ fn peer_ips<'a>( ) -> BTreeSet { let mut ips = BTreeSet::new(); for slice in slices { - if slice.address_type.eq_ignore_ascii_case("IPv6") != want_ipv6 { + if !matches_address_family(&slice.address_type, want_ipv6) { continue; } // Replica-sync membership follows EndpointSlice membership, not traffic @@ -317,8 +463,17 @@ fn peer_ips<'a>( #[cfg(test)] mod tests { + use std::sync::atomic::Ordering; + use super::*; + use axum::{ + Json, Router, extract::State, http::StatusCode, response::IntoResponse, routing::get, + }; use k8s_openapi::api::discovery::v1::{Endpoint, EndpointConditions, EndpointPort}; + use std::sync::atomic::AtomicUsize; + use std::time::Duration; + use tokio::net::TcpListener; + use tokio::sync::Notify; fn slice_with(ips: &[&str], terminating: bool, address_type: &str) -> EndpointSlice { EndpointSlice { @@ -349,6 +504,29 @@ mod tests { slice } + fn slice_with_peer_ports(replica_sync: i32, selection_http: i32) -> EndpointSlice { + let mut slice = slice_with(&["10.0.0.1"], false, "IPv4"); + slice.metadata.name = Some("epp-peers-abc".to_string()); + slice.ports = Some(vec![ + EndpointPort { + name: Some(REPLICA_AGG_PORT_NAME.to_string()), + port: Some(replica_sync), + ..Default::default() + }, + EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: Some(selection_http), + ..Default::default() + }, + ]); + slice + } + + fn parse_named_port(slices: &[EndpointSlice], port_name: &str) -> Result { + let slices: Vec<_> = slices.iter().collect(); + named_tcp_port(&slices, port_name) + } + #[test] fn peer_ips_keeps_non_terminating() { let slices = [slice_with(&["10.0.0.1", "10.0.0.2"], false, "IPv4")]; @@ -412,6 +590,15 @@ mod tests { assert!(v6.contains("fd00::1")); } + #[test] + fn peer_ips_rejects_fqdn_addresses() { + let slices = [slice_with(&["epp.example.test"], false, "FQDN")]; + assert!(peer_ips(slices.iter(), false).is_empty()); + assert!( + recovery_peer_urls(&store_from_slices(slices.to_vec()), "10.0.0.9", 9093).is_empty() + ); + } + #[test] fn authority_brackets_ipv6_only() { assert_eq!(authority("10.0.0.1", 9092), "10.0.0.1:9092"); @@ -424,13 +611,33 @@ mod tests { slice_with_replica_port(Some(9092)), slice_with_replica_port(Some(9092)), ]; - assert_eq!(replica_sync_port(slices.iter()).unwrap(), 9092); + assert_eq!( + parse_named_port(&slices, REPLICA_AGG_PORT_NAME).unwrap(), + 9092 + ); + } + + #[test] + fn resolves_selection_http_named_port() { + let slices = [ + slice_with_peer_ports(9092, 9093), + slice_with_peer_ports(9092, 9093), + ]; + assert_eq!( + peer_ports(slices.iter()).unwrap(), + PeerPorts { + replica_sync: 9092, + selection_http: 9093, + } + ); } #[test] fn rejects_missing_replica_agg_named_port() { let slices = [slice_with(&["10.0.0.1"], false, "IPv4")]; - let error = replica_sync_port(slices.iter()).unwrap_err().to_string(); + let error = parse_named_port(&slices, REPLICA_AGG_PORT_NAME) + .unwrap_err() + .to_string(); assert!(error.contains(REPLICA_AGG_PORT_NAME)); } @@ -440,7 +647,9 @@ mod tests { slice_with_replica_port(Some(9092)), slice_with_replica_port(Some(9093)), ]; - let error = replica_sync_port(slices.iter()).unwrap_err().to_string(); + let error = parse_named_port(&slices, REPLICA_AGG_PORT_NAME) + .unwrap_err() + .to_string(); assert!(error.contains("inconsistent ports")); } @@ -460,11 +669,19 @@ mod tests { fn accepts_absent_or_tcp_replica_agg_protocol() { // Absent protocol defaults to TCP in Kubernetes; explicit TCP is fine. assert_eq!( - replica_sync_port([slice_with_replica_port_protocol(None)].iter()).unwrap(), + parse_named_port( + &[slice_with_replica_port_protocol(None)], + REPLICA_AGG_PORT_NAME + ) + .unwrap(), 9092 ); assert_eq!( - replica_sync_port([slice_with_replica_port_protocol(Some("TCP"))].iter()).unwrap(), + parse_named_port( + &[slice_with_replica_port_protocol(Some("TCP"))], + REPLICA_AGG_PORT_NAME + ) + .unwrap(), 9092 ); } @@ -475,12 +692,90 @@ mod tests { // tcp://, so treating it as valid would be a silent transport mismatch. // With no TCP match left, resolution fails with the "does not expose" // error naming the port. - let error = replica_sync_port([slice_with_replica_port_protocol(Some("UDP"))].iter()) - .unwrap_err() - .to_string(); + let error = parse_named_port( + &[slice_with_replica_port_protocol(Some("UDP"))], + REPLICA_AGG_PORT_NAME, + ) + .unwrap_err() + .to_string(); assert!(error.contains(REPLICA_AGG_PORT_NAME)); } + #[test] + fn rejects_invalid_named_tcp_ports() { + let cases = [ + ( + "missing number", + vec![EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: None, + ..Default::default() + }], + ), + ( + "zero", + vec![EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: Some(0), + ..Default::default() + }], + ), + ( + "out of range", + vec![EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: Some(65_536), + ..Default::default() + }], + ), + ( + "udp", + vec![EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: Some(9093), + protocol: Some("UDP".to_string()), + ..Default::default() + }], + ), + ( + "duplicate", + vec![ + EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: Some(9093), + ..Default::default() + }, + EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: Some(9094), + ..Default::default() + }, + ], + ), + ]; + + for (name, ports) in cases { + let mut slice = slice_with(&["10.0.0.1"], false, "IPv4"); + slice.metadata.name = Some(name.to_string()); + slice.ports = Some(ports); + assert!( + parse_named_port(&[slice], SELECTION_HTTP_PORT_NAME).is_err(), + "case {name} must fail" + ); + } + + let slices = [ + slice_with_peer_ports(9092, 9093), + slice_with_peer_ports(9092, 9094), + ]; + assert!( + parse_named_port(&slices, SELECTION_HTTP_PORT_NAME) + .unwrap_err() + .to_string() + .contains("inconsistent ports") + ); + } + fn free_tcp_port() -> u16 { std::net::TcpListener::bind("127.0.0.1:0") .unwrap() @@ -505,6 +800,428 @@ mod tests { store } + fn recovery_slice(ip: &str, ready: Option, serving: Option) -> EndpointSlice { + EndpointSlice { + address_type: "IPv4".to_string(), + endpoints: vec![Endpoint { + addresses: vec![ip.to_string()], + conditions: Some(EndpointConditions { + ready, + serving, + ..Default::default() + }), + ..Default::default() + }], + ports: Some(vec![ + EndpointPort { + name: Some(REPLICA_AGG_PORT_NAME.to_string()), + port: Some(9092), + ..Default::default() + }, + EndpointPort { + name: Some(SELECTION_HTTP_PORT_NAME.to_string()), + port: Some(9093), + ..Default::default() + }, + ]), + ..Default::default() + } + } + + fn store_and_writer( + slices: Vec, + ) -> ( + Store, + kube::runtime::reflector::store::Writer, + ) { + use kube::runtime::watcher; + + let mut writer = kube::runtime::reflector::store::Writer::::default(); + let store = writer.as_reader(); + writer.apply_watcher_event(&watcher::Event::Init); + for (index, mut slice) in slices.into_iter().enumerate() { + slice + .metadata + .name + .get_or_insert_with(|| format!("epp-peers-{index}")); + writer.apply_watcher_event(&watcher::Event::InitApply(slice)); + } + writer.apply_watcher_event(&watcher::Event::InitDone); + (store, writer) + } + + fn start_recovery( + service: Arc, + store: Store, + selection_http_port: u16, + changes_rx: watch::Receiver, + cancel: CancellationToken, + ) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + let mut known = BTreeSet::new(); + let mut changes_rx = changes_rx; + recover_initial_index( + &service, + &store, + 9092, + selection_http_port, + "127.0.0.9", + &mut known, + &mut changes_rx, + &cancel, + Duration::from_millis(10), + Duration::from_millis(40), + ) + .await + }) + } + + async fn recovery_service() -> Arc { + use dynamo_kv_router::config::KvRouterConfig; + use dynamo_kv_router::services::selection::SelectionServiceBuilder; + + Arc::new( + SelectionServiceBuilder::new(KvRouterConfig::default()) + .indexer_threads(1) + .build() + .await + .expect("build selection service"), + ) + } + + #[derive(Clone)] + struct DumpGate { + requested: Arc, + release: Arc, + } + + async fn gated_dump(State(gate): State) -> Json { + gate.requested.notify_one(); + gate.release.notified().await; + Json(serde_json::json!({})) + } + + #[derive(Clone)] + struct FlakyDump { + first_failed: Arc, + attempts: Arc, + } + + async fn flaky_dump(State(state): State) -> impl IntoResponse { + if state.attempts.fetch_add(1, Ordering::SeqCst) == 0 { + state.first_failed.notify_one(); + (StatusCode::SERVICE_UNAVAILABLE, "not ready").into_response() + } else { + Json(serde_json::json!({})).into_response() + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn initial_recovery_bootstraps_without_peer() { + let service = recovery_service().await; + let (store, _writer) = store_and_writer(Vec::new()); + let (_changes_tx, mut changes_rx) = watch::channel(0u64); + let cancel = CancellationToken::new(); + let mut known = BTreeSet::new(); + + recover_initial_index( + &service, + &store, + 9092, + 9093, + "127.0.0.9", + &mut known, + &mut changes_rx, + &cancel, + Duration::from_millis(10), + Duration::from_millis(40), + ) + .await + .expect("empty peer set must bootstrap immediately"); + assert!(known.is_empty()); + + service.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn recovery_uses_selection_http_port_from_endpoint_slice() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + axum::serve( + listener, + Router::new().route("/dump", get(|| async { Json(serde_json::json!({})) })), + ) + .await + }); + let mut slice = recovery_slice("127.0.0.1", Some(true), Some(true)); + slice.ports.as_mut().unwrap()[1].port = Some(i32::from(port)); + let ports = peer_ports([&slice].into_iter()).expect("resolve peer ports"); + let (store, _writer) = store_and_writer(vec![slice]); + let (_changes_tx, changes_rx) = watch::channel(0u64); + let service = recovery_service().await; + let cancel = CancellationToken::new(); + + start_recovery( + service.clone(), + store, + ports.selection_http, + changes_rx, + cancel.clone(), + ) + .await + .expect("recovery task joins") + .expect("recovery must use the EndpointSlice HTTP port"); + + cancel.cancel(); + server.abort(); + service.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn recovery_retries_unchanged_peer_until_reachable() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let state = FlakyDump { + first_failed: Arc::new(Notify::new()), + attempts: Arc::new(AtomicUsize::new(0)), + }; + let server = tokio::spawn({ + let state = state.clone(); + async move { + axum::serve( + listener, + Router::new() + .route("/dump", get(flaky_dump)) + .with_state(state), + ) + .await + } + }); + let (store, _writer) = + store_and_writer(vec![recovery_slice("127.0.0.1", Some(true), Some(true))]); + let (_changes_tx, changes_rx) = watch::channel(0u64); + let service = recovery_service().await; + let cancel = CancellationToken::new(); + let task = start_recovery(service.clone(), store, port, changes_rx, cancel.clone()); + + tokio::time::timeout(Duration::from_secs(3), state.first_failed.notified()) + .await + .expect("first recovery request must fail"); + tokio::time::timeout(Duration::from_secs(3), task) + .await + .expect("unchanged peer must be retried") + .expect("recovery task joins") + .expect("second recovery succeeds"); + assert_eq!(state.attempts.load(Ordering::SeqCst), 2); + + cancel.cancel(); + server.abort(); + service.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn peer_change_cancels_inflight_recovery_and_uses_new_peer() { + use kube::runtime::watcher; + + struct DropSignal(Arc); + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.notify_one(); + } + } + + let port = 9093; + let old_ip = "192.0.2.10"; + let new_ip = "192.0.2.11"; + let old_slice = recovery_slice(old_ip, Some(true), Some(true)); + let (store, mut writer) = store_and_writer(vec![old_slice]); + let (changes_tx, changes_rx) = watch::channel(0u64); + let service = recovery_service().await; + let cancel = CancellationToken::new(); + + let first_started = Arc::new(Notify::new()); + let first_dropped = Arc::new(Notify::new()); + let attempts = Arc::new(std::sync::Mutex::new(Vec::>::new())); + let attempt_number = Arc::new(AtomicUsize::new(0)); + let task = tokio::spawn({ + let service = service.clone(); + let cancel = cancel.clone(); + let first_started = first_started.clone(); + let first_dropped = first_dropped.clone(); + let attempts = attempts.clone(); + let attempt_number = attempt_number.clone(); + async move { + let mut known = BTreeSet::new(); + let mut changes_rx = changes_rx; + recover_initial_index_with_attempt( + &service, + &store, + 9092, + port, + "192.0.2.99", + &mut known, + &mut changes_rx, + &cancel, + Duration::from_millis(10), + Duration::from_millis(40), + move |_service, peers| { + let peers = peers.to_vec(); + attempts.lock().unwrap().push(peers.clone()); + let number = attempt_number.fetch_add(1, Ordering::SeqCst); + let first_started = first_started.clone(); + let first_dropped = first_dropped.clone(); + Box::pin(async move { + if number == 0 { + let _drop_signal = DropSignal(first_dropped); + first_started.notify_one(); + std::future::pending::<()>().await; + unreachable!("the old recovery attempt must be cancelled"); + } + Ok(peers == vec![format!("http://{new_ip}:{port}")]) + }) + }, + ) + .await + } + }); + + tokio::time::timeout(Duration::from_secs(1), first_started.notified()) + .await + .expect("old recovery request must be in flight"); + let mut replacement = recovery_slice(new_ip, Some(true), Some(true)); + replacement.metadata.name = Some("epp-peers-0".to_string()); + writer.apply_watcher_event(&watcher::Event::Apply(replacement)); + changes_tx.send(1).unwrap(); + + tokio::time::timeout(Duration::from_secs(1), first_dropped.notified()) + .await + .expect("peer change must drop the old recovery future"); + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("peer change must not wait for old HTTP timeout") + .expect("recovery task joins") + .expect("new peer dump completes recovery"); + assert_eq!( + *attempts.lock().unwrap(), + vec![ + vec![format!("http://{old_ip}:{port}")], + vec![format!("http://{new_ip}:{port}")], + ] + ); + + cancel.cancel(); + service.shutdown().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn peers_disappearing_during_recovery_bootstraps() { + use kube::runtime::watcher; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let gate = DumpGate { + requested: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + }; + let server = tokio::spawn({ + let gate = gate.clone(); + async move { + axum::serve( + listener, + Router::new() + .route("/dump", get(gated_dump)) + .with_state(gate), + ) + .await + } + }); + let mut old_slice = recovery_slice("127.0.0.1", Some(true), Some(true)); + old_slice.metadata.name = Some("epp-peers-0".to_string()); + let (store, mut writer) = store_and_writer(vec![old_slice.clone()]); + let (changes_tx, changes_rx) = watch::channel(0u64); + let service = recovery_service().await; + let cancel = CancellationToken::new(); + let task = start_recovery(service.clone(), store, port, changes_rx, cancel.clone()); + + tokio::time::timeout(Duration::from_secs(3), gate.requested.notified()) + .await + .expect("old recovery request must be in flight"); + writer.apply_watcher_event(&watcher::Event::Delete(old_slice)); + changes_tx.send(1).unwrap(); + + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("empty peer set must bootstrap without waiting for old request") + .expect("recovery task joins") + .expect("empty peer set bootstraps"); + + cancel.cancel(); + server.abort(); + service.shutdown().await; + } + + #[test] + fn recovery_candidate_order_does_not_change_replica_membership() { + let slice = EndpointSlice { + address_type: "IPv4".to_string(), + endpoints: vec![ + Endpoint { + addresses: vec!["10.0.0.2".to_string()], + conditions: Some(EndpointConditions { + ready: Some(false), + serving: Some(false), + ..Default::default() + }), + ..Default::default() + }, + Endpoint { + addresses: vec!["10.0.0.3".to_string()], + conditions: Some(EndpointConditions { + ready: Some(true), + ..Default::default() + }), + ..Default::default() + }, + ], + ..Default::default() + }; + let (store, _writer) = store_and_writer(vec![slice.clone()]); + + assert_eq!( + peer_ips([&slice].into_iter(), false), + BTreeSet::from(["10.0.0.2".to_string(), "10.0.0.3".to_string()]) + ); + assert_eq!( + recovery_peer_urls(&store, "10.0.0.9", 9093), + vec![ + "http://10.0.0.3:9093".to_string(), + "http://10.0.0.2:9093".to_string(), + ] + ); + } + + #[test] + fn recovery_peer_urls_bracket_ipv6() { + let slice = EndpointSlice { + address_type: "IPv6".to_string(), + endpoints: vec![Endpoint { + addresses: vec!["fd00::2".to_string()], + conditions: Some(EndpointConditions { + ready: Some(true), + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + }; + let (store, _writer) = store_and_writer(vec![slice]); + assert_eq!( + recovery_peer_urls(&store, "fd00::1", 9093), + vec!["http://[fd00::2]:9093".to_string()] + ); + } + /// End-to-end at the reconcile boundary: a sibling that enters termination /// while still `serving` is draining in-flight ext-proc streams and will emit /// final `PrefillComplete`/`Free` events over replica sync. `reconcile_once` diff --git a/deploy/inference-gateway/ext-proc/src/peer_http.rs b/deploy/inference-gateway/ext-proc/src/peer_http.rs new file mode 100644 index 000000000000..9a7dafd4f923 --- /dev/null +++ b/deploy/inference-gateway/ext-proc/src/peer_http.rs @@ -0,0 +1,142 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Startup KV-index dump endpoint for sibling EPP replicas. + +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use axum::{Json, Router, extract::State, routing::get}; +use tokio::net::TcpListener; +use tokio_util::sync::CancellationToken; + +use dynamo_kv_router::services::selection::SelectionService; + +/// Bind the dump listener before returning so peer recovery cannot race server startup. +pub(crate) async fn spawn( + service: Arc, + port: u16, + pod_ip: IpAddr, + cancel: CancellationToken, +) -> Result<()> { + let address = listener_addr(pod_ip, port); + let listener = TcpListener::bind(address) + .await + .with_context(|| format!("binding EPP peer HTTP server on {address}"))?; + let app = Router::new().route("/dump", get(dump)).with_state(service); + + tokio::spawn(async move { + if let Err(error) = axum::serve(listener, app) + .with_graceful_shutdown(cancel.cancelled_owned()) + .await + { + tracing::error!(%error, port, "EPP peer HTTP server exited"); + } + }); + Ok(()) +} + +fn listener_addr(pod_ip: IpAddr, port: u16) -> SocketAddr { + match pod_ip { + IpAddr::V4(_) => SocketAddr::from((Ipv4Addr::UNSPECIFIED, port)), + IpAddr::V6(_) => SocketAddr::from((Ipv6Addr::UNSPECIFIED, port)), + } +} + +async fn dump(State(service): State>) -> Json { + Json(service.indexer_snapshot().await) +} + +#[cfg(test)] +mod tests { + use super::*; + use dynamo_kv_router::config::KvRouterConfig; + use dynamo_kv_router::services::selection::SelectionServiceBuilder; + + async fn service() -> Arc { + Arc::new( + SelectionServiceBuilder::new(KvRouterConfig::default()) + .indexer_threads(1) + .build() + .await + .expect("build selection service"), + ) + } + + fn free_tcp_port() -> u16 { + std::net::TcpListener::bind("127.0.0.1:0") + .expect("reserve port") + .local_addr() + .expect("read local address") + .port() + } + + #[tokio::test] + async fn dump_listener_is_bound_before_spawn_returns() { + let service = service().await; + let cancel = CancellationToken::new(); + let port = free_tcp_port(); + + spawn( + service.clone(), + port, + "127.0.0.1".parse().unwrap(), + cancel.clone(), + ) + .await + .expect("spawn peer HTTP server"); + tokio::net::TcpStream::connect(("127.0.0.1", port)) + .await + .expect("listener must already be bound"); + + cancel.cancel(); + service.shutdown().await; + } + + #[tokio::test] + async fn dump_endpoint_matches_selection_service_snapshot() { + let service = service().await; + let cancel = CancellationToken::new(); + let port = free_tcp_port(); + spawn( + service.clone(), + port, + "127.0.0.1".parse().unwrap(), + cancel.clone(), + ) + .await + .expect("spawn peer HTTP server"); + + let response: serde_json::Value = reqwest::get(format!("http://127.0.0.1:{port}/dump")) + .await + .expect("request dump") + .json() + .await + .expect("decode dump"); + assert_eq!(response, service.indexer_snapshot().await); + + cancel.cancel(); + service.shutdown().await; + } + + #[test] + fn listener_addr_matches_pod_ip_family() { + assert_eq!( + listener_addr("192.0.2.1".parse().unwrap(), 9093), + "0.0.0.0:9093".parse().unwrap() + ); + assert_eq!( + listener_addr("2001:db8::1".parse().unwrap(), 9093), + "[::]:9093".parse().unwrap() + ); + } + + #[tokio::test] + async fn ipv6_pod_binds_an_ipv6_listener() { + let listener = TcpListener::bind(listener_addr("::1".parse().unwrap(), 0)) + .await + .expect("IPv6 loopback listener should bind"); + assert!(listener.local_addr().unwrap().is_ipv6()); + } +} diff --git a/deploy/inference-gateway/ext-proc/src/runner.rs b/deploy/inference-gateway/ext-proc/src/runner.rs index e764d0fa0702..b499c49d9053 100644 --- a/deploy/inference-gateway/ext-proc/src/runner.rs +++ b/deploy/inference-gateway/ext-proc/src/runner.rs @@ -345,14 +345,13 @@ async fn run_inner( kv_router_config, factory, } => { - let service = - Selector::build_selection_service_with_worker_selection_policy_factory( - &selector_cfg, - *kv_router_config, - factory, - ) - .await?; - crate::EppRouter::from_selection_service(selector_cfg, service).await? + let selector = Selector::new_with_worker_selection_policy_factory( + &selector_cfg, + *kv_router_config, + factory, + ) + .await?; + crate::EppRouter::from_built_selector(selector_cfg, selector).await? } })) } => router?, @@ -408,7 +407,7 @@ async fn serve( // Continuously mirror readiness onto the health status. `is_ready()` is a // *live* signal that can flip both ways — standalone discovery clears it when // the InferencePool is deleted/invalid (nothing routable), and in replicated - // mode it stays false until the peer set finishes its initial sync. A latch + // mode it stays false until peer discovery and KV-index bootstrap finish. A latch // (set SERVING once) would strand those states, so a background task tracks // transitions and moves the health status in lock-step, dropping out of // SERVING when readiness drops and recovering when it returns. Health starts diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 05153523e292..be79fea510ac 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -11,7 +11,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::sync::atomic::AtomicBool; -use anyhow::{Result, anyhow}; +use anyhow::{Context, Result, anyhow}; use dynamo_kv_router::config::{KvRouterConfig, kv_router_config_from_dynamo_env}; use dynamo_kv_router::protocols::RoutingConstraints; @@ -24,6 +24,7 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use crate::epp_standalone_config::EppStandaloneConfig; +use crate::peer_discovery::PeerPorts; const DEFAULT_ROUTING_GROUP: &str = "default"; @@ -81,9 +82,8 @@ pub struct Selector { /// `Drop` tears down its core + replica-sync tasks. cancel: CancellationToken, reconcile_state: Mutex, - /// Peer-discovery readiness in replicated mode: `None` when replication is - /// disabled (single replica, always ready), or `Some(flag)` that latches - /// `true` once the initial peer-set sync completes. ANDed into EPP health. + /// Replication-bootstrap readiness: initial peer discovery plus KV-index + /// recovery, or authoritative no-peer bootstrap. Latched once initialized. peer_ready: Option>, } @@ -99,6 +99,37 @@ struct ReconcileState { tracked_worker_ids: HashSet, } +struct ReplicationConfig { + service_name: String, + ports: PeerPorts, +} + +struct StartupCancellation { + cancel: CancellationToken, + armed: bool, +} + +impl StartupCancellation { + fn new(cancel: CancellationToken) -> Self { + Self { + cancel, + armed: true, + } + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for StartupCancellation { + fn drop(&mut self) { + if self.armed { + self.cancel.cancel(); + } + } +} + impl Selector { fn validate_queueing_requirements( cfg: &EppStandaloneConfig, @@ -119,27 +150,43 @@ impl Selector { Self::new_with_kv_router_config(cfg, kv_router_config_from_dynamo_env()).await } - /// Build a selection service using the custom policy compiled into this EPP image. - pub(crate) async fn build_selection_service_with_worker_selection_policy_factory( + /// Build a selector using the custom policy compiled into this EPP image. + pub(crate) async fn new_with_worker_selection_policy_factory( cfg: &EppStandaloneConfig, kv_router_config: KvRouterConfig, factory: WorkerSelectionPolicyFactory, - ) -> Result { - Self::build_selection_service(cfg, kv_router_config, Some(factory)).await + ) -> Result { + Self::new_with_optional_factory(cfg, kv_router_config, Some(factory)).await } async fn new_with_kv_router_config( cfg: &EppStandaloneConfig, kv_router_config: KvRouterConfig, ) -> Result { - let service = Self::build_selection_service(cfg, kv_router_config, None).await?; - Self::from_service(cfg, service).await + Self::new_with_optional_factory(cfg, kv_router_config, None).await + } + + async fn new_with_optional_factory( + cfg: &EppStandaloneConfig, + kv_router_config: KvRouterConfig, + factory: Option, + ) -> Result { + let replication = Self::replication(cfg).await?; + let service = Self::build_selection_service( + cfg, + kv_router_config, + factory, + replication.as_ref().map(|config| config.ports.replica_sync), + ) + .await?; + Self::from_service_with_replication(cfg, Arc::new(service), replication).await } async fn build_selection_service( cfg: &EppStandaloneConfig, kv_router_config: KvRouterConfig, factory: Option, + replica_sync_port: Option, ) -> Result { // If queueing is enabled, we need to validate that the max_num_batched_tokens is set. // Done once at startup to avoid validating on every reconcile. @@ -151,10 +198,9 @@ impl Selector { let mut builder = SelectionServiceBuilder::new(kv_router_config) .indexer_threads(cfg.selector_threads) .resolved_worker_selection_policy_factory(factory); - let replication = Self::replication(cfg).await?; - if let Some((_, peer_sync_port)) = &replication { - builder = builder.replica_sync(*peer_sync_port, Vec::new()); + if let Some(peer_sync_port) = replica_sync_port { + builder = builder.replica_sync(peer_sync_port, Vec::new()); } builder @@ -173,26 +219,43 @@ impl Selector { .queueing_enabled(&cfg.model_name) .map_err(|e| anyhow!("resolving router policy for model {}: {e}", cfg.model_name))?; Self::validate_queueing_requirements(cfg, queueing_enabled)?; - let replication = match &cfg.peer_service { - Some(name) => Some(( - name.clone(), - service.replica_sync_port().ok_or_else(|| { - anyhow!( - "DYN_EPP_PEER_SERVICE requires a prebuilt SelectionService with replica sync enabled" - ) - })?, - )), - None => None, + let prebuilt_replica_sync_port = if cfg.peer_service.is_some() { + let replica_sync_port = service.replica_sync_port().ok_or_else(|| { + anyhow!( + "DYN_EPP_PEER_SERVICE requires a prebuilt SelectionService with replica sync enabled" + ) + })?; + if !service.list_workers(None, None).is_empty() { + anyhow::bail!( + "replicated prebuilt SelectionService must have an empty worker catalog before \ + KV-index recovery" + ); + } + Some(replica_sync_port) + } else { + None }; + let replication = Self::replication(cfg).await?; + if let Some(replication) = &replication { + let replica_sync_port = prebuilt_replica_sync_port + .ok_or_else(|| anyhow!("replicated prebuilt SelectionService was not validated"))?; + if replica_sync_port != replication.ports.replica_sync { + anyhow::bail!( + "prebuilt SelectionService replica-sync port {replica_sync_port} does not match \ + EndpointSlice port {}", + replication.ports.replica_sync + ); + } + } Self::from_service_with_replication(cfg, service, replication).await } - async fn replication(cfg: &EppStandaloneConfig) -> Result> { + async fn replication(cfg: &EppStandaloneConfig) -> Result> { match &cfg.peer_service { - Some(name) => Ok(Some(( - name.clone(), - crate::peer_discovery::resolve_replica_sync_port(&cfg.namespace, name).await?, - ))), + Some(name) => Ok(Some(ReplicationConfig { + service_name: name.clone(), + ports: crate::peer_discovery::resolve_peer_ports(&cfg.namespace, name).await?, + })), None => Ok(None), } } @@ -200,11 +263,12 @@ impl Selector { async fn from_service_with_replication( cfg: &EppStandaloneConfig, service: Arc, - replication: Option<(String, u16)>, + replication: Option, ) -> Result { let cancel = CancellationToken::new(); + let mut startup = StartupCancellation::new(cancel.clone()); - let peer_ready = if let Some((service_name, peer_sync_port)) = replication { + let peer_ready = if let Some(replication) = replication { // In replicated mode, we need to exclude ourselves from the peer set which requires the POD_IP let self_ip = std::env::var("POD_IP") .ok() @@ -216,14 +280,24 @@ impl Selector { via the downward API (fieldRef status.podIP) so this replica can \ exclude itself from its peer set" ) - })?; + })? + .parse::() + .context("POD_IP must be a valid IPv4 or IPv6 address")?; + crate::peer_http::spawn( + service.clone(), + replication.ports.selection_http, + self_ip, + cancel.clone(), + ) + .await?; Some( crate::peer_discovery::spawn( service.clone(), &cfg.namespace, - &service_name, - peer_sync_port, - self_ip, + &replication.service_name, + replication.ports.replica_sync, + replication.ports.selection_http, + self_ip.to_string(), cancel.clone(), ) .await?, @@ -232,6 +306,8 @@ impl Selector { None }; + startup.disarm(); + tracing::info!( replicated = peer_ready.is_some(), "Initialized in-process selection service" @@ -474,6 +550,14 @@ models: } } + fn free_tcp_port() -> u16 { + std::net::TcpListener::bind("127.0.0.1:0") + .expect("reserve test port") + .local_addr() + .expect("read test port") + .port() + } + /// Minimal single-replica config (no peer service, so no cluster access). /// `max_num_batched_tokens` is set so `Selector::new` never fails its /// fast-fail check regardless of the ambient router policy. @@ -580,6 +664,7 @@ models: Box::new(FirstEligiblePicker), ) })), + None, ) .await .expect("custom selection service should build"); @@ -904,6 +989,31 @@ models: ); } + #[tokio::test] + async fn prebuilt_service_rejects_existing_workers_before_recovery() { + let service = SelectionServiceBuilder::new(KvRouterConfig::default()) + .indexer_threads(1) + .replica_sync(free_tcp_port(), Vec::new()) + .build() + .await + .expect("replica-sync selection service should build"); + service + .upsert_worker(Selector::worker_request(&incomplete_registration(1))) + .await + .expect("incomplete worker should still enter the catalog"); + + let mut cfg = test_config(); + cfg.peer_service = Some("does-not-exist".to_string()); + let error = Selector::from_service(&cfg, service) + .await + .err() + .expect("prebuilt service with workers must be rejected"); + assert!( + error.to_string().contains("empty worker catalog"), + "{error}" + ); + } + #[tokio::test] async fn duplicate_worker_ids_are_rejected_before_reconciliation() { let selector = Selector::new(&test_config()) diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index 0606c3393526..34e99cd49a7f 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -379,8 +379,8 @@ The model name, block size, event port, and maximum batched-token values in the coherent configuration. Change them together when adapting another model or vLLM deployment. Create the replacement `dynamo-epp` Service with gRPC port `9002` and, for two replicas, the named -`replica-agg` port. The complete ServiceAccount, RBAC, Deployment, and Service definitions are -available in the +`replica-agg` and `selection-http` ports. The complete ServiceAccount, RBAC, Deployment, and Service +definitions are available in the [`agg.yaml` EPP resources](https://github.com/ai-dynamo/dynamo/blob/main/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml). ### Connect the InferencePool @@ -460,9 +460,10 @@ show no selection, the route is bypassing the `InferencePool`. ## EPP Replication Each EPP replica has an in-process selector and KV index. When `DYN_EPP_PEER_SERVICE` is set, replicas -watch their own Service's `EndpointSlice` resources and discover its named TCP `replica-agg` port. -They synchronize admission, prefill-complete, and free events so active-load accounting converges -across replicas. +watch their own Service's `EndpointSlice` resources and discover its named TCP `replica-agg` and +`selection-http` ports. The `replica-agg` port synchronizes admission, prefill-complete, and free +events so active-load accounting converges across replicas. The internal `selection-http` port exposes +`GET /dump` for KV-index recovery during startup. ```mermaid flowchart LR @@ -474,14 +475,28 @@ flowchart LR Workers["Ready vLLM pods"] -. "KV events :5557
independent index warm-up" .-> ReplicaA Workers -. "KV events :5557
independent index warm-up" .-> ReplicaB ReplicaA <-->|"replica-agg :9092
admission, prefill-complete, free"| ReplicaB + ReplicaA <-->|"selection-http :9093
startup GET /dump"| ReplicaB ``` -Replica synchronization does not copy the full KV index to a new EPP. A new replica warms its index -from live worker events and optional replay. For consistency details, see +Before a joining or restarted EPP starts worker discovery, it recovers its KV index from a sibling's +`GET /dump` endpoint. The EPP remains `NOT_SERVING` until recovery succeeds. If no sibling exists, +the first EPP follows the normal worker-event and optional replay bootstrap path. If discovered peers +are temporarily unreachable, recovery retries with backoff and retries immediately when peer membership +changes. + +The `agg.yaml` example uses a rolling-update strategy with `maxUnavailable: 0` and `maxSurge: 100%`. +During the first upgrade from an EPP image without `selection-http`, this lets two new Pods start +together and recover from each other's already-bound dump endpoints. The surge can temporarily double +the EPP Pod count, so make sure the cluster has capacity for it. Later rollouts recover from an existing +ready replica. + +Recovery restores KV-index placement only. It does not restore active reservations or provide an atomic +snapshot-plus-live-event handoff. Replica lifecycle synchronization and worker replay continue to converge +state after startup. For consistency details, see [Standalone Selection Service](../../developer-guide/knowledge-base/modular-components/router/standalone-selection.md). For a single EPP replica, set `spec.replicas: 1` and remove `DYN_EPP_PEER_SERVICE`, `POD_IP`, and the -`replica-agg` Service port. +`replica-agg` and `selection-http` container and Service ports. ## Standalone EPP Configuration Reference From 6d1b4298deb888eccb8ac91a4a5fb544d9227274 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Tue, 18 Aug 2026 18:50:49 +0800 Subject: [PATCH 02/17] fix(kv-router): reject empty peer index dumps Signed-off-by: Peter Pan --- lib/kv-router/src/services/indexer/recovery.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index 48e71dea1f44..61a16a4e4794 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -84,5 +84,22 @@ async fn try_recover_from_peer( } tracing::info!(total_events, "applied dump events from peer"); + require_events(total_events)?; Ok(()) } + +fn require_events(total_events: usize) -> Result<()> { + anyhow::ensure!(total_events > 0, "peer dump contained no index events"); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_peer_dump_is_not_a_successful_recovery() { + assert!(require_events(0).is_err()); + assert!(require_events(1).is_ok()); + } +} From 01041ab0ecc3f17dc37d5a94af8d6308aa9d4dcf Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Tue, 18 Aug 2026 18:55:22 +0800 Subject: [PATCH 03/17] docs(epp): document peer recovery failure handling Signed-off-by: Peter Pan --- .../kv-aware-routing/vanilla-vllm-onramp.mdx | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index 34e99cd49a7f..9c13d81d26d0 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -484,11 +484,20 @@ the first EPP follows the normal worker-event and optional replay bootstrap path are temporarily unreachable, recovery retries with backoff and retries immediately when peer membership changes. +If recovery remains `NOT_SERVING`, check that the peer Service resolves the named ports and that each +ready peer's `selection-http` address is reachable from the EPP Pod. Recovery intentionally retries +without an attempt limit: serving with an empty or stale index is less safe than remaining out of +service. Alert on a prolonged `NOT_SERVING` period and inspect EndpointSlices, NetworkPolicies, and +the `/dump` port before restarting replicas. + The `agg.yaml` example uses a rolling-update strategy with `maxUnavailable: 0` and `maxSurge: 100%`. During the first upgrade from an EPP image without `selection-http`, this lets two new Pods start together and recover from each other's already-bound dump endpoints. The surge can temporarily double the EPP Pod count, so make sure the cluster has capacity for it. Later rollouts recover from an existing -ready replica. +ready replica. The example assumes the Service's EndpointSlices expose a consistent pair of named +ports during rollout; do not mix an image that lacks `selection-http` into the same peer Service without +first applying the new manifest and planning the transition. If a mixed-version EndpointSlice omits the +port, the new EPP remains out of service until the peer set is upgraded and the ports converge. Recovery restores KV-index placement only. It does not restore active reservations or provide an atomic snapshot-plus-live-event handoff. Replica lifecycle synchronization and worker replay continue to converge From 4bb2b12b77738b7cfd7a7d08f1d7f67d0ff00f9e Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Wed, 19 Aug 2026 17:51:36 +0800 Subject: [PATCH 04/17] fix(epp): recover only from serving peers; accept empty dumps A replicated cold start or idle rollout deadlocked: every replica treated its not-ready sibling as a recovery candidate and rejected the sibling's empty /dump via require_events(0), retrying with bounded backoff forever, so no replica ever became Ready. Restrict recovery candidates to already-serving peers (not-ready siblings cannot hold a KV index: worker KV listeners start only after recovery completes), and treat an empty dump from a serving peer as a valid no-op recovery. Fixes a serious technical problem: multi-replica EPP deployments can now boot in cold-start and idle-rollout scenarios instead of hanging before readiness. Signed-off-by: Peter Pan --- .../ext-proc/src/peer_discovery.rs | 41 +++++++++---------- .../src/services/indexer/recovery.rs | 23 +++-------- 2 files changed, 26 insertions(+), 38 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index ee8c177a3ef8..49b8a9e14593 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -375,40 +375,40 @@ fn live_peer_ips(store: &Store, self_ip: &str) -> BTreeSet { ips } -/// Recovery prefers an already-serving sibling, but falls back to not-ready -/// siblings so same-generation replicas can bootstrap during a full surge. +/// Recovery targets only already-serving siblings (EndpointSlice `ready` or +/// `serving`). Not-ready siblings cannot contribute a meaningful index: this +/// replica starts worker KV listeners only *after* recovery completes, so a +/// not-ready sibling's index is empty by construction. Treating those as +/// recovery candidates turns a cold start (all replicas empty, none serving) +/// into a mutual-recovery deadlock — each replica rejects the other's empty +/// dump via `require_events(0)` and retries forever, so no replica ever +/// becomes Ready. With only serving peers as candidates, an empty set means +/// "no eligible peer" and bootstraps an empty index immediately. fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { let want_ipv6 = is_ipv6(self_ip); - let mut preferred = BTreeSet::new(); - let mut fallback = BTreeSet::new(); + let mut peers = BTreeSet::new(); for slice in store.state() { if !matches_address_family(&slice.address_type, want_ipv6) { continue; } for endpoint in &slice.endpoints { - let is_preferred = endpoint.conditions.as_ref().is_some_and(|conditions| { + let is_serving = endpoint.conditions.as_ref().is_some_and(|conditions| { conditions.ready == Some(true) || conditions.serving == Some(true) }); + if !is_serving { + continue; + } for address in &endpoint.addresses { - if address.is_empty() || address == self_ip { - continue; - } - if is_preferred { - preferred.insert(address.clone()); - } else { - fallback.insert(address.clone()); + if !address.is_empty() && address != self_ip { + peers.insert(address.clone()); } } } } - for address in &preferred { - fallback.remove(address); - } - preferred + peers .into_iter() - .chain(fallback) .map(|ip| format!("http://{}", authority(&ip, port))) .collect() } @@ -1192,12 +1192,11 @@ mod tests { peer_ips([&slice].into_iter(), false), BTreeSet::from(["10.0.0.2".to_string(), "10.0.0.3".to_string()]) ); + // Recovery candidates exclude the not-ready sibling (10.0.0.2): a + // not-ready replica has no KV index yet, so it cannot bootstrap a peer. assert_eq!( recovery_peer_urls(&store, "10.0.0.9", 9093), - vec![ - "http://10.0.0.3:9093".to_string(), - "http://10.0.0.2:9093".to_string(), - ] + vec!["http://10.0.0.3:9093".to_string()] ); } diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index 61a16a4e4794..3119d045b214 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -83,23 +83,12 @@ async fn try_recover_from_peer( } } + // An empty dump is a valid recovery. Recovery candidates are restricted to + // already-serving peers (see `recovery_peer_urls`), so a zero-event dump + // means the serving peer genuinely holds no KV index yet (idle cluster), + // not a transient race. Rejecting it would deadlock cold starts and idle + // rollouts: the joining replica would wait forever for events that no + // serving peer holds. tracing::info!(total_events, "applied dump events from peer"); - require_events(total_events)?; Ok(()) } - -fn require_events(total_events: usize) -> Result<()> { - anyhow::ensure!(total_events > 0, "peer dump contained no index events"); - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn empty_peer_dump_is_not_a_successful_recovery() { - assert!(require_events(0).is_err()); - assert!(require_events(1).is_ok()); - } -} From 7aa496dbbd20908f89ff3e9b1fea2a08289d8082 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Wed, 19 Aug 2026 18:16:31 +0800 Subject: [PATCH 05/17] fix(kv-router): make peer recovery dump fetch bounded and configurable The peer /dump snapshot can be tens of MB on a busy deployment, but recovery used a hard-coded 10s HTTP timeout and read the full body into memory with no size guard: a large index either tripped the timeout (and then the retry loop) or buffered unbounded memory on both sides. - Timeout is now configurable via DYN_EPP_RECOVERY_HTTP_TIMEOUT_MS and the default is raised from 10s to 30s to fit large snapshots. - A Content-Length pre-check (DYN_EPP_RECOVERY_MAX_DUMP_BYTES, default 512 MiB) fails fast with a clear error before the body is read. - Documented both knobs in the onramp env reference. Full snapshot-over-HTTP remains the design; pagination/streaming and incremental (seq-based) recovery are follow-ups. Signed-off-by: Peter Pan --- .../kv-aware-routing/vanilla-vllm-onramp.mdx | 2 + .../src/services/indexer/recovery.rs | 68 ++++++++++++++++++- 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index 9c13d81d26d0..67c21bb6c5cf 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -527,6 +527,8 @@ For a single EPP replica, set `spec.replicas: 1` and remove `DYN_EPP_PEER_SERVIC | `DYN_EPP_TOKENIZER_MAX_RESPONSE_BYTES` | No | Maximum render response size; defaults to `16777216`. | | `DYN_EPP_MAX_INFLIGHT_REQUESTS` | No | Concurrent EPP request guardrail; defaults to `1024`, and excess requests receive `503`. | | `DYN_EPP_PEER_SERVICE` | No | EPP Service used to discover sibling replicas. | +| `DYN_EPP_RECOVERY_HTTP_TIMEOUT_MS` | No | Timeout for one peer `/dump` recovery fetch (request + body); defaults to `30000` ms. Raise it when a large KV index snapshot would otherwise exceed the deadline. | +| `DYN_EPP_RECOVERY_MAX_DUMP_BYTES` | No | Maximum accepted `/dump` snapshot body; defaults to `536870912` (512 MiB). Recovery fails fast with a clear error when a peer snapshot exceeds this, instead of buffering unbounded memory. | | `POD_IP` | With peer service | EPP pod IP used to exclude the local replica from its peer set. | ## Scope and Limitations diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index 3119d045b214..f00c99c527b8 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use std::collections::HashMap; +use std::time::Duration; use anyhow::{Context, Result}; use serde::Deserialize; @@ -11,19 +12,48 @@ use crate::protocols::RouterEvent; use super::registry::WorkerRegistry; +/// Timeout for one peer `/dump` fetch (HTTP request + body transfer). A full +/// KV-index snapshot can be tens of MB on a busy deployment; 10s only fits +/// small indexes once serialization and parsing are counted. Configurable via +/// `DYN_EPP_RECOVERY_HTTP_TIMEOUT_MS` for larger clusters. +const DEFAULT_RECOVERY_HTTP_TIMEOUT_MS: u64 = 30_000; +const RECOVERY_HTTP_TIMEOUT_ENV: &str = "DYN_EPP_RECOVERY_HTTP_TIMEOUT_MS"; +/// Safety ceiling on the accepted `/dump` response body. The whole snapshot is +/// materialized in memory on both sides today (streaming is a follow-up), so +/// an unbounded body either OOMs or trips the timeout and then the retry loop. +/// Fail fast with a clear error instead. Configurable via +/// `DYN_EPP_RECOVERY_MAX_DUMP_BYTES`. +const DEFAULT_MAX_DUMP_BYTES: u64 = 512 * 1024 * 1024; +const MAX_DUMP_BYTES_ENV: &str = "DYN_EPP_RECOVERY_MAX_DUMP_BYTES"; + #[derive(Deserialize)] struct DumpEntry { block_size: u32, events: Vec, } +fn parse_u64(value: Option, default: u64) -> u64 { + value + .as_deref() + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(default) +} + +fn env_u64(key: &str, default: u64) -> u64 { + parse_u64(std::env::var(key).ok(), default) +} + pub async fn recover_from_peers(peers: &[String], registry: &WorkerRegistry) -> Result { + let timeout = Duration::from_millis(env_u64( + RECOVERY_HTTP_TIMEOUT_ENV, + DEFAULT_RECOVERY_HTTP_TIMEOUT_MS, + )); let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(10)) + .timeout(timeout) .build() .context("failed to build HTTP client")?; - tokio::time::sleep(std::time::Duration::from_secs(1)).await; + tokio::time::sleep(Duration::from_secs(1)).await; for peer_url in peers { match try_recover_from_peer(&client, peer_url, registry).await { @@ -58,6 +88,20 @@ async fn try_recover_from_peer( anyhow::bail!("peer returned status {}", resp.status()); } + // Fail fast on an oversized snapshot before reading the body: the dump is + // materialized fully in memory on both sides, so a large body either OOMs + // or trips the request timeout and the retry loop. A clear error is more + // actionable than either. + let max_dump_bytes = env_u64(MAX_DUMP_BYTES_ENV, DEFAULT_MAX_DUMP_BYTES); + if let Some(len) = resp.content_length() + && len > max_dump_bytes + { + anyhow::bail!( + "peer dump is too large: {len} bytes exceeds limit {max_dump_bytes} \ + (raise {MAX_DUMP_BYTES_ENV} to accept larger snapshots)" + ); + } + let dump: HashMap = resp.json().await.context("failed to parse dump response")?; let mut total_events = 0usize; @@ -92,3 +136,23 @@ async fn try_recover_from_peer( tracing::info!(total_events, "applied dump events from peer"); Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_u64_falls_back_on_missing_or_invalid() { + assert_eq!(parse_u64(None, 30_000), 30_000); + assert_eq!(parse_u64(Some("not-a-number".to_string()), 30_000), 30_000); + assert_eq!(parse_u64(Some(" 123 ".to_string()), 30_000), 123); + assert_eq!(parse_u64(Some("0".to_string()), 30_000), 0); + } + + #[test] + fn parse_u64_accepts_zero() { + // Zero is a valid explicit value (e.g. disabling a cap); it must not be + // treated as a parse failure that falls back to the default. + assert_eq!(parse_u64(Some("0".to_string()), 512), 0); + } +} From de9fbea33beec1042cfbd05bf20fedf093978ccb Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Thu, 20 Aug 2026 10:45:26 +0800 Subject: [PATCH 06/17] fix(epp): address peer-recovery review feedback - Keep an in-flight dump transfer when unrelated EndpointSlice churn fires; restart only when the recovery candidate set actually changes. - Require recovery sources to be serving and not terminating (ready fallback only when serving is absent); reject ready-but-not-serving and draining peers. - Shuffle recovery candidates per bootstrap so joining replicas do not all start from the lowest-IP peer. - Pass the accepted dump budget to the peer and return 413 before writing an over-budget snapshot; define 0 as unbounded. - Make selection-http optional: a Service without the dump port degrades to no-recovery instead of failing startup. - Gate /dump on local recovery completion (503 until then) so a stale-ready sibling cannot hand out an empty snapshot. - Drop the incorrect maxSurge:100% first-upgrade rationale from the onramp manifest and docs. Signed-off-by: Peter Pan --- .../ext-proc/examples/onramp/agg.yaml | 7 - .../ext-proc/src/peer_discovery.rs | 220 +++++++++++++++--- .../ext-proc/src/peer_http.rs | 118 +++++++++- .../ext-proc/src/selector.rs | 62 +++-- .../kv-aware-routing/vanilla-vllm-onramp.mdx | 13 +- .../src/services/indexer/recovery.rs | 25 +- 6 files changed, 371 insertions(+), 74 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml index 2244a4f798f6..86ebddb63639 100644 --- a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml +++ b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml @@ -213,13 +213,6 @@ metadata: app: dynamo-epp spec: replicas: 2 - # During the first upgrade from an image without the peer dump endpoint, two - # new EPP Pods must start together so they can recover from each other. - strategy: - type: RollingUpdate - rollingUpdate: - maxUnavailable: 0 - maxSurge: 100% selector: matchLabels: app: dynamo-epp diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index 49b8a9e14593..e52b330c0c06 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -7,13 +7,15 @@ //! in-process [`SelectionService`] as sibling EPP replicas join or leave. use std::collections::BTreeSet; +use std::collections::hash_map::RandomState; use std::future::Future; +use std::hash::{BuildHasher, Hash, Hasher}; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::AtomicBool; +use std::sync::atomic::{AtomicBool, Ordering}; use anyhow::{Context, Result}; -use k8s_openapi::api::discovery::v1::EndpointSlice; +use k8s_openapi::api::discovery::v1::{Endpoint, EndpointSlice}; use tokio::sync::watch; use tokio_util::sync::CancellationToken; @@ -34,7 +36,11 @@ const MAX_RECOVERY_BACKOFF: std::time::Duration = std::time::Duration::from_secs #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct PeerPorts { pub(crate) replica_sync: u16, - pub(crate) selection_http: u16, + /// `None` when the Service does not declare the `selection-http` port (an + /// image-only upgrade from a deployment that predates the dump endpoint). + /// Peer KV-index recovery is then disabled — the replica bootstraps empty + /// — rather than failing startup. + pub(crate) selection_http: Option, } type Store = kube::runtime::reflector::Store; @@ -63,10 +69,67 @@ fn peer_ports<'a>(slices: impl Iterator) -> Result = slices.collect(); Ok(PeerPorts { replica_sync: named_tcp_port(&slices, REPLICA_AGG_PORT_NAME)?, - selection_http: named_tcp_port(&slices, SELECTION_HTTP_PORT_NAME)?, + // `selection-http` is optional: a deployment upgraded before the dump + // endpoint existed declares only `replica-agg`. A missing dump port + // degrades to "no recovery" (bootstrap empty) instead of failing. + selection_http: optional_named_tcp_port(&slices, SELECTION_HTTP_PORT_NAME)?, }) } +/// Like [`named_tcp_port`], but returns `Ok(None)` when no EndpointSlice +/// exposes the named port. Slices that omit the port are skipped; conflicting +/// values or invalid ports still error. +fn optional_named_tcp_port(slices: &[&EndpointSlice], port_name: &str) -> Result> { + let mut resolved = BTreeSet::new(); + + for slice in slices { + let slice_name = slice.metadata.name.as_deref().unwrap_or(""); + let mut matches = slice + .ports + .as_deref() + .unwrap_or_default() + .iter() + .filter(|port| port.name.as_deref() == Some(port_name)); + let Some(endpoint_port) = matches.next() else { + continue; + }; + anyhow::ensure!( + matches.next().is_none(), + "EndpointSlice {slice_name} exposes named port {port_name:?} more than once" + ); + anyhow::ensure!( + endpoint_port + .protocol + .as_deref() + .is_none_or(|protocol| protocol.eq_ignore_ascii_case("TCP")), + "EndpointSlice {slice_name} named port {port_name:?} must use TCP" + ); + let raw_port = endpoint_port.port.with_context(|| { + format!("EndpointSlice {slice_name} named port {port_name:?} has no port number") + })?; + let port = u16::try_from(raw_port).with_context(|| { + format!( + "EndpointSlice {slice_name} named port {port_name:?} has invalid port {raw_port}" + ) + })?; + anyhow::ensure!( + port > 0, + "named port {port_name:?} must be greater than zero" + ); + resolved.insert(port); + } + + anyhow::ensure!(!slices.is_empty(), "peer Service has no EndpointSlices"); + if resolved.is_empty() { + return Ok(None); + } + anyhow::ensure!( + resolved.len() == 1, + "named port {port_name:?} resolves to inconsistent ports {resolved:?}" + ); + Ok(resolved.first().copied()) +} + fn named_tcp_port(slices: &[&EndpointSlice], port_name: &str) -> Result { let mut resolved = BTreeSet::new(); @@ -134,7 +197,8 @@ pub async fn spawn( selection_http_port: u16, self_ip: String, cancel: CancellationToken, -) -> Result> { + recovered: Arc, +) -> Result<()> { use futures::StreamExt; use kube::{Api, Client, runtime::WatchStreamExt, runtime::reflector, runtime::watcher}; @@ -219,7 +283,9 @@ pub async fn spawn( ) .await?; - let peer_ready = Arc::new(AtomicBool::new(true)); + // Recovery/bootstrap finished: the /dump endpoint may now serve a + // non-empty snapshot (see `peer_http` gating on this flag). + recovered.store(true, Ordering::Release); tracing::info!("EPP peer discovery and KV-index bootstrap complete"); tokio::spawn(async move { @@ -235,7 +301,7 @@ pub async fn spawn( reconcile_once(&service, &store, sync_port, &self_ip, &mut known).await; } }); - Ok(peer_ready) + Ok(()) } #[allow(clippy::too_many_arguments)] @@ -286,7 +352,7 @@ where { let mut backoff = initial_backoff; - loop { + 'attempt: loop { reconcile_once(service, store, sync_port, self_ip, known).await; let peers = recovery_peer_urls(store, self_ip, selection_http_port); if peers.is_empty() { @@ -296,17 +362,29 @@ where let attempt = recover(service, &peers); tokio::pin!(attempt); - let result = tokio::select! { - biased; - _ = cancel.cancelled() => { - anyhow::bail!("EPP peer discovery cancelled during KV-index recovery") - } - changed = changes_rx.changed() => { - changed.context("EPP peer EndpointSlice watch ended during KV-index recovery")?; - backoff = initial_backoff; - continue; + + // Await the attempt, restarting only when the candidate set actually + // changes (or becomes empty). Unrelated EndpointSlice churn (readiness + // flips, metadata/zone updates) must not discard in-flight progress: a + // large dump under churn would otherwise be dropped repeatedly, leaving + // a partially applied snapshot that the next attempt re-applies. + let result = loop { + tokio::select! { + biased; + _ = cancel.cancelled() => { + anyhow::bail!("EPP peer discovery cancelled during KV-index recovery") + } + changed = changes_rx.changed() => { + changed.context("EPP peer EndpointSlice watch ended during KV-index recovery")?; + if recovery_peer_urls(store, self_ip, selection_http_port) != peers { + // Candidate set changed: restart with the new set. + backoff = initial_backoff; + continue 'attempt; + } + // Unrelated churn: keep awaiting the same attempt. + } + result = &mut attempt => break result, } - result = &mut attempt => result, }; match result { @@ -375,15 +453,13 @@ fn live_peer_ips(store: &Store, self_ip: &str) -> BTreeSet { ips } -/// Recovery targets only already-serving siblings (EndpointSlice `ready` or -/// `serving`). Not-ready siblings cannot contribute a meaningful index: this +/// Recovery targets only siblings that are actively serving and not +/// terminating. Not-ready siblings cannot contribute a meaningful index: this /// replica starts worker KV listeners only *after* recovery completes, so a /// not-ready sibling's index is empty by construction. Treating those as /// recovery candidates turns a cold start (all replicas empty, none serving) -/// into a mutual-recovery deadlock — each replica rejects the other's empty -/// dump via `require_events(0)` and retries forever, so no replica ever -/// becomes Ready. With only serving peers as candidates, an empty set means -/// "no eligible peer" and bootstraps an empty index immediately. +/// into a mutual-recovery deadlock. An empty candidate set means "no eligible +/// peer" and bootstraps an empty index immediately. fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { let want_ipv6 = is_ipv6(self_ip); let mut peers = BTreeSet::new(); @@ -393,10 +469,7 @@ fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { continue; } for endpoint in &slice.endpoints { - let is_serving = endpoint.conditions.as_ref().is_some_and(|conditions| { - conditions.ready == Some(true) || conditions.serving == Some(true) - }); - if !is_serving { + if !is_eligible_recovery_endpoint(endpoint) { continue; } for address in &endpoint.addresses { @@ -407,10 +480,40 @@ fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { } } - peers + // Shuffle so N simultaneously-joining replicas do not all deterministically + // pick the lowest-IP peer (BTreeSet order), concentrating dump work on one + // serving EPP. RandomState is seeded per process, so each bootstrap gets a + // different order; serial fallback through the shuffled list is retained. + let hasher = RandomState::new(); + let mut urls: Vec = peers .into_iter() .map(|ip| format!("http://{}", authority(&ip, port))) - .collect() + .collect(); + urls.sort_by_cached_key(|url| { + let mut h = hasher.build_hasher(); + url.hash(&mut h); + h.finish() + }); + urls +} + +/// A peer is an eligible recovery source only when it is actively serving and +/// not terminating. `ready` alone is insufficient in both directions: with +/// `publishNotReadyAddresses` an endpoint can be `ready=true, serving=false` +/// (a live-but-cold replica that would return an empty dump), while a draining +/// pod is `serving=true, terminating=true` and can vanish mid-transfer. When +/// the `serving` condition is absent (legacy slices), fall back to `ready`. +fn is_eligible_recovery_endpoint(endpoint: &Endpoint) -> bool { + let Some(conditions) = endpoint.conditions.as_ref() else { + return false; + }; + if conditions.terminating == Some(true) { + return false; + } + match conditions.serving { + Some(serving) => serving, + None => conditions.ready == Some(true), + } } /// Format `host:port`, bracketing IPv6 literals (`fd00::1` -> `[fd00::1]`) so the @@ -627,7 +730,7 @@ mod tests { peer_ports(slices.iter()).unwrap(), PeerPorts { replica_sync: 9092, - selection_http: 9093, + selection_http: Some(9093), } ); } @@ -965,7 +1068,7 @@ mod tests { start_recovery( service.clone(), store, - ports.selection_http, + ports.selection_http.unwrap(), changes_rx, cancel.clone(), ) @@ -1200,6 +1303,59 @@ mod tests { ); } + #[test] + fn recovery_excludes_ready_but_not_serving_peer() { + // `publishNotReadyAddresses` can yield ready=true, serving=false for a + // live-but-cold replica: it would return an empty dump, so it must not + // be a recovery source. + let slice = EndpointSlice { + address_type: "IPv4".to_string(), + endpoints: vec![Endpoint { + addresses: vec!["10.0.0.2".to_string()], + conditions: Some(EndpointConditions { + ready: Some(true), + serving: Some(false), + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + }; + let (store, _writer) = store_and_writer(vec![slice]); + assert!(recovery_peer_urls(&store, "10.0.0.9", 9093).is_empty()); + } + + #[test] + fn recovery_excludes_terminating_serving_peer() { + // A draining pod is serving=true, terminating=true and can vanish + // mid-transfer, so it must not be a recovery source. + let slice = EndpointSlice { + address_type: "IPv4".to_string(), + endpoints: vec![Endpoint { + addresses: vec!["10.0.0.2".to_string()], + conditions: Some(EndpointConditions { + serving: Some(true), + terminating: Some(true), + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + }; + let (store, _writer) = store_and_writer(vec![slice]); + assert!(recovery_peer_urls(&store, "10.0.0.9", 9093).is_empty()); + } + + #[test] + fn selection_http_port_is_optional() { + // A Service without `selection-http` (a deployment predating the dump + // endpoint) resolves with selection_http=None instead of failing. + let slice = slice_with_replica_port(Some(9092)); + let ports = peer_ports([&slice].into_iter()).expect("resolve ports"); + assert_eq!(ports.replica_sync, 9092); + assert_eq!(ports.selection_http, None); + } + #[test] fn recovery_peer_urls_bracket_ipv6() { let slice = EndpointSlice { diff --git a/deploy/inference-gateway/ext-proc/src/peer_http.rs b/deploy/inference-gateway/ext-proc/src/peer_http.rs index 9a7dafd4f923..5c9df89f2dda 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_http.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_http.rs @@ -5,26 +5,48 @@ use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use anyhow::{Context, Result}; -use axum::{Json, Router, extract::State, routing::get}; +use axum::{ + Router, + extract::{Query, State}, + http::StatusCode, + response::{IntoResponse, Response}, + routing::get, +}; use tokio::net::TcpListener; use tokio_util::sync::CancellationToken; use dynamo_kv_router::services::selection::SelectionService; +#[derive(Clone)] +struct AppState { + service: Arc, + recovered: Arc, +} + +#[derive(serde::Deserialize)] +struct DumpQuery { + /// Caller's accepted snapshot budget in bytes. The peer rejects over-budget + /// snapshots with 413 *before* writing the body. Absent or `0` = unbounded. + max_bytes: Option, +} + /// Bind the dump listener before returning so peer recovery cannot race server startup. pub(crate) async fn spawn( service: Arc, port: u16, pod_ip: IpAddr, cancel: CancellationToken, + recovered: Arc, ) -> Result<()> { let address = listener_addr(pod_ip, port); let listener = TcpListener::bind(address) .await .with_context(|| format!("binding EPP peer HTTP server on {address}"))?; - let app = Router::new().route("/dump", get(dump)).with_state(service); + let state = AppState { service, recovered }; + let app = Router::new().route("/dump", get(dump)).with_state(state); tokio::spawn(async move { if let Err(error) = axum::serve(listener, app) @@ -44,8 +66,45 @@ fn listener_addr(pod_ip: IpAddr, port: u16) -> SocketAddr { } } -async fn dump(State(service): State>) -> Json { - Json(service.indexer_snapshot().await) +async fn dump(State(state): State, Query(query): Query) -> Response { + // Do not serve a snapshot until local recovery/bootstrap has finished: the + // index is empty during recovery, and an early /dump could let a sibling + // latch onto an empty index while a warm one exists elsewhere. + if !state.recovered.load(Ordering::Acquire) { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "peer KV index not yet recovered", + ) + .into_response(); + } + let snapshot = state.service.indexer_snapshot().await; + let bytes = match serde_json::to_vec(&snapshot) { + Ok(bytes) => bytes, + Err(error) => { + tracing::warn!(%error, "Failed to serialize peer KV-index snapshot"); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + "snapshot serialization failed", + ) + .into_response(); + } + }; + if let Some(max) = query.max_bytes + && max > 0 + && bytes.len() as u64 > max + { + return ( + StatusCode::PAYLOAD_TOO_LARGE, + "peer KV index snapshot exceeds max_bytes", + ) + .into_response(); + } + ( + StatusCode::OK, + [(axum::http::header::CONTENT_TYPE, "application/json")], + bytes, + ) + .into_response() } #[cfg(test)] @@ -83,6 +142,7 @@ mod tests { port, "127.0.0.1".parse().unwrap(), cancel.clone(), + Arc::new(AtomicBool::new(true)), ) .await .expect("spawn peer HTTP server"); @@ -104,6 +164,7 @@ mod tests { port, "127.0.0.1".parse().unwrap(), cancel.clone(), + Arc::new(AtomicBool::new(true)), ) .await .expect("spawn peer HTTP server"); @@ -120,6 +181,55 @@ mod tests { service.shutdown().await; } + #[tokio::test] + async fn dump_returns_503_until_recovered() { + let service = service().await; + let cancel = CancellationToken::new(); + let port = free_tcp_port(); + spawn( + service.clone(), + port, + "127.0.0.1".parse().unwrap(), + cancel.clone(), + Arc::new(AtomicBool::new(false)), + ) + .await + .expect("spawn peer HTTP server"); + + let resp = reqwest::get(format!("http://127.0.0.1:{port}/dump")) + .await + .expect("request dump"); + assert_eq!(resp.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE); + + cancel.cancel(); + service.shutdown().await; + } + + #[tokio::test] + async fn dump_rejects_over_budget_snapshot() { + let service = service().await; + let cancel = CancellationToken::new(); + let port = free_tcp_port(); + spawn( + service.clone(), + port, + "127.0.0.1".parse().unwrap(), + cancel.clone(), + Arc::new(AtomicBool::new(true)), + ) + .await + .expect("spawn peer HTTP server"); + + // max_bytes=1 is smaller than any serialized snapshot ("{}" is 2 bytes). + let resp = reqwest::get(format!("http://127.0.0.1:{port}/dump?max_bytes=1")) + .await + .expect("request dump"); + assert_eq!(resp.status(), reqwest::StatusCode::PAYLOAD_TOO_LARGE); + + cancel.cancel(); + service.shutdown().await; + } + #[test] fn listener_addr_matches_pod_ip_family() { assert_eq!( diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index be79fea510ac..079b5d085e11 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -283,25 +283,49 @@ impl Selector { })? .parse::() .context("POD_IP must be a valid IPv4 or IPv6 address")?; - crate::peer_http::spawn( - service.clone(), - replication.ports.selection_http, - self_ip, - cancel.clone(), - ) - .await?; - Some( - crate::peer_discovery::spawn( - service.clone(), - &cfg.namespace, - &replication.service_name, - replication.ports.replica_sync, - replication.ports.selection_http, - self_ip.to_string(), - cancel.clone(), - ) - .await?, - ) + + match replication.ports.selection_http { + None => { + // The Service does not declare `selection-http` (an + // image-only upgrade from a deployment that predates the + // dump endpoint). Degrade to no-recovery instead of failing + // startup: the replica bootstraps empty and serves. + tracing::warn!( + service = %replication.service_name, + "selection-http port not found; peer KV-index recovery disabled \ + (bootstrapping empty). Update the peer Service to add the \ + selection-http named port." + ); + None + } + Some(selection_http_port) => { + // Shared flag gating /dump: the endpoint answers 503 until + // this replica's own recovery/bootstrap finishes, so a + // stale-ready sibling cannot hand out an empty snapshot + // mid-recovery. + let recovered = Arc::new(AtomicBool::new(false)); + crate::peer_http::spawn( + service.clone(), + selection_http_port, + self_ip, + cancel.clone(), + recovered.clone(), + ) + .await?; + crate::peer_discovery::spawn( + service.clone(), + &cfg.namespace, + &replication.service_name, + replication.ports.replica_sync, + selection_http_port, + self_ip.to_string(), + cancel.clone(), + recovered.clone(), + ) + .await?; + Some(recovered) + } + } } else { None }; diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index 67c21bb6c5cf..1e4295b0d862 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -490,14 +490,11 @@ without an attempt limit: serving with an empty or stale index is less safe than service. Alert on a prolonged `NOT_SERVING` period and inspect EndpointSlices, NetworkPolicies, and the `/dump` port before restarting replicas. -The `agg.yaml` example uses a rolling-update strategy with `maxUnavailable: 0` and `maxSurge: 100%`. -During the first upgrade from an EPP image without `selection-http`, this lets two new Pods start -together and recover from each other's already-bound dump endpoints. The surge can temporarily double -the EPP Pod count, so make sure the cluster has capacity for it. Later rollouts recover from an existing -ready replica. The example assumes the Service's EndpointSlices expose a consistent pair of named -ports during rollout; do not mix an image that lacks `selection-http` into the same peer Service without -first applying the new manifest and planning the transition. If a mixed-version EndpointSlice omits the -port, the new EPP remains out of service until the peer set is upgraded and the ports converge. +The example uses the Deployment's default rolling-update strategy. A recovery source must already be +serving (`serving=true`, not terminating), so a replica recovers only from an existing ready peer — +never from a peer that is itself still bootstrapping. For the first upgrade from an EPP image without +`selection-http`, apply the manifest (which adds the named port) before upgrading the image; an +image-only upgrade degrades to no-recovery (the replica bootstraps empty) rather than failing startup. Recovery restores KV-index placement only. It does not restore active reservations or provide an atomic snapshot-plus-live-event handoff. Replica lifecycle synchronization and worker replay continue to converge diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index f00c99c527b8..647d0c9a3c57 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -75,7 +75,16 @@ async fn try_recover_from_peer( peer_url: &str, registry: &WorkerRegistry, ) -> Result<()> { - let dump_url = format!("{peer_url}/dump"); + // Pass the accepted budget to the peer so it can reject an over-budget + // snapshot with 413 *before* serializing/transmitting it. `0` means no + // budget (unbounded). The Content-Length check below remains as receiver + // defense in depth. + let max_dump_bytes = env_u64(MAX_DUMP_BYTES_ENV, DEFAULT_MAX_DUMP_BYTES); + let dump_url = if max_dump_bytes > 0 { + format!("{peer_url}/dump?max_bytes={max_dump_bytes}") + } else { + format!("{peer_url}/dump") + }; tracing::info!(url = %dump_url, "fetching dump from peer"); let resp = client @@ -90,10 +99,11 @@ async fn try_recover_from_peer( // Fail fast on an oversized snapshot before reading the body: the dump is // materialized fully in memory on both sides, so a large body either OOMs - // or trips the request timeout and the retry loop. A clear error is more - // actionable than either. - let max_dump_bytes = env_u64(MAX_DUMP_BYTES_ENV, DEFAULT_MAX_DUMP_BYTES); + // or trips the request timeout and the retry loop. The peer already rejects + // over-budget bodies with 413, so this only fires for a peer without the + // budget-aware endpoint. if let Some(len) = resp.content_length() + && max_dump_bytes > 0 && len > max_dump_bytes { anyhow::bail!( @@ -119,6 +129,13 @@ async fn try_recover_from_peer( // peer's dump land in the matching lower-tier slot rather than the // device primary. The peer side retags lower-tier events in // `Indexer::dump_events`, so the `storage_tier` here is correct. + // + // Re-application is idempotent by construction: the radix index + // keys blocks by `tokens_hash`, so applying the same (or an + // overlapping) peer dump again is a no-op for blocks already + // present. A cancelled-then-retried attempt can therefore leave a + // partially applied snapshot that the next attempt safely re-applies + // on top of, without inflating or duplicating residency state. indexer .apply_event_routed(event) .await From 5cd033fc7e71b7d994b366089f1302226286f5bf Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Thu, 20 Aug 2026 10:51:06 +0800 Subject: [PATCH 07/17] style(epp): satisfy clippy on peer-discovery review changes - Use BuildHasher::hash_one for the recovery-candidate shuffle. - Allow too_many_arguments on spawn after the recovered-flag parameter. Signed-off-by: Peter Pan --- deploy/inference-gateway/ext-proc/src/peer_discovery.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index e52b330c0c06..8b17add50bfb 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -9,7 +9,7 @@ use std::collections::BTreeSet; use std::collections::hash_map::RandomState; use std::future::Future; -use std::hash::{BuildHasher, Hash, Hasher}; +use std::hash::BuildHasher; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -189,6 +189,7 @@ fn named_tcp_port(slices: &[&EndpointSlice], port_name: &str) -> Result { /// /// This call does not return until initial KV-index recovery succeeds or the /// authoritative sibling set is empty. The dump server must already be bound. +#[allow(clippy::too_many_arguments)] pub async fn spawn( service: Arc, namespace: &str, @@ -489,11 +490,7 @@ fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { .into_iter() .map(|ip| format!("http://{}", authority(&ip, port))) .collect(); - urls.sort_by_cached_key(|url| { - let mut h = hasher.build_hasher(); - url.hash(&mut h); - h.finish() - }); + urls.sort_by_cached_key(|url| hasher.hash_one(url)); urls } From 733f694eb3ffaa22064f9f92acffd8821915a6a4 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Thu, 20 Aug 2026 12:36:50 +0800 Subject: [PATCH 08/17] fix(epp): address CodeRabbit review on peer recovery - peer_discovery: only reset the backoff when the recovery candidate set actually changes; unrelated EndpointSlice churn no longer resets it, keeping exponential backoff effective during rolling updates. - agg.yaml: add a NetworkPolicy restricting replica-agg and selection-http (/dump) ingress to sibling EPP pods; the GAIE gateway and kubelet probes keep reaching gRPC and grpc-health. - peer_http: return 500 instead of a 200 {"error": ...} body when an indexer dump fails, so the recovery consumer never parses a shape it cannot deserialize. - recovery: read the /dump body as a bounded stream (chunked responses without Content-Length can no longer buffer without bound); zero cap stays disabled, with coverage for the no-Content-Length path. Signed-off-by: Peter Pan --- .../ext-proc/examples/onramp/agg.yaml | 36 ++++++ .../ext-proc/src/peer_discovery.rs | 9 +- .../ext-proc/src/peer_http.rs | 52 ++++++++ .../src/services/indexer/recovery.rs | 121 +++++++++++++++++- 4 files changed, 215 insertions(+), 3 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml index 86ebddb63639..582602471574 100644 --- a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml +++ b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml @@ -319,6 +319,42 @@ spec: port: 9093 targetPort: selection-http --- +# `selection-http` serves the peer KV-index snapshot (`GET /dump`) without +# authentication, so it is reachable only from sibling EPP pods. The same +# peer-only restriction applies to `replica-agg` lifecycle sync. The GAIE +# gateway still reaches gRPC (9002) and the kubelet still probes grpc-health +# (9003) from any source; KV events flow *out* of the EPP to workers, so no +# inbound rule is needed for them. Clusters without a NetworkPolicy CNI +# simply ignore this object. +apiVersion: networking.k8s.io/v1 +kind: NetworkPolicy +metadata: + name: dynamo-epp + labels: + app: dynamo-epp +spec: + podSelector: + matchLabels: + app: dynamo-epp + policyTypes: + - Ingress + ingress: + - from: + - podSelector: + matchLabels: + app: dynamo-epp + ports: + - protocol: TCP + port: replica-agg + - protocol: TCP + port: selection-http + - from: [] # any source: GAIE gateway + kubelet probes + ports: + - protocol: TCP + port: grpc + - protocol: TCP + port: grpc-health +--- # InferencePool selects the raw vLLM pods and delegates endpoint selection to # the EPP. The EPP reads this object (read-only) to learn the pod selector and # target port — the same object the gateway routes to. diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index 8b17add50bfb..5188b6fed762 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -407,7 +407,14 @@ where } changed = changes_rx.changed() => { changed.context("EPP peer EndpointSlice watch ended during KV-index recovery")?; - backoff = initial_backoff; + // Only a genuine candidate-set change earns an immediate retry + // with the initial backoff. Unrelated churn (readiness flips, + // metadata/zone updates) must not reset the backoff either: + // under a rolling update it would keep the retry loop hot at + // the initial backoff instead of letting it grow. + if recovery_peer_urls(store, self_ip, selection_http_port) != peers { + backoff = initial_backoff; + } } _ = tokio::time::sleep(backoff) => { backoff = backoff.saturating_mul(2).min(max_backoff); diff --git a/deploy/inference-gateway/ext-proc/src/peer_http.rs b/deploy/inference-gateway/ext-proc/src/peer_http.rs index 5c9df89f2dda..a8268ff62b20 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_http.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_http.rs @@ -66,6 +66,17 @@ fn listener_addr(pod_ip: IpAddr, port: u16) -> SocketAddr { } } +/// True when the indexer snapshot contains a per-model `{"error": ...}` +/// entry, which happens when one model's indexer failed to dump. The +/// recovery consumer deserializes each entry as `DumpEntry` (`block_size` +/// plus `events`), so an error entry would make the whole body unparseable; +/// the dump handler fails such snapshots with a non-success status. +fn snapshot_has_failed_dump(snapshot: &serde_json::Value) -> bool { + snapshot + .as_object() + .is_some_and(|entries| entries.values().any(|entry| entry.get("error").is_some())) +} + async fn dump(State(state): State, Query(query): Query) -> Response { // Do not serve a snapshot until local recovery/bootstrap has finished: the // index is empty during recovery, and an early /dump could let a sibling @@ -78,6 +89,19 @@ async fn dump(State(state): State, Query(query): Query) -> .into_response(); } let snapshot = state.service.indexer_snapshot().await; + // A per-model indexer dump failure is surfaced as an `{"error": ...}` + // entry inside the snapshot. The recovery consumer expects `DumpEntry` + // values and cannot parse that shape, so fail the whole dump with a + // non-success status instead of returning a 200 body the consumer would + // reject (and then fall back to an empty index on). + if snapshot_has_failed_dump(&snapshot) { + tracing::warn!("Peer KV-index snapshot contains an indexer dump failure"); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + "peer KV-index snapshot generation failed", + ) + .into_response(); + } let bytes = match serde_json::to_vec(&snapshot) { Ok(bytes) => bytes, Err(error) => { @@ -242,6 +266,34 @@ mod tests { ); } + #[test] + fn snapshot_with_failed_indexer_dump_is_detected() { + // A per-model indexer dump failure surfaces as `{"error": ...}`; the + // recovery consumer expects `DumpEntry` values, so such snapshots must + // be rejected as a whole. + let failed = serde_json::json!({ + "model:default": {"error": "indexer dump failed"}, + }); + assert!(snapshot_has_failed_dump(&failed)); + + let mixed = serde_json::json!({ + "model:a": {"block_size": 16, "events": []}, + "model:b": {"error": "boom"}, + }); + assert!(snapshot_has_failed_dump(&mixed)); + } + + #[test] + fn snapshot_without_failed_indexer_dump_is_accepted() { + let healthy = serde_json::json!({ + "model:default": {"block_size": 16, "events": []}, + }); + assert!(!snapshot_has_failed_dump(&healthy)); + + let empty = serde_json::json!({}); + assert!(!snapshot_has_failed_dump(&empty)); + } + #[tokio::test] async fn ipv6_pod_binds_an_ipv6_listener() { let listener = TcpListener::bind(listener_addr("::1".parse().unwrap(), 0)) diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index 647d0c9a3c57..3c307b95901e 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -87,7 +87,7 @@ async fn try_recover_from_peer( }; tracing::info!(url = %dump_url, "fetching dump from peer"); - let resp = client + let mut resp = client .get(&dump_url) .send() .await @@ -112,8 +112,13 @@ async fn try_recover_from_peer( ); } + // Read the body as a bounded stream: a chunked response without + // Content-Length must not buffer without bound (resp.json() would). + // `0` disables the cap, consistent with the documented behavior. + let body = read_dump_body(&mut resp, max_dump_bytes).await?; + let dump: HashMap = - resp.json().await.context("failed to parse dump response")?; + serde_json::from_slice(&body).context("failed to parse dump response")?; let mut total_events = 0usize; for (map_key, entry) in dump { let (model_name, routing_group) = map_key @@ -154,6 +159,30 @@ async fn try_recover_from_peer( Ok(()) } +/// Read the `/dump` response body as a bounded stream, failing as soon as the +/// configured cap is exceeded so a chunked response without `Content-Length` +/// cannot buffer without bound. `max_dump_bytes == 0` disables the cap, +/// consistent with the documented behavior. +async fn read_dump_body(resp: &mut reqwest::Response, max_dump_bytes: u64) -> Result> { + let mut body = Vec::new(); + while let Some(chunk) = resp + .chunk() + .await + .context("failed to read dump response body")? + { + if max_dump_bytes > 0 + && (body.len() as u64).saturating_add(chunk.len() as u64) > max_dump_bytes + { + anyhow::bail!( + "peer dump is too large: exceeds limit {max_dump_bytes} bytes \ + (raise {MAX_DUMP_BYTES_ENV} to accept larger snapshots)" + ); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + #[cfg(test)] mod tests { use super::*; @@ -172,4 +201,92 @@ mod tests { // treated as a parse failure that falls back to the default. assert_eq!(parse_u64(Some("0".to_string()), 512), 0); } + + /// Serve one `chunked` HTTP response with no `Content-Length` header, the + /// shape `read_dump_body` must bound without relying on the header. + async fn serve_chunked_without_content_length(body: &[u8]) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test server"); + let addr = listener.local_addr().expect("read test addr"); + let body = body.to_vec(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accept test client"); + // Drain the request line + headers. + let mut buf = [0u8; 4096]; + let _ = socket.read(&mut buf).await.expect("read request"); + // Transfer-Encoding: chunked, deliberately no Content-Length. + let mut resp = String::from( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nTransfer-Encoding: chunked\r\n\r\n", + ); + for chunk in body.chunks(7) { + resp.push_str(&format!("{:x}\r\n", chunk.len())); + resp.push_str(&String::from_utf8_lossy(chunk)); + resp.push_str("\r\n"); + } + resp.push_str("0\r\n\r\n"); + socket + .write_all(resp.as_bytes()) + .await + .expect("write response"); + }); + let url = format!("http://{addr}/dump"); + // Detach: the server task finishes once the client has consumed the + // body; the test runtime aborts it when the test ends. + drop(server); + url + } + + #[tokio::test] + async fn read_dump_body_bounds_chunked_response_without_content_length() { + // Small body, cap not hit: bounded read succeeds and returns the body. + let url = + serve_chunked_without_content_length(br#"{"m:default":{"block_size":16,"events":[]}}"#) + .await; + let client = reqwest::Client::new(); + let mut resp = client.get(&url).send().await.expect("GET dump"); + assert!( + resp.content_length().is_none(), + "fixture must not declare Content-Length" + ); + let body = read_dump_body(&mut resp, 1024) + .await + .expect("bounded read under cap"); + assert_eq!( + body, + br#"{"m:default":{"block_size":16,"events":[]}}"#.to_vec() + ); + } + + #[tokio::test] + async fn read_dump_body_rejects_oversized_chunked_response_without_content_length() { + // Body larger than the cap, no Content-Length: the bounded reader must + // fail instead of buffering without bound (resp.json() would). + let url = serve_chunked_without_content_length(&[b'x'; 64]).await; + let client = reqwest::Client::new(); + let mut resp = client.get(&url).send().await.expect("GET dump"); + assert!(resp.content_length().is_none()); + let err = read_dump_body(&mut resp, 32) + .await + .expect_err("cap must be enforced without Content-Length"); + assert!( + err.to_string().contains("peer dump is too large"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn read_dump_body_zero_cap_is_unbounded() { + // max_dump_bytes == 0 disables the cap, even for a body larger than + // any plausible default and with no Content-Length header. + let url = serve_chunked_without_content_length(&[b'x'; 512]).await; + let client = reqwest::Client::new(); + let mut resp = client.get(&url).send().await.expect("GET dump"); + let body = read_dump_body(&mut resp, 0) + .await + .expect("zero disables the cap"); + assert_eq!(body.len(), 512); + } } From 47022ee4c6f1a997794ed038af8582a5e5df347e Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Thu, 20 Aug 2026 21:32:07 +0800 Subject: [PATCH 09/17] feat(epp): stream peer KV-index dump as NDJSON MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The EPP's 9093 /dump endpoint is new in this PR and serves only EPP peer recovery, so it uses a single format: NDJSON, one StreamDumpRecord per line. The receiver applies each event as it arrives and drops it, so the whole snapshot is never buffered on either side — no max-bytes budget knob is needed (it was a magic number users could not set correctly; a serving peer already proves the index fits, since replicas are homogeneous). The pre-existing JSON /dump + recover_from_peers path is untouched: it serves standalone indexer/selection/python P2P recovery in other processes. EPP recovery goes through a dedicated recover_indexer_from_peers_streaming. - kv-router: dump_registry_records + StreamDumpRecord (structured, per-event), recover_from_peers_streaming (line-by-line apply, idempotent on retry) - peer_http: /dump streams NDJSON via Body::from_stream - peer_discovery: EPP recovery uses the streaming entry point; test dumps now serve an empty NDJSON body - tests: streaming recovery (empty/multi/truncated), NDJSON endpoint Signed-off-by: Peter Pan --- .../ext-proc/src/peer_discovery.rs | 19 +- .../ext-proc/src/peer_http.rs | 172 +++++---------- .../src/services/indexer/recovery.rs | 203 ++++++++++++++++++ lib/kv-router/src/services/indexer/server.rs | 55 +++++ .../src/services/selection/core/mod.rs | 13 ++ .../src/services/selection/service.rs | 19 ++ 6 files changed, 358 insertions(+), 123 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index 5188b6fed762..3ee0a103dd84 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -329,7 +329,7 @@ async fn recover_initial_index( cancel, initial_backoff, max_backoff, - |service, peers| Box::pin(service.recover_indexer_from_peers(peers)), + |service, peers| Box::pin(service.recover_indexer_from_peers_streaming(peers)), ) .await } @@ -573,9 +573,7 @@ mod tests { use std::sync::atomic::Ordering; use super::*; - use axum::{ - Json, Router, extract::State, http::StatusCode, response::IntoResponse, routing::get, - }; + use axum::{Router, extract::State, http::StatusCode, response::IntoResponse, routing::get}; use k8s_openapi::api::discovery::v1::{Endpoint, EndpointConditions, EndpointPort}; use std::sync::atomic::AtomicUsize; use std::time::Duration; @@ -1002,10 +1000,11 @@ mod tests { release: Arc, } - async fn gated_dump(State(gate): State) -> Json { + async fn gated_dump(State(gate): State) -> axum::response::Response { gate.requested.notify_one(); gate.release.notified().await; - Json(serde_json::json!({})) + // Empty body = an empty streaming (NDJSON) dump, a valid recovery. + axum::response::Response::new(axum::body::Body::empty()) } #[derive(Clone)] @@ -1019,7 +1018,8 @@ mod tests { state.first_failed.notify_one(); (StatusCode::SERVICE_UNAVAILABLE, "not ready").into_response() } else { - Json(serde_json::json!({})).into_response() + // Empty body = an empty streaming (NDJSON) dump, a valid recovery. + axum::response::Response::new(axum::body::Body::empty()).into_response() } } @@ -1057,7 +1057,10 @@ mod tests { let server = tokio::spawn(async move { axum::serve( listener, - Router::new().route("/dump", get(|| async { Json(serde_json::json!({})) })), + Router::new().route( + "/dump", + get(|| async { axum::response::Response::new(axum::body::Body::empty()) }), + ), ) .await }); diff --git a/deploy/inference-gateway/ext-proc/src/peer_http.rs b/deploy/inference-gateway/ext-proc/src/peer_http.rs index a8268ff62b20..608cd54cec52 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_http.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_http.rs @@ -10,29 +10,28 @@ use std::sync::atomic::{AtomicBool, Ordering}; use anyhow::{Context, Result}; use axum::{ Router, - extract::{Query, State}, - http::StatusCode, + body::Body, + extract::State, + http::{StatusCode, header::CONTENT_TYPE}, response::{IntoResponse, Response}, routing::get, }; +use futures::StreamExt; use tokio::net::TcpListener; use tokio_util::sync::CancellationToken; +use dynamo_kv_router::services::indexer::server::StreamDumpRecord; use dynamo_kv_router::services::selection::SelectionService; +/// Media type of the streaming NDJSON dump (one [`StreamDumpRecord`] per line). +const STREAM_DUMP_MEDIA_TYPE: &str = "application/x-ndjson"; + #[derive(Clone)] struct AppState { service: Arc, recovered: Arc, } -#[derive(serde::Deserialize)] -struct DumpQuery { - /// Caller's accepted snapshot budget in bytes. The peer rejects over-budget - /// snapshots with 413 *before* writing the body. Absent or `0` = unbounded. - max_bytes: Option, -} - /// Bind the dump listener before returning so peer recovery cannot race server startup. pub(crate) async fn spawn( service: Arc, @@ -66,18 +65,7 @@ fn listener_addr(pod_ip: IpAddr, port: u16) -> SocketAddr { } } -/// True when the indexer snapshot contains a per-model `{"error": ...}` -/// entry, which happens when one model's indexer failed to dump. The -/// recovery consumer deserializes each entry as `DumpEntry` (`block_size` -/// plus `events`), so an error entry would make the whole body unparseable; -/// the dump handler fails such snapshots with a non-success status. -fn snapshot_has_failed_dump(snapshot: &serde_json::Value) -> bool { - snapshot - .as_object() - .is_some_and(|entries| entries.values().any(|entry| entry.get("error").is_some())) -} - -async fn dump(State(state): State, Query(query): Query) -> Response { +async fn dump(State(state): State) -> Response { // Do not serve a snapshot until local recovery/bootstrap has finished: the // index is empty during recovery, and an early /dump could let a sibling // latch onto an empty index while a warm one exists elsewhere. @@ -88,47 +76,42 @@ async fn dump(State(state): State, Query(query): Query) -> ) .into_response(); } - let snapshot = state.service.indexer_snapshot().await; - // A per-model indexer dump failure is surfaced as an `{"error": ...}` - // entry inside the snapshot. The recovery consumer expects `DumpEntry` - // values and cannot parse that shape, so fail the whole dump with a - // non-success status instead of returning a 200 body the consumer would - // reject (and then fall back to an empty index on). - if snapshot_has_failed_dump(&snapshot) { - tracing::warn!("Peer KV-index snapshot contains an indexer dump failure"); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - "peer KV-index snapshot generation failed", - ) - .into_response(); - } - let bytes = match serde_json::to_vec(&snapshot) { - Ok(bytes) => bytes, + + let records = match state.service.indexer_stream_records().await { + Ok(records) => records, Err(error) => { - tracing::warn!(%error, "Failed to serialize peer KV-index snapshot"); + tracing::warn!(%error, "Failed to collect peer KV-index dump records"); return ( StatusCode::INTERNAL_SERVER_ERROR, - "snapshot serialization failed", + "peer KV-index snapshot generation failed", ) .into_response(); } }; - if let Some(max) = query.max_bytes - && max > 0 - && bytes.len() as u64 > max - { - return ( - StatusCode::PAYLOAD_TOO_LARGE, - "peer KV index snapshot exceeds max_bytes", - ) - .into_response(); - } - ( - StatusCode::OK, - [(axum::http::header::CONTENT_TYPE, "application/json")], - bytes, - ) - .into_response() + + // Stream NDJSON, one record per line. The receiver applies each event and + // drops it, so the whole snapshot is never buffered on either side and no + // max-bytes budget is needed. + let stream = futures::stream::iter(records).map(|record: StreamDumpRecord| { + match serde_json::to_vec(&record) { + Ok(mut bytes) => { + bytes.push(b'\n'); + Ok::<_, std::convert::Infallible>(bytes) + } + Err(error) => { + tracing::warn!(%error, "Failed to serialize dump record"); + // Signal a mid-stream failure with an empty frame; the receiver + // treats a truncated record as an error and retries. + Ok::<_, std::convert::Infallible>(Vec::new()) + } + } + }); + + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, STREAM_DUMP_MEDIA_TYPE) + .body(Body::from_stream(stream)) + .expect("static response builder") } #[cfg(test)] @@ -179,7 +162,7 @@ mod tests { } #[tokio::test] - async fn dump_endpoint_matches_selection_service_snapshot() { + async fn streaming_dump_emits_ndjson_records() { let service = service().await; let cancel = CancellationToken::new(); let port = free_tcp_port(); @@ -193,13 +176,25 @@ mod tests { .await .expect("spawn peer HTTP server"); - let response: serde_json::Value = reqwest::get(format!("http://127.0.0.1:{port}/dump")) - .await - .expect("request dump") - .json() + let resp = reqwest::get(format!("http://127.0.0.1:{port}/dump")) .await - .expect("decode dump"); - assert_eq!(response, service.indexer_snapshot().await); + .expect("request dump"); + assert_eq!(resp.status(), reqwest::StatusCode::OK); + assert_eq!( + resp.headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()), + Some("application/x-ndjson") + ); + let body = resp.text().await.expect("read body"); + // An empty index yields an empty stream; every non-empty line must be a + // parseable StreamDumpRecord. + for line in body.lines() { + if line.trim().is_empty() { + continue; + } + serde_json::from_str::(line).expect("valid NDJSON record"); + } cancel.cancel(); service.shutdown().await; @@ -229,31 +224,6 @@ mod tests { service.shutdown().await; } - #[tokio::test] - async fn dump_rejects_over_budget_snapshot() { - let service = service().await; - let cancel = CancellationToken::new(); - let port = free_tcp_port(); - spawn( - service.clone(), - port, - "127.0.0.1".parse().unwrap(), - cancel.clone(), - Arc::new(AtomicBool::new(true)), - ) - .await - .expect("spawn peer HTTP server"); - - // max_bytes=1 is smaller than any serialized snapshot ("{}" is 2 bytes). - let resp = reqwest::get(format!("http://127.0.0.1:{port}/dump?max_bytes=1")) - .await - .expect("request dump"); - assert_eq!(resp.status(), reqwest::StatusCode::PAYLOAD_TOO_LARGE); - - cancel.cancel(); - service.shutdown().await; - } - #[test] fn listener_addr_matches_pod_ip_family() { assert_eq!( @@ -266,34 +236,6 @@ mod tests { ); } - #[test] - fn snapshot_with_failed_indexer_dump_is_detected() { - // A per-model indexer dump failure surfaces as `{"error": ...}`; the - // recovery consumer expects `DumpEntry` values, so such snapshots must - // be rejected as a whole. - let failed = serde_json::json!({ - "model:default": {"error": "indexer dump failed"}, - }); - assert!(snapshot_has_failed_dump(&failed)); - - let mixed = serde_json::json!({ - "model:a": {"block_size": 16, "events": []}, - "model:b": {"error": "boom"}, - }); - assert!(snapshot_has_failed_dump(&mixed)); - } - - #[test] - fn snapshot_without_failed_indexer_dump_is_accepted() { - let healthy = serde_json::json!({ - "model:default": {"block_size": 16, "events": []}, - }); - assert!(!snapshot_has_failed_dump(&healthy)); - - let empty = serde_json::json!({}); - assert!(!snapshot_has_failed_dump(&empty)); - } - #[tokio::test] async fn ipv6_pod_binds_an_ipv6_listener() { let listener = TcpListener::bind(listener_addr("::1".parse().unwrap(), 0)) diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index 3c307b95901e..e0255c51e27d 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -11,6 +11,7 @@ use crate::identity::RoutingPartitionId; use crate::protocols::RouterEvent; use super::registry::WorkerRegistry; +use super::server::StreamDumpRecord; /// Timeout for one peer `/dump` fetch (HTTP request + body transfer). A full /// KV-index snapshot can be tens of MB on a busy deployment; 10s only fits @@ -70,6 +71,119 @@ pub async fn recover_from_peers(peers: &[String], registry: &WorkerRegistry) -> Ok(false) } +/// Streaming peer recovery over an NDJSON `/dump` (one [`StreamDumpRecord`] per +/// line): each event is applied as it arrives, so the whole snapshot is never +/// buffered. No `max_bytes` budget is needed — peak memory is the index plus one +/// line, not the index plus the full JSON body. This is the EPP peer-recovery +/// protocol; the JSON [`recover_from_peers`] path stays for standalone +/// indexer/selection/python peers. +pub async fn recover_from_peers_streaming( + peers: &[String], + registry: &WorkerRegistry, +) -> Result { + let timeout = Duration::from_millis(env_u64( + RECOVERY_HTTP_TIMEOUT_ENV, + DEFAULT_RECOVERY_HTTP_TIMEOUT_MS, + )); + let client = reqwest::Client::builder() + .timeout(timeout) + .build() + .context("failed to build HTTP client")?; + + tokio::time::sleep(Duration::from_secs(1)).await; + + for peer_url in peers { + match try_recover_from_peer_streaming(&client, peer_url, registry).await { + Ok(()) => { + tracing::info!(peer = %peer_url, "streaming recovery from peer succeeded"); + return Ok(true); + } + Err(e) => { + tracing::warn!(peer = %peer_url, error = %e, "streaming recovery from peer failed, trying next"); + } + } + } + + Ok(false) +} + +async fn try_recover_from_peer_streaming( + client: &reqwest::Client, + peer_url: &str, + registry: &WorkerRegistry, +) -> Result<()> { + let dump_url = format!("{peer_url}/dump"); + tracing::info!(url = %dump_url, "streaming dump from peer"); + + let mut resp = client + .get(&dump_url) + .send() + .await + .context("HTTP request failed")?; + + if !resp.status().is_success() { + anyhow::bail!("peer returned status {}", resp.status()); + } + + let mut total_events = 0usize; + let mut line = String::new(); + // Read the body chunk-by-chunk and split on newlines. A chunk may contain + // several records or only part of one; `line` carries the partial record + // across chunks. Each record is applied and dropped immediately. + while let Some(chunk) = resp + .chunk() + .await + .context("failed to read dump response body")? + { + line.push_str(&String::from_utf8_lossy(&chunk)); + while let Some(pos) = line.find('\n') { + let record_line = line[..pos].to_string(); + line.drain(..=pos); + if record_line.is_empty() { + continue; + } + let record: StreamDumpRecord = serde_json::from_str(&record_line) + .with_context(|| format!("failed to parse dump record: {record_line}"))?; + apply_dump_record(registry, record).await?; + total_events += 1; + } + } + // A well-formed stream ends with a trailing newline; a leftover partial + // line means the stream was truncated mid-record. + if !line.is_empty() { + anyhow::bail!("peer dump stream ended mid-record"); + } + + // An empty dump is a valid recovery. Recovery candidates are restricted to + // already-serving peers (see `recovery_peer_urls`), so a zero-event dump + // means the serving peer genuinely holds no KV index yet (idle cluster), + // not a transient race. Rejecting it would deadlock cold starts and idle + // rollouts. + tracing::info!(total_events, "applied streamed dump events from peer"); + Ok(()) +} + +async fn apply_dump_record(registry: &WorkerRegistry, record: StreamDumpRecord) -> Result<()> { + let (model_name, routing_group) = record + .key + .split_once(':') + .ok_or_else(|| anyhow::anyhow!("invalid dump key format: {}", record.key))?; + + let key = RoutingPartitionId::new(model_name, routing_group); + let indexer = registry.get_or_create_indexer(key, record.block_size); + + // Re-application is idempotent by construction: the radix index keys blocks + // by `tokens_hash`, so applying the same (or an overlapping) peer dump again + // is a no-op for blocks already present. A cancelled-then-retried attempt + // can therefore leave a partially applied snapshot that the next attempt + // safely re-applies on top of. + indexer + .apply_event_routed(record.event) + .await + .context("peer recovery event was rejected by the local indexer")?; + Ok(()) +} + async fn try_recover_from_peer( client: &reqwest::Client, peer_url: &str, @@ -289,4 +403,93 @@ mod tests { .expect("zero disables the cap"); assert_eq!(body.len(), 512); } + + /// Serve a chunked HTTP response with the given Content-Type and body (no + /// Content-Length header), for exercising the streaming dump reader. + async fn serve_ndjson(body: &[u8]) -> String { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test server"); + let addr = listener.local_addr().expect("read test addr"); + let body = body.to_vec(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accept test client"); + let mut buf = [0u8; 4096]; + let _ = socket.read(&mut buf).await.expect("read request"); + let mut resp = String::from( + "HTTP/1.1 200 OK\r\nContent-Type: application/x-ndjson\r\nTransfer-Encoding: chunked\r\n\r\n", + ); + for chunk in body.chunks(7) { + resp.push_str(&format!("{:x}\r\n", chunk.len())); + resp.push_str(&String::from_utf8_lossy(chunk)); + resp.push_str("\r\n"); + } + resp.push_str("0\r\n\r\n"); + socket + .write_all(resp.as_bytes()) + .await + .expect("write response"); + }); + let url = format!("http://{addr}/dump"); + drop(server); + url + } + + #[tokio::test] + async fn streaming_recovery_applies_empty_dump() { + // A streaming dump with zero records (just a trailing newline) is a + // valid recovery, matching the non-streaming "empty dump is success". + let url = serve_ndjson(b"\n").await; + let client = reqwest::Client::new(); + let registry = WorkerRegistry::new(1); + try_recover_from_peer_streaming(&client, &url, ®istry) + .await + .expect("empty streamed dump must succeed"); + } + + #[tokio::test] + async fn streaming_recovery_rejects_truncated_record() { + // A stream that ends mid-record (no trailing newline) must fail instead + // of silently applying a partial event. + let url = serve_ndjson(br#"{"key":"m:default","block_size":16,"event":{"#).await; + let client = reqwest::Client::new(); + let registry = WorkerRegistry::new(1); + let err = try_recover_from_peer_streaming(&client, &url, ®istry) + .await + .expect_err("truncated stream must fail"); + assert!( + err.to_string().contains("mid-record"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn streaming_recovery_applies_multiple_records() { + // A well-formed NDJSON stream applies each record line-by-line. Empty + // blocks make each event a no-op but must parse and apply without error. + let record = |worker_id: u64| { + serde_json::json!({ + "key": "m:default", + "block_size": 16, + "event": { + "worker_id": worker_id, + "storage_tier": "device", + "event": {"event_id": 0, "data": {"stored": {"parent_hash": null, "blocks": []}}}, + } + }) + }; + let mut body = serde_json::to_vec(&record(1)).unwrap(); + body.push(b'\n'); + body.extend(serde_json::to_vec(&record(2)).unwrap()); + body.push(b'\n'); + + let url = serve_ndjson(&body).await; + let client = reqwest::Client::new(); + let registry = WorkerRegistry::new(1); + try_recover_from_peer_streaming(&client, &url, ®istry) + .await + .expect("well-formed stream must apply"); + } } diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index b07ee2f07405..a9a5e0f079b3 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -517,6 +517,61 @@ pub(crate) async fn dump_registry(registry: &WorkerRegistry) -> serde_json::Valu serde_json::json!(result) } +/// One NDJSON line of a streaming KV-index dump: a single `RouterEvent` with the +/// routing partition and block size it belongs to. The streaming dump emits one +/// of these per line so the receiver can apply each event and drop it instead of +/// buffering the whole snapshot. `key` is `"model:routing_group"`. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct StreamDumpRecord { + pub key: String, + pub block_size: u32, + pub event: crate::protocols::RouterEvent, +} + +/// Collect every indexer's events as structured records for a streaming dump. +/// +/// Structured counterpart to [`dump_registry`]: the same snapshot, but one +/// [`StreamDumpRecord`] per event so a caller can serialize and send records one +/// at a time instead of building one JSON object. A failed indexer dump aborts +/// the whole stream with an error. +pub(crate) async fn dump_registry_records( + registry: &WorkerRegistry, +) -> Result, String> { + let all = registry.all_indexers_with_block_size(); + let mut handles = Vec::with_capacity(all.len()); + + for (key, indexer, block_size) in all { + handles.push(tokio::spawn(async move { + let events = indexer.dump_events().await; + (key, events, block_size) + })); + } + + let mut records = Vec::new(); + for handle in handles { + match handle.await { + Ok((key, Ok(events), block_size)) => { + let map_key = format!("{}:{}", key.model_name, key.routing_group); + for event in events { + records.push(StreamDumpRecord { + key: map_key.clone(), + block_size, + event, + }); + } + } + Ok((key, Err(e), _)) => { + let map_key = format!("{}:{}", key.model_name, key.routing_group); + return Err(format!("indexer {map_key} dump failed: {e}")); + } + Err(e) => { + return Err(format!("dump task join error: {e}")); + } + } + } + Ok(records) +} + async fn handle_health() -> StatusCode { StatusCode::OK } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 134ac45579b3..0357412087e3 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -273,6 +273,13 @@ impl SelectionCore { recovery::recover_from_peers(peers, &self.indexer_registry).await } + pub(crate) async fn recover_indexer_from_peers_streaming( + &self, + peers: &[String], + ) -> anyhow::Result { + recovery::recover_from_peers_streaming(peers, &self.indexer_registry).await + } + pub(crate) fn signal_indexer_ready(&self) { self.indexer_registry.signal_ready(); } @@ -281,6 +288,12 @@ impl SelectionCore { crate::services::indexer::server::dump_registry(&self.indexer_registry).await } + pub(crate) async fn dump_indexer_records( + &self, + ) -> Result, String> { + crate::services::indexer::server::dump_registry_records(&self.indexer_registry).await + } + pub(crate) fn dispatch_replica_event(&self, envelope: ScopedReplicaEvent) { let (key, block_size, event) = envelope.into_parts(); if self diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 37a5afcdd453..13fa5078c46e 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -382,10 +382,29 @@ impl SelectionService { self.core.dump_indexer_events().await } + /// Structured snapshot of the local KV index for a streaming peer dump: one + /// record per event. The caller serializes records one at a time instead of + /// building a single JSON object in memory. + pub async fn indexer_stream_records( + &self, + ) -> Result, String> { + self.core.dump_indexer_records().await + } + pub async fn recover_indexer_from_peers(&self, peers: &[String]) -> anyhow::Result { self.core.recover_indexer_from_peers(peers).await } + /// EPP peer recovery over the streaming NDJSON `/dump` format (one event per + /// line, applied as it arrives). Distinct from [`Self::recover_indexer_from_peers`], + /// which uses the single-JSON format for standalone indexer/selection peers. + pub async fn recover_indexer_from_peers_streaming( + &self, + peers: &[String], + ) -> anyhow::Result { + self.core.recover_indexer_from_peers_streaming(peers).await + } + pub async fn cancelled(&self) { self.cancel_token.cancelled().await; } From 327d173267e2de342fba254f5402e20d4d930027 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Thu, 20 Aug 2026 22:14:41 +0800 Subject: [PATCH 10/17] feat(epp): expose KV-index recovery outcome as a metric Addresses review P1: a silent full-index loss (all replicas restart at once and every one bootstraps empty) and a silently-degraded rollout (peer Service lacks selection-http, so recovery is disabled with only a warn) are both invisible to operators. Log levels alone cannot distinguish a normal first deploy / rolling upgrade from an actual index loss. Add dynamo_epp_kv_recovery_state{state=recovered|empty_bootstrap|recovery_disabled}, set once at startup by the existing metrics infrastructure (9090, prometheus crate). Alerts should key on the *transition* (a replica that previously exported recovered flipping to empty_bootstrap = full-index loss) or a *stuck* non-recovered state past the upgrade window, not the bare state. Also lift the empty-bootstrap log from info to warn with an explicit 'EMPTY KV index' message so it stands out in logs. Signed-off-by: Peter Pan --- .../inference-gateway/ext-proc/src/metrics.rs | 52 ++++++++++++++++++- .../ext-proc/src/peer_discovery.rs | 12 ++++- .../ext-proc/src/selector.rs | 1 + 3 files changed, 62 insertions(+), 3 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/metrics.rs b/deploy/inference-gateway/ext-proc/src/metrics.rs index eca99706049d..c803b9d217ea 100644 --- a/deploy/inference-gateway/ext-proc/src/metrics.rs +++ b/deploy/inference-gateway/ext-proc/src/metrics.rs @@ -18,7 +18,9 @@ use axum::{ routing::get, }; use dynamo_llm::http::service::metrics::generate_log_buckets; -use prometheus::{Encoder, HistogramOpts, HistogramVec, Registry, TEXT_FORMAT, TextEncoder}; +use prometheus::{ + Encoder, HistogramOpts, HistogramVec, IntGaugeVec, Opts, Registry, TEXT_FORMAT, TextEncoder, +}; /// Port the `/metrics` endpoint binds to unless `DYN_EPP_METRICS_PORT` says /// otherwise. Distinct from the ext_proc gRPC port (9002) and the health port @@ -89,6 +91,54 @@ pub fn observe_cached_tokens(cached_tokens: u64) { .observe(cached_tokens as f64); } +/// Possible startup KV-index recovery outcomes for one EPP replica. The gauge +/// reports which state is current (1 for it, 0 for the others), so a scrape can +/// distinguish "recovered from a peer" from "bootstrapped empty" and from +/// "recovery disabled" — the difference between a normal cold start and a +/// silent full-index loss. +pub const KV_RECOVERY_RECOVERED: &str = "recovered"; +pub const KV_RECOVERY_EMPTY_BOOTSTRAP: &str = "empty_bootstrap"; +pub const KV_RECOVERY_DISABLED: &str = "recovery_disabled"; + +/// Current startup KV-index recovery outcome for this replica. Set once at +/// startup. A `1` appears on exactly one `state` label; the others are `0`. +/// +/// Alerts consume the *transition* (a replica that previously exported +/// `recovered` flipping to `empty_bootstrap` means a full-index loss, whereas +/// `empty_bootstrap` from first deployment is normal) or a *stuck* `disabled` / +/// `empty_bootstrap` lasting past the expected upgrade window — not the bare +/// state, which is expected during first deploy and rolling upgrades. +static KV_RECOVERY_STATE: LazyLock = LazyLock::new(|| { + let gauge = IntGaugeVec::new( + Opts::new( + "dynamo_epp_kv_recovery_state", + "Startup KV-index peer-recovery outcome for this EPP replica: \ + recovered (restored from a peer dump), empty_bootstrap (no eligible \ + serving peer), or recovery_disabled (peer Service lacks the \ + selection-http port). Exactly one label is 1.", + ), + &["state"], + ) + .expect("kv_recovery_state gauge options are statically valid"); + REGISTRY + .register(Box::new(gauge.clone())) + .expect("kv_recovery_state is the only registrant of its name"); + gauge +}); + +/// Mark this replica's startup KV-index recovery outcome. +pub fn set_kv_recovery_state(state: &str) { + for candidate in [ + KV_RECOVERY_RECOVERED, + KV_RECOVERY_EMPTY_BOOTSTRAP, + KV_RECOVERY_DISABLED, + ] { + KV_RECOVERY_STATE + .with_label_values(&[candidate]) + .set(i64::from(candidate == state)); + } +} + /// Serve `/metrics` until the process exits. pub async fn serve(port: u16) -> anyhow::Result<()> { let app = AxumRouter::new().route("/metrics", get(render)); diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index 3ee0a103dd84..84957557cda0 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -357,7 +357,12 @@ where reconcile_once(service, store, sync_port, self_ip, known).await; let peers = recovery_peer_urls(store, self_ip, selection_http_port); if peers.is_empty() { - tracing::info!("No sibling EPP peers found; bootstrapping an empty KV index"); + crate::metrics::set_kv_recovery_state(crate::metrics::KV_RECOVERY_EMPTY_BOOTSTRAP); + tracing::warn!( + "No serving sibling EPP peer found; bootstrapping an EMPTY KV index \ + (normal on first deploy, but a full-index loss if the cluster was \ + already serving — check expected replica count / recovery history)" + ); return Ok(()); } @@ -389,7 +394,10 @@ where }; match result { - Ok(true) => return Ok(()), + Ok(true) => { + crate::metrics::set_kv_recovery_state(crate::metrics::KV_RECOVERY_RECOVERED); + return Ok(()); + } Ok(false) => tracing::warn!( retry_ms = backoff.as_millis(), "No reachable EPP peer dump; retrying KV-index recovery" diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 079b5d085e11..57420f306688 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -290,6 +290,7 @@ impl Selector { // image-only upgrade from a deployment that predates the // dump endpoint). Degrade to no-recovery instead of failing // startup: the replica bootstraps empty and serves. + crate::metrics::set_kv_recovery_state(crate::metrics::KV_RECOVERY_DISABLED); tracing::warn!( service = %replication.service_name, "selection-http port not found; peer KV-index recovery disabled \ From c6615c4cacc12a5940e0d031aa6388555a2d5e32 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Thu, 20 Aug 2026 22:15:30 +0800 Subject: [PATCH 11/17] docs(kv-router): note the two dump protocols coexist intentionally Signed-off-by: Peter Pan --- lib/kv-router/src/services/indexer/recovery.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index e0255c51e27d..129c68538521 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -1,6 +1,22 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +//! Peer KV-index recovery over each peer's `/dump` endpoint. +//! +//! Two dump protocols coexist on purpose: +//! +//! - [`recover_from_peers`] consumes the legacy single-JSON-object format. It is +//! shared by standalone indexer / selection / python P2P recovery, which talk +//! to those processes' own `/dump` endpoints — pre-existing and untouched. +//! - [`recover_from_peers_streaming`] consumes the NDJSON format served by the +//! EPP's peer `/dump` (one `StreamDumpRecord` per line, applied as it +//! arrives). This is the EPP peer-recovery protocol introduced alongside it. +//! +//! Both are intentionally kept: changing the shared JSON path would ripple into +//! the standalone deployments, and the EPP path needs streaming to avoid +//! buffering a whole snapshot. Keep idempotency and retry semantics in sync +//! when touching either; a future convergence is possible but not required. + use std::collections::HashMap; use std::time::Duration; From 757a97682ae4264509da92f7938a377d7d85dc67 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Fri, 21 Aug 2026 13:56:13 +0800 Subject: [PATCH 12/17] refactor(kv-router): drop the NDJSON streaming dump protocol Remove the second core /dump wire contract introduced for EPP peer recovery (StreamDumpRecord, dump_registry_records, recover_from_peers_streaming) and its selection-service surface (indexer_stream_records, recover_indexer_from_peers_streaming). The EPP reuses the established single-JSON path (recover_from_peers + dump_registry) instead, so there is exactly one recovery mechanism to maintain. Streaming was added to bound memory on large snapshots, but the typical snapshot is MB-scale (not the projected hundreds of MB), and the reviewer prefers one wire contract over a second one with a single consumer. Signed-off-by: Peter Pan --- .../src/services/indexer/recovery.rs | 222 +----------------- lib/kv-router/src/services/indexer/server.rs | 55 ----- .../src/services/selection/core/mod.rs | 13 - .../src/services/selection/service.rs | 19 -- 4 files changed, 6 insertions(+), 303 deletions(-) diff --git a/lib/kv-router/src/services/indexer/recovery.rs b/lib/kv-router/src/services/indexer/recovery.rs index 129c68538521..25ed56af018e 100644 --- a/lib/kv-router/src/services/indexer/recovery.rs +++ b/lib/kv-router/src/services/indexer/recovery.rs @@ -3,19 +3,12 @@ //! Peer KV-index recovery over each peer's `/dump` endpoint. //! -//! Two dump protocols coexist on purpose: -//! -//! - [`recover_from_peers`] consumes the legacy single-JSON-object format. It is -//! shared by standalone indexer / selection / python P2P recovery, which talk -//! to those processes' own `/dump` endpoints — pre-existing and untouched. -//! - [`recover_from_peers_streaming`] consumes the NDJSON format served by the -//! EPP's peer `/dump` (one `StreamDumpRecord` per line, applied as it -//! arrives). This is the EPP peer-recovery protocol introduced alongside it. -//! -//! Both are intentionally kept: changing the shared JSON path would ripple into -//! the standalone deployments, and the EPP path needs streaming to avoid -//! buffering a whole snapshot. Keep idempotency and retry semantics in sync -//! when touching either; a future convergence is possible but not required. +//! [`recover_from_peers`] consumes the single-JSON-object format served by the +//! standalone indexer / selection / python P2P recovery and by the EPP's peer +//! `/dump`. The EPP reuses this shared path rather than introducing a second +//! wire contract: the receiver buffers the body with a bounded read +//! ([`read_dump_body`]) and applies each event idempotently (the radix index +//! keys blocks by `tokens_hash`). use std::collections::HashMap; use std::time::Duration; @@ -27,7 +20,6 @@ use crate::identity::RoutingPartitionId; use crate::protocols::RouterEvent; use super::registry::WorkerRegistry; -use super::server::StreamDumpRecord; /// Timeout for one peer `/dump` fetch (HTTP request + body transfer). A full /// KV-index snapshot can be tens of MB on a busy deployment; 10s only fits @@ -87,119 +79,6 @@ pub async fn recover_from_peers(peers: &[String], registry: &WorkerRegistry) -> Ok(false) } -/// Streaming peer recovery over an NDJSON `/dump` (one [`StreamDumpRecord`] per -/// line): each event is applied as it arrives, so the whole snapshot is never -/// buffered. No `max_bytes` budget is needed — peak memory is the index plus one -/// line, not the index plus the full JSON body. This is the EPP peer-recovery -/// protocol; the JSON [`recover_from_peers`] path stays for standalone -/// indexer/selection/python peers. -pub async fn recover_from_peers_streaming( - peers: &[String], - registry: &WorkerRegistry, -) -> Result { - let timeout = Duration::from_millis(env_u64( - RECOVERY_HTTP_TIMEOUT_ENV, - DEFAULT_RECOVERY_HTTP_TIMEOUT_MS, - )); - let client = reqwest::Client::builder() - .timeout(timeout) - .build() - .context("failed to build HTTP client")?; - - tokio::time::sleep(Duration::from_secs(1)).await; - - for peer_url in peers { - match try_recover_from_peer_streaming(&client, peer_url, registry).await { - Ok(()) => { - tracing::info!(peer = %peer_url, "streaming recovery from peer succeeded"); - return Ok(true); - } - Err(e) => { - tracing::warn!(peer = %peer_url, error = %e, "streaming recovery from peer failed, trying next"); - } - } - } - - Ok(false) -} - -async fn try_recover_from_peer_streaming( - client: &reqwest::Client, - peer_url: &str, - registry: &WorkerRegistry, -) -> Result<()> { - let dump_url = format!("{peer_url}/dump"); - tracing::info!(url = %dump_url, "streaming dump from peer"); - - let mut resp = client - .get(&dump_url) - .send() - .await - .context("HTTP request failed")?; - - if !resp.status().is_success() { - anyhow::bail!("peer returned status {}", resp.status()); - } - - let mut total_events = 0usize; - let mut line = String::new(); - // Read the body chunk-by-chunk and split on newlines. A chunk may contain - // several records or only part of one; `line` carries the partial record - // across chunks. Each record is applied and dropped immediately. - while let Some(chunk) = resp - .chunk() - .await - .context("failed to read dump response body")? - { - line.push_str(&String::from_utf8_lossy(&chunk)); - while let Some(pos) = line.find('\n') { - let record_line = line[..pos].to_string(); - line.drain(..=pos); - if record_line.is_empty() { - continue; - } - let record: StreamDumpRecord = serde_json::from_str(&record_line) - .with_context(|| format!("failed to parse dump record: {record_line}"))?; - apply_dump_record(registry, record).await?; - total_events += 1; - } - } - // A well-formed stream ends with a trailing newline; a leftover partial - // line means the stream was truncated mid-record. - if !line.is_empty() { - anyhow::bail!("peer dump stream ended mid-record"); - } - - // An empty dump is a valid recovery. Recovery candidates are restricted to - // already-serving peers (see `recovery_peer_urls`), so a zero-event dump - // means the serving peer genuinely holds no KV index yet (idle cluster), - // not a transient race. Rejecting it would deadlock cold starts and idle - // rollouts. - tracing::info!(total_events, "applied streamed dump events from peer"); - Ok(()) -} - -async fn apply_dump_record(registry: &WorkerRegistry, record: StreamDumpRecord) -> Result<()> { - let (model_name, routing_group) = record - .key - .split_once(':') - .ok_or_else(|| anyhow::anyhow!("invalid dump key format: {}", record.key))?; - - let key = RoutingPartitionId::new(model_name, routing_group); - let indexer = registry.get_or_create_indexer(key, record.block_size); - - // Re-application is idempotent by construction: the radix index keys blocks - // by `tokens_hash`, so applying the same (or an overlapping) peer dump again - // is a no-op for blocks already present. A cancelled-then-retried attempt - // can therefore leave a partially applied snapshot that the next attempt - // safely re-applies on top of. - indexer - .apply_event_routed(record.event) - .await - .context("peer recovery event was rejected by the local indexer")?; - Ok(()) -} - async fn try_recover_from_peer( client: &reqwest::Client, peer_url: &str, @@ -419,93 +298,4 @@ mod tests { .expect("zero disables the cap"); assert_eq!(body.len(), 512); } - - /// Serve a chunked HTTP response with the given Content-Type and body (no - /// Content-Length header), for exercising the streaming dump reader. - async fn serve_ndjson(body: &[u8]) -> String { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind test server"); - let addr = listener.local_addr().expect("read test addr"); - let body = body.to_vec(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accept test client"); - let mut buf = [0u8; 4096]; - let _ = socket.read(&mut buf).await.expect("read request"); - let mut resp = String::from( - "HTTP/1.1 200 OK\r\nContent-Type: application/x-ndjson\r\nTransfer-Encoding: chunked\r\n\r\n", - ); - for chunk in body.chunks(7) { - resp.push_str(&format!("{:x}\r\n", chunk.len())); - resp.push_str(&String::from_utf8_lossy(chunk)); - resp.push_str("\r\n"); - } - resp.push_str("0\r\n\r\n"); - socket - .write_all(resp.as_bytes()) - .await - .expect("write response"); - }); - let url = format!("http://{addr}/dump"); - drop(server); - url - } - - #[tokio::test] - async fn streaming_recovery_applies_empty_dump() { - // A streaming dump with zero records (just a trailing newline) is a - // valid recovery, matching the non-streaming "empty dump is success". - let url = serve_ndjson(b"\n").await; - let client = reqwest::Client::new(); - let registry = WorkerRegistry::new(1); - try_recover_from_peer_streaming(&client, &url, ®istry) - .await - .expect("empty streamed dump must succeed"); - } - - #[tokio::test] - async fn streaming_recovery_rejects_truncated_record() { - // A stream that ends mid-record (no trailing newline) must fail instead - // of silently applying a partial event. - let url = serve_ndjson(br#"{"key":"m:default","block_size":16,"event":{"#).await; - let client = reqwest::Client::new(); - let registry = WorkerRegistry::new(1); - let err = try_recover_from_peer_streaming(&client, &url, ®istry) - .await - .expect_err("truncated stream must fail"); - assert!( - err.to_string().contains("mid-record"), - "unexpected error: {err}" - ); - } - - #[tokio::test] - async fn streaming_recovery_applies_multiple_records() { - // A well-formed NDJSON stream applies each record line-by-line. Empty - // blocks make each event a no-op but must parse and apply without error. - let record = |worker_id: u64| { - serde_json::json!({ - "key": "m:default", - "block_size": 16, - "event": { - "worker_id": worker_id, - "storage_tier": "device", - "event": {"event_id": 0, "data": {"stored": {"parent_hash": null, "blocks": []}}}, - } - }) - }; - let mut body = serde_json::to_vec(&record(1)).unwrap(); - body.push(b'\n'); - body.extend(serde_json::to_vec(&record(2)).unwrap()); - body.push(b'\n'); - - let url = serve_ndjson(&body).await; - let client = reqwest::Client::new(); - let registry = WorkerRegistry::new(1); - try_recover_from_peer_streaming(&client, &url, ®istry) - .await - .expect("well-formed stream must apply"); - } } diff --git a/lib/kv-router/src/services/indexer/server.rs b/lib/kv-router/src/services/indexer/server.rs index a9a5e0f079b3..b07ee2f07405 100644 --- a/lib/kv-router/src/services/indexer/server.rs +++ b/lib/kv-router/src/services/indexer/server.rs @@ -517,61 +517,6 @@ pub(crate) async fn dump_registry(registry: &WorkerRegistry) -> serde_json::Valu serde_json::json!(result) } -/// One NDJSON line of a streaming KV-index dump: a single `RouterEvent` with the -/// routing partition and block size it belongs to. The streaming dump emits one -/// of these per line so the receiver can apply each event and drop it instead of -/// buffering the whole snapshot. `key` is `"model:routing_group"`. -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct StreamDumpRecord { - pub key: String, - pub block_size: u32, - pub event: crate::protocols::RouterEvent, -} - -/// Collect every indexer's events as structured records for a streaming dump. -/// -/// Structured counterpart to [`dump_registry`]: the same snapshot, but one -/// [`StreamDumpRecord`] per event so a caller can serialize and send records one -/// at a time instead of building one JSON object. A failed indexer dump aborts -/// the whole stream with an error. -pub(crate) async fn dump_registry_records( - registry: &WorkerRegistry, -) -> Result, String> { - let all = registry.all_indexers_with_block_size(); - let mut handles = Vec::with_capacity(all.len()); - - for (key, indexer, block_size) in all { - handles.push(tokio::spawn(async move { - let events = indexer.dump_events().await; - (key, events, block_size) - })); - } - - let mut records = Vec::new(); - for handle in handles { - match handle.await { - Ok((key, Ok(events), block_size)) => { - let map_key = format!("{}:{}", key.model_name, key.routing_group); - for event in events { - records.push(StreamDumpRecord { - key: map_key.clone(), - block_size, - event, - }); - } - } - Ok((key, Err(e), _)) => { - let map_key = format!("{}:{}", key.model_name, key.routing_group); - return Err(format!("indexer {map_key} dump failed: {e}")); - } - Err(e) => { - return Err(format!("dump task join error: {e}")); - } - } - } - Ok(records) -} - async fn handle_health() -> StatusCode { StatusCode::OK } diff --git a/lib/kv-router/src/services/selection/core/mod.rs b/lib/kv-router/src/services/selection/core/mod.rs index 0357412087e3..134ac45579b3 100644 --- a/lib/kv-router/src/services/selection/core/mod.rs +++ b/lib/kv-router/src/services/selection/core/mod.rs @@ -273,13 +273,6 @@ impl SelectionCore { recovery::recover_from_peers(peers, &self.indexer_registry).await } - pub(crate) async fn recover_indexer_from_peers_streaming( - &self, - peers: &[String], - ) -> anyhow::Result { - recovery::recover_from_peers_streaming(peers, &self.indexer_registry).await - } - pub(crate) fn signal_indexer_ready(&self) { self.indexer_registry.signal_ready(); } @@ -288,12 +281,6 @@ impl SelectionCore { crate::services::indexer::server::dump_registry(&self.indexer_registry).await } - pub(crate) async fn dump_indexer_records( - &self, - ) -> Result, String> { - crate::services::indexer::server::dump_registry_records(&self.indexer_registry).await - } - pub(crate) fn dispatch_replica_event(&self, envelope: ScopedReplicaEvent) { let (key, block_size, event) = envelope.into_parts(); if self diff --git a/lib/kv-router/src/services/selection/service.rs b/lib/kv-router/src/services/selection/service.rs index 13fa5078c46e..37a5afcdd453 100644 --- a/lib/kv-router/src/services/selection/service.rs +++ b/lib/kv-router/src/services/selection/service.rs @@ -382,29 +382,10 @@ impl SelectionService { self.core.dump_indexer_events().await } - /// Structured snapshot of the local KV index for a streaming peer dump: one - /// record per event. The caller serializes records one at a time instead of - /// building a single JSON object in memory. - pub async fn indexer_stream_records( - &self, - ) -> Result, String> { - self.core.dump_indexer_records().await - } - pub async fn recover_indexer_from_peers(&self, peers: &[String]) -> anyhow::Result { self.core.recover_indexer_from_peers(peers).await } - /// EPP peer recovery over the streaming NDJSON `/dump` format (one event per - /// line, applied as it arrives). Distinct from [`Self::recover_indexer_from_peers`], - /// which uses the single-JSON format for standalone indexer/selection peers. - pub async fn recover_indexer_from_peers_streaming( - &self, - peers: &[String], - ) -> anyhow::Result { - self.core.recover_indexer_from_peers_streaming(peers).await - } - pub async fn cancelled(&self) { self.cancel_token.cancelled().await; } From 4b9ad7baf56f4d47a5c765f4ce9a9e6bd3af2c33 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Fri, 21 Aug 2026 13:56:20 +0800 Subject: [PATCH 13/17] refactor(epp): reuse the JSON peer dump; subscribe to workers before recovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - /dump serves the single-JSON snapshot (dump_registry shape) instead of NDJSON; the recovery consumer uses the established recover_from_peers. - Subscribe-first ordering: peer KV-index recovery is deferred out of selector construction and runs (start_peer_recovery) after the topology adapter registers workers, so the ZMQ KV-event listeners are already buffering live events when the dump is transferred — the snapshot and the live stream overlap idempotently instead of leaving a gap (the correctness gap called out in review: the EPP dumped before subscribing, and the dump carried no per-worker cursor). - Drop the 'empty worker catalog before recovery' guard: workers may (and now must) be registered before recovery. - is_ready already ANDs the recovered flag, so the EPP stays NOT_READY and /dump stays 503 until recovery completes. Signed-off-by: Peter Pan --- .../ext-proc/src/epp_router.rs | 30 ++-- .../ext-proc/src/peer_http.rs | 82 ++++----- .../ext-proc/src/selector.rs | 164 ++++++++++++------ 3 files changed, 156 insertions(+), 120 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/epp_router.rs b/deploy/inference-gateway/ext-proc/src/epp_router.rs index 4130bd0453f6..7d43b36879fb 100644 --- a/deploy/inference-gateway/ext-proc/src/epp_router.rs +++ b/deploy/inference-gateway/ext-proc/src/epp_router.rs @@ -73,13 +73,7 @@ impl EppRouter { ) -> Result { let selector = Arc::new(selector); let (renderer, reflector, reflector_ready) = Self::dependencies(&cfg).await?; - Ok(Self::from_selector_parts( - cfg, - renderer, - reflector, - reflector_ready, - selector, - )) + Self::from_selector_parts(cfg, renderer, reflector, reflector_ready, selector).await } /// Assemble a custom EPP image around a prebuilt selection service. @@ -89,13 +83,7 @@ impl EppRouter { ) -> Result { let selector = Arc::new(Selector::from_service(&cfg, service).await?); let (renderer, reflector, reflector_ready) = Self::dependencies(&cfg).await?; - Ok(Self::from_selector_parts( - cfg, - renderer, - reflector, - reflector_ready, - selector, - )) + Self::from_selector_parts(cfg, renderer, reflector, reflector_ready, selector).await } async fn dependencies( @@ -113,23 +101,29 @@ impl EppRouter { Ok((renderer, reflector, reflector_ready)) } - fn from_selector_parts( + async fn from_selector_parts( cfg: EppStandaloneConfig, renderer: VllmRenderClient, reflector: Arc, reflector_ready: Arc, selector: Arc, - ) -> Self { + ) -> Result { let peer_ready = selector.peer_ready(); let defaults = RegistrationDefaults::from_config(&cfg); let adapter = TopologyAdapter::spawn(reflector.as_ref().clone(), selector.clone(), defaults); + // Subscribe-first: the topology adapter registers workers (their ZMQ + // KV-event listeners begin buffering live events) before the peer dump, + // so recovery overlaps the live event stream instead of leaving a gap. + // `is_ready` stays false (and /dump stays 503) until recovery completes. + selector.start_peer_recovery().await?; + // Readiness is driven solely by the live pod+pool signal (see `is_ready`); // we do not block startup on a schedulable worker. A valid, empty pool is // ready immediately and returns 503 per-request until capacity appears. - Self { + Ok(Self { renderer, reflector, selector, @@ -138,7 +132,7 @@ impl EppRouter { peer_ready, model_name: cfg.model_name, inflight: Arc::new(Semaphore::new(cfg.max_inflight_requests)), - } + }) } /// Overall EPP readiness: worker discovery is ready and replicated mode has diff --git a/deploy/inference-gateway/ext-proc/src/peer_http.rs b/deploy/inference-gateway/ext-proc/src/peer_http.rs index 608cd54cec52..ae3e5a52d70f 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_http.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_http.rs @@ -9,23 +9,17 @@ use std::sync::atomic::{AtomicBool, Ordering}; use anyhow::{Context, Result}; use axum::{ - Router, - body::Body, + Json, Router, extract::State, - http::{StatusCode, header::CONTENT_TYPE}, + http::StatusCode, response::{IntoResponse, Response}, routing::get, }; -use futures::StreamExt; use tokio::net::TcpListener; use tokio_util::sync::CancellationToken; -use dynamo_kv_router::services::indexer::server::StreamDumpRecord; use dynamo_kv_router::services::selection::SelectionService; -/// Media type of the streaming NDJSON dump (one [`StreamDumpRecord`] per line). -const STREAM_DUMP_MEDIA_TYPE: &str = "application/x-ndjson"; - #[derive(Clone)] struct AppState { service: Arc, @@ -77,41 +71,25 @@ async fn dump(State(state): State) -> Response { .into_response(); } - let records = match state.service.indexer_stream_records().await { - Ok(records) => records, - Err(error) => { - tracing::warn!(%error, "Failed to collect peer KV-index dump records"); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - "peer KV-index snapshot generation failed", - ) - .into_response(); - } - }; - - // Stream NDJSON, one record per line. The receiver applies each event and - // drops it, so the whole snapshot is never buffered on either side and no - // max-bytes budget is needed. - let stream = futures::stream::iter(records).map(|record: StreamDumpRecord| { - match serde_json::to_vec(&record) { - Ok(mut bytes) => { - bytes.push(b'\n'); - Ok::<_, std::convert::Infallible>(bytes) - } - Err(error) => { - tracing::warn!(%error, "Failed to serialize dump record"); - // Signal a mid-stream failure with an empty frame; the receiver - // treats a truncated record as an error and retries. - Ok::<_, std::convert::Infallible>(Vec::new()) - } - } - }); + let snapshot = state.service.indexer_snapshot().await; + // `dump_registry` embeds per-model `{"error": ...}` entries when an indexer + // dump fails; surface those as a whole-snapshot 500 so the recovery consumer + // does not deserialize an incompatible entry (and fall back to empty state). + let failed = snapshot + .as_object() + .into_iter() + .flatten() + .any(|(_key, entry)| entry.get("error").is_some()); + if failed { + tracing::warn!("peer KV-index snapshot contains a failed indexer dump"); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + "peer KV-index snapshot generation failed", + ) + .into_response(); + } - Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, STREAM_DUMP_MEDIA_TYPE) - .body(Body::from_stream(stream)) - .expect("static response builder") + (StatusCode::OK, Json(snapshot)).into_response() } #[cfg(test)] @@ -162,7 +140,7 @@ mod tests { } #[tokio::test] - async fn streaming_dump_emits_ndjson_records() { + async fn dump_returns_json_snapshot() { let service = service().await; let cancel = CancellationToken::new(); let port = free_tcp_port(); @@ -184,17 +162,17 @@ mod tests { resp.headers() .get(reqwest::header::CONTENT_TYPE) .and_then(|v| v.to_str().ok()), - Some("application/x-ndjson") + Some("application/json") ); let body = resp.text().await.expect("read body"); - // An empty index yields an empty stream; every non-empty line must be a - // parseable StreamDumpRecord. - for line in body.lines() { - if line.trim().is_empty() { - continue; - } - serde_json::from_str::(line).expect("valid NDJSON record"); - } + // An empty index yields an empty snapshot object; it must parse as the + // `HashMap` shape `recover_from_peers` consumes. + let snapshot: serde_json::Value = + serde_json::from_str(&body).expect("dump body must be valid JSON"); + assert!( + snapshot.is_object(), + "snapshot must be a JSON object, got: {snapshot}" + ); cancel.cancel(); service.shutdown().await; diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 57420f306688..c92d4dcea5aa 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -8,6 +8,7 @@ //! [`crate::peer_discovery`]. use std::collections::{HashMap, HashSet}; +use std::net::IpAddr; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -85,6 +86,8 @@ pub struct Selector { /// Replication-bootstrap readiness: initial peer discovery plus KV-index /// recovery, or authoritative no-peer bootstrap. Latched once initialized. peer_ready: Option>, + /// Deferred peer KV-index recovery parameters (see [`PeerRecovery`]). + peer_recovery: Option, } /// Local bookkeeping for desired-state reconciliation. @@ -104,6 +107,21 @@ struct ReplicationConfig { ports: PeerPorts, } +/// Deferred peer KV-index recovery. Parameters are resolved at selector +/// construction, but the recovery itself is deliberately NOT run there: +/// workers must register (and their ZMQ KV-event listeners subscribe) first, +/// so the peer dump merges with already-buffered live events instead of +/// leaving a gap. The EPP router starts worker registration, then calls +/// [`Selector::start_peer_recovery`]. +struct PeerRecovery { + namespace: String, + service_name: String, + ports: PeerPorts, + self_ip: IpAddr, + /// Shared with the `/dump` endpoint: 503 until recovery/bootstrap done. + recovered: Arc, +} + struct StartupCancellation { cancel: CancellationToken, armed: bool, @@ -225,12 +243,10 @@ impl Selector { "DYN_EPP_PEER_SERVICE requires a prebuilt SelectionService with replica sync enabled" ) })?; - if !service.list_workers(None, None).is_empty() { - anyhow::bail!( - "replicated prebuilt SelectionService must have an empty worker catalog before \ - KV-index recovery" - ); - } + // Workers may already be registered here: recovery runs only after + // the topology adapter starts (subscribe-first, see + // `start_peer_recovery`), so an empty catalog is no longer a + // precondition. Some(replica_sync_port) } else { None @@ -268,6 +284,7 @@ impl Selector { let cancel = CancellationToken::new(); let mut startup = StartupCancellation::new(cancel.clone()); + let mut peer_recovery = None; let peer_ready = if let Some(replication) = replication { // In replicated mode, we need to exclude ourselves from the peer set which requires the POD_IP let self_ip = std::env::var("POD_IP") @@ -313,17 +330,18 @@ impl Selector { recovered.clone(), ) .await?; - crate::peer_discovery::spawn( - service.clone(), - &cfg.namespace, - &replication.service_name, - replication.ports.replica_sync, - selection_http_port, - self_ip.to_string(), - cancel.clone(), - recovered.clone(), - ) - .await?; + // KV-index recovery itself is deferred: worker registration + // (and with it the ZMQ KV-event subscription) must start + // first so the dump overlaps the live event stream instead + // of leaving a gap. `start_peer_recovery` runs it once the + // EPP router has started the topology adapter. + peer_recovery = Some(PeerRecovery { + namespace: cfg.namespace.clone(), + service_name: replication.service_name, + ports: replication.ports, + self_ip, + recovered: recovered.clone(), + }); Some(recovered) } } @@ -343,9 +361,46 @@ impl Selector { cancel, reconcile_state: Mutex::new(ReconcileState::default()), peer_ready, + peer_recovery, }) } + /// Run peer KV-index recovery now that worker registration has started. + /// + /// Subscribe-first ordering: the topology adapter registers workers (their + /// ZMQ KV-event listeners begin buffering) before this runs, so the peer + /// dump covers past history and the buffered events cover everything after + /// it — only an overlap remains, absorbed idempotently. Blocks until + /// recovery succeeds or no peer exists (empty bootstrap). No-op when + /// replication is disabled or the Service lacks `selection-http`. + pub(crate) async fn start_peer_recovery(&self) -> Result<()> { + let Some(recovery) = &self.peer_recovery else { + return Ok(()); + }; + let selection_http_port = recovery + .ports + .selection_http + .expect("guarded by construction"); + + // Give the topology adapter a moment to register workers before the + // dump: the snapshot then overlaps the live event stream rather than + // racing ahead of it. Bounded — a slow cold start must not hang the + // replica forever; the peer dump is still the best available state. + wait_for_registered_worker(&self.service, &self.cancel).await?; + + crate::peer_discovery::spawn( + self.service.clone(), + &recovery.namespace, + &recovery.service_name, + recovery.ports.replica_sync, + selection_http_port, + recovery.self_ip.to_string(), + self.cancel.clone(), + recovery.recovered.clone(), + ) + .await + } + pub fn peer_ready(&self) -> Option> { self.peer_ready.clone() } @@ -505,6 +560,48 @@ impl Selector { } } +/// Bounded wait for at least one registered worker before the peer dump, so +/// the snapshot overlaps the already-buffered live event stream (subscribe- +/// first). A cold start that never registers a worker in time proceeds anyway: +/// the peer dump is then the best available state. +const PEER_RECOVERY_WORKER_WAIT: std::time::Duration = std::time::Duration::from_secs(15); + +async fn wait_for_registered_worker( + service: &SelectionService, + cancel: &CancellationToken, +) -> Result<()> { + if !service.list_workers(None, None).is_empty() { + return Ok(()); + } + tracing::info!( + "Waiting up to {PEER_RECOVERY_WORKER_WAIT:?} for a registered worker before peer \ + KV-index recovery (subscribe-first)" + ); + let deadline = tokio::time::Instant::now() + PEER_RECOVERY_WORKER_WAIT; + loop { + tokio::select! { + biased; + _ = cancel.cancelled() => { + anyhow::bail!( + "EPP startup cancelled while waiting for workers before peer recovery" + ) + } + _ = tokio::time::sleep_until(deadline) => { + tracing::warn!( + "No worker registered within {PEER_RECOVERY_WORKER_WAIT:?}; proceeding with \ + peer KV-index recovery anyway (best effort)" + ); + return Ok(()); + } + _ = tokio::time::sleep(std::time::Duration::from_millis(250)) => { + if !service.list_workers(None, None).is_empty() { + return Ok(()); + } + } + } + } +} + impl Drop for Selector { fn drop(&mut self) { // Stop the peer-discovery watch; the service's own Drop stops the core, @@ -575,14 +672,6 @@ models: } } - fn free_tcp_port() -> u16 { - std::net::TcpListener::bind("127.0.0.1:0") - .expect("reserve test port") - .local_addr() - .expect("read test port") - .port() - } - /// Minimal single-replica config (no peer service, so no cluster access). /// `max_num_batched_tokens` is set so `Selector::new` never fails its /// fast-fail check regardless of the ambient router policy. @@ -1014,31 +1103,6 @@ models: ); } - #[tokio::test] - async fn prebuilt_service_rejects_existing_workers_before_recovery() { - let service = SelectionServiceBuilder::new(KvRouterConfig::default()) - .indexer_threads(1) - .replica_sync(free_tcp_port(), Vec::new()) - .build() - .await - .expect("replica-sync selection service should build"); - service - .upsert_worker(Selector::worker_request(&incomplete_registration(1))) - .await - .expect("incomplete worker should still enter the catalog"); - - let mut cfg = test_config(); - cfg.peer_service = Some("does-not-exist".to_string()); - let error = Selector::from_service(&cfg, service) - .await - .err() - .expect("prebuilt service with workers must be rejected"); - assert!( - error.to_string().contains("empty worker catalog"), - "{error}" - ); - } - #[tokio::test] async fn duplicate_worker_ids_are_rejected_before_reconciliation() { let selector = Selector::new(&test_config()) From f5234caf718cd449fd51a665f43e5760cae361b1 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Fri, 21 Aug 2026 13:56:26 +0800 Subject: [PATCH 14/17] fix(epp): keep active peer recovery through EndpointSlice churn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The restart guard compared shuffled peer vectors, and recovery_peer_urls re-randomizes on every call, so any >1-peer churn looked like a membership change and dropped the in-flight dump. Split the concerns: - recovery_peer_set: deterministic BTreeSet for change detection (BTreeSet ordering is stable, so equality detects only real membership changes). - shuffled_peer_urls: one attempt-scoped random priority per cycle; shuffling lives in the attempt path, never in a comparison. The loop now tries one peer per attempt with a pending/tried model: churn reconciles only the pending candidates (drop unattempted peers that are no longer eligible, add newly eligible ones) and never cancels an active request — it may complete even after its source leaves the slice. A fresh randomized cycle starts only after the eligible set is exhausted/backed off. Tests: unchanged churn, a join during recovery, removal of an unattempted peer, and removal of the active peer without cancelling its request. Signed-off-by: Peter Pan --- .../ext-proc/src/peer_discovery.rs | 490 ++++++++++++------ 1 file changed, 344 insertions(+), 146 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs index 84957557cda0..a44ae0ebd743 100644 --- a/deploy/inference-gateway/ext-proc/src/peer_discovery.rs +++ b/deploy/inference-gateway/ext-proc/src/peer_discovery.rs @@ -329,7 +329,7 @@ async fn recover_initial_index( cancel, initial_backoff, max_backoff, - |service, peers| Box::pin(service.recover_indexer_from_peers_streaming(peers)), + |service, peers| Box::pin(service.recover_indexer_from_peers(peers)), ) .await } @@ -353,10 +353,14 @@ where { let mut backoff = initial_backoff; - 'attempt: loop { + loop { reconcile_once(service, store, sync_port, self_ip, known).await; - let peers = recovery_peer_urls(store, self_ip, selection_http_port); - if peers.is_empty() { + + // Deterministic eligible set for change detection; the shuffled order is + // derived once per cycle (attempt-scoped priority), never compared, so + // an unrelated EndpointSlice update cannot look like a membership change. + let eligible = recovery_peer_set(store, self_ip, selection_http_port); + if eligible.is_empty() { crate::metrics::set_kv_recovery_state(crate::metrics::KV_RECOVERY_EMPTY_BOOTSTRAP); tracing::warn!( "No serving sibling EPP peer found; bootstrapping an EMPTY KV index \ @@ -366,49 +370,65 @@ where return Ok(()); } - let attempt = recover(service, &peers); - tokio::pin!(attempt); - - // Await the attempt, restarting only when the candidate set actually - // changes (or becomes empty). Unrelated EndpointSlice churn (readiness - // flips, metadata/zone updates) must not discard in-flight progress: a - // large dump under churn would otherwise be dropped repeatedly, leaving - // a partially applied snapshot that the next attempt re-applies. - let result = loop { - tokio::select! { - biased; - _ = cancel.cancelled() => { - anyhow::bail!("EPP peer discovery cancelled during KV-index recovery") - } - changed = changes_rx.changed() => { - changed.context("EPP peer EndpointSlice watch ended during KV-index recovery")?; - if recovery_peer_urls(store, self_ip, selection_http_port) != peers { - // Candidate set changed: restart with the new set. - backoff = initial_backoff; - continue 'attempt; + let mut priority = shuffled_peer_urls(&eligible); + let mut tried: BTreeSet = BTreeSet::new(); + + // Try one peer per attempt. EndpointSlice churn never cancels an + // in-flight request — the dump may still complete after its source + // leaves the slice — so churn only reconciles the *pending* candidates: + // drop unattempted peers that are no longer eligible, add newly eligible + // ones. Peers that failed stay in `tried` until the cycle is exhausted. + 'attempt: loop { + let Some(peer) = priority.iter().find(|p| !tried.contains(*p)).cloned() else { + break 'attempt; + }; + + let peers = [peer.clone()]; + let attempt = recover(service, &peers); + tokio::pin!(attempt); + + let result = loop { + tokio::select! { + biased; + _ = cancel.cancelled() => { + anyhow::bail!("EPP peer discovery cancelled during KV-index recovery") } - // Unrelated churn: keep awaiting the same attempt. + changed = changes_rx.changed() => { + changed.context("EPP peer EndpointSlice watch ended during KV-index recovery")?; + let current = recovery_peer_set(store, self_ip, selection_http_port); + // Keep the active request (even if its own source left + // the slice); reconcile only the unattempted pending set. + priority.retain(|p| tried.contains(p) || current.contains(p)); + for newly in current.difference(&tried) { + if !priority.contains(newly) { + priority.push(newly.clone()); + } + } + } + result = &mut attempt => break result, } - result = &mut attempt => break result, - } - }; + }; - match result { - Ok(true) => { - crate::metrics::set_kv_recovery_state(crate::metrics::KV_RECOVERY_RECOVERED); - return Ok(()); + match result { + Ok(true) => { + crate::metrics::set_kv_recovery_state(crate::metrics::KV_RECOVERY_RECOVERED); + return Ok(()); + } + Ok(false) => tracing::warn!( + peer = %peer, + "No reachable EPP peer dump; trying next recovery candidate" + ), + Err(error) => tracing::warn!( + peer = %peer, + %error, + "EPP peer KV-index recovery failed; trying next recovery candidate" + ), } - Ok(false) => tracing::warn!( - retry_ms = backoff.as_millis(), - "No reachable EPP peer dump; retrying KV-index recovery" - ), - Err(error) => tracing::warn!( - %error, - retry_ms = backoff.as_millis(), - "EPP peer KV-index recovery failed; retrying" - ), + tried.insert(peer); } + // The current eligible set is exhausted (or emptied by churn): back off, + // then start a fresh cycle with a fresh shuffle. tokio::select! { _ = cancel.cancelled() => { anyhow::bail!("EPP peer discovery cancelled during KV-index recovery") @@ -420,7 +440,7 @@ where // metadata/zone updates) must not reset the backoff either: // under a rolling update it would keep the retry loop hot at // the initial backoff instead of letting it grow. - if recovery_peer_urls(store, self_ip, selection_http_port) != peers { + if recovery_peer_set(store, self_ip, selection_http_port) != eligible { backoff = initial_backoff; } } @@ -476,7 +496,11 @@ fn live_peer_ips(store: &Store, self_ip: &str) -> BTreeSet { /// recovery candidates turns a cold start (all replicas empty, none serving) /// into a mutual-recovery deadlock. An empty candidate set means "no eligible /// peer" and bootstraps an empty index immediately. -fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { +/// Deterministic eligible peer URL set. Used for change detection and as the +/// input to an attempt-scoped shuffle. Unlike a shuffled vector, the BTreeSet +/// order is stable, so equality across calls detects only real membership +/// changes — never a re-randomization. +fn recovery_peer_set(store: &Store, self_ip: &str, port: u16) -> BTreeSet { let want_ipv6 = is_ipv6(self_ip); let mut peers = BTreeSet::new(); @@ -490,21 +514,23 @@ fn recovery_peer_urls(store: &Store, self_ip: &str, port: u16) -> Vec { } for address in &endpoint.addresses { if !address.is_empty() && address != self_ip { - peers.insert(address.clone()); + peers.insert(format!("http://{}", authority(address, port))); } } } } + peers +} - // Shuffle so N simultaneously-joining replicas do not all deterministically - // pick the lowest-IP peer (BTreeSet order), concentrating dump work on one - // serving EPP. RandomState is seeded per process, so each bootstrap gets a - // different order; serial fallback through the shuffled list is retained. +/// One attempt-scoped random priority order over the eligible set. Shuffling +/// happens only when a new recovery cycle starts, never inside a comparison: +/// N simultaneously-joining replicas must not all deterministically pick the +/// lowest-IP peer (BTreeSet order), concentrating dump work on one serving EPP. +/// `RandomState` is seeded per process, so each bootstrap gets a different +/// order; serial fallback through the shuffled order is retained. +fn shuffled_peer_urls(eligible: &BTreeSet) -> Vec { let hasher = RandomState::new(); - let mut urls: Vec = peers - .into_iter() - .map(|ip| format!("http://{}", authority(&ip, port))) - .collect(); + let mut urls: Vec = eligible.iter().cloned().collect(); urls.sort_by_cached_key(|url| hasher.hash_one(url)); urls } @@ -708,7 +734,7 @@ mod tests { let slices = [slice_with(&["epp.example.test"], false, "FQDN")]; assert!(peer_ips(slices.iter(), false).is_empty()); assert!( - recovery_peer_urls(&store_from_slices(slices.to_vec()), "10.0.0.9", 9093).is_empty() + recovery_peer_set(&store_from_slices(slices.to_vec()), "10.0.0.9", 9093).is_empty() ); } @@ -1011,8 +1037,8 @@ mod tests { async fn gated_dump(State(gate): State) -> axum::response::Response { gate.requested.notify_one(); gate.release.notified().await; - // Empty body = an empty streaming (NDJSON) dump, a valid recovery. - axum::response::Response::new(axum::body::Body::empty()) + // An empty JSON object = an empty snapshot, a valid recovery. + axum::response::Response::new(axum::body::Body::from("{}")) } #[derive(Clone)] @@ -1026,8 +1052,8 @@ mod tests { state.first_failed.notify_one(); (StatusCode::SERVICE_UNAVAILABLE, "not ready").into_response() } else { - // Empty body = an empty streaming (NDJSON) dump, a valid recovery. - axum::response::Response::new(axum::body::Body::empty()).into_response() + // An empty JSON object = an empty snapshot, a valid recovery. + axum::response::Response::new(axum::body::Body::from("{}")).into_response() } } @@ -1067,7 +1093,7 @@ mod tests { listener, Router::new().route( "/dump", - get(|| async { axum::response::Response::new(axum::body::Body::empty()) }), + get(|| async { axum::response::Response::new(axum::body::Body::from("{}")) }), ), ) .await @@ -1138,102 +1164,271 @@ mod tests { service.shutdown().await; } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn peer_change_cancels_inflight_recovery_and_uses_new_peer() { - use kube::runtime::watcher; + /// Shared driver for the churn-behavior tests: drives + /// `recover_initial_index_with_attempt` with a scripted recover closure. + struct ChurnHarness { + task: tokio::task::JoinHandle>, + cancel: CancellationToken, + service: Arc, + changes_tx: watch::Sender, + writer: kube::runtime::reflector::store::Writer, + attempts: Arc>>>, + first_started: Arc, + release: Arc, + first_dropped: Arc, + } + + impl ChurnHarness { + /// Start recovery with one slice containing `initial_ips`, all + /// eligible. The first attempt is held until `release`. + async fn start(initial_ips: &[&str], port: u16, first_attempt_outcome: bool) -> Self { + use kube::runtime::watcher; - struct DropSignal(Arc); - impl Drop for DropSignal { - fn drop(&mut self) { - self.0.notify_one(); + let mut slices: Vec = initial_ips + .iter() + .map(|ip| recovery_slice(ip, Some(true), Some(true))) + .collect(); + // The reflector keys slices by name; name them by IP so later + // Apply/Delete events for the same peer hit the same object. + for (ip, slice) in initial_ips.iter().zip(slices.iter_mut()) { + slice.metadata.name = Some(format!("peer-{}", ip.replace('.', "-"))); } - } + let mut writer = kube::runtime::reflector::store::Writer::::default(); + let store = writer.as_reader(); + writer.apply_watcher_event(&watcher::Event::Init); + for slice in slices { + writer.apply_watcher_event(&watcher::Event::InitApply(slice)); + } + writer.apply_watcher_event(&watcher::Event::InitDone); - let port = 9093; - let old_ip = "192.0.2.10"; - let new_ip = "192.0.2.11"; - let old_slice = recovery_slice(old_ip, Some(true), Some(true)); - let (store, mut writer) = store_and_writer(vec![old_slice]); - let (changes_tx, changes_rx) = watch::channel(0u64); - let service = recovery_service().await; - let cancel = CancellationToken::new(); + let (changes_tx, changes_rx) = watch::channel(0u64); + let service = recovery_service().await; + let cancel = CancellationToken::new(); - let first_started = Arc::new(Notify::new()); - let first_dropped = Arc::new(Notify::new()); - let attempts = Arc::new(std::sync::Mutex::new(Vec::>::new())); - let attempt_number = Arc::new(AtomicUsize::new(0)); - let task = tokio::spawn({ - let service = service.clone(); - let cancel = cancel.clone(); - let first_started = first_started.clone(); - let first_dropped = first_dropped.clone(); - let attempts = attempts.clone(); - let attempt_number = attempt_number.clone(); - async move { - let mut known = BTreeSet::new(); - let mut changes_rx = changes_rx; - recover_initial_index_with_attempt( - &service, - &store, - 9092, - port, - "192.0.2.99", - &mut known, - &mut changes_rx, - &cancel, - Duration::from_millis(10), - Duration::from_millis(40), - move |_service, peers| { - let peers = peers.to_vec(); - attempts.lock().unwrap().push(peers.clone()); - let number = attempt_number.fetch_add(1, Ordering::SeqCst); - let first_started = first_started.clone(); - let first_dropped = first_dropped.clone(); - Box::pin(async move { - if number == 0 { + let attempts = Arc::new(std::sync::Mutex::new(Vec::>::new())); + let first_started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let first_dropped = Arc::new(Notify::new()); + + struct DropSignal(Arc); + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.notify_one(); + } + } + + let attempt_no = Arc::new(AtomicUsize::new(0)); + let task = { + let service = service.clone(); + let cancel = cancel.clone(); + let attempts = attempts.clone(); + let attempt_no = attempt_no.clone(); + let first_started = first_started.clone(); + let release = release.clone(); + let first_dropped = first_dropped.clone(); + tokio::spawn(async move { + let mut known = BTreeSet::new(); + let mut changes_rx = changes_rx; + recover_initial_index_with_attempt( + &service, + &store, + 9092, + port, + "192.0.2.99", + &mut known, + &mut changes_rx, + &cancel, + Duration::from_millis(10), + Duration::from_millis(40), + move |_service, peers| { + let peers = peers.to_vec(); + attempts.lock().unwrap().push(peers.clone()); + let number = attempt_no.fetch_add(1, Ordering::SeqCst); + let first_started = first_started.clone(); + let release = release.clone(); + let first_dropped = first_dropped.clone(); + Box::pin(async move { let _drop_signal = DropSignal(first_dropped); - first_started.notify_one(); - std::future::pending::<()>().await; - unreachable!("the old recovery attempt must be cancelled"); - } - Ok(peers == vec![format!("http://{new_ip}:{port}")]) - }) - }, - ) - .await + if number == 0 { + first_started.notify_one(); + // Hold the first attempt until the test + // releases it. The test chooses its outcome: + // `Ok(false)` moves on to the next pending + // peer, `Ok(true)` completes recovery. + release.notified().await; + Ok(first_attempt_outcome) + } else { + // Later attempts succeed: recovery completes. + Ok(true) + } + }) + }, + ) + .await + }) + }; + + Self { + task, + cancel, + service, + changes_tx, + writer, + attempts, + first_started, + release, + first_dropped, } - }); + } - tokio::time::timeout(Duration::from_secs(1), first_started.notified()) - .await - .expect("old recovery request must be in flight"); - let mut replacement = recovery_slice(new_ip, Some(true), Some(true)); - replacement.metadata.name = Some("epp-peers-0".to_string()); - writer.apply_watcher_event(&watcher::Event::Apply(replacement)); - changes_tx.send(1).unwrap(); + async fn apply(&mut self, ip: &str, eligible: bool) { + let mut slice = recovery_slice(ip, Some(true), Some(eligible)); + slice.metadata.name = Some(format!("peer-{}", ip.replace('.', "-"))); + self.writer + .apply_watcher_event(&kube::runtime::watcher::Event::Apply(slice)); + self.changes_tx.send(1).unwrap(); + } + + async fn remove(&mut self, ip: &str) { + let mut slice = recovery_slice(ip, Some(true), Some(true)); + slice.metadata.name = Some(format!("peer-{}", ip.replace('.', "-"))); + self.writer + .apply_watcher_event(&kube::runtime::watcher::Event::Delete(slice)); + self.changes_tx.send(1).unwrap(); + } + + async fn assert_active_not_dropped(&self) { + // The in-flight attempt must survive the churn: no Drop. + let dropped = self.first_dropped.clone(); + let result = tokio::time::timeout(Duration::from_millis(200), dropped.notified()).await; + assert!( + result.is_err(), + "in-flight recovery attempt must NOT be cancelled by EndpointSlice churn" + ); + } + + async fn finish(self) { + self.release.notify_one(); + tokio::time::timeout(Duration::from_secs(3), self.task) + .await + .expect("recovery must complete") + .expect("recovery task joins") + .expect("recovery must succeed"); + self.cancel.cancel(); + self.service.shutdown().await; + } + } - tokio::time::timeout(Duration::from_secs(1), first_dropped.notified()) + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn unchanged_churn_keeps_inflight_attempt() { + let port = 9093; + let mut harness = ChurnHarness::start(&["192.0.2.10"], port, true).await; + tokio::time::timeout(Duration::from_secs(1), harness.first_started.notified()) .await - .expect("peer change must drop the old recovery future"); - tokio::time::timeout(Duration::from_secs(1), task) + .expect("first recovery attempt must be in flight"); + + // Unrelated churn: same eligible set, metadata-only change (the reflector + // bumps the change channel on every Apply). + let mut slice = recovery_slice("192.0.2.10", Some(true), Some(true)); + slice.metadata.name = Some("epp-peers-churn".to_string()); + slice.metadata.annotations = Some( + [("note".to_string(), "churn".to_string())] + .into_iter() + .collect(), + ); + harness + .writer + .apply_watcher_event(&kube::runtime::watcher::Event::Apply(slice)); + harness.changes_tx.send(1).unwrap(); + + harness.assert_active_not_dropped().await; + let attempts = harness.attempts.clone(); + harness.finish().await; + assert_eq!( + *attempts.lock().unwrap(), + vec![vec![format!("http://192.0.2.10:{port}")]], + "unrelated churn must not restart the attempt" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn join_during_recovery_keeps_active_then_tries_new_peer() { + let port = 9093; + let mut harness = ChurnHarness::start(&["192.0.2.10"], port, false).await; + tokio::time::timeout(Duration::from_secs(1), harness.first_started.notified()) .await - .expect("peer change must not wait for old HTTP timeout") - .expect("recovery task joins") - .expect("new peer dump completes recovery"); + .expect("first recovery attempt must be in flight"); + + // A new serving peer joins mid-recovery: the active request must not be + // cancelled; once it fails, the new peer is tried next. + harness.apply("192.0.2.11", true).await; + harness.assert_active_not_dropped().await; + let attempts = harness.attempts.clone(); + harness.finish().await; assert_eq!( *attempts.lock().unwrap(), vec![ - vec![format!("http://{old_ip}:{port}")], - vec![format!("http://{new_ip}:{port}")], - ] + vec![format!("http://192.0.2.10:{port}")], + vec![format!("http://192.0.2.11:{port}")], + ], + "a join must not cancel the active attempt; the new peer is tried after it fails" ); + } - cancel.cancel(); - service.shutdown().await; + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn removal_of_unattempted_peer_drops_it_from_pending() { + let port = 9093; + let mut harness = ChurnHarness::start(&["192.0.2.10", "192.0.2.11"], port, false).await; + tokio::time::timeout(Duration::from_secs(1), harness.first_started.notified()) + .await + .expect("first recovery attempt must be in flight"); + + // The first attempt is on whichever peer the attempt-scoped shuffle put + // first; the OTHER peer is unattempted. Remove that one mid-flight: it + // must be dropped from pending and never tried. + let first_peer = { + let attempts = harness.attempts.lock().unwrap(); + attempts.last().expect("first attempt recorded")[0].clone() + }; + let unattempted = if first_peer.contains("192.0.2.10") { + "192.0.2.11" + } else { + "192.0.2.10" + }; + harness.remove(unattempted).await; + harness.assert_active_not_dropped().await; + let attempts = harness.attempts.clone(); + harness.finish().await; + let attempts = attempts.lock().unwrap().clone(); + assert!( + attempts.iter().all(|a| a == &vec![first_peer.clone()]), + "the removed peer must never be tried, got: {attempts:?}" + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn peers_disappearing_during_recovery_bootstraps() { + async fn removal_of_active_peer_does_not_cancel_its_request() { + let port = 9093; + let mut harness = ChurnHarness::start(&["192.0.2.10"], port, true).await; + tokio::time::timeout(Duration::from_secs(1), harness.first_started.notified()) + .await + .expect("first recovery attempt must be in flight"); + + // The active peer leaves the slice mid-transfer: the request is kept + // and may still complete successfully. + harness.remove("192.0.2.10").await; + harness.assert_active_not_dropped().await; + let attempts = harness.attempts.clone(); + harness.finish().await; + assert_eq!( + *attempts.lock().unwrap(), + vec![vec![format!("http://192.0.2.10:{port}")]], + "removal of the active peer must not cancel its request" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn peers_disappearing_during_recovery_keep_active_request() { use kube::runtime::watcher; let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -1268,11 +1463,14 @@ mod tests { writer.apply_watcher_event(&watcher::Event::Delete(old_slice)); changes_tx.send(1).unwrap(); - tokio::time::timeout(Duration::from_secs(1), task) + // The in-flight request is kept even though every peer is gone; only + // once it completes (empty snapshot = success) does recovery finish. + gate.release.notify_one(); + tokio::time::timeout(Duration::from_secs(3), task) .await - .expect("empty peer set must bootstrap without waiting for old request") + .expect("recovery must finish after the kept request completes") .expect("recovery task joins") - .expect("empty peer set bootstraps"); + .expect("kept request completes recovery"); cancel.cancel(); server.abort(); @@ -1313,8 +1511,8 @@ mod tests { // Recovery candidates exclude the not-ready sibling (10.0.0.2): a // not-ready replica has no KV index yet, so it cannot bootstrap a peer. assert_eq!( - recovery_peer_urls(&store, "10.0.0.9", 9093), - vec!["http://10.0.0.3:9093".to_string()] + recovery_peer_set(&store, "10.0.0.9", 9093), + BTreeSet::from(["http://10.0.0.3:9093".to_string()]) ); } @@ -1337,7 +1535,7 @@ mod tests { ..Default::default() }; let (store, _writer) = store_and_writer(vec![slice]); - assert!(recovery_peer_urls(&store, "10.0.0.9", 9093).is_empty()); + assert!(recovery_peer_set(&store, "10.0.0.9", 9093).is_empty()); } #[test] @@ -1358,7 +1556,7 @@ mod tests { ..Default::default() }; let (store, _writer) = store_and_writer(vec![slice]); - assert!(recovery_peer_urls(&store, "10.0.0.9", 9093).is_empty()); + assert!(recovery_peer_set(&store, "10.0.0.9", 9093).is_empty()); } #[test] @@ -1372,7 +1570,7 @@ mod tests { } #[test] - fn recovery_peer_urls_bracket_ipv6() { + fn recovery_peer_set_brackets_ipv6() { let slice = EndpointSlice { address_type: "IPv6".to_string(), endpoints: vec![Endpoint { @@ -1387,8 +1585,8 @@ mod tests { }; let (store, _writer) = store_and_writer(vec![slice]); assert_eq!( - recovery_peer_urls(&store, "fd00::1", 9093), - vec!["http://[fd00::2]:9093".to_string()] + recovery_peer_set(&store, "fd00::1", 9093), + BTreeSet::from(["http://[fd00::2]:9093".to_string()]) ); } From da813f9e3dd5ff2f31e068db4dc65ef539156d2a Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Fri, 21 Aug 2026 13:56:29 +0800 Subject: [PATCH 15/17] chore(epp): split the peer Service from the request Service in the onramp Give the replica plane its own Service (dynamo-epp-peer: replica-agg + selection-http) and point DYN_EPP_PEER_SERVICE at it; the request Service (dynamo-epp) carries only gRPC. The EPP resolves the peer ports and watches EndpointSlices for the peer Service, making the peer plane and its port contract explicit without prescribing cluster ingress policy. The example keeps a minimal NetworkPolicy on the peer ports (the /dump endpoint is unauthenticated); operators with their own baseline may drop it. Signed-off-by: Peter Pan --- .../ext-proc/examples/onramp/agg.yaml | 45 +++++++++++++------ .../kv-aware-routing/vanilla-vllm-onramp.mdx | 10 +++-- 2 files changed, 37 insertions(+), 18 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml index 582602471574..10635d9ff977 100644 --- a/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml +++ b/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml @@ -258,12 +258,13 @@ spec: # the vLLM engine's --max-num-batched-tokens. - name: DYN_EPP_MAX_NUM_BATCHED_TOKENS value: "8192" - # Replication: watch THIS Deployment's own Service - # (EndpointSlices) to find sibling EPP replicas and sync active load - # over its required named replica-agg port. Omit this env to run a - # single fully-local replica (and set replicas: 1). + # Replication: watch the peer Service (EndpointSlices) to find + # sibling EPP replicas and sync active load over its required named + # replica-agg port, and to recover the KV index over selection-http. + # Omit this env to run a single fully-local replica (and set + # replicas: 1). - name: DYN_EPP_PEER_SERVICE - value: dynamo-epp + value: dynamo-epp-peer # The reflector watches pods in the EPP's own namespace. - name: POD_NAMESPACE valueFrom: @@ -298,8 +299,11 @@ spec: drop: - ALL --- -# GAIE calls the gRPC port. EPP replicas use `replica-agg` for lifecycle sync -# and `selection-http` for startup KV-index recovery. +# GAIE calls the gRPC port on the request Service; EPP replicas talk to each +# other over the dedicated peer Service below. Keeping the peer plane on its +# own Service makes the EndpointSlice discovery (and its port contract) +# explicit: the EPP reads `replica-agg` / `selection-http` from +# `dynamo-epp-peer`, and the request Service carries only gRPC. apiVersion: v1 kind: Service metadata: @@ -312,6 +316,18 @@ spec: appProtocol: http2 port: 9002 targetPort: grpc +--- +# Peer Service: `replica-agg` for lifecycle replica sync, `selection-http` for +# startup KV-index recovery. The EPP watches THIS Service's EndpointSlices for +# sibling replicas and resolves the named peer ports from its spec. +apiVersion: v1 +kind: Service +metadata: + name: dynamo-epp-peer +spec: + selector: + app: dynamo-epp + ports: - name: replica-agg port: 9092 targetPort: replica-agg @@ -319,13 +335,14 @@ spec: port: 9093 targetPort: selection-http --- -# `selection-http` serves the peer KV-index snapshot (`GET /dump`) without -# authentication, so it is reachable only from sibling EPP pods. The same -# peer-only restriction applies to `replica-agg` lifecycle sync. The GAIE -# gateway still reaches gRPC (9002) and the kubelet still probes grpc-health -# (9003) from any source; KV events flow *out* of the EPP to workers, so no -# inbound rule is needed for them. Clusters without a NetworkPolicy CNI -# simply ignore this object. +# The peer Service's `selection-http` serves the KV-index snapshot +# (`GET /dump`) without authentication, so it is reachable only from sibling +# EPP pods; the same peer-only restriction applies to `replica-agg` lifecycle +# sync. The GAIE gateway still reaches gRPC (9002) and the kubelet still probes +# grpc-health (9003) from any source; KV events flow *out* of the EPP to +# workers, so no inbound rule is needed for them. Clusters without a +# NetworkPolicy CNI simply ignore this object; operators with their own network +# baseline may drop or replace it. apiVersion: networking.k8s.io/v1 kind: NetworkPolicy metadata: diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx index 1e4295b0d862..ec58df5c5e06 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx +++ b/docs/fern/pages/kubernetes/kv-aware-routing/vanilla-vllm-onramp.mdx @@ -365,7 +365,7 @@ spec: - name: DYN_EPP_MAX_NUM_BATCHED_TOKENS value: "8192" - name: DYN_EPP_PEER_SERVICE - value: dynamo-epp + value: dynamo-epp-peer - name: POD_IP valueFrom: fieldRef: @@ -378,9 +378,11 @@ configuration. The model name, block size, event port, and maximum batched-token values in the example are one coherent configuration. Change them together when adapting another model or vLLM deployment. -Create the replacement `dynamo-epp` Service with gRPC port `9002` and, for two replicas, the named -`replica-agg` and `selection-http` ports. The complete ServiceAccount, RBAC, Deployment, and Service -definitions are available in the +Create the replacement request Service `dynamo-epp` with gRPC port `9002`, plus a dedicated peer +Service `dynamo-epp-peer` that carries the two replica-plane ports (`replica-agg` for lifecycle sync +and `selection-http` for startup KV-index recovery); the EPP resolves those named ports from +`dynamo-epp-peer`. The complete ServiceAccount, RBAC, Deployment, and Service definitions are +available in the [`agg.yaml` EPP resources](https://github.com/ai-dynamo/dynamo/blob/main/deploy/inference-gateway/ext-proc/examples/onramp/agg.yaml). ### Connect the InferencePool From 667ab67a7950c93a6ef4f0553113c7af2299e191 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Fri, 21 Aug 2026 15:02:28 +0800 Subject: [PATCH 16/17] chore(epp): bound the pre-recovery worker wait to a short best-effort margin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 15s wait was a magic number: it only ever fires when no worker is registered yet, and beyond the topology adapter's first reconcile (~1s in the common restart case) waiting longer cannot improve the dump overlap — a pure cold start with no peers would just stall startup before bootstrapping empty. Shorten to 5s and document that the wait is a best-effort margin, not a correctness guarantee. Signed-off-by: Peter Pan --- deploy/inference-gateway/ext-proc/src/selector.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index c92d4dcea5aa..65d57c11858d 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -562,9 +562,12 @@ impl Selector { /// Bounded wait for at least one registered worker before the peer dump, so /// the snapshot overlaps the already-buffered live event stream (subscribe- -/// first). A cold start that never registers a worker in time proceeds anyway: -/// the peer dump is then the best available state. -const PEER_RECOVERY_WORKER_WAIT: std::time::Duration = std::time::Duration::from_secs(15); +/// first). This is a best-effort margin for the topology adapter's first +/// reconcile — in the common restart case workers register within ~1s and the +/// wait is a no-op. A cold start that never registers a worker in time +/// proceeds anyway: the peer dump is then the best available state, and a +/// longer wait cannot improve it. +const PEER_RECOVERY_WORKER_WAIT: std::time::Duration = std::time::Duration::from_secs(5); async fn wait_for_registered_worker( service: &SelectionService, From d8da747aa92af575e54f2ae066d6ee0c42119865 Mon Sep 17 00:00:00 2001 From: Peter Pan Date: Fri, 21 Aug 2026 17:30:07 +0800 Subject: [PATCH 17/17] docs(epp): state the pre-recovery worker-wait limitation explicitly The 5s bound is a best-effort magic-number timeout, not a correctness guarantee: a worker registering after the window still leaves a gap, and the precise fix is a deterministic first-reconcile signal rather than a wall clock. Document that instead of implying the wait closes the gap. Signed-off-by: Peter Pan --- .../inference-gateway/ext-proc/src/selector.rs | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/deploy/inference-gateway/ext-proc/src/selector.rs b/deploy/inference-gateway/ext-proc/src/selector.rs index 65d57c11858d..545ae0515b77 100644 --- a/deploy/inference-gateway/ext-proc/src/selector.rs +++ b/deploy/inference-gateway/ext-proc/src/selector.rs @@ -562,11 +562,18 @@ impl Selector { /// Bounded wait for at least one registered worker before the peer dump, so /// the snapshot overlaps the already-buffered live event stream (subscribe- -/// first). This is a best-effort margin for the topology adapter's first -/// reconcile — in the common restart case workers register within ~1s and the -/// wait is a no-op. A cold start that never registers a worker in time -/// proceeds anyway: the peer dump is then the best available state, and a -/// longer wait cannot improve it. +/// first). +/// +/// LIMITATION: this is a best-effort timeout, not a correctness guarantee, and +/// the `5s` value is a magic number — there is no measured basis for how long +/// the topology adapter's first reconcile takes (in the common restart case +/// workers register within ~1s and the wait is a no-op). If a worker registers +/// only after this window, the dump still runs ahead of the live subscription +/// and leaves a gap; a longer wait would not help because the worker simply +/// was not there to subscribe to. The precise fix would be to await a +/// deterministic "first reconcile complete" signal from the topology adapter +/// instead of a wall-clock timeout; this bound exists only so a slow cold start +/// does not hang forever. const PEER_RECOVERY_WORKER_WAIT: std::time::Duration = std::time::Duration::from_secs(5); async fn wait_for_registered_worker(