From bb6dcfd5c59c868b82f3230b150c9de17b1ff190 Mon Sep 17 00:00:00 2001 From: Shang-Pin Date: Wed, 16 Sep 2026 23:24:18 +0000 Subject: [PATCH] standalone-indexer: subscribe to every data-parallel rank of a discovered pod Pod discovery registered each engine pod once, as dp_rank 0 on tcp://:. vLLM offsets the KV-event ZMQ port and the /kv_recover port by data_parallel_rank, so on a --data-parallel-size N engine ranks 1..N-1 publish to ports nobody subscribes to and everything they cache is invisible to /query while the engine still hits it. Measured on deepseek-ai/DeepSeek-V4.1-Flash (DP=2): every prefill the scheduler placed on rank 1 (about half) showed 0 in the standalone indexer and in rank 0's /kv_recover dump, and all its blocks in rank 1's dump; the KV ladder read actual 0.78 > h24 0.65 > perfect 0.63 > reality 0.59. Add --watch-dp-size (default 1, so single-rank engines are unchanged). The watcher registers one listener per rank under the pod's instance: dp_rank r, tcp://:, http://:. register() rejects a rank that is already present, so a pod whose registration fails part-way is deregistered before the retry instead of being left half-subscribed forever. Co-Authored-By: Claude Fable 5.1 --- lib/bindings/python/rust/llm/kv.rs | 12 ++ lib/kv-router/src/standalone_indexer/docs.md | 11 ++ lib/kv-router/src/standalone_indexer/mod.rs | 11 +- .../src/standalone_indexer/pod_watcher.rs | 140 +++++++++++++++--- 4 files changed, 151 insertions(+), 23 deletions(-) diff --git a/lib/bindings/python/rust/llm/kv.rs b/lib/bindings/python/rust/llm/kv.rs index 85a033527092..3bec45699ea2 100644 --- a/lib/bindings/python/rust/llm/kv.rs +++ b/lib/bindings/python/rust/llm/kv.rs @@ -118,6 +118,14 @@ struct KvIndexerCli { #[arg(long)] watch_recover_port: Option, + /// Data-parallel ranks per discovered engine pod (vLLM + /// `--data-parallel-size`). Rank r publishes KV events on + /// `--watch-zmq-port + r` and serves recovery on `--watch-recover-port + r`; + /// every rank is subscribed under the pod's instance. Leaving this at 1 on + /// a DP engine silently indexes only rank 0's cache. + #[arg(long, default_value_t = 1)] + watch_dp_size: u32, + /// Model name whose engine pods to discover, e.g. "openai/gpt-oss-120b". /// Unless --watch-label overrides it, the pod watch uses the selector /// `di/model_name=` (the stable label the backend stamps @@ -161,11 +169,15 @@ where let block_size = cli.block_size.ok_or_else(|| { anyhow::anyhow!("--block-size is required when --watch-namespace is set") })?; + if cli.watch_dp_size == 0 { + anyhow::bail!("--watch-dp-size must be at least 1"); + } Some(KubeDiscoveryConfig { namespace, label_selector, zmq_port: cli.watch_zmq_port, recover_port: cli.watch_recover_port, + dp_size: cli.watch_dp_size, model_name: cli.watch_model_name.unwrap_or_else(|| cli.model_name.clone()), tenant_id: cli.tenant_id.clone(), block_size, diff --git a/lib/kv-router/src/standalone_indexer/docs.md b/lib/kv-router/src/standalone_indexer/docs.md index 71c50a687f69..aa9537c6d05b 100644 --- a/lib/kv-router/src/standalone_indexer/docs.md +++ b/lib/kv-router/src/standalone_indexer/docs.md @@ -26,6 +26,17 @@ If no `recover_endpoint` is configured, gaps are logged and the dropped batches are lost. Implementation lives in `listener.rs` (`recover_gap`, `apply_recover_response`, `apply_recovered_events`). +## Data-parallel engines (`--watch-dp-size`) + +A vLLM engine with `--data-parallel-size N` runs N ranks per pod, each with its +own KV cache and its own event stream: rank `r` publishes on `zmq_port + r` and +serves `/kv_recover` on `kv_recover_port + r`. Pod discovery registers one +listener per rank under the pod's instance (`--watch-dp-size N`, default 1), +with `dp_rank = r` and the per-rank endpoints, so `/workers` shows N +`listeners` per pod. With the default on a DP engine only rank 0's cache is +indexed and every prefill scheduled on another rank is invisible to `/query` +while the engine still hits it (`pod_watcher.rs`, `rank_endpoints`). + ## Audit logging (`--enable-logging`) Pass `--enable-logging` to `python -m dynamo.indexer` to turn on verbose audit diff --git a/lib/kv-router/src/standalone_indexer/mod.rs b/lib/kv-router/src/standalone_indexer/mod.rs index ee71fd60c16a..8a03ecf151c2 100644 --- a/lib/kv-router/src/standalone_indexer/mod.rs +++ b/lib/kv-router/src/standalone_indexer/mod.rs @@ -124,12 +124,17 @@ pub struct KubeDiscoveryConfig { /// from the model name (`di/model_name=`, spans every /// engine_hash) or supplied raw, e.g. `engine_hash=d4b7a85131172ca6`. pub label_selector: String, - /// ZMQ KV-event port the engines publish on (e.g. 5557). + /// ZMQ KV-event port the engines publish on (e.g. 5557). Data-parallel + /// rank `r` publishes on `zmq_port + r`. pub zmq_port: u16, /// Optional HTTP port serving `GET /kv_recover` on the engines, used for - /// per-worker gap recovery. `http://:` is the base - /// URL the indexer queries on a detected gap. + /// per-worker gap recovery. `http://:` is the + /// base URL the indexer queries for rank `r` on a detected gap. pub recover_port: Option, + /// Data-parallel ranks per engine pod. Each rank has its own KV cache and + /// its own event stream, so every rank is registered as a listener of the + /// same instance. Engines without data parallelism run one rank. + pub dp_size: u32, /// Model name discovered pods are registered under. pub model_name: String, /// Tenant id discovered pods are registered under. diff --git a/lib/kv-router/src/standalone_indexer/pod_watcher.rs b/lib/kv-router/src/standalone_indexer/pod_watcher.rs index 376cf9051af2..ada1a33fb022 100644 --- a/lib/kv-router/src/standalone_indexer/pod_watcher.rs +++ b/lib/kv-router/src/standalone_indexer/pod_watcher.rs @@ -119,6 +119,7 @@ pub fn spawn_pod_watcher( namespace = %config.namespace, label_selector = %config.label_selector, zmq_port = config.zmq_port, + dp_size = config.dp_size, model_name = %config.model_name, block_size = config.block_size, "Starting Kubernetes pod watcher" @@ -259,30 +260,35 @@ async fn reconcile( let instance_id = instance_id_for(name); let ip = &pod.ip; let tenant_id = pod.tenant_id(&config.tenant_id); - let endpoint = format!("tcp://{ip}:{}", config.zmq_port); - let recover_endpoint = config - .recover_port - .map(|port| format!("http://{ip}:{port}")); - match registry - .register( - instance_id, - endpoint, - 0, // dp_rank: single-rank engines - config.model_name.clone(), - tenant_id.clone(), - config.block_size, - recover_endpoint, - Some(name.clone()), - ) - .await - { - Ok(()) => { + let mut failed: Option<(u32, anyhow::Error)> = None; + for rank in rank_endpoints(ip, config) { + if let Err(error) = registry + .register( + instance_id, + rank.endpoint, + rank.dp_rank, + config.model_name.clone(), + tenant_id.clone(), + config.block_size, + rank.recover_endpoint, + Some(name.clone()), + ) + .await + { + failed = Some((rank.dp_rank, error)); + break; + } + } + + match failed { + None => { tracing::info!( pod = %name, ip = %ip, instance_id, tenant_id = %tenant_id, + dp_size = config.dp_size, "Subscribed to engine pod" ); subscribed.insert( @@ -294,19 +300,57 @@ async fn reconcile( }, ); } - Err(error) => { - // Not recorded, so it is retried on the next reconcile. + Some((dp_rank, error)) => { + // Not recorded, so it is retried on the next reconcile. Drop + // any ranks that did register: register() rejects a rank that + // is already present, so a partial instance would never + // converge on retry. tracing::warn!( pod = %name, ip = %ip, + dp_rank, error = %error, "Failed to register engine pod; will retry" ); + if let Err(error) = registry + .deregister(instance_id, &config.model_name, &tenant_id) + .await + { + tracing::debug!( + pod = %name, + error = %error, + "Cleanup deregister was a no-op" + ); + } } } } } +/// One listener to register for a pod: its data-parallel rank and the +/// per-rank event and recovery endpoints. +#[derive(Debug, PartialEq)] +struct RankEndpoint { + dp_rank: u32, + endpoint: String, + recover_endpoint: Option, +} + +/// Endpoints for every data-parallel rank of a pod. vLLM offsets both the ZMQ +/// KV-event port and the `/kv_recover` port by the rank, so rank `r` of a pod +/// at `ip` publishes on `zmq_port + r` and recovers on `recover_port + r`. +fn rank_endpoints(ip: &str, config: &KubeDiscoveryConfig) -> Vec { + (0..config.dp_size.max(1)) + .map(|dp_rank| RankEndpoint { + dp_rank, + endpoint: format!("tcp://{ip}:{}", u32::from(config.zmq_port) + dp_rank), + recover_endpoint: config + .recover_port + .map(|port| format!("http://{ip}:{}", u32::from(port) + dp_rank)), + }) + .collect() +} + /// Derive a stable `WorkerId` from the pod name so the same pod always maps to /// the same registry entry across MODIFIED events. fn instance_id_for(pod_name: &str) -> WorkerId { @@ -456,4 +500,60 @@ mod tests { assert_eq!(instance_id_for("engine-abc"), instance_id_for("engine-abc")); assert_ne!(instance_id_for("engine-abc"), instance_id_for("engine-xyz")); } + + fn discovery(dp_size: u32, recover_port: Option) -> KubeDiscoveryConfig { + KubeDiscoveryConfig { + namespace: "deepinfra".to_string(), + label_selector: "di/model_name=m".to_string(), + zmq_port: 5557, + recover_port, + dp_size, + model_name: "m".to_string(), + tenant_id: "default".to_string(), + block_size: 128, + } + } + + #[test] + fn single_rank_engine_registers_rank_zero_on_base_ports() { + assert_eq!( + rank_endpoints("10.0.0.1", &discovery(1, Some(5559))), + vec![RankEndpoint { + dp_rank: 0, + endpoint: "tcp://10.0.0.1:5557".to_string(), + recover_endpoint: Some("http://10.0.0.1:5559".to_string()), + }] + ); + } + + #[test] + fn data_parallel_ranks_offset_event_and_recovery_ports() { + assert_eq!( + rank_endpoints("10.0.0.1", &discovery(2, Some(5559))), + vec![ + RankEndpoint { + dp_rank: 0, + endpoint: "tcp://10.0.0.1:5557".to_string(), + recover_endpoint: Some("http://10.0.0.1:5559".to_string()), + }, + RankEndpoint { + dp_rank: 1, + endpoint: "tcp://10.0.0.1:5558".to_string(), + recover_endpoint: Some("http://10.0.0.1:5560".to_string()), + }, + ] + ); + } + + #[test] + fn no_recover_port_means_no_recovery_endpoint_on_any_rank() { + let ranks = rank_endpoints("10.0.0.1", &discovery(2, None)); + assert_eq!(ranks.len(), 2); + assert!(ranks.iter().all(|r| r.recover_endpoint.is_none())); + } + + #[test] + fn zero_dp_size_still_registers_rank_zero() { + assert_eq!(rank_endpoints("10.0.0.1", &discovery(0, None)).len(), 1); + } }