From 09d47e5250117c3fc4d2958ab91ed416abf9d607 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Tue, 15 Sep 2026 17:49:30 -0700 Subject: [PATCH 01/17] feat(sglang): relay node-local KV events from headless sidecars Signed-off-by: jain-ria --- Cargo.lock | 6 + lib/bindings/python/Cargo.lock | 2 + lib/bindings/python/rust/backend.rs | 6 +- lib/sidecar/sglang/Cargo.toml | 6 + lib/sidecar/sglang/README.md | 53 +++- lib/sidecar/sglang/src/args.rs | 29 +- lib/sidecar/sglang/src/context.rs | 281 +++++++++++++++++ lib/sidecar/sglang/src/engine.rs | 116 ++++++- lib/sidecar/sglang/src/headless.rs | 411 +++++++++++++++++++++++++ lib/sidecar/sglang/src/lib.rs | 4 + lib/sidecar/sglang/src/main.rs | 11 +- lib/sidecar/sglang/src/startup.rs | 74 +++++ lib/sidecar/sglang/tests/executable.rs | 15 + 13 files changed, 1001 insertions(+), 13 deletions(-) create mode 100644 lib/sidecar/sglang/src/context.rs create mode 100644 lib/sidecar/sglang/src/headless.rs create mode 100644 lib/sidecar/sglang/src/startup.rs diff --git a/Cargo.lock b/Cargo.lock index f640db17e0a6..cc2c91f1120c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3055,13 +3055,19 @@ dependencies = [ "async-trait", "clap", "dynamo-backend-common", + "dynamo-kv-router", + "dynamo-llm", "dynamo-runtime", "dynamo-sidecar-common", "futures", "prost 0.13.5", "rand 0.9.4", "reqwest 0.12.28", + "rmp-serde", + "serde", "serde_json", + "tempfile", + "tmq", "tokio", "tokio-util", "tonic 0.13.1", diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 269910dc2c5f..2d77b2289563 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2552,12 +2552,14 @@ dependencies = [ "async-trait", "clap", "dynamo-backend-common", + "dynamo-llm", "dynamo-runtime", "dynamo-sidecar-common", "futures", "prost 0.13.5", "rand 0.9.4", "reqwest 0.12.28", + "serde", "serde_json", "tokio", "tokio-util", diff --git a/lib/bindings/python/rust/backend.rs b/lib/bindings/python/rust/backend.rs index 6031f47cc8fd..902be61b10aa 100644 --- a/lib/bindings/python/rust/backend.rs +++ b/lib/bindings/python/rust/backend.rs @@ -101,11 +101,11 @@ fn sglang_sidecar_argv(argv: Vec) -> Vec { #[pyo3(signature = (argv=None))] fn _run_sglang_sidecar(py: Python<'_>, argv: Option>) -> PyResult<()> { let cli_argv = sglang_sidecar_argv(argv.unwrap_or_default()); - let (engine, config) = py - .allow_threads(move || dynamo_sglang_sidecar::SglangSidecarEngine::try_from_args(cli_argv)) + let sidecar = py + .allow_threads(move || dynamo_sglang_sidecar::SglangSidecar::try_from_args(cli_argv)) .map_err(sidecar_startup_to_pyerr)?; - py.allow_threads(move || dynamo_backend_common::run(Arc::new(engine), config)) + py.allow_threads(move || sidecar.run()) .map_err(|err| pyo3::exceptions::PyRuntimeError::new_err(err.to_string())) } diff --git a/lib/sidecar/sglang/Cargo.toml b/lib/sidecar/sglang/Cargo.toml index 596192eecfe9..b86d48cc03b5 100644 --- a/lib/sidecar/sglang/Cargo.toml +++ b/lib/sidecar/sglang/Cargo.toml @@ -21,6 +21,7 @@ path = "src/main.rs" [dependencies] dynamo-backend-common = { workspace = true } +dynamo-llm = { path = "../../llm", default-features = false } dynamo-runtime = { workspace = true } dynamo-sidecar-common = { workspace = true } @@ -31,6 +32,7 @@ clap = { version = "4", features = ["derive", "env"] } futures = { workspace = true } rand = { workspace = true } reqwest = { workspace = true } +serde = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } tokio-util = { workspace = true } @@ -45,3 +47,7 @@ tonic-build = { workspace = true } [dev-dependencies] dynamo-backend-common = { workspace = true, features = ["testing"] } +dynamo-kv-router = { workspace = true } +rmp-serde = { workspace = true } +tempfile = { workspace = true } +tmq = { workspace = true } diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index da2c0330b66a..72643c6eb483 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -57,7 +57,58 @@ python3 -m sglang.launch_server \ The entry point configures Dynamo logging when `main()` runs, then calls the private `dynamo._core.backend._run_sglang_sidecar(argv)` binding. The binding prepends the executable name expected by clap, releases the GIL, and runs the -same unified worker lifecycle as the standalone executable. +same mode dispatcher as the standalone executable. + +## Node-local KV sidecars (multinode) + +This opt-in path requires the version-1 launch context from +[SGLang #39659](https://github.com/sgl-project/sglang/pull/39659). It does not +require follower gRPC servers. Enable `--sidecar-scope local-telemetry` alongside +`--sidecar dynamo.sglang.sidecar` in SGLang. Without a launch context, the +existing gRPC-only startup and multinode KV-discovery guard are unchanged. + +SGLang supplies `SGLANG_SIDECAR_CONTEXT` when it launches each sidecar: + +- Node 0 uses `mode: "full"`: normal request serving and global DP registration, + but subscribes only to the explicitly listed local KV sources. +- A follower with local KV publishers uses `mode: "telemetry"`: no gRPC client, + model registration, or inference request endpoint. It relays its own local + events directly to Dynamo. +- A TP-only follower without a local KV publisher needs no sidecar. SGLang + decides whether to launch one; Dynamo does not infer ownership from node rank. + +Both modes consume `kv_event_sources` exactly as supplied, including global +`dp_rank`, connectable `endpoint`, `topic`, and logical `block_size`. Dynamo +does not reconstruct rank ranges or ports from `dp_size`. An empty full-mode +source list stays empty. + +The leader advertises the existing `sglang_worker_group_id` derived from the +shared `dist_init_addr`, plus its KV block size and locally owned ranks in +runtime metadata. Followers use Dynamo's runtime-config discovery watch to +find that group's serving worker. Each follower retains its own publisher +identity, while its events carry the **leader's worker ID and the source's +global DP rank**. This associates cache locations with a routable target; it +does not send the events through the leader sidecar. + +All sidecars in a group must use the same Dynamo namespace, component, and +endpoint. For PD disaggregation, pass `DYN_COMPONENT=prefill` on every prefill +node and `DYN_COMPONENT=decode` on every decode node (or the corresponding +`--component` sidecar argument). A headless sidecar has no gRPC discovery from +which to infer the PD role. Separate prefill/decode engine groups need distinct +`dist_init_addr` values, as in the in-process integration. + +Followers may start before the leader; discovery is cancellable and bounded +by `--leader-discovery-timeout-secs` (default 1800). Missing, ambiguous, or +incompatible leader metadata produces an explicit startup error. If a running +follower loses its leader or observes a changed publishing identity/config, +it exits: restart the distributed engine group, not just the sidecar. This +matches SGLang's managed-sidecar failure policy and avoids continuing under a +stale serving-worker identity. + +This change adds **KV events only**, not forward-pass metrics or scheduler +load publishing. It reuses the existing advisory ZMQ relay: it does not add a +readiness handshake, lossless startup, or independent sidecar restart/replay. +`replay_endpoint` is accepted in the context but is not consumed by this relay. ## Deploy on Kubernetes (quick start) diff --git a/lib/sidecar/sglang/src/args.rs b/lib/sidecar/sglang/src/args.rs index 153bec1549a7..356002474828 100644 --- a/lib/sidecar/sglang/src/args.rs +++ b/lib/sidecar/sglang/src/args.rs @@ -3,7 +3,24 @@ //! Command-line arguments for the SGLang sidecar. -use dynamo_sidecar_common::SidecarArgs; +use dynamo_backend_common::CommonArgs; +use dynamo_sidecar_common::GrpcTransportArgs; + +use crate::context::SidecarContext; + +/// SGLang alone supports a telemetry-only launch without a gRPC endpoint. +#[derive(clap::Args, Debug, Clone)] +pub struct SglangSidecarArgs { + #[command(flatten)] + pub common: CommonArgs, + + /// Engine gRPC endpoint, required in full mode. Falls back to SGLANG_GRPC_ENDPOINT. + #[arg(long, env = "DYN_SIDECAR_GRPC_ENDPOINT")] + pub grpc_endpoint: Option, + + #[command(flatten)] + pub grpc: GrpcTransportArgs, +} /// Parsed sidecar arguments. #[derive(clap::Parser, Debug, Clone)] @@ -13,7 +30,15 @@ use dynamo_sidecar_common::SidecarArgs; )] pub struct Args { #[command(flatten)] - pub sidecar: SidecarArgs, + pub sidecar: SglangSidecarArgs, + + /// Versioned node-local context supplied by SGLang's managed sidecar launcher. + #[arg(long, env = "SGLANG_SIDECAR_CONTEXT")] + pub sidecar_context: Option, + + /// Maximum wait for a matching leader registration in telemetry mode. + #[arg(long, default_value_t = 1800, value_parser = clap::value_parser!(u64).range(1..))] + pub leader_discovery_timeout_secs: u64, /// Reachable host that decode workers use to connect to a prefill worker's /// SGLang disaggregation bootstrap port. By default this is derived from diff --git a/lib/sidecar/sglang/src/context.rs b/lib/sidecar/sglang/src/context.rs new file mode 100644 index 000000000000..93025f03b855 --- /dev/null +++ b/lib/sidecar/sglang/src/context.rs @@ -0,0 +1,281 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! SGLang's versioned, engine-owned managed-sidecar launch contract. + +use std::collections::HashSet; +use std::str::FromStr; + +use anyhow::{Context, Result, bail, ensure}; +use serde::Deserialize; + +pub(crate) const WORKER_GROUP_KEY: &str = "sglang_worker_group_id"; +pub(crate) const KV_CONFIG_KEY: &str = "sglang_sidecar_kv_events"; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum SidecarMode { + Full, + Telemetry, +} + +/// Matches sglang.srt.entrypoints.sidecar_context.SidecarContext, version 1. +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SidecarContext { + pub version: u32, + pub mode: SidecarMode, + pub node_rank: u32, + pub nnodes: u32, + pub dp_size: u32, + pub dist_init_addr: Option, + pub kv_event_sources: Vec, +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct KvEventSource { + pub dp_rank: u32, + pub endpoint: String, + pub topic: String, + pub block_size: u32, + /// Accepted for wire compatibility; the existing Dynamo relay is live-only. + pub replay_endpoint: Option, +} + +impl FromStr for SidecarContext { + type Err = String; + + fn from_str(raw: &str) -> Result { + let context: Self = serde_json::from_str(raw) + .map_err(|error| format!("invalid SGLANG_SIDECAR_CONTEXT: {error}"))?; + context.validate().map_err(|error| error.to_string())?; + Ok(context) + } +} + +impl SidecarContext { + fn validate(&self) -> Result<()> { + ensure!( + self.version == 1, + "unsupported SGLANG_SIDECAR_CONTEXT version {}", + self.version + ); + ensure!( + self.nnodes > 0 && self.node_rank < self.nnodes, + "invalid sidecar node topology" + ); + ensure!(self.dp_size > 0, "sidecar dp_size must be positive"); + ensure!( + (self.mode == SidecarMode::Full) == (self.node_rank == 0), + "full mode requires node_rank=0; telemetry mode requires a follower node" + ); + ensure!( + self.mode != SidecarMode::Telemetry || !self.kv_event_sources.is_empty(), + "telemetry mode requires at least one local KV source" + ); + ensure!( + self.nnodes == 1 + || self + .dist_init_addr + .as_ref() + .is_some_and(|addr| !addr.trim().is_empty()), + "multinode sidecars require dist_init_addr for leader matching" + ); + let mut ranks = HashSet::new(); + let mut endpoints = HashSet::new(); + for source in &self.kv_event_sources { + ensure!( + source.dp_rank < self.dp_size, + "KV source rank {} is outside dp_size {}", + source.dp_rank, + self.dp_size + ); + ensure!( + ranks.insert(source.dp_rank), + "duplicate local KV source rank {}", + source.dp_rank + ); + ensure!( + endpoints.insert(&source.endpoint), + "duplicate local KV source endpoint {}", + source.endpoint + ); + ensure!( + source.block_size > 0, + "KV source block_size must be positive" + ); + validate_endpoint(&source.endpoint)?; + if let Some(endpoint) = &source.replay_endpoint { + validate_endpoint(endpoint)?; + } + } + Ok(()) + } + + pub(crate) fn validate_registration( + &self, + dp_size: u32, + block_size: Option, + ) -> Result<()> { + ensure!( + self.dp_size == dp_size, + "sidecar context dp_size does not match leader registration" + ); + for source in &self.kv_event_sources { + ensure!( + Some(source.block_size) == block_size, + "KV source block_size does not match leader registration" + ); + } + Ok(()) + } + + pub(crate) fn worker_group_id(&self) -> Result> { + if self.nnodes == 1 { + return Ok(None); + } + let raw = self + .dist_init_addr + .as_deref() + .context("missing dist_init_addr")? + .trim(); + let url = if raw.contains("://") { + raw.to_owned() + } else { + format!("tcp://{raw}") + }; + let address = url::Url::parse(&url).context("invalid dist_init_addr")?; + ensure!( + address.scheme() == "tcp", + "dist_init_addr must be a TCP address" + ); + validate_endpoint(&url)?; + // Use the in-process group-key format: resolve the shared rendezvous + // address, not a per-node source or gRPC address. + let resolved = address.socket_addrs(|| None)?; + let address = resolved + .first() + .context("dist_init_addr resolved to no addresses")?; + Ok(Some(format!("dist_init:tcp://{address}"))) + } +} + +fn validate_endpoint(endpoint: &str) -> Result<()> { + if let Some(path) = endpoint.strip_prefix("ipc://") { + ensure!( + (path.starts_with('/') || path.starts_with('@')) + && path.len() > 1 + && !path.contains('\0'), + "IPC source must have an absolute or abstract socket path" + ); + return Ok(()); + } + let url = url::Url::parse(endpoint).context("invalid KV source endpoint")?; + let host = url + .host_str() + .context("KV source endpoint requires a host")?; + ensure!( + url.scheme() == "tcp" && url.port().is_some_and(|port| port > 0), + "KV source must use tcp://HOST:PORT or ipc://PATH" + ); + ensure!( + !matches!(host, "*" | "0.0.0.0" | "[::]" | "::"), + "KV source endpoint must be dialable, not a wildcard bind address" + ); + if !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + { + bail!("KV source TCP endpoint must contain only host and port"); + } + Ok(()) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use serde_json::{Value, json}; + + pub(crate) fn context_json(mode: &str) -> Value { + json!({"version":1,"mode":mode,"node_rank":if mode == "full" {0} else {1}, + "nnodes":2,"dp_size":8,"dist_init_addr":"tcp://127.0.0.1:2345", + "kv_event_sources":[{"dp_rank":4,"endpoint":"tcp://127.0.0.1:5561", + "topic":"","block_size":64,"replay_endpoint":null}]}) + } + + #[test] + fn accepts_upstream_contract_and_global_not_local_rank() { + let context: SidecarContext = context_json("telemetry").to_string().parse().unwrap(); + assert_eq!(context.kv_event_sources[0].dp_rank, 4); + assert_eq!( + context.worker_group_id().unwrap().as_deref(), + Some("dist_init:tcp://127.0.0.1:2345") + ); + context.validate_registration(8, Some(64)).unwrap(); + assert!(context.validate_registration(4, Some(64)).is_err()); + assert!(context.validate_registration(8, Some(32)).is_err()); + } + + #[test] + fn validates_context_before_runtime_or_grpc_startup() { + for (field, value) in [ + ("version", json!(2)), + ("mode", json!("unknown")), + ("node_rank", json!(0)), + ("node_rank", json!(2)), + ("dp_size", json!(0)), + ("dist_init_addr", Value::Null), + ("kv_event_sources", json!([])), + ] { + let mut raw = context_json("telemetry"); + raw[field] = value; + assert!(raw.to_string().parse::().is_err(), "{raw}"); + } + } + + #[test] + fn rejects_conflicting_or_undialable_sources() { + for (field, value) in [ + ("dp_rank", json!(8)), + ("dp_rank", json!(-1)), + ("block_size", json!(0)), + ("endpoint", json!("tcp://*:5561")), + ("endpoint", json!("inproc://events")), + ("endpoint", json!("ipc://relative")), + ] { + let mut raw = context_json("telemetry"); + raw["kv_event_sources"][0][field] = value; + assert!(raw.to_string().parse::().is_err(), "{raw}"); + } + let mut raw = context_json("telemetry"); + let source = raw["kv_event_sources"][0].clone(); + raw["kv_event_sources"].as_array_mut().unwrap().push(source); + assert!(raw.to_string().parse::().is_err()); + raw["kv_event_sources"][1]["dp_rank"] = json!(5); + assert!(raw.to_string().parse::().is_err()); + } + + #[test] + fn leader_may_have_no_local_publishers() { + let mut raw = context_json("full"); + raw["kv_event_sources"] = json!([]); + let context: SidecarContext = raw.to_string().parse().unwrap(); + context.validate_registration(8, Some(64)).unwrap(); + assert!(context.kv_event_sources.is_empty()); + } + + #[test] + fn ipv6_and_ipc_are_supported() { + let mut raw = context_json("telemetry"); + raw["dist_init_addr"] = json!("tcp://[::1]:2345"); + raw["kv_event_sources"][0]["endpoint"] = json!("ipc:///tmp/sglang-kv-test.sock"); + let context: SidecarContext = raw.to_string().parse().unwrap(); + assert_eq!( + context.worker_group_id().unwrap().as_deref(), + Some("dist_init:tcp://[::1]:2345") + ); + } +} diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index 5357e6c33ff2..c07bb822478a 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -22,6 +22,7 @@ use tokio_util::sync::CancellationToken; use crate::args::Args; use crate::client::{self, Client, Discovery, Pool}; +use crate::context::{KV_CONFIG_KEY, SidecarContext, SidecarMode, WORKER_GROUP_KEY}; use crate::native_http::{self, NativeHttp}; use crate::proto as pb; use crate::protocol::{ @@ -35,6 +36,8 @@ pub struct SglangSidecarEngine { disaggregation_mode: DisaggregationMode, bootstrap_host: Option, bootstrap_port: Option, + local_context: Option, + worker_group_id: Option, state: OnceCell, cancel: CancellationToken, } @@ -69,14 +72,39 @@ impl SglangSidecarEngine { Self::from_parsed(args).map_err(Into::into) } - fn from_parsed(args: Args) -> Result<(Self, WorkerConfig), DynamoError> { + pub(crate) fn from_parsed(args: Args) -> Result<(Self, WorkerConfig), DynamoError> { + if args + .sidecar_context + .as_ref() + .is_some_and(|ctx| ctx.mode != SidecarMode::Full) + { + return Err(client::invalid_arg( + "telemetry context requires the headless sidecar entry point", + )); + } if args.sidecar.common.route_to_encoder { return Err(client::invalid_arg( "route-to-encoder is not supported by the SGLang sidecar", )); } - let endpoint = args.sidecar.grpc_endpoint; + let endpoint = args + .sidecar + .grpc_endpoint + .or_else(|| std::env::var("SGLANG_GRPC_ENDPOINT").ok()) + .ok_or_else(|| { + client::invalid_arg( + "full sidecar mode requires --grpc-endpoint or SGLANG_GRPC_ENDPOINT", + ) + })?; + let endpoint = GrpcEndpoint::parse(&endpoint, "--grpc-endpoint")?; + let worker_group_id = args + .sidecar_context + .as_ref() + .map(SidecarContext::worker_group_id) + .transpose() + .map_err(|error| client::invalid_arg(error.to_string()))? + .flatten(); let transport = args.sidecar.grpc.config(); let discovery = bootstrap_discover(&endpoint, &transport)?; let disaggregation_mode = discovery_mode(&discovery)?; @@ -135,6 +163,8 @@ impl SglangSidecarEngine { disaggregation_mode, bootstrap_host, bootstrap_port, + local_context: args.sidecar_context, + worker_group_id, state: OnceCell::new(), cancel: CancellationToken::new(), }, @@ -220,7 +250,12 @@ impl LLMEngine for SglangSidecarEngine { .runtime_data .insert("sglang_generate".into(), true.into()); } - let kv_event_sources = discover_kv_event_sources(&discovery, &config, &self.endpoint)?; + let kv_event_sources = match &self.local_context { + Some(context) => { + local_kv_event_sources(context, self.worker_group_id.as_deref(), &mut config)? + } + None => discover_kv_event_sources(&discovery, &config, &self.endpoint)?, + }; let connection_count = pool.len(); let kv_event_source_count = kv_event_sources.len(); self.state @@ -613,6 +648,38 @@ fn is_routable_host(host: &str) -> bool { .unwrap_or(true) } +fn local_kv_event_sources( + context: &SidecarContext, + worker_group_id: Option<&str>, + config: &mut EngineConfig, +) -> Result, DynamoError> { + let llm = config + .llm + .as_ref() + .ok_or_else(|| client::invalid_arg("KV sources require LLM registration"))?; + context + .validate_registration(llm.data_parallel_size.unwrap_or(1), llm.kv_cache_block_size) + .map_err(|error| client::invalid_arg(error.to_string()))?; + if let Some(group_id) = worker_group_id { + config + .runtime_data + .insert(WORKER_GROUP_KEY.into(), group_id.into()); + } + config.runtime_data.insert(KV_CONFIG_KEY.into(), serde_json::json!({ + "block_size": llm.kv_cache_block_size, + "local_dp_ranks": context.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), + })); + Ok(context + .kv_event_sources + .iter() + .map(|source| DiscoveredKvEventSource { + endpoint: source.endpoint.clone(), + topic: source.topic.clone(), + dp_rank: source.dp_rank, + }) + .collect()) +} + fn discover_kv_event_sources( discovery: &Discovery, engine_config: &EngineConfig, @@ -1264,4 +1331,47 @@ mod tests { ] ); } + + #[test] + fn local_context_overrides_global_sources_without_shrinking_registration() { + use crate::context::{ + KV_CONFIG_KEY, SidecarContext, WORKER_GROUP_KEY, tests::context_json, + }; + + let discovery = discovery(json!({ + "page_size": 64, "dp_size": 8, "nnodes": 2, + "kv_events": { + "publisher": "zmq", "endpoint_host": "*", + "endpoint_port_base": 5557, "topic": "kv", "block_size": 64, + "dp_size": 8 + } + })); + let mut config = + build_engine_config(&discovery, DisaggregationMode::Aggregated, None, None).unwrap(); + let mut context: SidecarContext = context_json("full").to_string().parse().unwrap(); + let sources = super::local_kv_event_sources(&context, Some("group"), &mut config).unwrap(); + assert_eq!(sources.len(), context.kv_event_sources.len()); + assert_eq!(sources[0].endpoint, context.kv_event_sources[0].endpoint); + assert_eq!(sources[0].dp_rank, context.kv_event_sources[0].dp_rank); + assert_eq!(config.llm.as_ref().unwrap().data_parallel_size, Some(8)); + assert_eq!(config.runtime_data[WORKER_GROUP_KEY], "group"); + assert_eq!(config.runtime_data[KV_CONFIG_KEY]["block_size"], 64); + + // An explicit empty source list must not fall back to all global ranks. + context.kv_event_sources.clear(); + assert!( + super::local_kv_event_sources(&context, Some("group"), &mut config) + .unwrap() + .is_empty() + ); + assert_eq!( + config.runtime_data[KV_CONFIG_KEY]["local_dp_ranks"], + json!([]) + ); + assert_eq!(config.llm.as_ref().unwrap().data_parallel_size, Some(8)); + + // Engines without the new launch contract retain the old guard. + let endpoint = GrpcEndpoint::parse("http://worker.example:30001", "test").unwrap(); + assert!(discover_kv_event_sources(&discovery, &config, &endpoint).is_err()); + } } diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs new file mode 100644 index 000000000000..5c16b0845b9c --- /dev/null +++ b/lib/sidecar/sglang/src/headless.rs @@ -0,0 +1,411 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Node-local KV relaying. No engine, model registration, or request ingress. + +use std::collections::HashMap; +use std::time::Duration; + +use anyhow::{Context, Result, bail, ensure}; +use dynamo_backend_common::{CommonArgs, DynamoError}; +use dynamo_llm::discovery::{RuntimeConfigWatch, runtime_config_watch}; +use dynamo_llm::kv_router::publisher::{KvEventPublisher, KvEventSourceConfig}; +use dynamo_llm::local_model::runtime_config::ModelRuntimeConfig; +use dynamo_runtime::component::Endpoint; +use dynamo_runtime::distributed::{DistributedConfig, DistributedRuntime}; +use dynamo_runtime::{Runtime, logging}; +use serde::Deserialize; +use tokio_util::sync::CancellationToken; + +use crate::args::Args; +use crate::client; +use crate::context::{KV_CONFIG_KEY, SidecarContext, WORKER_GROUP_KEY}; + +pub(crate) struct HeadlessSidecar { + context: SidecarContext, + group_id: String, + common: CommonArgs, + discovery_timeout: Duration, +} + +#[derive(Clone, Debug, Deserialize)] +struct LeaderKvConfig { + block_size: u32, + local_dp_ranks: Vec, +} + +impl HeadlessSidecar { + pub(crate) fn from_args(args: Args) -> Result { + let context = args + .sidecar_context + .expect("dispatcher checked telemetry context"); + let group_id = context + .worker_group_id() + .map_err(|error| client::invalid_arg(error.to_string()))? + .ok_or_else(|| client::invalid_arg("telemetry mode requires a multinode group"))?; + if args.sidecar.common.route_to_encoder || args.sidecar.common.enable_rl { + return Err(client::invalid_arg( + "telemetry mode cannot register encoder or RL request routes", + )); + } + Ok(Self { + context, + group_id, + common: args.sidecar.common, + discovery_timeout: Duration::from_secs(args.leader_discovery_timeout_secs), + }) + } + + pub(crate) fn run(self) -> Result<()> { + logging::init(); + let runtime = Runtime::from_settings()?; + runtime.secondary().block_on(async { + let shutdown = CancellationToken::new(); + let result = async { + // Install signal listeners before discovery so a follower waiting + // for its leader can still shut down immediately. + let mut terminate = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; + let mut interrupt = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())?; + tokio::select! { + result = self.run_inner(runtime.clone(), shutdown.clone()) => result, + _ = terminate.recv() => Ok(()), + _ = interrupt.recv() => Ok(()), + _ = runtime.primary_token().cancelled_owned() => Ok(()), + } + } + .await; + shutdown.cancel(); + runtime.shutdown(); + result + }) + } + + async fn run_inner(&self, runtime: Runtime, shutdown: CancellationToken) -> Result<()> { + let drt = DistributedRuntime::new(runtime, DistributedConfig::from_settings()).await?; + let endpoint = drt + .namespace(&self.common.namespace)? + .component(&self.common.component)? + .endpoint(&self.common.endpoint); + tracing::info!(node_rank = self.context.node_rank, group = %self.group_id, + endpoint = %endpoint.id(), "Waiting for SGLang leader for local KV publishing"); + let mut configs = runtime_config_watch(&endpoint, shutdown.clone()).await?; + let (worker_id, config) = tokio::time::timeout(self.discovery_timeout, + wait_for_leader(&mut configs, &self.group_id, &shutdown)).await + .context("timed out waiting for SGLang leader; check namespace, component, endpoint, and dist_init_addr")??; + let leader = validate_leader(&self.context, &config)?; + let _publishers = start_publishers(&endpoint, &self.context, worker_id, &config, &leader)?; + tracing::info!(node_rank = self.context.node_rank, worker_id, group = %self.group_id, + local_ranks = ?self.context.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), + "SGLang headless sidecar publishing local KV events"); + + loop { + tokio::select! { + _ = shutdown.cancelled() => return Ok(()), + result = configs.changed() => { + result.context("leader discovery watch closed")?; + check_leader_unchanged(&configs.borrow_and_update(), &self.group_id, worker_id, &config)?; + } + } + } + // Publisher Drop cancels local subscriptions; the follower never owns + // the leader's serving registration. Independent sidecar restart is not + // supported by SGLang's managed lifecycle. + } +} + +fn check_leader_unchanged( + configs: &HashMap, + group_id: &str, + worker_id: u64, + original: &ModelRuntimeConfig, +) -> Result<()> { + let (current_id, current) = matching_leader(configs, group_id)? + .context("SGLang leader disappeared; restart the distributed engine instance")?; + ensure!( + current_id == worker_id + && current.data_parallel_start_rank == original.data_parallel_start_rank + && current.data_parallel_size == original.data_parallel_size + && current.runtime_data.get(KV_CONFIG_KEY) == original.runtime_data.get(KV_CONFIG_KEY) + && current.enable_local_indexer == original.enable_local_indexer + && current.kv_state_endpoint == original.kv_state_endpoint, + "SGLang leader identity or KV publishing configuration changed; restart the distributed engine instance" + ); + Ok(()) +} + +fn matching_leader( + configs: &HashMap, + group_id: &str, +) -> Result> { + let mut matches = configs.iter().filter(|(_, config)| { + config + .runtime_data + .get(WORKER_GROUP_KEY) + .and_then(|value| value.as_str()) + == Some(group_id) + }); + let first = matches.next(); + ensure!( + matches.next().is_none(), + "multiple SGLang leaders registered for group {group_id}" + ); + Ok(first.map(|(id, config)| (*id, config.clone()))) +} + +async fn wait_for_leader( + configs: &mut RuntimeConfigWatch, + group_id: &str, + shutdown: &CancellationToken, +) -> Result<(u64, ModelRuntimeConfig)> { + loop { + if let Some(leader) = matching_leader(&configs.borrow_and_update(), group_id)? { + return Ok(leader); + } + tokio::select! { + _ = shutdown.cancelled() => bail!("leader lookup cancelled"), + result = configs.changed() => result.context("leader discovery watch closed")?, + } + } +} + +fn validate_leader( + context: &SidecarContext, + config: &ModelRuntimeConfig, +) -> Result { + ensure!( + config.data_parallel_start_rank == 0, + "SGLang leader must serve the complete global DP range" + ); + let metadata = config + .runtime_data + .get(KV_CONFIG_KEY) + .context("leader does not advertise node-local sidecar KV metadata")?; + let leader: LeaderKvConfig = + serde_json::from_value(metadata.clone()).context("invalid leader KV metadata")?; + context.validate_registration(config.data_parallel_size, Some(leader.block_size))?; + for source in &context.kv_event_sources { + ensure!( + !leader.local_dp_ranks.contains(&source.dp_rank), + "local KV rank {} is already published by the leader", + source.dp_rank + ); + } + Ok(leader) +} + +fn start_publishers( + endpoint: &Endpoint, + context: &SidecarContext, + worker_id: u64, + config: &ModelRuntimeConfig, + leader: &LeaderKvConfig, +) -> Result> { + context + .kv_event_sources + .iter() + .map(|source| { + KvEventPublisher::new_with_local_indexer_and_worker_id_at( + endpoint.clone(), + config.effective_kv_state_endpoint(&endpoint.id()), + Some(worker_id), + leader.block_size, + Some(KvEventSourceConfig::Zmq { + endpoint: source.endpoint.clone(), + topic: source.topic.clone(), + // Same SGLang normalization contract as the full worker relay. + image_token_id: None, + video_token_id: None, + }), + config.enable_local_indexer, + source.dp_rank, + None, + ) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::context::tests::context_json; + use serde_json::json; + use tokio::sync::watch; + + fn leader(group: &str) -> ModelRuntimeConfig { + ModelRuntimeConfig { + data_parallel_start_rank: 0, + data_parallel_size: 8, + runtime_data: HashMap::from([ + (WORKER_GROUP_KEY.into(), json!(group)), + ( + KV_CONFIG_KEY.into(), + json!({"block_size":64,"local_dp_ranks":[0,1,2,3]}), + ), + ]), + ..Default::default() + } + } + + #[test] + fn matches_exact_group_not_arbitrary_worker() { + let mut configs = HashMap::from([(42, leader("A")), (99, leader("B"))]); + assert_eq!(matching_leader(&configs, "A").unwrap().unwrap().0, 42); + assert!(matching_leader(&configs, "C").unwrap().is_none()); + configs.insert(100, leader("A")); + assert!(matching_leader(&configs, "A").is_err()); + } + + #[test] + fn leader_restart_or_topology_change_requires_group_restart() { + let original = leader("A"); + let mut configs = HashMap::from([(42, original.clone())]); + check_leader_unchanged(&configs, "A", 42, &original).unwrap(); + // Changes to unrelated workers must not interrupt this group's relay. + configs.insert(99, leader("B")); + check_leader_unchanged(&configs, "A", 42, &original).unwrap(); + configs.remove(&42); + assert!(check_leader_unchanged(&configs, "A", 42, &original).is_err()); + configs.insert(43, original.clone()); + assert!(check_leader_unchanged(&configs, "A", 42, &original).is_err()); + configs.remove(&43); + configs.insert(42, original.clone()); + configs.get_mut(&42).unwrap().enable_local_indexer = !original.enable_local_indexer; + assert!(check_leader_unchanged(&configs, "A", 42, &original).is_err()); + } + + #[test] + fn rejects_incompatible_or_duplicate_rank_ownership() { + let context: SidecarContext = context_json("telemetry").to_string().parse().unwrap(); + let mut config = leader("A"); + validate_leader(&context, &config).unwrap(); + config.data_parallel_size = 4; + assert!(validate_leader(&context, &config).is_err()); + config.data_parallel_size = 8; + config.runtime_data.get_mut(KV_CONFIG_KEY).unwrap()["local_dp_ranks"] = json!([4]); + assert!(validate_leader(&context, &config).is_err()); + } + + #[tokio::test] + async fn follower_can_start_before_leader() { + let (tx, mut rx) = watch::channel(HashMap::new()); + let cancel = CancellationToken::new(); + let lookup = wait_for_leader(&mut rx, "A", &cancel); + let publish = async { + tokio::task::yield_now().await; + tx.send(HashMap::from([(42, leader("A"))])).unwrap(); + }; + let (found, ()) = tokio::join!(lookup, publish); + assert_eq!(found.unwrap().0, 42); + } + + #[tokio::test] + async fn shutdown_interrupts_a_pending_lookup() { + let (_tx, mut rx) = watch::channel(HashMap::new()); + let cancel = CancellationToken::new(); + cancel.cancel(); + assert!(wait_for_leader(&mut rx, "A", &cancel).await.is_err()); + } + + #[tokio::test] + async fn local_zmq_events_keep_global_rank_and_leader_identity_without_serving() { + use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent}; + use dynamo_runtime::discovery::{DiscoveryQuery, EventSourceQuery}; + use dynamo_runtime::transports::event_plane::EventSubscriber; + use futures::SinkExt; + + let runtime = Runtime::from_current().unwrap(); + let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) + .await + .unwrap(); + let endpoint = drt + .namespace("headless-kv-test") + .unwrap() + .component("backend") + .unwrap() + .endpoint("generate"); + let socket_dir = tempfile::tempdir().unwrap(); + let source_address = format!("ipc://{}/kv.sock", socket_dir.path().display()); + let zmq_context = tmq::Context::new(); + let mut source = tmq::publish::publish(&zmq_context) + .set_linger(0) + .bind(&source_address) + .unwrap(); + let mut raw = context_json("telemetry"); + raw["kv_event_sources"][0]["endpoint"] = json!(source_address); + let context: SidecarContext = raw.to_string().parse().unwrap(); + let config = leader("A"); + let metadata = validate_leader(&context, &config).unwrap(); + let mut subscriber = EventSubscriber::for_endpoint(&endpoint, KV_EVENT_SUBJECT) + .await + .unwrap() + .typed::>(); + let publishers = start_publishers(&endpoint, &context, 42, &config, &metadata).unwrap(); + + // Repeat until the real ZMQ subscriptions are connected; no fixed sleep. + let received = tokio::time::timeout(Duration::from_secs(5), async { + let mut ticks = tokio::time::interval(Duration::from_millis(50)); + let mut seq = 0_u64; + loop { + tokio::select! { + batch = subscriber.next() => { + let (envelope, events) = batch.unwrap().unwrap(); + if let Some(event) = events.first() { + break (envelope, event.clone()); + } + } + _ = ticks.tick() => { + seq += 1; + let payload = rmp_serde::to_vec_named(&json!([ + 0.0, [{"type":"BlockRemoved", "block_hashes":[42]}], 4 + ])).unwrap(); + source.send(vec![Vec::new(), seq.to_be_bytes().to_vec(), payload]).await.unwrap(); + } + } + } + }).await.expect("headless relay should forward local events"); + assert_eq!(received.1.worker_id, 42); + assert_eq!(received.1.event.dp_rank, 4); + assert_ne!(received.0.publisher_id, 42); + assert!( + drt.discovery() + .list(DiscoveryQuery::Endpoint { + namespace: "headless-kv-test".into(), + component: "backend".into(), + endpoint: "generate".into(), + }) + .await + .unwrap() + .is_empty(), + "telemetry must not register an inference worker" + ); + let source_query = DiscoveryQuery::EventSources(EventSourceQuery::endpoint_topic( + endpoint.id(), + KV_EVENT_SUBJECT, + )); + assert_eq!( + drt.discovery() + .list(source_query.clone()) + .await + .unwrap() + .len(), + 1 + ); + drop(publishers); + tokio::time::timeout(Duration::from_secs(5), async { + while !drt + .discovery() + .list(source_query.clone()) + .await + .unwrap() + .is_empty() + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("dropping a relay must unregister only its local source"); + drt.shutdown(); + } +} diff --git a/lib/sidecar/sglang/src/lib.rs b/lib/sidecar/sglang/src/lib.rs index 7e8968fc3207..17d7bd183530 100644 --- a/lib/sidecar/sglang/src/lib.rs +++ b/lib/sidecar/sglang/src/lib.rs @@ -14,8 +14,11 @@ pub mod args; pub mod client; +pub mod context; pub mod engine; +mod headless; mod native_http; +mod startup; /// Generated SGLang gRPC types, temporarily exposed for the Mocker server /// until SGLang publishes its upstream protocol package. @@ -24,3 +27,4 @@ pub mod proto; mod protocol; pub use engine::SglangSidecarEngine; +pub use startup::SglangSidecar; diff --git a/lib/sidecar/sglang/src/main.rs b/lib/sidecar/sglang/src/main.rs index 6aa72b6f5063..e94b16b01e29 100644 --- a/lib/sidecar/sglang/src/main.rs +++ b/lib/sidecar/sglang/src/main.rs @@ -1,9 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::sync::Arc; - fn main() -> anyhow::Result<()> { - let (engine, config) = dynamo_sglang_sidecar::SglangSidecarEngine::from_args(None)?; - dynamo_backend_common::run(Arc::new(engine), config) + let sidecar = + match dynamo_sglang_sidecar::SglangSidecar::try_from_args(std::env::args().collect()) { + Ok(sidecar) => sidecar, + Err(dynamo_sidecar_common::SidecarStartupError::Cli(error)) => error.exit(), + Err(error) => return Err(error.into()), + }; + sidecar.run() } diff --git a/lib/sidecar/sglang/src/startup.rs b/lib/sidecar/sglang/src/startup.rs new file mode 100644 index 000000000000..2483be7a202f --- /dev/null +++ b/lib/sidecar/sglang/src/startup.rs @@ -0,0 +1,74 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +use clap::Parser; +use dynamo_backend_common::WorkerConfig; +use dynamo_sidecar_common::SidecarStartupError; + +use crate::args::Args; +use crate::context::SidecarMode; +use crate::engine::SglangSidecarEngine; +use crate::headless::HeadlessSidecar; + +/// Common startup dispatcher for the managed Python module and Rust executable. +pub struct SglangSidecar(Mode); + +enum Mode { + Full(Box<(SglangSidecarEngine, WorkerConfig)>), + Telemetry(Box), +} + +impl SglangSidecar { + pub fn try_from_args(argv: Vec) -> Result { + let args = Args::try_parse_from(argv)?; + let mode = if args + .sidecar_context + .as_ref() + .is_some_and(|context| context.mode == SidecarMode::Telemetry) + { + Mode::Telemetry(Box::new(HeadlessSidecar::from_args(args)?)) + } else { + Mode::Full(Box::new(SglangSidecarEngine::from_parsed(args)?)) + }; + Ok(Self(mode)) + } + + pub fn run(self) -> anyhow::Result<()> { + match self.0 { + Mode::Full(full) => { + let (engine, config) = *full; + dynamo_backend_common::run(Arc::new(engine), config) + } + Mode::Telemetry(headless) => headless.run(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::context::tests::context_json; + + #[test] + fn telemetry_needs_no_grpc_and_does_not_parse_an_inherited_endpoint() { + let sidecar = SglangSidecar::try_from_args(vec![ + "sidecar".into(), + "--sidecar-context".into(), + context_json("telemetry").to_string(), + "--grpc-endpoint".into(), + "not-a-grpc-address".into(), + ]) + .unwrap(); + assert!(matches!(sidecar.0, Mode::Telemetry(_))); + } + + #[test] + fn help_retains_structured_cli_exit() { + let result = SglangSidecar::try_from_args(vec!["sidecar".into(), "--help".into()]); + assert!( + matches!(result, Err(SidecarStartupError::Cli(error)) if error.kind() == clap::error::ErrorKind::DisplayHelp) + ); + } +} diff --git a/lib/sidecar/sglang/tests/executable.rs b/lib/sidecar/sglang/tests/executable.rs index d2530990194c..7f687a8374e8 100644 --- a/lib/sidecar/sglang/tests/executable.rs +++ b/lib/sidecar/sglang/tests/executable.rs @@ -27,7 +27,22 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { "DYN_SIDECAR_GRPC_RETRY_INTERVAL_SECS", "--grpc-startup-deadline-secs", "DYN_SIDECAR_GRPC_STARTUP_DEADLINE_SECS", + "--sidecar-context", + "SGLANG_SIDECAR_CONTEXT", + "--leader-discovery-timeout-secs", ] { assert!(stdout.contains(expected), "help omits {expected}"); } } + +#[test] +fn invalid_launch_context_fails_before_grpc_or_discovery() { + let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) + .env("SGLANG_SIDECAR_CONTEXT", "not-json") + .env_remove("SGLANG_GRPC_ENDPOINT") + .env_remove("DYN_SIDECAR_GRPC_ENDPOINT") + .output() + .unwrap(); + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr).contains("invalid SGLANG_SIDECAR_CONTEXT")); +} From 7dd5ea76100e8459b2d914329f061babaa8e179c Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 16 Sep 2026 11:41:47 -0700 Subject: [PATCH 02/17] refactor(sglang): launch sidecars directly from their context mode Signed-off-by: jain-ria --- lib/bindings/python/rust/backend.rs | 11 ++--- lib/sidecar/sglang/src/lib.rs | 63 +++++++++++++++++++++++- lib/sidecar/sglang/src/main.rs | 12 ++--- lib/sidecar/sglang/src/startup.rs | 74 ----------------------------- 4 files changed, 71 insertions(+), 89 deletions(-) delete mode 100644 lib/sidecar/sglang/src/startup.rs diff --git a/lib/bindings/python/rust/backend.rs b/lib/bindings/python/rust/backend.rs index 902be61b10aa..294e0e890db9 100644 --- a/lib/bindings/python/rust/backend.rs +++ b/lib/bindings/python/rust/backend.rs @@ -101,12 +101,11 @@ fn sglang_sidecar_argv(argv: Vec) -> Vec { #[pyo3(signature = (argv=None))] fn _run_sglang_sidecar(py: Python<'_>, argv: Option>) -> PyResult<()> { let cli_argv = sglang_sidecar_argv(argv.unwrap_or_default()); - let sidecar = py - .allow_threads(move || dynamo_sglang_sidecar::SglangSidecar::try_from_args(cli_argv)) - .map_err(sidecar_startup_to_pyerr)?; - - py.allow_threads(move || sidecar.run()) - .map_err(|err| pyo3::exceptions::PyRuntimeError::new_err(err.to_string())) + py.allow_threads(move || dynamo_sglang_sidecar::run(cli_argv)) + .map_err(|error| match error.downcast::() { + Ok(error) => sidecar_startup_to_pyerr(error), + Err(error) => pyo3::exceptions::PyRuntimeError::new_err(error.to_string()), + }) } const VLLM_SIDECAR_PROGRAM_NAME: &str = "dynamo-vllm-sidecar"; diff --git a/lib/sidecar/sglang/src/lib.rs b/lib/sidecar/sglang/src/lib.rs index 17d7bd183530..7464b62767d4 100644 --- a/lib/sidecar/sglang/src/lib.rs +++ b/lib/sidecar/sglang/src/lib.rs @@ -12,13 +12,21 @@ //! The crate never depends on `sglang` or any engine crate — only //! `dynamo-backend-common`, `tonic`/`prost`, `clap`, and tokio. +use std::sync::Arc; + +use clap::Parser; +use dynamo_sidecar_common::SidecarStartupError; + +use args::Args; +use context::SidecarMode; +use headless::HeadlessSidecar; + pub mod args; pub mod client; pub mod context; pub mod engine; mod headless; mod native_http; -mod startup; /// Generated SGLang gRPC types, temporarily exposed for the Mocker server /// until SGLang publishes its upstream protocol package. @@ -27,4 +35,55 @@ pub mod proto; mod protocol; pub use engine::SglangSidecarEngine; -pub use startup::SglangSidecar; + +/// Parse and run the sidecar for both the Python launcher and Rust executable. +/// Startup errors retain their type so callers can preserve CLI exit codes and +/// distinguish invalid configuration from runtime failures. +pub fn run(argv: Vec) -> anyhow::Result<()> { + let args = Args::try_parse_from(argv).map_err(SidecarStartupError::from)?; + match args + .sidecar_context + .as_ref() + .map_or(SidecarMode::Full, |context| context.mode) + { + SidecarMode::Full => { + let (engine, config) = + SglangSidecarEngine::from_parsed(args).map_err(SidecarStartupError::from)?; + dynamo_backend_common::run(Arc::new(engine), config) + } + SidecarMode::Telemetry => HeadlessSidecar::from_args(args) + .map_err(SidecarStartupError::from)? + .run(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::context::tests::context_json; + + #[test] + fn telemetry_uses_headless_validation_before_grpc() { + let error = run(vec![ + "sidecar".into(), + "--sidecar-context".into(), + context_json("telemetry").to_string(), + "--grpc-endpoint".into(), + "not-a-grpc-address".into(), + // Stop at headless validation, before starting a runtime. + "--route-to-encoder".into(), + ]) + .unwrap_err(); + let error = error.downcast::().unwrap(); + assert!(matches!(error, SidecarStartupError::Dynamo(ref error) + if error.to_string().contains("telemetry mode cannot register encoder or RL request routes"))); + } + + #[test] + fn help_retains_structured_cli_exit() { + let error = run(vec!["sidecar".into(), "--help".into()]).unwrap_err(); + let error = error.downcast::().unwrap(); + assert!(matches!(error, SidecarStartupError::Cli(error) + if error.kind() == clap::error::ErrorKind::DisplayHelp)); + } +} diff --git a/lib/sidecar/sglang/src/main.rs b/lib/sidecar/sglang/src/main.rs index e94b16b01e29..c0efb9a0c9b7 100644 --- a/lib/sidecar/sglang/src/main.rs +++ b/lib/sidecar/sglang/src/main.rs @@ -2,11 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 fn main() -> anyhow::Result<()> { - let sidecar = - match dynamo_sglang_sidecar::SglangSidecar::try_from_args(std::env::args().collect()) { - Ok(sidecar) => sidecar, - Err(dynamo_sidecar_common::SidecarStartupError::Cli(error)) => error.exit(), - Err(error) => return Err(error.into()), - }; - sidecar.run() + dynamo_sglang_sidecar::run(std::env::args().collect()).inspect_err(|error| { + if let Some(dynamo_sidecar_common::SidecarStartupError::Cli(cli)) = error.downcast_ref() { + cli.exit(); + } + }) } diff --git a/lib/sidecar/sglang/src/startup.rs b/lib/sidecar/sglang/src/startup.rs deleted file mode 100644 index 2483be7a202f..000000000000 --- a/lib/sidecar/sglang/src/startup.rs +++ /dev/null @@ -1,74 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::sync::Arc; - -use clap::Parser; -use dynamo_backend_common::WorkerConfig; -use dynamo_sidecar_common::SidecarStartupError; - -use crate::args::Args; -use crate::context::SidecarMode; -use crate::engine::SglangSidecarEngine; -use crate::headless::HeadlessSidecar; - -/// Common startup dispatcher for the managed Python module and Rust executable. -pub struct SglangSidecar(Mode); - -enum Mode { - Full(Box<(SglangSidecarEngine, WorkerConfig)>), - Telemetry(Box), -} - -impl SglangSidecar { - pub fn try_from_args(argv: Vec) -> Result { - let args = Args::try_parse_from(argv)?; - let mode = if args - .sidecar_context - .as_ref() - .is_some_and(|context| context.mode == SidecarMode::Telemetry) - { - Mode::Telemetry(Box::new(HeadlessSidecar::from_args(args)?)) - } else { - Mode::Full(Box::new(SglangSidecarEngine::from_parsed(args)?)) - }; - Ok(Self(mode)) - } - - pub fn run(self) -> anyhow::Result<()> { - match self.0 { - Mode::Full(full) => { - let (engine, config) = *full; - dynamo_backend_common::run(Arc::new(engine), config) - } - Mode::Telemetry(headless) => headless.run(), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::context::tests::context_json; - - #[test] - fn telemetry_needs_no_grpc_and_does_not_parse_an_inherited_endpoint() { - let sidecar = SglangSidecar::try_from_args(vec![ - "sidecar".into(), - "--sidecar-context".into(), - context_json("telemetry").to_string(), - "--grpc-endpoint".into(), - "not-a-grpc-address".into(), - ]) - .unwrap(); - assert!(matches!(sidecar.0, Mode::Telemetry(_))); - } - - #[test] - fn help_retains_structured_cli_exit() { - let result = SglangSidecar::try_from_args(vec!["sidecar".into(), "--help".into()]); - assert!( - matches!(result, Err(SidecarStartupError::Cli(error)) if error.kind() == clap::error::ErrorKind::DisplayHelp) - ); - } -} From 9e46ddba6b87c9dde9e889e6fb48496f5e97be2b Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 16 Sep 2026 14:52:48 -0700 Subject: [PATCH 03/17] refactor(sglang): configure telemetry-only sidecars over gRPC Signed-off-by: jain-ria --- lib/sidecar/sglang/README.md | 100 ++++-- .../sglang/launch/multinode_kv_router.sh | 116 ++++++ lib/sidecar/sglang/src/args.rs | 29 +- lib/sidecar/sglang/src/client.rs | 327 ++++++++++++++++- lib/sidecar/sglang/src/context.rs | 281 --------------- lib/sidecar/sglang/src/engine.rs | 140 ++++---- lib/sidecar/sglang/src/headless.rs | 331 ++++++++++++++++-- lib/sidecar/sglang/src/lib.rs | 25 +- lib/sidecar/sglang/tests/executable.rs | 25 +- 9 files changed, 899 insertions(+), 475 deletions(-) create mode 100644 lib/sidecar/sglang/launch/multinode_kv_router.sh delete mode 100644 lib/sidecar/sglang/src/context.rs diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index 72643c6eb483..359f93a941fc 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -38,7 +38,7 @@ The sidecar discovers the model and tokenizer paths, served model name, parser d SGLang remains the source of truth for the worker's aggregated, prefill, or decode role. The inherited `--disaggregation-mode` option and `DYN_DISAGGREGATION_MODE` environment variable have no effect in this sidecar. The SGLang sidecar rejects `--route-to-encoder` because its native protocol does not support encoder workers. Disaggregated workers continue to register under their fixed role components; aggregated workers honor `--component` or `DYN_COMPONENT`. -The sidecar opens eight gRPC connections by default. Override the pool size with `--grpc-connections` or `DYN_SIDECAR_GRPC_CONNECTIONS`. +The full sidecar opens eight gRPC connections by default. Override the pool size with `--grpc-connections` or `DYN_SIDECAR_GRPC_CONNECTIONS`. Telemetry-only mode uses one metadata connection. Connection startup uses a 30-second timeout per attempt, a one-second retry and readiness interval, and a 30-minute deadline for establishing the full connection pool. Override them with `--grpc-connect-attempt-timeout-secs`, `--grpc-retry-interval-secs`, and `--grpc-startup-deadline-secs`, or with the corresponding `DYN_SIDECAR_GRPC_*` environment variables. @@ -61,26 +61,36 @@ same mode dispatcher as the standalone executable. ## Node-local KV sidecars (multinode) -This opt-in path requires the version-1 launch context from -[SGLang #39659](https://github.com/sgl-project/sglang/pull/39659). It does not -require follower gRPC servers. Enable `--sidecar-scope local-telemetry` alongside -`--sidecar dynamo.sglang.sidecar` in SGLang. Without a launch context, the -existing gRPC-only startup and multinode KV-discovery guard are unchanged. +This opt-in path requires an SGLang build containing +[SGLang #39659](https://github.com/sgl-project/sglang/pull/39659) +(metadata contract at `c6c49ac1d2721b93becdfc2be7cab46ddf58c0e8`). It exposes +node-local `kv_event_sources` through `GetServerInfo`, including on followers +started with `--grpc-port`. Follower gRPC servers provide metadata only, not +an inference API. -SGLang supplies `SGLANG_SIDECAR_CONTEXT` when it launches each sidecar: +Launch the engines and sidecars separately, with a sidecar colocated with each +node that owns KV publishers: -- Node 0 uses `mode: "full"`: normal request serving and global DP registration, - but subscribes only to the explicitly listed local KV sources. -- A follower with local KV publishers uses `mode: "telemetry"`: no gRPC client, - model registration, or inference request endpoint. It relays its own local - events directly to Dynamo. -- A TP-only follower without a local KV publisher needs no sidecar. SGLang - decides whether to launch one; Dynamo does not infer ownership from node rank. +```bash +# Node 0: normal request serving and global DP registration. +python3 -m dynamo.sglang.sidecar --grpc-endpoint http://127.0.0.1:30001 + +# Follower node: local KV events only; no model or inference registration. +python3 -m dynamo.sglang.sidecar --grpc-endpoint http://127.0.0.1:30001 --telemetry-only +``` -Both modes consume `kv_event_sources` exactly as supplied, including global -`dp_rank`, connectable `endpoint`, `topic`, and logical `block_size`. Dynamo -does not reconstruct rank ranges or ports from `dp_size`. An empty full-mode -source list stays empty. +Both modes read the local engine's gRPC metadata. Telemetry-only mode calls +`GetServerInfo` without model discovery, inference, or health-check RPCs. +SGLang-managed follower sidecar launching is not implemented by this change; +the existing single-node `--sidecar dynamo.sglang.sidecar` path is unchanged. + +Both modes consume `kv_event_sources` exactly as supplied: global `dp_rank`, +connectable `endpoint`, `topic`, and logical `block_size`. Dynamo does not +reconstruct local rank ranges or ports from `dp_size`. An explicit empty list +stays empty; a telemetry-only sidecar rejects it. A TP-only follower without a +local KV publisher should run just the engine, not a sidecar. Older leader +engines without this field retain legacy discovery, including its multinode +DP guard. The leader advertises the existing `sglang_worker_group_id` derived from the shared `dist_init_addr`, plus its KV block size and locally owned ranks in @@ -90,25 +100,57 @@ identity, while its events carry the **leader's worker ID and the source's global DP rank**. This associates cache locations with a routable target; it does not send the events through the leader sidecar. -All sidecars in a group must use the same Dynamo namespace, component, and -endpoint. For PD disaggregation, pass `DYN_COMPONENT=prefill` on every prefill -node and `DYN_COMPONENT=decode` on every decode node (or the corresponding -`--component` sidecar argument). A headless sidecar has no gRPC discovery from -which to infer the PD role. Separate prefill/decode engine groups need distinct -`dist_init_addr` values, as in the in-process integration. +All sidecars in a group must use the same Dynamo namespace and endpoint. +Both modes discover the prefill/decode role from SGLang and use its fixed +component; aggregated workers must also share the same configured component. +Separate prefill/decode engine groups need distinct `dist_init_addr` values, +as in the in-process integration. Followers may start before the leader; discovery is cancellable and bounded by `--leader-discovery-timeout-secs` (default 1800). Missing, ambiguous, or incompatible leader metadata produces an explicit startup error. If a running -follower loses its leader or observes a changed publishing identity/config, -it exits: restart the distributed engine group, not just the sidecar. This -matches SGLang's managed-sidecar failure policy and avoids continuing under a -stale serving-worker identity. +follower loses its leader or local engine, or observes changed publishing +metadata, it exits. The external launcher must coordinate restarting the +distributed engine group and its sidecars, rather than restarting a sidecar +independently with stale router state. This change adds **KV events only**, not forward-pass metrics or scheduler load publishing. It reuses the existing advisory ZMQ relay: it does not add a readiness handshake, lossless startup, or independent sidecar restart/replay. -`replay_endpoint` is accepted in the context but is not consumed by this relay. +`replay_endpoint` in engine metadata is not consumed by this relay. + +### Externally managed multinode example + +[launch/multinode_kv_router.sh](launch/multinode_kv_router.sh) starts one local +engine and one sidecar. Run it once per node, using the same model, topology, +namespace, and shared Dynamo discovery/event services (`ETCD_ENDPOINTS` and +`NATS_SERVER`). Start one Dynamo frontend separately: + +```bash +python3 -m dynamo.frontend --router-mode kv + +# Node 0 (replace 10.0.0.10 with its reachable address). +NODE_RANK=0 DIST_INIT_ADDR=10.0.0.10:5000 \ + bash lib/sidecar/sglang/launch/multinode_kv_router.sh + +# Node 1, in a separate terminal on that host. +NODE_RANK=1 DIST_INIT_ADDR=10.0.0.10:5000 \ + bash lib/sidecar/sglang/launch/multinode_kv_router.sh +``` + +The default is two nodes with one GPU each: global TP=2, attention DP=2. +SGLang multinode DP requires `--enable-dp-attention`; it is not an arbitrary +collection of independent full-model replicas. For two eight-GPU nodes with +four attention-DP ranks per node, set `TP_SIZE=16 DP_SIZE=8` on both hosts. +The example is limited to layouts where every node owns a publisher. + +For PD disaggregation, run a separate group with `ROLE=prefill` and another +with `ROLE=decode`, using distinct rendezvous addresses. Set +`SGLANG_BOOTSTRAP_HOST` to the prefill leader's reachable address on its node 0. +The script passes the role and NIXL transfer configuration to SGLang; sidecars +discover that role over gRPC. See `--help` for per-node port overrides and pass +additional engine/network options as script arguments. This example has not +been validated on a multinode GPU deployment. ## Deploy on Kubernetes (quick start) diff --git a/lib/sidecar/sglang/launch/multinode_kv_router.sh b/lib/sidecar/sglang/launch/multinode_kv_router.sh new file mode 100644 index 000000000000..9f13893d4a79 --- /dev/null +++ b/lib/sidecar/sglang/launch/multinode_kv_router.sh @@ -0,0 +1,116 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# One node of an externally managed SGLang attention-DP group with local KV events. +# Requires SGLang #39659 and this Dynamo sidecar on every node. Start the frontend separately. +# Defaults: two nodes, one GPU per node, global TP=2 / attention DP=2. + +set -e + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" +export DYNAMO_HOME="${DYNAMO_HOME:-$(readlink -f "$SCRIPT_DIR/../../../..")}" +# shellcheck disable=SC1091 # Resolved relative to this script at runtime. +source "$DYNAMO_HOME/examples/common/gpu_utils.sh" +# shellcheck disable=SC1091 # Resolved relative to this script at runtime. +source "$DYNAMO_HOME/examples/common/launch_utils.sh" + +if [[ "${1:-}" == --help || "${1:-}" == -h ]]; then + echo "Usage: NODE_RANK= DIST_INIT_ADDR= bash $0 [SGLang options...]" + echo "Run once per node; use the same Dynamo namespace, discovery and event services." + echo "Required: NODE_RANK, DIST_INIT_ADDR" + echo "Topology: NNODES=2, TP_SIZE=NNODES, DP_SIZE=NNODES (global sizes; PP=1)" + echo " DP_SIZE must be divisible by NNODES so every node owns publishers." + echo "Model: MODEL=Qwen/Qwen3-0.6B, MAX_MODEL_LEN=4096, MAX_CONCURRENT_SEQS=32" + echo "Role: ROLE=aggregated|prefill|decode (default: aggregated)" + echo "PD: SGLANG_BOOTSTRAP_HOST= on prefill node 0" + echo "Ports: SGLANG_HTTP_PORT=30000, SGLANG_GRPC_PORT=30001, SGLANG_KV_EVENT_PORT=5557" + echo " SGLANG_DISAGGREGATION_BOOTSTRAP_PORT=8998, DYN_SYSTEM_PORT=8081" + echo " DYN_HTTP_PORT=8000 is the separately launched frontend's port." + echo "Other: SGLANG_PYTHON=python3, SGLANG_HOST=0.0.0.0, SGLANG_PAGE_SIZE=64" + echo " CUDA_VISIBLE_DEVICES is inherited. Extra arguments go to SGLang only." + exit 0 +fi + +: "${NODE_RANK:?Set NODE_RANK to the local node rank within the engine group}" +: "${DIST_INIT_ADDR:?Set DIST_INIT_ADDR to the reachable engine leader host:port}" +NNODES="${NNODES:-2}" +TP_SIZE="${TP_SIZE:-$NNODES}" +DP_SIZE="${DP_SIZE:-$NNODES}" +for size in "$NNODES" "$TP_SIZE" "$DP_SIZE"; do + if [[ ! "$size" =~ ^[1-9][0-9]*$ ]]; then + echo "NNODES, TP_SIZE and DP_SIZE must be positive integers" >&2 + exit 1 + fi +done +if [[ ! "$NODE_RANK" =~ ^(0|[1-9][0-9]*)$ ]] || (( NODE_RANK >= NNODES )); then + echo "NODE_RANK must be in [0, NNODES)" >&2 + exit 1 +fi +if (( NNODES < 2 || TP_SIZE % DP_SIZE != 0 || DP_SIZE % NNODES != 0 )); then + echo "This example requires NNODES >= 2, TP_SIZE divisible by DP_SIZE, and DP_SIZE divisible by NNODES" >&2 + exit 1 +fi + +MODEL="${MODEL:-Qwen/Qwen3-0.6B}" +ROLE="${ROLE:-aggregated}" +SGLANG_PYTHON="${SGLANG_PYTHON:-python3}" +SGLANG_HOST="${SGLANG_HOST:-0.0.0.0}" +SGLANG_HTTP_PORT="${SGLANG_HTTP_PORT:-30000}" +SGLANG_GRPC_PORT="${SGLANG_GRPC_PORT:-30001}" +SGLANG_KV_EVENT_PORT="${SGLANG_KV_EVENT_PORT:-5557}" +SGLANG_PAGE_SIZE="${SGLANG_PAGE_SIZE:-64}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" +MAX_CONCURRENT_SEQS="${MAX_CONCURRENT_SEQS:-32}" +HTTP_PORT="${DYN_HTTP_PORT:-8000}" +GPU_MEM_ARGS=$(build_sglang_gpu_mem_args) +ENGINE_ROLE_ARGS=() +SIDECAR_ARGS=() +case "$SGLANG_HOST" in + 0.0.0.0) SIDECAR_HOST=127.0.0.1 ;; + ::) SIDECAR_HOST='[::1]' ;; + *:*) SIDECAR_HOST="[$SGLANG_HOST]" ;; + *) SIDECAR_HOST="$SGLANG_HOST" ;; +esac + +case "$ROLE" in + aggregated) ;; + prefill|decode) + ENGINE_ROLE_ARGS=(--disaggregation-mode "$ROLE" --disaggregation-transfer-backend nixl + --disaggregation-bootstrap-port "${SGLANG_DISAGGREGATION_BOOTSTRAP_PORT:-8998}") + if [[ "$ROLE" == prefill && "$NODE_RANK" == 0 ]]; then + : "${SGLANG_BOOTSTRAP_HOST:?Set SGLANG_BOOTSTRAP_HOST to the reachable prefill leader address}" + SIDECAR_ARGS+=(--bootstrap-host "$SGLANG_BOOTSTRAP_HOST") + fi + ;; + *) echo "ROLE must be aggregated, prefill or decode" >&2; exit 1 ;; +esac +if (( NODE_RANK > 0 )); then + SIDECAR_ARGS+=(--telemetry-only) +fi + +print_launch_banner --no-curl "SGLang multinode KV sidecar ($ROLE, node $NODE_RANK)" "$MODEL" "$HTTP_PORT" \ + "Topology: NNODES=$NNODES, TP_SIZE=$TP_SIZE, DP_SIZE=$DP_SIZE (attention DP)" \ + "Rendezvous: $DIST_INIT_ADDR; local gRPC: $SGLANG_GRPC_PORT" \ + "Start the frontend separately. Restart all nodes and sidecars together after failure." + +trap dynamo_exit_trap EXIT + +# Sidecars retry gRPC startup; no sleep or inference health check is needed on followers. +# shellcheck disable=SC2086 # GPU_MEM_ARGS intentionally expands into multiple flags. +"$SGLANG_PYTHON" -m sglang.launch_server \ + --model-path "$MODEL" \ + --host "$SGLANG_HOST" --port "$SGLANG_HTTP_PORT" --grpc-port "$SGLANG_GRPC_PORT" \ + --incremental-streaming-output \ + --nnodes "$NNODES" --node-rank "$NODE_RANK" --dist-init-addr "$DIST_INIT_ADDR" \ + --tp-size "$TP_SIZE" --dp-size "$DP_SIZE" --enable-dp-attention \ + --kv-events-config "{\"publisher\":\"zmq\",\"endpoint\":\"tcp://127.0.0.1:${SGLANG_KV_EVENT_PORT}\",\"topic\":\"\"}" \ + --page-size "$SGLANG_PAGE_SIZE" --context-length "$MAX_MODEL_LEN" \ + --max-running-requests "$MAX_CONCURRENT_SEQS" \ + "${ENGINE_ROLE_ARGS[@]}" $GPU_MEM_ARGS "$@" & + +DYN_SYSTEM_PORT="${DYN_SYSTEM_PORT:-8081}" \ + python3 -m dynamo.sglang.sidecar \ + --grpc-endpoint "http://$SIDECAR_HOST:$SGLANG_GRPC_PORT" "${SIDECAR_ARGS[@]}" & + +wait_any_exit diff --git a/lib/sidecar/sglang/src/args.rs b/lib/sidecar/sglang/src/args.rs index 356002474828..14376db1306a 100644 --- a/lib/sidecar/sglang/src/args.rs +++ b/lib/sidecar/sglang/src/args.rs @@ -3,18 +3,18 @@ //! Command-line arguments for the SGLang sidecar. -use dynamo_backend_common::CommonArgs; -use dynamo_sidecar_common::GrpcTransportArgs; +use dynamo_backend_common::{CommonArgs, DynamoError}; +use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportArgs}; -use crate::context::SidecarContext; +use crate::client; -/// SGLang alone supports a telemetry-only launch without a gRPC endpoint. +/// Both serving and telemetry-only sidecars discover their local engine over gRPC. #[derive(clap::Args, Debug, Clone)] pub struct SglangSidecarArgs { #[command(flatten)] pub common: CommonArgs, - /// Engine gRPC endpoint, required in full mode. Falls back to SGLANG_GRPC_ENDPOINT. + /// Local engine gRPC endpoint. Falls back to SGLANG_GRPC_ENDPOINT. #[arg(long, env = "DYN_SIDECAR_GRPC_ENDPOINT")] pub grpc_endpoint: Option, @@ -22,6 +22,19 @@ pub struct SglangSidecarArgs { pub grpc: GrpcTransportArgs, } +impl SglangSidecarArgs { + pub(crate) fn resolve_grpc_endpoint(&self) -> Result { + let endpoint = self + .grpc_endpoint + .clone() + .or_else(|| std::env::var("SGLANG_GRPC_ENDPOINT").ok()) + .ok_or_else(|| { + client::invalid_arg("sidecar requires --grpc-endpoint or SGLANG_GRPC_ENDPOINT") + })?; + GrpcEndpoint::parse(&endpoint, "--grpc-endpoint") + } +} + /// Parsed sidecar arguments. #[derive(clap::Parser, Debug, Clone)] #[command( @@ -32,9 +45,9 @@ pub struct Args { #[command(flatten)] pub sidecar: SglangSidecarArgs, - /// Versioned node-local context supplied by SGLang's managed sidecar launcher. - #[arg(long, env = "SGLANG_SIDECAR_CONTEXT")] - pub sidecar_context: Option, + /// Relay a follower node's local KV events without registering a request endpoint. + #[arg(long)] + pub telemetry_only: bool, /// Maximum wait for a matching leader registration in telemetry mode. #[arg(long, default_value_t = 1800, value_parser = clap::value_parser!(u64).range(1..))] diff --git a/lib/sidecar/sglang/src/client.rs b/lib/sidecar/sglang/src/client.rs index cf2b830422a8..7424b70a1bc3 100644 --- a/lib/sidecar/sglang/src/client.rs +++ b/lib/sidecar/sglang/src/client.rs @@ -3,12 +3,15 @@ //! Thin sidecar client for SGLang's native `sglang.runtime.v1.SglangService`. +use std::collections::HashSet; use std::future::Future; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; -use dynamo_backend_common::{BackendError, DynamoError, ErrorType}; +use anyhow::{Context, bail, ensure}; +use dynamo_backend_common::{BackendError, DisaggregationMode, DynamoError, ErrorType}; use dynamo_sidecar_common::{DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcEndpoint, GrpcTransportConfig}; +use serde::Deserialize; use serde_json::Value; use tokio::time::{Instant, timeout_at}; use tonic::transport::{Channel, Endpoint}; @@ -18,6 +21,169 @@ use crate::proto::sglang_service_client::SglangServiceClient; pub type Client = SglangServiceClient; +pub(crate) const WORKER_GROUP_KEY: &str = "sglang_worker_group_id"; +pub(crate) const KV_CONFIG_KEY: &str = "sglang_sidecar_kv_events"; + +/// Node-local KV publishers reported by the engine's GetServerInfo RPC. +/// Other server fields (and unused source fields such as replay_endpoint) are +/// ignored so the metadata-only and full servers share the same wire contract. +#[derive(Clone, Debug, Deserialize, PartialEq, Eq)] +pub(crate) struct NodeMetadata { + pub node_rank: u32, + pub nnodes: u32, + pub dp_size: u32, + pub dist_init_addr: Option, + pub kv_event_sources: Vec, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Eq)] +pub(crate) struct LocalKvEventSource { + pub dp_rank: u32, + pub endpoint: String, + pub topic: String, + pub block_size: u32, +} + +impl NodeMetadata { + /// An absent source list permits legacy leader discovery; an explicit empty + /// list is authoritative and must never fall back to all global DP ranks. + pub(crate) fn from_server_info(server_info: &Value) -> anyhow::Result> { + ensure!( + server_info.is_object(), + "GetServerInfo must contain a JSON object" + ); + if server_info.get("kv_event_sources").is_none() { + return Ok(None); + } + let metadata: Self = serde_json::from_value(server_info.clone()) + .context("invalid GetServerInfo node-local KV metadata")?; + metadata.validate()?; + Ok(Some(metadata)) + } + + fn validate(&self) -> anyhow::Result<()> { + ensure!( + self.nnodes > 0 && self.node_rank < self.nnodes, + "invalid GetServerInfo node topology" + ); + ensure!(self.dp_size > 0, "GetServerInfo dp_size must be positive"); + ensure!( + self.nnodes == 1 + || self + .dist_init_addr + .as_ref() + .is_some_and(|addr| !addr.trim().is_empty()), + "multinode sidecars require dist_init_addr for leader matching" + ); + let mut ranks = HashSet::new(); + let mut endpoints = HashSet::new(); + for source in &self.kv_event_sources { + ensure!( + source.dp_rank < self.dp_size, + "KV source rank {} is outside dp_size {}", + source.dp_rank, + self.dp_size + ); + ensure!( + ranks.insert(source.dp_rank), + "duplicate local KV source rank {}", + source.dp_rank + ); + ensure!( + endpoints.insert(&source.endpoint), + "duplicate local KV source endpoint {}", + source.endpoint + ); + ensure!( + source.block_size > 0, + "KV source block_size must be positive" + ); + validate_endpoint(&source.endpoint)?; + } + Ok(()) + } + + pub(crate) fn validate_registration( + &self, + dp_size: u32, + block_size: Option, + ) -> anyhow::Result<()> { + ensure!( + self.dp_size == dp_size, + "GetServerInfo dp_size does not match leader registration" + ); + for source in &self.kv_event_sources { + ensure!( + Some(source.block_size) == block_size, + "KV source block_size does not match leader registration" + ); + } + Ok(()) + } + + pub(crate) fn worker_group_id(&self) -> anyhow::Result> { + if self.nnodes == 1 { + return Ok(None); + } + let raw = self + .dist_init_addr + .as_deref() + .context("missing dist_init_addr")? + .trim(); + let url = if raw.contains("://") { + raw.to_owned() + } else { + format!("tcp://{raw}") + }; + let address = url::Url::parse(&url).context("invalid dist_init_addr")?; + ensure!( + address.scheme() == "tcp", + "dist_init_addr must be a TCP address" + ); + validate_endpoint(&url)?; + // Match the in-process group key using the shared rendezvous address, + // not the local source or gRPC address. + let resolved = address.socket_addrs(|| None)?; + let address = resolved + .first() + .context("dist_init_addr resolved to no addresses")?; + Ok(Some(format!("dist_init:tcp://{address}"))) + } +} + +fn validate_endpoint(endpoint: &str) -> anyhow::Result<()> { + if let Some(path) = endpoint.strip_prefix("ipc://") { + ensure!( + (path.starts_with('/') || path.starts_with('@')) + && path.len() > 1 + && !path.contains('\0'), + "IPC source must have an absolute or abstract socket path" + ); + return Ok(()); + } + let url = url::Url::parse(endpoint).context("invalid KV source endpoint")?; + let host = url + .host_str() + .context("KV source endpoint requires a host")?; + ensure!( + url.scheme() == "tcp" && url.port().is_some_and(|port| port > 0), + "KV source must use tcp://HOST:PORT or ipc://PATH" + ); + ensure!( + !matches!(host, "*" | "0.0.0.0" | "[::]" | "::"), + "KV source endpoint must be dialable, not a wildcard bind address" + ); + if !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + { + bail!("KV source TCP endpoint must contain only host and port"); + } + Ok(()) +} + /// Metadata exposed by SGLang's model/server discovery RPCs. #[derive(Clone, Debug)] pub struct Discovery { @@ -119,6 +285,12 @@ impl Pool { } pub async fn discover(client: &mut Client, deadline: Instant) -> Result { + let server_info = get_server_info(client, deadline).await?; + if json_u32(&server_info, "node_rank").is_some_and(|rank| rank > 0) { + return Err(invalid_arg( + "full sidecar requires node_rank=0; use --telemetry-only for a follower with local KV sources", + )); + } let model = rpc_with_deadline( "GetModelInfo", deadline, @@ -126,13 +298,6 @@ pub async fn discover(client: &mut Client, deadline: Instant) -> Result Result Result { + let server = rpc_with_deadline( + "GetServerInfo", + deadline, + client.get_server_info(pb::GetServerInfoRequest {}), + ) + .await? + .into_inner(); + parse_json_object("GetServerInfo.json_info", &server.json_info) } pub async fn health_check(client: &mut Client, deadline: Instant) -> Result { @@ -180,11 +361,10 @@ where fn parse_discovery( model: pb::GetModelInfoResponse, - server: pb::GetServerInfoResponse, + server_info: Value, models: Vec, ) -> Result { let model_info = parse_json_object("GetModelInfo.json_info", &model.json_info)?; - let server_info = parse_json_object("GetServerInfo.json_info", &server.json_info)?; let model_path = if model.model_path.trim().is_empty() { model_info .get("model_path") @@ -237,6 +417,21 @@ fn parse_discovery( }) } +pub(crate) fn discovery_mode(server_info: &Value) -> Result { + match server_info + .get("disaggregation_mode") + .and_then(Value::as_str) + .unwrap_or("null") + { + "null" | "agg" | "aggregated" => Ok(DisaggregationMode::Aggregated), + "prefill" => Ok(DisaggregationMode::Prefill), + "decode" => Ok(DisaggregationMode::Decode), + mode => Err(protocol_error(format!( + "unsupported SGLang disaggregation_mode `{mode}`" + ))), + } +} + fn parse_json_object(label: &str, raw: &str) -> Result { let value: Value = serde_json::from_str(raw) .map_err(|err| protocol_error(format!("invalid {label}: {err}")))?; @@ -315,7 +510,10 @@ mod tests { use tokio::time::Instant; use tonic::transport::Endpoint; - use super::{client_from_channel, discover, json_u32, json_u64, parse_discovery}; + use super::{ + NodeMetadata, client_from_channel, discover, discovery_mode, json_u32, json_u64, + parse_discovery, + }; use crate::proto as pb; #[test] @@ -333,9 +531,7 @@ mod tests { model_path: "model-repo".to_string(), json_info: json!({"tokenizer_path": "tokenizer-repo"}).to_string(), }, - pb::GetServerInfoResponse { - json_info: json!({}).to_string(), - }, + json!({}), Vec::new(), ) .unwrap(); @@ -343,6 +539,107 @@ mod tests { assert_eq!(discovery.tokenizer_path, "tokenizer-repo"); } + fn node_metadata_json() -> serde_json::Value { + json!({ + "node_rank": 1, "nnodes": 2, "dp_size": 8, + "dist_init_addr": "127.0.0.1:2345", + "kv_event_sources": [{ + "dp_rank": 4, "endpoint": "tcp://127.0.0.1:5561", + "topic": "", "block_size": 64 + }] + }) + } + + #[test] + fn parses_node_local_sources_and_ignores_unrelated_server_fields() { + let mut raw = node_metadata_json(); + raw["model_path"] = json!("model-repo"); + // Live-only relaying does not interpret the engine's optional replay field. + raw["kv_event_sources"][0]["replay_endpoint"] = json!("unused"); + let metadata = NodeMetadata::from_server_info(&raw).unwrap().unwrap(); + assert_eq!(metadata.kv_event_sources[0].dp_rank, 4); + assert_eq!( + metadata.worker_group_id().unwrap().as_deref(), + Some("dist_init:tcp://127.0.0.1:2345") + ); + metadata.validate_registration(8, Some(64)).unwrap(); + assert!(metadata.validate_registration(4, Some(64)).is_err()); + assert!(metadata.validate_registration(8, Some(32)).is_err()); + } + + #[test] + fn local_sources_distinguish_absent_empty_and_invalid_metadata() { + assert!( + NodeMetadata::from_server_info(&json!({})) + .unwrap() + .is_none() + ); + let mut raw = node_metadata_json(); + raw["kv_event_sources"] = json!([]); + assert!( + NodeMetadata::from_server_info(&raw) + .unwrap() + .unwrap() + .kv_event_sources + .is_empty() + ); + raw["kv_event_sources"] = serde_json::Value::Null; + assert!(NodeMetadata::from_server_info(&raw).is_err()); + } + + #[test] + fn rejects_duplicate_rank_or_endpoint() { + let mut raw = node_metadata_json(); + let mut source = raw["kv_event_sources"][0].clone(); + source["endpoint"] = json!("tcp://127.0.0.1:5562"); + raw["kv_event_sources"].as_array_mut().unwrap().push(source); + assert!(NodeMetadata::from_server_info(&raw).is_err()); + raw["kv_event_sources"][1]["dp_rank"] = json!(5); + raw["kv_event_sources"][1]["endpoint"] = json!("tcp://127.0.0.1:5561"); + assert!(NodeMetadata::from_server_info(&raw).is_err()); + } + + #[test] + fn source_rank_and_block_size_must_be_valid() { + let mut raw = node_metadata_json(); + raw["kv_event_sources"][0]["dp_rank"] = json!(8); + assert!(NodeMetadata::from_server_info(&raw).is_err()); + raw["kv_event_sources"][0]["dp_rank"] = json!(4); + raw["kv_event_sources"][0]["block_size"] = json!(0); + assert!(NodeMetadata::from_server_info(&raw).is_err()); + } + + #[test] + fn normalizes_ipv6_group_id_and_accepts_bound_ipc_sources() { + let mut raw = node_metadata_json(); + raw["dist_init_addr"] = json!("tcp://[::1]:2345"); + raw["kv_event_sources"][0]["endpoint"] = json!("ipc:///engine/kv-events"); + let metadata = NodeMetadata::from_server_info(&raw).unwrap().unwrap(); + assert_eq!( + metadata.worker_group_id().unwrap().as_deref(), + Some("dist_init:tcp://[::1]:2345") + ); + } + + #[test] + fn discovery_mode_reads_only_server_metadata() { + use dynamo_backend_common::DisaggregationMode; + + assert_eq!( + discovery_mode(&json!({})).unwrap(), + DisaggregationMode::Aggregated + ); + assert_eq!( + discovery_mode(&json!({"disaggregation_mode": "prefill"})).unwrap(), + DisaggregationMode::Prefill + ); + assert_eq!( + discovery_mode(&json!({"disaggregation_mode": "decode"})).unwrap(), + DisaggregationMode::Decode + ); + assert!(discovery_mode(&json!({"disaggregation_mode": "unknown"})).is_err()); + } + #[tokio::test] async fn discovery_deadline_bounds_a_half_open_peer() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/lib/sidecar/sglang/src/context.rs b/lib/sidecar/sglang/src/context.rs deleted file mode 100644 index 93025f03b855..000000000000 --- a/lib/sidecar/sglang/src/context.rs +++ /dev/null @@ -1,281 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! SGLang's versioned, engine-owned managed-sidecar launch contract. - -use std::collections::HashSet; -use std::str::FromStr; - -use anyhow::{Context, Result, bail, ensure}; -use serde::Deserialize; - -pub(crate) const WORKER_GROUP_KEY: &str = "sglang_worker_group_id"; -pub(crate) const KV_CONFIG_KEY: &str = "sglang_sidecar_kv_events"; - -#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] -pub enum SidecarMode { - Full, - Telemetry, -} - -/// Matches sglang.srt.entrypoints.sidecar_context.SidecarContext, version 1. -#[derive(Clone, Debug, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct SidecarContext { - pub version: u32, - pub mode: SidecarMode, - pub node_rank: u32, - pub nnodes: u32, - pub dp_size: u32, - pub dist_init_addr: Option, - pub kv_event_sources: Vec, -} - -#[derive(Clone, Debug, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct KvEventSource { - pub dp_rank: u32, - pub endpoint: String, - pub topic: String, - pub block_size: u32, - /// Accepted for wire compatibility; the existing Dynamo relay is live-only. - pub replay_endpoint: Option, -} - -impl FromStr for SidecarContext { - type Err = String; - - fn from_str(raw: &str) -> Result { - let context: Self = serde_json::from_str(raw) - .map_err(|error| format!("invalid SGLANG_SIDECAR_CONTEXT: {error}"))?; - context.validate().map_err(|error| error.to_string())?; - Ok(context) - } -} - -impl SidecarContext { - fn validate(&self) -> Result<()> { - ensure!( - self.version == 1, - "unsupported SGLANG_SIDECAR_CONTEXT version {}", - self.version - ); - ensure!( - self.nnodes > 0 && self.node_rank < self.nnodes, - "invalid sidecar node topology" - ); - ensure!(self.dp_size > 0, "sidecar dp_size must be positive"); - ensure!( - (self.mode == SidecarMode::Full) == (self.node_rank == 0), - "full mode requires node_rank=0; telemetry mode requires a follower node" - ); - ensure!( - self.mode != SidecarMode::Telemetry || !self.kv_event_sources.is_empty(), - "telemetry mode requires at least one local KV source" - ); - ensure!( - self.nnodes == 1 - || self - .dist_init_addr - .as_ref() - .is_some_and(|addr| !addr.trim().is_empty()), - "multinode sidecars require dist_init_addr for leader matching" - ); - let mut ranks = HashSet::new(); - let mut endpoints = HashSet::new(); - for source in &self.kv_event_sources { - ensure!( - source.dp_rank < self.dp_size, - "KV source rank {} is outside dp_size {}", - source.dp_rank, - self.dp_size - ); - ensure!( - ranks.insert(source.dp_rank), - "duplicate local KV source rank {}", - source.dp_rank - ); - ensure!( - endpoints.insert(&source.endpoint), - "duplicate local KV source endpoint {}", - source.endpoint - ); - ensure!( - source.block_size > 0, - "KV source block_size must be positive" - ); - validate_endpoint(&source.endpoint)?; - if let Some(endpoint) = &source.replay_endpoint { - validate_endpoint(endpoint)?; - } - } - Ok(()) - } - - pub(crate) fn validate_registration( - &self, - dp_size: u32, - block_size: Option, - ) -> Result<()> { - ensure!( - self.dp_size == dp_size, - "sidecar context dp_size does not match leader registration" - ); - for source in &self.kv_event_sources { - ensure!( - Some(source.block_size) == block_size, - "KV source block_size does not match leader registration" - ); - } - Ok(()) - } - - pub(crate) fn worker_group_id(&self) -> Result> { - if self.nnodes == 1 { - return Ok(None); - } - let raw = self - .dist_init_addr - .as_deref() - .context("missing dist_init_addr")? - .trim(); - let url = if raw.contains("://") { - raw.to_owned() - } else { - format!("tcp://{raw}") - }; - let address = url::Url::parse(&url).context("invalid dist_init_addr")?; - ensure!( - address.scheme() == "tcp", - "dist_init_addr must be a TCP address" - ); - validate_endpoint(&url)?; - // Use the in-process group-key format: resolve the shared rendezvous - // address, not a per-node source or gRPC address. - let resolved = address.socket_addrs(|| None)?; - let address = resolved - .first() - .context("dist_init_addr resolved to no addresses")?; - Ok(Some(format!("dist_init:tcp://{address}"))) - } -} - -fn validate_endpoint(endpoint: &str) -> Result<()> { - if let Some(path) = endpoint.strip_prefix("ipc://") { - ensure!( - (path.starts_with('/') || path.starts_with('@')) - && path.len() > 1 - && !path.contains('\0'), - "IPC source must have an absolute or abstract socket path" - ); - return Ok(()); - } - let url = url::Url::parse(endpoint).context("invalid KV source endpoint")?; - let host = url - .host_str() - .context("KV source endpoint requires a host")?; - ensure!( - url.scheme() == "tcp" && url.port().is_some_and(|port| port > 0), - "KV source must use tcp://HOST:PORT or ipc://PATH" - ); - ensure!( - !matches!(host, "*" | "0.0.0.0" | "[::]" | "::"), - "KV source endpoint must be dialable, not a wildcard bind address" - ); - if !url.username().is_empty() - || url.password().is_some() - || url.query().is_some() - || url.fragment().is_some() - || !matches!(url.path(), "" | "/") - { - bail!("KV source TCP endpoint must contain only host and port"); - } - Ok(()) -} - -#[cfg(test)] -pub(crate) mod tests { - use super::*; - use serde_json::{Value, json}; - - pub(crate) fn context_json(mode: &str) -> Value { - json!({"version":1,"mode":mode,"node_rank":if mode == "full" {0} else {1}, - "nnodes":2,"dp_size":8,"dist_init_addr":"tcp://127.0.0.1:2345", - "kv_event_sources":[{"dp_rank":4,"endpoint":"tcp://127.0.0.1:5561", - "topic":"","block_size":64,"replay_endpoint":null}]}) - } - - #[test] - fn accepts_upstream_contract_and_global_not_local_rank() { - let context: SidecarContext = context_json("telemetry").to_string().parse().unwrap(); - assert_eq!(context.kv_event_sources[0].dp_rank, 4); - assert_eq!( - context.worker_group_id().unwrap().as_deref(), - Some("dist_init:tcp://127.0.0.1:2345") - ); - context.validate_registration(8, Some(64)).unwrap(); - assert!(context.validate_registration(4, Some(64)).is_err()); - assert!(context.validate_registration(8, Some(32)).is_err()); - } - - #[test] - fn validates_context_before_runtime_or_grpc_startup() { - for (field, value) in [ - ("version", json!(2)), - ("mode", json!("unknown")), - ("node_rank", json!(0)), - ("node_rank", json!(2)), - ("dp_size", json!(0)), - ("dist_init_addr", Value::Null), - ("kv_event_sources", json!([])), - ] { - let mut raw = context_json("telemetry"); - raw[field] = value; - assert!(raw.to_string().parse::().is_err(), "{raw}"); - } - } - - #[test] - fn rejects_conflicting_or_undialable_sources() { - for (field, value) in [ - ("dp_rank", json!(8)), - ("dp_rank", json!(-1)), - ("block_size", json!(0)), - ("endpoint", json!("tcp://*:5561")), - ("endpoint", json!("inproc://events")), - ("endpoint", json!("ipc://relative")), - ] { - let mut raw = context_json("telemetry"); - raw["kv_event_sources"][0][field] = value; - assert!(raw.to_string().parse::().is_err(), "{raw}"); - } - let mut raw = context_json("telemetry"); - let source = raw["kv_event_sources"][0].clone(); - raw["kv_event_sources"].as_array_mut().unwrap().push(source); - assert!(raw.to_string().parse::().is_err()); - raw["kv_event_sources"][1]["dp_rank"] = json!(5); - assert!(raw.to_string().parse::().is_err()); - } - - #[test] - fn leader_may_have_no_local_publishers() { - let mut raw = context_json("full"); - raw["kv_event_sources"] = json!([]); - let context: SidecarContext = raw.to_string().parse().unwrap(); - context.validate_registration(8, Some(64)).unwrap(); - assert!(context.kv_event_sources.is_empty()); - } - - #[test] - fn ipv6_and_ipc_are_supported() { - let mut raw = context_json("telemetry"); - raw["dist_init_addr"] = json!("tcp://[::1]:2345"); - raw["kv_event_sources"][0]["endpoint"] = json!("ipc:///tmp/sglang-kv-test.sock"); - let context: SidecarContext = raw.to_string().parse().unwrap(); - assert_eq!( - context.worker_group_id().unwrap().as_deref(), - Some("dist_init:tcp://[::1]:2345") - ); - } -} diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index c07bb822478a..0d7d848eb510 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -21,8 +21,7 @@ use tokio::time::Instant; use tokio_util::sync::CancellationToken; use crate::args::Args; -use crate::client::{self, Client, Discovery, Pool}; -use crate::context::{KV_CONFIG_KEY, SidecarContext, SidecarMode, WORKER_GROUP_KEY}; +use crate::client::{self, Client, Discovery, KV_CONFIG_KEY, NodeMetadata, Pool, WORKER_GROUP_KEY}; use crate::native_http::{self, NativeHttp}; use crate::proto as pb; use crate::protocol::{ @@ -36,8 +35,6 @@ pub struct SglangSidecarEngine { disaggregation_mode: DisaggregationMode, bootstrap_host: Option, bootstrap_port: Option, - local_context: Option, - worker_group_id: Option, state: OnceCell, cancel: CancellationToken, } @@ -73,13 +70,9 @@ impl SglangSidecarEngine { } pub(crate) fn from_parsed(args: Args) -> Result<(Self, WorkerConfig), DynamoError> { - if args - .sidecar_context - .as_ref() - .is_some_and(|ctx| ctx.mode != SidecarMode::Full) - { + if args.telemetry_only { return Err(client::invalid_arg( - "telemetry context requires the headless sidecar entry point", + "--telemetry-only requires the headless sidecar entry point", )); } if args.sidecar.common.route_to_encoder { @@ -88,26 +81,10 @@ impl SglangSidecarEngine { )); } - let endpoint = args - .sidecar - .grpc_endpoint - .or_else(|| std::env::var("SGLANG_GRPC_ENDPOINT").ok()) - .ok_or_else(|| { - client::invalid_arg( - "full sidecar mode requires --grpc-endpoint or SGLANG_GRPC_ENDPOINT", - ) - })?; - let endpoint = GrpcEndpoint::parse(&endpoint, "--grpc-endpoint")?; - let worker_group_id = args - .sidecar_context - .as_ref() - .map(SidecarContext::worker_group_id) - .transpose() - .map_err(|error| client::invalid_arg(error.to_string()))? - .flatten(); + let endpoint = args.sidecar.resolve_grpc_endpoint()?; let transport = args.sidecar.grpc.config(); let discovery = bootstrap_discover(&endpoint, &transport)?; - let disaggregation_mode = discovery_mode(&discovery)?; + let disaggregation_mode = client::discovery_mode(&discovery.server_info)?; let bootstrap_host = if disaggregation_mode.is_prefill() { resolve_bootstrap_host( args.bootstrap_host.as_deref(), @@ -163,8 +140,6 @@ impl SglangSidecarEngine { disaggregation_mode, bootstrap_host, bootstrap_port, - local_context: args.sidecar_context, - worker_group_id, state: OnceCell::new(), cancel: CancellationToken::new(), }, @@ -209,7 +184,7 @@ impl LLMEngine for SglangSidecarEngine { let mut control = pool.control_client(); self.await_ready(&mut control, deadline).await?; let discovery = client::discover(&mut control, deadline).await?; - let observed_mode = discovery_mode(&discovery)?; + let observed_mode = client::discovery_mode(&discovery.server_info)?; if observed_mode != self.disaggregation_mode { return Err(client::invalid_arg(format!( "SGLang role changed since bootstrap: registered as {:?}, now reports {:?}", @@ -250,12 +225,7 @@ impl LLMEngine for SglangSidecarEngine { .runtime_data .insert("sglang_generate".into(), true.into()); } - let kv_event_sources = match &self.local_context { - Some(context) => { - local_kv_event_sources(context, self.worker_group_id.as_deref(), &mut config)? - } - None => discover_kv_event_sources(&discovery, &config, &self.endpoint)?, - }; + let kv_event_sources = discover_kv_event_sources(&discovery, &mut config, &self.endpoint)?; let connection_count = pool.len(); let kv_event_source_count = kv_event_sources.len(); self.state @@ -529,22 +499,6 @@ fn bootstrap_discover( }) } -fn discovery_mode(discovery: &Discovery) -> Result { - match discovery - .server_info - .get("disaggregation_mode") - .and_then(Value::as_str) - .unwrap_or("null") - { - "null" | "agg" | "aggregated" => Ok(DisaggregationMode::Aggregated), - "prefill" => Ok(DisaggregationMode::Prefill), - "decode" => Ok(DisaggregationMode::Decode), - mode => Err(client::protocol_error(format!( - "unsupported SGLang disaggregation_mode `{mode}`" - ))), - } -} - fn discovery_string(value: &Value, key: &str) -> Option { value .get(key) @@ -649,27 +603,29 @@ fn is_routable_host(host: &str) -> bool { } fn local_kv_event_sources( - context: &SidecarContext, - worker_group_id: Option<&str>, + metadata: &NodeMetadata, config: &mut EngineConfig, ) -> Result, DynamoError> { let llm = config .llm .as_ref() .ok_or_else(|| client::invalid_arg("KV sources require LLM registration"))?; - context + metadata .validate_registration(llm.data_parallel_size.unwrap_or(1), llm.kv_cache_block_size) .map_err(|error| client::invalid_arg(error.to_string()))?; - if let Some(group_id) = worker_group_id { + if let Some(group_id) = metadata + .worker_group_id() + .map_err(|error| client::invalid_arg(error.to_string()))? + { config .runtime_data .insert(WORKER_GROUP_KEY.into(), group_id.into()); } config.runtime_data.insert(KV_CONFIG_KEY.into(), serde_json::json!({ "block_size": llm.kv_cache_block_size, - "local_dp_ranks": context.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), + "local_dp_ranks": metadata.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), })); - Ok(context + Ok(metadata .kv_event_sources .iter() .map(|source| DiscoveredKvEventSource { @@ -682,9 +638,16 @@ fn local_kv_event_sources( fn discover_kv_event_sources( discovery: &Discovery, - engine_config: &EngineConfig, + engine_config: &mut EngineConfig, grpc_endpoint: &GrpcEndpoint, ) -> Result, DynamoError> { + // An explicit list is authoritative, including an empty list. Only older + // engines without this field may use the legacy single-node descriptor. + if let Some(metadata) = NodeMetadata::from_server_info(&discovery.server_info) + .map_err(|error| client::protocol_error(error.to_string()))? + { + return local_kv_event_sources(&metadata, engine_config); + } let Some(descriptor) = discovery.server_info.get("kv_events") else { return Ok(Vec::new()); }; @@ -1309,11 +1272,11 @@ mod tests { "dp_size": 2 } })); - let config = + let mut config = build_engine_config(&discovery, DisaggregationMode::Aggregated, None, None).unwrap(); let endpoint = GrpcEndpoint::parse("http://worker.example:30001", "test").unwrap(); - let sources = discover_kv_event_sources(&discovery, &config, &endpoint).unwrap(); + let sources = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap(); assert_eq!( sources, @@ -1333,13 +1296,18 @@ mod tests { } #[test] - fn local_context_overrides_global_sources_without_shrinking_registration() { - use crate::context::{ - KV_CONFIG_KEY, SidecarContext, WORKER_GROUP_KEY, tests::context_json, - }; - - let discovery = discovery(json!({ - "page_size": 64, "dp_size": 8, "nnodes": 2, + fn local_metadata_overrides_global_sources_without_shrinking_registration() { + use crate::client::{KV_CONFIG_KEY, WORKER_GROUP_KEY}; + + let mut discovery = discovery(json!({ + "page_size": 64, "dp_size": 8, "nnodes": 2, "node_rank": 0, + "dist_init_addr": "127.0.0.1:2345", + "kv_event_sources": [ + {"dp_rank":0,"endpoint":"tcp://127.0.0.1:5557","topic":"kv","block_size":64}, + {"dp_rank":1,"endpoint":"tcp://127.0.0.1:5558","topic":"kv","block_size":64}, + {"dp_rank":2,"endpoint":"tcp://127.0.0.1:5559","topic":"kv","block_size":64}, + {"dp_rank":3,"endpoint":"tcp://127.0.0.1:5560","topic":"kv","block_size":64} + ], "kv_events": { "publisher": "zmq", "endpoint_host": "*", "endpoint_port_base": 5557, "topic": "kv", "block_size": 64, @@ -1348,19 +1316,27 @@ mod tests { })); let mut config = build_engine_config(&discovery, DisaggregationMode::Aggregated, None, None).unwrap(); - let mut context: SidecarContext = context_json("full").to_string().parse().unwrap(); - let sources = super::local_kv_event_sources(&context, Some("group"), &mut config).unwrap(); - assert_eq!(sources.len(), context.kv_event_sources.len()); - assert_eq!(sources[0].endpoint, context.kv_event_sources[0].endpoint); - assert_eq!(sources[0].dp_rank, context.kv_event_sources[0].dp_rank); + let endpoint = GrpcEndpoint::parse("http://127.0.0.1:30001", "test").unwrap(); + let sources = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap(); + assert_eq!( + sources + .iter() + .map(|source| source.dp_rank) + .collect::>(), + vec![0, 1, 2, 3] + ); + assert_eq!(sources[0].endpoint, "tcp://127.0.0.1:5557"); assert_eq!(config.llm.as_ref().unwrap().data_parallel_size, Some(8)); - assert_eq!(config.runtime_data[WORKER_GROUP_KEY], "group"); + assert_eq!( + config.runtime_data[WORKER_GROUP_KEY], + "dist_init:tcp://127.0.0.1:2345" + ); assert_eq!(config.runtime_data[KV_CONFIG_KEY]["block_size"], 64); // An explicit empty source list must not fall back to all global ranks. - context.kv_event_sources.clear(); + discovery.server_info["kv_event_sources"] = json!([]); assert!( - super::local_kv_event_sources(&context, Some("group"), &mut config) + discover_kv_event_sources(&discovery, &mut config, &endpoint) .unwrap() .is_empty() ); @@ -1370,8 +1346,12 @@ mod tests { ); assert_eq!(config.llm.as_ref().unwrap().data_parallel_size, Some(8)); - // Engines without the new launch contract retain the old guard. - let endpoint = GrpcEndpoint::parse("http://worker.example:30001", "test").unwrap(); - assert!(discover_kv_event_sources(&discovery, &config, &endpoint).is_err()); + // Engines without the new metadata field retain the old multinode guard. + discovery + .server_info + .as_object_mut() + .unwrap() + .remove("kv_event_sources"); + assert!(discover_kv_event_sources(&discovery, &mut config, &endpoint).is_err()); } } diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index 5c16b0845b9c..ccf3de864235 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -1,29 +1,34 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Node-local KV relaying. No engine, model registration, or request ingress. +//! Node-local KV relaying using the follower engine's metadata-only gRPC server. +//! No model registration or request ingress. use std::collections::HashMap; use std::time::Duration; use anyhow::{Context, Result, bail, ensure}; -use dynamo_backend_common::{CommonArgs, DynamoError}; +use dynamo_backend_common::{CommonArgs, DisaggregationMode, DynamoError}; use dynamo_llm::discovery::{RuntimeConfigWatch, runtime_config_watch}; use dynamo_llm::kv_router::publisher::{KvEventPublisher, KvEventSourceConfig}; use dynamo_llm::local_model::runtime_config::ModelRuntimeConfig; use dynamo_runtime::component::Endpoint; use dynamo_runtime::distributed::{DistributedConfig, DistributedRuntime}; use dynamo_runtime::{Runtime, logging}; +use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportConfig}; use serde::Deserialize; +use serde_json::Value; +use tokio::time::Instant; use tokio_util::sync::CancellationToken; use crate::args::Args; -use crate::client; -use crate::context::{KV_CONFIG_KEY, SidecarContext, WORKER_GROUP_KEY}; +use crate::client::{self, KV_CONFIG_KEY, NodeMetadata, WORKER_GROUP_KEY}; + +const METADATA_POLL_INTERVAL: Duration = Duration::from_secs(5); pub(crate) struct HeadlessSidecar { - context: SidecarContext, - group_id: String, + grpc_endpoint: GrpcEndpoint, + transport: GrpcTransportConfig, common: CommonArgs, discovery_timeout: Duration, } @@ -36,21 +41,14 @@ struct LeaderKvConfig { impl HeadlessSidecar { pub(crate) fn from_args(args: Args) -> Result { - let context = args - .sidecar_context - .expect("dispatcher checked telemetry context"); - let group_id = context - .worker_group_id() - .map_err(|error| client::invalid_arg(error.to_string()))? - .ok_or_else(|| client::invalid_arg("telemetry mode requires a multinode group"))?; if args.sidecar.common.route_to_encoder || args.sidecar.common.enable_rl { return Err(client::invalid_arg( "telemetry mode cannot register encoder or RL request routes", )); } Ok(Self { - context, - group_id, + grpc_endpoint: args.sidecar.resolve_grpc_endpoint()?, + transport: args.sidecar.grpc.config(), common: args.sidecar.common, discovery_timeout: Duration::from_secs(args.leader_discovery_timeout_secs), }) @@ -83,21 +81,61 @@ impl HeadlessSidecar { } async fn run_inner(&self, runtime: Runtime, shutdown: CancellationToken) -> Result<()> { + // Followers only implement GetServerInfo. Do not use the full engine's + // model discovery, HealthCheck, connection pool, or native HTTP client. + let (mut client, metadata, mode) = self.connect_local_engine().await?; + let group_id = metadata + .worker_group_id()? + .context("telemetry-only mode requires a multinode group")?; + + // Supervise the local engine while waiting for the leader too. Engines + // and sidecars are managed externally and require a coordinated restart + // after a failure; the live KV stream has no replay/recovery here. + tokio::select! { + _ = shutdown.cancelled() => Ok(()), + result = monitor_local_engine(&mut client, &metadata, mode, self.transport.connect_attempt_timeout) => result, + result = self.relay(runtime, shutdown.clone(), &metadata, mode, &group_id) => result, + } + } + + async fn connect_local_engine( + &self, + ) -> Result<(client::Client, NodeMetadata, DisaggregationMode)> { + let deadline = Instant::now() + self.transport.startup_deadline; + let mut client = client::connect(&self.grpc_endpoint, &self.transport, deadline).await?; + let info = client::get_server_info(&mut client, deadline).await?; + let (metadata, mode) = follower_metadata(&info)?; + Ok((client, metadata, mode)) + } + + async fn relay( + &self, + runtime: Runtime, + shutdown: CancellationToken, + metadata: &NodeMetadata, + mode: DisaggregationMode, + group_id: &str, + ) -> Result<()> { let drt = DistributedRuntime::new(runtime, DistributedConfig::from_settings()).await?; + let component = if mode == DisaggregationMode::Aggregated { + &self.common.component + } else { + mode.discovery_component() + }; let endpoint = drt .namespace(&self.common.namespace)? - .component(&self.common.component)? + .component(component)? .endpoint(&self.common.endpoint); - tracing::info!(node_rank = self.context.node_rank, group = %self.group_id, + tracing::info!(node_rank = metadata.node_rank, group = group_id, endpoint = %endpoint.id(), "Waiting for SGLang leader for local KV publishing"); let mut configs = runtime_config_watch(&endpoint, shutdown.clone()).await?; let (worker_id, config) = tokio::time::timeout(self.discovery_timeout, - wait_for_leader(&mut configs, &self.group_id, &shutdown)).await + wait_for_leader(&mut configs, group_id, &shutdown)).await .context("timed out waiting for SGLang leader; check namespace, component, endpoint, and dist_init_addr")??; - let leader = validate_leader(&self.context, &config)?; - let _publishers = start_publishers(&endpoint, &self.context, worker_id, &config, &leader)?; - tracing::info!(node_rank = self.context.node_rank, worker_id, group = %self.group_id, - local_ranks = ?self.context.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), + let leader = validate_leader(metadata, &config)?; + let _publishers = start_publishers(&endpoint, metadata, worker_id, &config, &leader)?; + tracing::info!(node_rank = metadata.node_rank, worker_id, group = group_id, + local_ranks = ?metadata.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), "SGLang headless sidecar publishing local KV events"); loop { @@ -105,16 +143,64 @@ impl HeadlessSidecar { _ = shutdown.cancelled() => return Ok(()), result = configs.changed() => { result.context("leader discovery watch closed")?; - check_leader_unchanged(&configs.borrow_and_update(), &self.group_id, worker_id, &config)?; + check_leader_unchanged(&configs.borrow_and_update(), group_id, worker_id, &config)?; } } } // Publisher Drop cancels local subscriptions; the follower never owns - // the leader's serving registration. Independent sidecar restart is not - // supported by SGLang's managed lifecycle. + // the leader's serving registration. + } +} + +fn follower_metadata(info: &Value) -> Result<(NodeMetadata, DisaggregationMode)> { + let metadata = NodeMetadata::from_server_info(info)? + .context("telemetry-only mode requires GetServerInfo.kv_event_sources; upgrade SGLang")?; + ensure!( + metadata.node_rank > 0 && metadata.nnodes > 1, + "--telemetry-only requires a follower node; run the leader sidecar without this flag" + ); + ensure!( + !metadata.kv_event_sources.is_empty(), + "telemetry-only node has no local KV sources; do not launch a sidecar on a TP-only follower without a publisher" + ); + Ok((metadata, client::discovery_mode(info)?)) +} + +async fn monitor_local_engine( + client: &mut client::Client, + original: &NodeMetadata, + mode: DisaggregationMode, + timeout: Duration, +) -> Result<()> { + let mut ticks = tokio::time::interval_at( + Instant::now() + METADATA_POLL_INTERVAL, + METADATA_POLL_INTERVAL, + ); + ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + ticks.tick().await; + check_local_engine(client, original, mode, timeout).await?; } } +async fn check_local_engine( + client: &mut client::Client, + original: &NodeMetadata, + mode: DisaggregationMode, + timeout: Duration, +) -> Result<()> { + let info = client::get_server_info(client, Instant::now() + timeout) + .await + .context("local SGLang metadata server unavailable; restart the distributed engine instance and its sidecars")?; + let (current, current_mode) = follower_metadata(&info) + .context("local SGLang metadata became invalid; restart the distributed engine instance and its sidecars")?; + ensure!( + ¤t == original && current_mode == mode, + "local SGLang topology or KV source configuration changed; restart the distributed engine instance and its sidecars" + ); + Ok(()) +} + fn check_leader_unchanged( configs: &HashMap, group_id: &str, @@ -170,10 +256,7 @@ async fn wait_for_leader( } } -fn validate_leader( - context: &SidecarContext, - config: &ModelRuntimeConfig, -) -> Result { +fn validate_leader(context: &NodeMetadata, config: &ModelRuntimeConfig) -> Result { ensure!( config.data_parallel_start_rank == 0, "SGLang leader must serve the complete global DP range" @@ -197,7 +280,7 @@ fn validate_leader( fn start_publishers( endpoint: &Endpoint, - context: &SidecarContext, + context: &NodeMetadata, worker_id: u64, config: &ModelRuntimeConfig, leader: &LeaderKvConfig, @@ -229,9 +312,189 @@ fn start_publishers( #[cfg(test)] mod tests { use super::*; - use crate::context::tests::context_json; + use std::convert::Infallible; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{Arc, Mutex}; + use std::task::{Context as TaskContext, Poll}; + + use clap::Parser; use serde_json::json; use tokio::sync::watch; + use tonic::body::Body; + use tonic::codegen::{BoxFuture, Service, http}; + + use crate::proto as pb; + + // One real gRPC method, matching the upstream follower server. Recording + // every URI makes accidental full-engine discovery/health RPCs visible. + #[derive(Clone)] + struct MetadataService { + info: Arc>, + calls: Arc>>, + unavailable: Arc, + } + + impl tonic::server::NamedService for MetadataService { + const NAME: &'static str = "sglang.runtime.v1.SglangService"; + } + + impl tonic::server::UnaryService for MetadataService { + type Response = pb::GetServerInfoResponse; + type Future = BoxFuture, tonic::Status>; + + fn call(&mut self, _: tonic::Request) -> Self::Future { + let json_info = self.info.lock().unwrap().to_string(); + let unavailable = self.unavailable.load(Ordering::Relaxed); + Box::pin(async move { + if unavailable { + return Err(tonic::Status::unavailable("local engine stopped")); + } + Ok(tonic::Response::new(pb::GetServerInfoResponse { + json_info, + })) + }) + } + } + + impl Service> for MetadataService { + type Response = http::Response; + type Error = Infallible; + type Future = BoxFuture; + + fn poll_ready(&mut self, _: &mut TaskContext<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + let path = request.uri().path().to_owned(); + self.calls.lock().unwrap().push(path.clone()); + let service = self.clone(); + Box::pin(async move { + if path != "/sglang.runtime.v1.SglangService/GetServerInfo" { + return Ok(tonic::Status::unimplemented("metadata-only follower").into_http()); + } + let mut grpc = tonic::server::Grpc::new(tonic::codec::ProstCodec::default()); + Ok(grpc.unary(service, request).await) + }) + } + } + + struct MetadataServer { + endpoint: String, + service: MetadataService, + task: tokio::task::JoinHandle<()>, + } + + impl MetadataServer { + async fn start() -> Self { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let service = MetadataService { + info: Arc::new(Mutex::new(follower_info())), + calls: Arc::default(), + unavailable: Arc::default(), + }; + let incoming = futures::stream::unfold(listener, |listener| async move { + let connection = listener.accept().await.map(|(stream, _)| stream); + Some((connection, listener)) + }); + let server = tonic::transport::Server::builder().add_service(service.clone()); + let task = + tokio::spawn(async move { server.serve_with_incoming(incoming).await.unwrap() }); + Self { + endpoint, + service, + task, + } + } + } + + impl Drop for MetadataServer { + fn drop(&mut self) { + self.task.abort(); + } + } + + fn follower_info() -> Value { + json!({"node_rank":1,"nnodes":2,"dp_size":8, + "dist_init_addr":"127.0.0.1:2345", "disaggregation_mode":"prefill", + "kv_event_sources":[{"dp_rank":4,"endpoint":"tcp://127.0.0.1:5561", + "topic":"","block_size":64}]}) + } + + #[test] + fn telemetry_requires_follower_metadata_and_local_sources() { + let mut info = follower_info(); + let (metadata, mode) = follower_metadata(&info).unwrap(); + assert_eq!(metadata.node_rank, 1); + assert_eq!(mode, DisaggregationMode::Prefill); + info["node_rank"] = json!(0); + assert!(follower_metadata(&info).is_err()); + info["node_rank"] = json!(1); + info["kv_event_sources"] = json!([]); + assert!(follower_metadata(&info).is_err()); + info.as_object_mut().unwrap().remove("kv_event_sources"); + assert!(follower_metadata(&info).is_err()); + } + + #[tokio::test] + async fn follower_discovers_and_monitors_using_only_server_info() { + let server = MetadataServer::start().await; + let args = Args::try_parse_from([ + "sidecar", + "--telemetry-only", + "--grpc-endpoint", + &server.endpoint, + "--grpc-startup-deadline-secs", + "5", + ]) + .unwrap(); + let sidecar = HeadlessSidecar::from_args(args).unwrap(); + let (mut client, original, mode) = sidecar.connect_local_engine().await.unwrap(); + assert_eq!(mode, DisaggregationMode::Prefill); + let timeout = Duration::from_secs(1); + // A misplaced full sidecar fails with an actionable mode error, without + // asking the follower for any of the unsupported inference RPCs. + let error = client::discover(&mut client, Instant::now() + timeout) + .await + .unwrap_err(); + assert!(error.to_string().contains("use --telemetry-only")); + check_local_engine(&mut client, &original, mode, timeout) + .await + .unwrap(); + // Unrelated live server fields do not invalidate the KV contract. + server.service.info.lock().unwrap()["uptime"] = json!(42); + check_local_engine(&mut client, &original, mode, timeout) + .await + .unwrap(); + server.service.info.lock().unwrap()["kv_event_sources"][0]["endpoint"] = + json!("tcp://127.0.0.1:5562"); + assert!( + check_local_engine(&mut client, &original, mode, timeout) + .await + .is_err() + ); + *server.service.info.lock().unwrap() = follower_info(); + server.service.info.lock().unwrap()["disaggregation_mode"] = json!("decode"); + assert!( + check_local_engine(&mut client, &original, mode, timeout) + .await + .is_err() + ); + assert_eq!( + server.service.calls.lock().unwrap().as_slice(), + vec!["/sglang.runtime.v1.SglangService/GetServerInfo"; 6] + ); + server.service.unavailable.store(true, Ordering::Relaxed); + let error = check_local_engine(&mut client, &original, mode, timeout) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("local SGLang metadata server unavailable") + ); + } fn leader(group: &str) -> ModelRuntimeConfig { ModelRuntimeConfig { @@ -277,7 +540,7 @@ mod tests { #[test] fn rejects_incompatible_or_duplicate_rank_ownership() { - let context: SidecarContext = context_json("telemetry").to_string().parse().unwrap(); + let (context, _) = follower_metadata(&follower_info()).unwrap(); let mut config = leader("A"); validate_leader(&context, &config).unwrap(); config.data_parallel_size = 4; @@ -332,9 +595,9 @@ mod tests { .set_linger(0) .bind(&source_address) .unwrap(); - let mut raw = context_json("telemetry"); + let mut raw = follower_info(); raw["kv_event_sources"][0]["endpoint"] = json!(source_address); - let context: SidecarContext = raw.to_string().parse().unwrap(); + let (context, _) = follower_metadata(&raw).unwrap(); let config = leader("A"); let metadata = validate_leader(&context, &config).unwrap(); let mut subscriber = EventSubscriber::for_endpoint(&endpoint, KV_EVENT_SUBJECT) diff --git a/lib/sidecar/sglang/src/lib.rs b/lib/sidecar/sglang/src/lib.rs index 7464b62767d4..60ef6be7347a 100644 --- a/lib/sidecar/sglang/src/lib.rs +++ b/lib/sidecar/sglang/src/lib.rs @@ -18,12 +18,10 @@ use clap::Parser; use dynamo_sidecar_common::SidecarStartupError; use args::Args; -use context::SidecarMode; use headless::HeadlessSidecar; pub mod args; pub mod client; -pub mod context; pub mod engine; mod headless; mod native_http; @@ -41,33 +39,26 @@ pub use engine::SglangSidecarEngine; /// distinguish invalid configuration from runtime failures. pub fn run(argv: Vec) -> anyhow::Result<()> { let args = Args::try_parse_from(argv).map_err(SidecarStartupError::from)?; - match args - .sidecar_context - .as_ref() - .map_or(SidecarMode::Full, |context| context.mode) - { - SidecarMode::Full => { - let (engine, config) = - SglangSidecarEngine::from_parsed(args).map_err(SidecarStartupError::from)?; - dynamo_backend_common::run(Arc::new(engine), config) - } - SidecarMode::Telemetry => HeadlessSidecar::from_args(args) + if args.telemetry_only { + HeadlessSidecar::from_args(args) .map_err(SidecarStartupError::from)? - .run(), + .run() + } else { + let (engine, config) = + SglangSidecarEngine::from_parsed(args).map_err(SidecarStartupError::from)?; + dynamo_backend_common::run(Arc::new(engine), config) } } #[cfg(test)] mod tests { use super::*; - use crate::context::tests::context_json; #[test] fn telemetry_uses_headless_validation_before_grpc() { let error = run(vec![ "sidecar".into(), - "--sidecar-context".into(), - context_json("telemetry").to_string(), + "--telemetry-only".into(), "--grpc-endpoint".into(), "not-a-grpc-address".into(), // Stop at headless validation, before starting a runtime. diff --git a/lib/sidecar/sglang/tests/executable.rs b/lib/sidecar/sglang/tests/executable.rs index 7f687a8374e8..7a722fb6b0e7 100644 --- a/lib/sidecar/sglang/tests/executable.rs +++ b/lib/sidecar/sglang/tests/executable.rs @@ -27,8 +27,7 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { "DYN_SIDECAR_GRPC_RETRY_INTERVAL_SECS", "--grpc-startup-deadline-secs", "DYN_SIDECAR_GRPC_STARTUP_DEADLINE_SECS", - "--sidecar-context", - "SGLANG_SIDECAR_CONTEXT", + "--telemetry-only", "--leader-discovery-timeout-secs", ] { assert!(stdout.contains(expected), "help omits {expected}"); @@ -36,13 +35,17 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { } #[test] -fn invalid_launch_context_fails_before_grpc_or_discovery() { - let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) - .env("SGLANG_SIDECAR_CONTEXT", "not-json") - .env_remove("SGLANG_GRPC_ENDPOINT") - .env_remove("DYN_SIDECAR_GRPC_ENDPOINT") - .output() - .unwrap(); - assert!(!output.status.success()); - assert!(String::from_utf8_lossy(&output.stderr).contains("invalid SGLANG_SIDECAR_CONTEXT")); +fn both_modes_require_a_local_grpc_endpoint() { + for mode in [None, Some("--telemetry-only")] { + let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) + .args(mode) + .env_remove("SGLANG_GRPC_ENDPOINT") + .env_remove("DYN_SIDECAR_GRPC_ENDPOINT") + .output() + .unwrap(); + assert!(!output.status.success()); + assert!( + String::from_utf8_lossy(&output.stderr).contains("sidecar requires --grpc-endpoint") + ); + } } From 0d41b47c28b0a092671f73e3db5bd5d4908ee8f5 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 16 Sep 2026 15:18:28 -0700 Subject: [PATCH 04/17] refactor(sglang): reuse shared sidecar arguments Signed-off-by: jain-ria --- lib/sidecar/sglang/src/args.rs | 34 ++------------------------ lib/sidecar/sglang/src/engine.rs | 2 +- lib/sidecar/sglang/src/headless.rs | 2 +- lib/sidecar/sglang/tests/executable.rs | 9 +++---- 4 files changed, 8 insertions(+), 39 deletions(-) diff --git a/lib/sidecar/sglang/src/args.rs b/lib/sidecar/sglang/src/args.rs index 14376db1306a..397c8be15709 100644 --- a/lib/sidecar/sglang/src/args.rs +++ b/lib/sidecar/sglang/src/args.rs @@ -3,37 +3,7 @@ //! Command-line arguments for the SGLang sidecar. -use dynamo_backend_common::{CommonArgs, DynamoError}; -use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportArgs}; - -use crate::client; - -/// Both serving and telemetry-only sidecars discover their local engine over gRPC. -#[derive(clap::Args, Debug, Clone)] -pub struct SglangSidecarArgs { - #[command(flatten)] - pub common: CommonArgs, - - /// Local engine gRPC endpoint. Falls back to SGLANG_GRPC_ENDPOINT. - #[arg(long, env = "DYN_SIDECAR_GRPC_ENDPOINT")] - pub grpc_endpoint: Option, - - #[command(flatten)] - pub grpc: GrpcTransportArgs, -} - -impl SglangSidecarArgs { - pub(crate) fn resolve_grpc_endpoint(&self) -> Result { - let endpoint = self - .grpc_endpoint - .clone() - .or_else(|| std::env::var("SGLANG_GRPC_ENDPOINT").ok()) - .ok_or_else(|| { - client::invalid_arg("sidecar requires --grpc-endpoint or SGLANG_GRPC_ENDPOINT") - })?; - GrpcEndpoint::parse(&endpoint, "--grpc-endpoint") - } -} +use dynamo_sidecar_common::SidecarArgs; /// Parsed sidecar arguments. #[derive(clap::Parser, Debug, Clone)] @@ -43,7 +13,7 @@ impl SglangSidecarArgs { )] pub struct Args { #[command(flatten)] - pub sidecar: SglangSidecarArgs, + pub sidecar: SidecarArgs, /// Relay a follower node's local KV events without registering a request endpoint. #[arg(long)] diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index 0d7d848eb510..b5783176ed2a 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -81,7 +81,7 @@ impl SglangSidecarEngine { )); } - let endpoint = args.sidecar.resolve_grpc_endpoint()?; + let endpoint = args.sidecar.grpc_endpoint; let transport = args.sidecar.grpc.config(); let discovery = bootstrap_discover(&endpoint, &transport)?; let disaggregation_mode = client::discovery_mode(&discovery.server_info)?; diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index ccf3de864235..f64bcf8e4b1f 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -47,7 +47,7 @@ impl HeadlessSidecar { )); } Ok(Self { - grpc_endpoint: args.sidecar.resolve_grpc_endpoint()?, + grpc_endpoint: args.sidecar.grpc_endpoint, transport: args.sidecar.grpc.config(), common: args.sidecar.common, discovery_timeout: Duration::from_secs(args.leader_discovery_timeout_secs), diff --git a/lib/sidecar/sglang/tests/executable.rs b/lib/sidecar/sglang/tests/executable.rs index 7a722fb6b0e7..a00353c84162 100644 --- a/lib/sidecar/sglang/tests/executable.rs +++ b/lib/sidecar/sglang/tests/executable.rs @@ -39,13 +39,12 @@ fn both_modes_require_a_local_grpc_endpoint() { for mode in [None, Some("--telemetry-only")] { let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) .args(mode) - .env_remove("SGLANG_GRPC_ENDPOINT") .env_remove("DYN_SIDECAR_GRPC_ENDPOINT") .output() .unwrap(); - assert!(!output.status.success()); - assert!( - String::from_utf8_lossy(&output.stderr).contains("sidecar requires --grpc-endpoint") - ); + assert_eq!(output.status.code(), Some(2)); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("the following required arguments were not provided")); + assert!(stderr.contains("--grpc-endpoint ")); } } From 35dcc9e0015badd38fe9dc87f9f0b60a8d599540 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 16 Sep 2026 16:57:13 -0700 Subject: [PATCH 05/17] docs(sglang): shorten multinode KV sidecar guidance Signed-off-by: jain-ria --- lib/sidecar/sglang/README.md | 106 +++++------------------------------ 1 file changed, 14 insertions(+), 92 deletions(-) diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index 359f93a941fc..f204c11edcb8 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -59,98 +59,20 @@ private `dynamo._core.backend._run_sglang_sidecar(argv)` binding. The binding prepends the executable name expected by clap, releases the GIL, and runs the same mode dispatcher as the standalone executable. -## Node-local KV sidecars (multinode) - -This opt-in path requires an SGLang build containing -[SGLang #39659](https://github.com/sgl-project/sglang/pull/39659) -(metadata contract at `c6c49ac1d2721b93becdfc2be7cab46ddf58c0e8`). It exposes -node-local `kv_event_sources` through `GetServerInfo`, including on followers -started with `--grpc-port`. Follower gRPC servers provide metadata only, not -an inference API. - -Launch the engines and sidecars separately, with a sidecar colocated with each -node that owns KV publishers: - -```bash -# Node 0: normal request serving and global DP registration. -python3 -m dynamo.sglang.sidecar --grpc-endpoint http://127.0.0.1:30001 - -# Follower node: local KV events only; no model or inference registration. -python3 -m dynamo.sglang.sidecar --grpc-endpoint http://127.0.0.1:30001 --telemetry-only -``` - -Both modes read the local engine's gRPC metadata. Telemetry-only mode calls -`GetServerInfo` without model discovery, inference, or health-check RPCs. -SGLang-managed follower sidecar launching is not implemented by this change; -the existing single-node `--sidecar dynamo.sglang.sidecar` path is unchanged. - -Both modes consume `kv_event_sources` exactly as supplied: global `dp_rank`, -connectable `endpoint`, `topic`, and logical `block_size`. Dynamo does not -reconstruct local rank ranges or ports from `dp_size`. An explicit empty list -stays empty; a telemetry-only sidecar rejects it. A TP-only follower without a -local KV publisher should run just the engine, not a sidecar. Older leader -engines without this field retain legacy discovery, including its multinode -DP guard. - -The leader advertises the existing `sglang_worker_group_id` derived from the -shared `dist_init_addr`, plus its KV block size and locally owned ranks in -runtime metadata. Followers use Dynamo's runtime-config discovery watch to -find that group's serving worker. Each follower retains its own publisher -identity, while its events carry the **leader's worker ID and the source's -global DP rank**. This associates cache locations with a routable target; it -does not send the events through the leader sidecar. - -All sidecars in a group must use the same Dynamo namespace and endpoint. -Both modes discover the prefill/decode role from SGLang and use its fixed -component; aggregated workers must also share the same configured component. -Separate prefill/decode engine groups need distinct `dist_init_addr` values, -as in the in-process integration. - -Followers may start before the leader; discovery is cancellable and bounded -by `--leader-discovery-timeout-secs` (default 1800). Missing, ambiguous, or -incompatible leader metadata produces an explicit startup error. If a running -follower loses its leader or local engine, or observes changed publishing -metadata, it exits. The external launcher must coordinate restarting the -distributed engine group and its sidecars, rather than restarting a sidecar -independently with stale router state. - -This change adds **KV events only**, not forward-pass metrics or scheduler -load publishing. It reuses the existing advisory ZMQ relay: it does not add a -readiness handshake, lossless startup, or independent sidecar restart/replay. -`replay_endpoint` in engine metadata is not consumed by this relay. - -### Externally managed multinode example - -[launch/multinode_kv_router.sh](launch/multinode_kv_router.sh) starts one local -engine and one sidecar. Run it once per node, using the same model, topology, -namespace, and shared Dynamo discovery/event services (`ETCD_ENDPOINTS` and -`NATS_SERVER`). Start one Dynamo frontend separately: - -```bash -python3 -m dynamo.frontend --router-mode kv - -# Node 0 (replace 10.0.0.10 with its reachable address). -NODE_RANK=0 DIST_INIT_ADDR=10.0.0.10:5000 \ - bash lib/sidecar/sglang/launch/multinode_kv_router.sh - -# Node 1, in a separate terminal on that host. -NODE_RANK=1 DIST_INIT_ADDR=10.0.0.10:5000 \ - bash lib/sidecar/sglang/launch/multinode_kv_router.sh -``` - -The default is two nodes with one GPU each: global TP=2, attention DP=2. -SGLang multinode DP requires `--enable-dp-attention`; it is not an arbitrary -collection of independent full-model replicas. For two eight-GPU nodes with -four attention-DP ranks per node, set `TP_SIZE=16 DP_SIZE=8` on both hosts. -The example is limited to layouts where every node owns a publisher. - -For PD disaggregation, run a separate group with `ROLE=prefill` and another -with `ROLE=decode`, using distinct rendezvous addresses. Set -`SGLANG_BOOTSTRAP_HOST` to the prefill leader's reachable address on its node 0. -The script passes the role and NIXL transfer configuration to SGLang; sidecars -discover that role over gRPC. See `--help` for per-node port overrides and pass -additional engine/network options as script arguments. This example has not -been validated on a multinode GPU deployment. +## Multinode DP and KV routing + +For multinode deployments with multiple data-parallel (DP) replicas and KV-aware +routing, run a full sidecar on the leader (`node_rank=0`) and a separate +`--telemetry-only` sidecar on each follower node that owns a DP scheduler/KV +publisher. Each DP rank has its own KV cache; follower sidecars publish those +local KV events to Dynamo without accepting inference requests. A TP-only +follower that holds part of a replica but has no local KV publisher does not +need a sidecar. + +Launch each sidecar separately with `--grpc-endpoint` pointing to its local +SGLang server. This requires an SGLang build exposing node-local +`kv_event_sources` through `GetServerInfo` on leaders and followers started +with `--grpc-port`. ## Deploy on Kubernetes (quick start) From d2061af56e2f07b3290e8663954f3cbb684fdfe2 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 16 Sep 2026 17:32:22 -0700 Subject: [PATCH 06/17] fix(sglang): bind local KV publishers and ready follower sidecars Signed-off-by: jain-ria --- lib/sidecar/sglang/launch/multinode_kv_router.sh | 3 ++- lib/sidecar/sglang/src/headless.rs | 16 ++++++++++++++-- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/lib/sidecar/sglang/launch/multinode_kv_router.sh b/lib/sidecar/sglang/launch/multinode_kv_router.sh index 9f13893d4a79..805c401ec6d7 100644 --- a/lib/sidecar/sglang/launch/multinode_kv_router.sh +++ b/lib/sidecar/sglang/launch/multinode_kv_router.sh @@ -97,6 +97,7 @@ print_launch_banner --no-curl "SGLang multinode KV sidecar ($ROLE, node $NODE_RA trap dynamo_exit_trap EXIT # Sidecars retry gRPC startup; no sleep or inference health check is needed on followers. +# SGLang needs a wildcard bind to advertise a node-local KV-event source. # shellcheck disable=SC2086 # GPU_MEM_ARGS intentionally expands into multiple flags. "$SGLANG_PYTHON" -m sglang.launch_server \ --model-path "$MODEL" \ @@ -104,7 +105,7 @@ trap dynamo_exit_trap EXIT --incremental-streaming-output \ --nnodes "$NNODES" --node-rank "$NODE_RANK" --dist-init-addr "$DIST_INIT_ADDR" \ --tp-size "$TP_SIZE" --dp-size "$DP_SIZE" --enable-dp-attention \ - --kv-events-config "{\"publisher\":\"zmq\",\"endpoint\":\"tcp://127.0.0.1:${SGLANG_KV_EVENT_PORT}\",\"topic\":\"\"}" \ + --kv-events-config "{\"publisher\":\"zmq\",\"endpoint\":\"tcp://*:${SGLANG_KV_EVENT_PORT}\",\"topic\":\"\"}" \ --page-size "$SGLANG_PAGE_SIZE" --context-length "$MAX_MODEL_LEN" \ --max-running-requests "$MAX_CONCURRENT_SEQS" \ "${ENGINE_ROLE_ARGS[@]}" $GPU_MEM_ARGS "$@" & diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index f64bcf8e4b1f..32dc28fe9f9c 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -13,7 +13,9 @@ use dynamo_llm::discovery::{RuntimeConfigWatch, runtime_config_watch}; use dynamo_llm::kv_router::publisher::{KvEventPublisher, KvEventSourceConfig}; use dynamo_llm::local_model::runtime_config::ModelRuntimeConfig; use dynamo_runtime::component::Endpoint; +use dynamo_runtime::config::HealthStatus; use dynamo_runtime::distributed::{DistributedConfig, DistributedRuntime}; +use dynamo_runtime::prelude::DistributedRuntimeProvider; use dynamo_runtime::{Runtime, logging}; use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportConfig}; use serde::Deserialize; @@ -285,7 +287,7 @@ fn start_publishers( config: &ModelRuntimeConfig, leader: &LeaderKvConfig, ) -> Result> { - context + let publishers = context .kv_event_sources .iter() .map(|source| { @@ -306,7 +308,15 @@ fn start_publishers( None, ) }) - .collect() + .collect::>>()?; + // There is no serving worker to mark this process ready. Publish readiness + // only after all local publishers have been created. + endpoint + .drt() + .system_health() + .lock() + .set_health_status(HealthStatus::Ready); + Ok(publishers) } #[cfg(test)] @@ -604,7 +614,9 @@ mod tests { .await .unwrap() .typed::>(); + assert!(!drt.system_health().lock().get_health_status().0); let publishers = start_publishers(&endpoint, &context, 42, &config, &metadata).unwrap(); + assert!(drt.system_health().lock().get_health_status().0); // Repeat until the real ZMQ subscriptions are connected; no fixed sleep. let received = tokio::time::timeout(Duration::from_secs(5), async { From 4d35ab404aa2b90b32d036905efd97bc4d1ab8fc Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 16 Sep 2026 23:40:05 -0700 Subject: [PATCH 07/17] fix(sglang): mark multinode KV router launcher executable Signed-off-by: jain-ria --- lib/sidecar/sglang/launch/multinode_kv_router.sh | 0 1 file changed, 0 insertions(+), 0 deletions(-) mode change 100644 => 100755 lib/sidecar/sglang/launch/multinode_kv_router.sh diff --git a/lib/sidecar/sglang/launch/multinode_kv_router.sh b/lib/sidecar/sglang/launch/multinode_kv_router.sh old mode 100644 new mode 100755 From 12ed73f0277d2c4d9c5a75d738dfcf2c4e101d36 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Thu, 17 Sep 2026 00:40:29 -0700 Subject: [PATCH 08/17] fix(sglang): restrict local KV events and stabilize group discovery Signed-off-by: jain-ria --- lib/sidecar/sglang/README.md | 10 +++++ .../sglang/launch/multinode_kv_router.sh | 4 +- lib/sidecar/sglang/src/client.rs | 37 ++++++++++++++++--- 3 files changed, 43 insertions(+), 8 deletions(-) diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index f204c11edcb8..6cd13fbb4902 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -74,6 +74,16 @@ SGLang server. This requires an SGLang build exposing node-local `kv_event_sources` through `GetServerInfo` on leaders and followers started with `--grpc-port`. +The `launch/multinode_kv_router.sh` example also requires SGLang's explicit +`"bind": true` KV-event option (included in +[SGLang #39659](https://github.com/sgl-project/sglang/pull/39659)). It binds each +rank's publisher to `127.0.0.1`, so the engine and its sidecar must share a +network namespace. KV events can include token IDs and request metadata; +loopback prevents remote access but does not isolate other processes in that +namespace. If overriding this to a non-loopback bind, restrict every publisher +port (`SGLANG_KV_EVENT_PORT + dp_rank`) to authorized consumers using a firewall +or network policy. + ## Deploy on Kubernetes (quick start) `deploy/agg.yaml` runs an aggregated deployment (a frontend plus one worker pod diff --git a/lib/sidecar/sglang/launch/multinode_kv_router.sh b/lib/sidecar/sglang/launch/multinode_kv_router.sh index 805c401ec6d7..6c7787b6699b 100755 --- a/lib/sidecar/sglang/launch/multinode_kv_router.sh +++ b/lib/sidecar/sglang/launch/multinode_kv_router.sh @@ -97,7 +97,7 @@ print_launch_banner --no-curl "SGLang multinode KV sidecar ($ROLE, node $NODE_RA trap dynamo_exit_trap EXIT # Sidecars retry gRPC startup; no sleep or inference health check is needed on followers. -# SGLang needs a wildcard bind to advertise a node-local KV-event source. +# Bind KV events only on loopback; engine and sidecar share a network namespace. # shellcheck disable=SC2086 # GPU_MEM_ARGS intentionally expands into multiple flags. "$SGLANG_PYTHON" -m sglang.launch_server \ --model-path "$MODEL" \ @@ -105,7 +105,7 @@ trap dynamo_exit_trap EXIT --incremental-streaming-output \ --nnodes "$NNODES" --node-rank "$NODE_RANK" --dist-init-addr "$DIST_INIT_ADDR" \ --tp-size "$TP_SIZE" --dp-size "$DP_SIZE" --enable-dp-attention \ - --kv-events-config "{\"publisher\":\"zmq\",\"endpoint\":\"tcp://*:${SGLANG_KV_EVENT_PORT}\",\"topic\":\"\"}" \ + --kv-events-config "{\"publisher\":\"zmq\",\"endpoint\":\"tcp://127.0.0.1:${SGLANG_KV_EVENT_PORT}\",\"bind\":true,\"topic\":\"\"}" \ --page-size "$SGLANG_PAGE_SIZE" --context-length "$MAX_MODEL_LEN" \ --max-running-requests "$MAX_CONCURRENT_SEQS" \ "${ENGINE_ROLE_ARGS[@]}" $GPU_MEM_ARGS "$@" & diff --git a/lib/sidecar/sglang/src/client.rs b/lib/sidecar/sglang/src/client.rs index 7424b70a1bc3..c394961e9354 100644 --- a/lib/sidecar/sglang/src/client.rs +++ b/lib/sidecar/sglang/src/client.rs @@ -5,6 +5,7 @@ use std::collections::HashSet; use std::future::Future; +use std::net::SocketAddr; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; @@ -143,14 +144,19 @@ impl NodeMetadata { validate_endpoint(&url)?; // Match the in-process group key using the shared rendezvous address, // not the local source or gRPC address. - let resolved = address.socket_addrs(|| None)?; - let address = resolved - .first() - .context("dist_init_addr resolved to no addresses")?; - Ok(Some(format!("dist_init:tcp://{address}"))) + worker_group_id_from_addresses(address.socket_addrs(|| None)?).map(Some) } } +fn worker_group_id_from_addresses(mut resolved: Vec) -> anyhow::Result { + // Nodes may receive the same DNS answers in different orders. + resolved.sort_unstable(); + let address = resolved + .first() + .context("dist_init_addr resolved to no addresses")?; + Ok(format!("dist_init:tcp://{address}")) +} + fn validate_endpoint(endpoint: &str) -> anyhow::Result<()> { if let Some(path) = endpoint.strip_prefix("ipc://") { ensure!( @@ -512,7 +518,7 @@ mod tests { use super::{ NodeMetadata, client_from_channel, discover, discovery_mode, json_u32, json_u64, - parse_discovery, + parse_discovery, worker_group_id_from_addresses, }; use crate::proto as pb; @@ -621,6 +627,25 @@ mod tests { ); } + #[test] + fn worker_group_id_is_independent_of_dns_answer_order() { + let addresses = vec![ + "[::1]:2345".parse().unwrap(), + "127.0.0.2:2345".parse().unwrap(), + "127.0.0.1:2345".parse().unwrap(), + ]; + let reversed = addresses.iter().copied().rev().collect(); + let expected = "dist_init:tcp://127.0.0.1:2345"; + assert_eq!(worker_group_id_from_addresses(addresses).unwrap(), expected); + assert_eq!(worker_group_id_from_addresses(reversed).unwrap(), expected); + assert_eq!( + worker_group_id_from_addresses(Vec::new()) + .unwrap_err() + .to_string(), + "dist_init_addr resolved to no addresses" + ); + } + #[test] fn discovery_mode_reads_only_server_metadata() { use dynamo_backend_common::DisaggregationMode; From 50cb6a1d0e8195bd75d18ea1efa0694dff4ea78b Mon Sep 17 00:00:00 2001 From: jain-ria Date: Thu, 17 Sep 2026 12:44:45 -0700 Subject: [PATCH 09/17] fix(sglang): gate frontend readiness on connected KV relays Preserve early leader registration for follower discovery, while holding KV-routing frontend readiness until every declared DP rank has an active relay and an established router subscription. Use the same contract for single-node and multinode SGLang sidecars, including leaders whose KV publishers are all remote. Advertise relay sources after native ingress and outbound connections are established. Track ZMQ handshakes with safe monitor ownership and hostname reconnect handling, and flush NATS subscriptions before activation. Exercise delayed follower startup through real HTTP and ZMQ. Verify early leader discovery, rejected requests before readiness, and each rank's first request event reaching the router without event retries or warmup traffic. Cover monitor teardown and hostname reconnects as well. Signed-off-by: jain-ria --- Cargo.lock | 1 + lib/llm/src/direct_zmq_sub_pool.rs | 47 +- lib/llm/src/discovery/model.rs | 27 +- lib/llm/src/discovery/model_manager.rs | 22 + lib/llm/src/discovery/watcher.rs | 11 + lib/llm/src/discovery/worker_set.rs | 15 + lib/llm/src/http/service.rs | 3 + lib/llm/src/http/service/health.rs | 2 +- .../http/service/kv_relay_readiness_tests.rs | 370 +++++++++++++++ lib/llm/src/kv_router.rs | 15 + .../kv_router/indexer/recovery/broker_zmq.rs | 7 +- .../kv_router/indexer/recovery/direct_zmq.rs | 2 +- .../kv_router/indexer/recovery/subscriber.rs | 207 ++++++++- .../indexer/recovery/worker_query.rs | 26 ++ .../kv_router/prefill_router/activation.rs | 33 ++ lib/llm/src/kv_router/prefill_router/mod.rs | 36 ++ lib/llm/src/kv_router/publisher/mod.rs | 59 ++- lib/llm/src/kv_router/publisher/tests.rs | 19 +- .../src/kv_router/publisher/zmq_listener.rs | 29 +- lib/llm/src/local_model/runtime_config.rs | 11 + lib/runtime/Cargo.toml | 1 + lib/runtime/src/distributed.rs | 6 +- lib/runtime/src/transports/event_plane/mod.rs | 52 ++- .../src/transports/event_plane/transport.rs | 5 + .../transports/event_plane/zmq_transport.rs | 436 +++++++++++++++++- lib/sidecar/sglang/README.md | 11 + lib/sidecar/sglang/src/engine.rs | 34 ++ lib/sidecar/sglang/src/headless.rs | 98 ++-- 28 files changed, 1504 insertions(+), 81 deletions(-) create mode 100644 lib/llm/src/http/service/kv_relay_readiness_tests.rs diff --git a/Cargo.lock b/Cargo.lock index cc2c91f1120c..5459bb870d40 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3024,6 +3024,7 @@ dependencies = [ "uuid", "validator", "xxhash-rust", + "zmq", ] [[package]] diff --git a/lib/llm/src/direct_zmq_sub_pool.rs b/lib/llm/src/direct_zmq_sub_pool.rs index 640ad700471d..60d3932f58f6 100644 --- a/lib/llm/src/direct_zmq_sub_pool.rs +++ b/lib/llm/src/direct_zmq_sub_pool.rs @@ -11,6 +11,7 @@ use std::{ }; use anyhow::Result; +use dynamo_runtime::transports::event_plane::zmq_transport::ZmqConnectionReadiness; use dynamo_runtime::transports::event_plane::{ Codec, DynamicZmqSubSocket, ValidatedEnvelope, ValidatedZmqSource, ZmqWireMessage, }; @@ -97,6 +98,7 @@ struct SocketGroup { command_tx: mpsc::UnboundedSender, cancel: CancellationToken, handle: JoinHandle<()>, + readiness: ZmqConnectionReadiness, } struct PoolInner { @@ -122,6 +124,7 @@ pub(crate) struct DirectZmqSubRegistration { publisher_id: u64, generation: u64, armed: bool, + readiness: ZmqConnectionReadiness, } pub(crate) enum DirectZmqSubConnection { @@ -187,6 +190,35 @@ impl Drop for DirectZmqSubRegistration { } impl DirectZmqSubPool { + /// Unlike socket construction, readiness requires a completed transport handshake. + pub(crate) async fn connect_ready( + &self, + publisher_id: u64, + endpoint: &str, + generation: u64, + ) -> Result { + if self.endpoints_per_sub == 1 { + anyhow::ensure!( + !self.inner.lock().closed, + "direct-ZMQ socket pool is closed" + ); + let source = + ValidatedZmqSource::connect_ready(endpoint, &self.topic, publisher_id, self.rcvhwm) + .await?; + anyhow::ensure!( + !self.inner.lock().closed, + "direct-ZMQ socket pool is closed" + ); + return Ok(DirectZmqSubConnection::Dedicated(source)); + } + let registration = self + .register_grouped(publisher_id, endpoint, generation) + .await?; + // The group task keeps forwarding existing peers while this peer connects. + registration.readiness.wait_connected(endpoint).await?; + Ok(DirectZmqSubConnection::Grouped(registration)) + } + pub(crate) fn new( topic: impl Into>, endpoints_per_sub: usize, @@ -242,7 +274,7 @@ impl DirectZmqSubPool { ) -> Result { let (sender, receiver) = mpsc::channel(self.rcvhwm as usize); let disconnected = CancellationToken::new(); - let (group_id, completion) = { + let (group_id, completion, readiness) = { let mut inner = self.inner.lock(); Self::reap_failed_groups(&mut inner); anyhow::ensure!(!inner.closed, "direct-ZMQ socket pool is closed"); @@ -286,12 +318,15 @@ impl DirectZmqSubPool { group.assignments.remove(&publisher_id); anyhow::bail!("direct-ZMQ socket group stopped"); } - (group_id, Some(completion)) + (group_id, Some(completion), group.readiness.clone()) } else { // libzmq connects asynchronously. Keep group creation serialized so // concurrent registrations cannot create excess transient sockets. - let socket = - DynamicZmqSubSocket::connect_with_rcvhwm(endpoint, &self.topic, self.rcvhwm)?; + let (socket, readiness) = DynamicZmqSubSocket::connect_with_readiness( + endpoint, + &self.topic, + self.rcvhwm, + )?; let group_id = inner.next_group_id; inner.next_group_id = inner.next_group_id.wrapping_add(1); let (command_tx, command_rx) = mpsc::unbounded_channel(); @@ -320,9 +355,10 @@ impl DirectZmqSubPool { command_tx, cancel, handle, + readiness: readiness.clone(), }, ); - (group_id, None) + (group_id, None, readiness) } }; let registration = DirectZmqSubRegistration { @@ -333,6 +369,7 @@ impl DirectZmqSubPool { receiver, disconnected, armed: true, + readiness, }; let Some(completion) = completion else { diff --git a/lib/llm/src/discovery/model.rs b/lib/llm/src/discovery/model.rs index bfc37dbcc178..6b98cb07cb1a 100644 --- a/lib/llm/src/discovery/model.rs +++ b/lib/llm/src/discovery/model.rs @@ -385,7 +385,7 @@ impl Model { warn_legacy_readiness_once(&self.name, wsets[0].namespace()); } NamespaceReadinessEval { - ready: eval.ready, + ready: eval.ready && wsets.iter().all(|ws| ws.kv_event_sources_ready()), has_legacy: eval.has_legacy, legacy_live_workers: eval.legacy_live_workers, present: eval.present, @@ -408,6 +408,12 @@ impl Model { self.first_ready_workers().is_some() } + pub(crate) fn kv_event_sources_ready(&self) -> bool { + self.worker_sets + .iter() + .all(|entry| entry.value().kv_event_sources_ready()) + } + /// Structured per-namespace worker readiness for this model — the data /// behind the `GET /v1/models/{model}/ready` observability endpoint. /// @@ -471,6 +477,8 @@ impl Model { } else { None } + } else if wsets.iter().any(|ws| !ws.kv_event_sources_ready()) { + Some("waiting for required KV event feeds".to_string()) } else if eval.has_legacy { Some("legacy worker(s) present but no live worker".to_string()) } else if !eval.ambiguous.is_empty() { @@ -1329,6 +1337,23 @@ mod tests { assert!(!model.is_workers_ready("dynamo")); } + #[test] + fn kv_relay_readiness_does_not_withdraw_serving_membership() { + let model = Model::new("llama".to_string()); + let (mut worker_set, _workers) = + ws_with_type("dynamo", "mdc", WorkerType::Aggregated, vec![], vec![7]); + let (ready_tx, ready_rx) = watch::channel(false); + Arc::get_mut(&mut worker_set).unwrap().kv_event_readiness = Some(ready_rx); + model.add_worker_set("dynamo".to_string(), worker_set.clone()); + assert_eq!(worker_set.worker_count(), 1); + assert!(!model.is_workers_ready("dynamo")); + ready_tx.send(true).unwrap(); + assert!(model.is_workers_ready("dynamo")); + ready_tx.send(false).unwrap(); + assert!(!model.is_workers_ready("dynamo")); + assert_eq!(worker_set.worker_count(), 1); + } + #[test] fn readiness_pd_pair_ready() { let model = Model::new("llama".to_string()); diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index f3aefdea39ad..84fde0ed9a34 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -157,6 +157,8 @@ pub(crate) struct RemovedDiscoveryGroup { /// /// Note: Don't implement Clone for this, put it in an Arc instead. pub struct ModelManager { + /// KV frontends must not report ready before any model and its feeds are usable. + require_ready_model: AtomicBool, /// Model name → Model (which contains WorkerSets with engines) models: DashMap>, @@ -215,6 +217,7 @@ impl Default for ModelManager { impl ModelManager { pub fn new() -> Self { Self { + require_ready_model: AtomicBool::new(false), models: DashMap::new(), catalog: ArcSwap::from_pointee(CommittedCatalog::default()), cards: DashMap::new(), @@ -1095,6 +1098,25 @@ impl ModelManager { .any(|model| model.is_ready_to_serve()) } + pub(crate) fn require_model_readiness(&self) { + self.require_ready_model.store(true, Ordering::Relaxed); + } + + pub(crate) fn frontend_routing_ready(&self) -> bool { + if !self.require_ready_model.load(Ordering::Relaxed) { + return true; + } + let catalog = self.catalog.load(); + catalog + .models + .values() + .any(|model| model.is_ready_to_serve()) + && catalog + .models + .values() + .all(|model| model.kv_event_sources_ready()) + } + pub fn model_display_names(&self) -> HashSet { self.catalog .load() diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index 6f7859e72260..ae3563446c8e 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -317,6 +317,11 @@ where require_typed_worker_role: bool, worker_selector_factory: WorkerSelectorFactory, ) -> Self { + if router_config.router_mode == RouterMode::KV + && router_config.kv_router_config.use_kv_events + { + model_manager.require_model_readiness(); + } Self { manager: model_manager, drt: runtime, @@ -438,6 +443,11 @@ where // are process-local, so preserve them when the MDC supplies the base config. let router_config = effective_router_config(card.router_config.as_ref(), &self.router_config); + if router_config.router_mode == RouterMode::KV + && router_config.kv_router_config.use_kv_events + { + self.manager.require_model_readiness(); + } let component = self .drt @@ -611,6 +621,7 @@ where Arc::get_mut(&mut chooser) .expect("new KV chooser must have one owner") .set_teardown_task_guard(allocator_trim.clone()); + worker_set.kv_event_readiness = chooser.kv_event_readiness(); Some(chooser) } else { None diff --git a/lib/llm/src/discovery/worker_set.rs b/lib/llm/src/discovery/worker_set.rs index 352a3cc708a2..b1004382ded8 100644 --- a/lib/llm/src/discovery/worker_set.rs +++ b/lib/llm/src/discovery/worker_set.rs @@ -234,6 +234,9 @@ pub struct WorkerSet { /// None for in-process models (http/grpc) which don't have a discovery client. instance_count_rx: Option>>, + /// Frontend-local KV subscription readiness, independent of serving registration. + pub(crate) kv_event_readiness: Option>, + /// Cancels background work created while materializing this WorkerSet. lifecycle_cancellation: Option, @@ -266,6 +269,7 @@ impl WorkerSet { prefill_router: None, encoder_router: None, instance_count_rx: None, + kv_event_readiness: None, lifecycle_cancellation: None, allocator_trim: None, allocator_trim_wrapped: false, @@ -433,6 +437,16 @@ impl WorkerSet { } } + pub(crate) fn kv_event_sources_ready(&self) -> bool { + self.kv_event_readiness + .as_ref() + .is_none_or(|ready| *ready.borrow()) + && self + .prefill_router + .as_ref() + .is_none_or(|router| router.kv_event_sources_ready()) + } + /// Store the instance watcher from the Client's discovery system. /// Must be called before the WorkerSet is wrapped in Arc. pub fn set_instance_watcher(&mut self, rx: watch::Receiver>) { @@ -510,6 +524,7 @@ impl WorkerSet { prefill_router: self.prefill_router.clone(), encoder_router: self.encoder_router.clone(), instance_count_rx: self.instance_count_rx.clone(), + kv_event_readiness: self.kv_event_readiness.clone(), lifecycle_cancellation: None, allocator_trim: None, allocator_trim_wrapped: false, diff --git a/lib/llm/src/http/service.rs b/lib/llm/src/http/service.rs index f4d2a5fbe838..4d77d16cb1ef 100644 --- a/lib/llm/src/http/service.rs +++ b/lib/llm/src/http/service.rs @@ -34,6 +34,9 @@ pub mod realtime; pub mod service_v2; pub mod sglang_generate; +#[cfg(test)] +mod kv_relay_readiness_tests; + pub use axum; pub use frontend_extension::{ FrontendExtensionContext, FrontendRouteExtension, FrontendRouteSet, diff --git a/lib/llm/src/http/service/health.rs b/lib/llm/src/http/service/health.rs index 1b6c1da8ee86..81441cbbf388 100644 --- a/lib/llm/src/http/service/health.rs +++ b/lib/llm/src/http/service/health.rs @@ -63,7 +63,7 @@ async fn live_handler( async fn health_handler( axum::extract::State(state): axum::extract::State>, ) -> impl IntoResponse { - if !state.is_ready() { + if !state.is_ready() || !state.manager().frontend_routing_ready() { return ( StatusCode::SERVICE_UNAVAILABLE, Json(json!({ diff --git a/lib/llm/src/http/service/kv_relay_readiness_tests.rs b/lib/llm/src/http/service/kv_relay_readiness_tests.rs new file mode 100644 index 000000000000..439e35ee0c51 --- /dev/null +++ b/lib/llm/src/http/service/kv_relay_readiness_tests.rs @@ -0,0 +1,370 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Exercise the HTTP readiness gate with real engine and router ZMQ feeds. + +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; +use std::time::Duration; + +use dynamo_kv_router::{ + config::KvRouterConfig, + protocols::{KV_EVENT_SUBJECT, WorkerWithDpRank}, +}; +use dynamo_runtime::{ + CancellationToken, Runtime, + discovery::{DiscoveryQuery, EventSourceQuery}, + distributed::{DistributedConfig, DistributedRuntime}, + pipeline::{ + AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, ResponseStream, SingleIn, + async_trait, + }, +}; +use futures::SinkExt; +use serde_json::json; +use tmq::AsZmqSocket; +use tokio::sync::Mutex; + +use super::service_v2::HttpService; +use crate::{ + discovery::WorkerSet, + kv_router::publisher::{KvEventPublisher, KvEventSourceConfig}, + local_model::{register_model_card, runtime_config::SGLANG_GENERATE_CAPABILITY}, + model_card::ModelDeploymentCard, + model_type::{ModelInput, ModelType}, + protocols::{ + Annotated, + common::llm_backend::{LLMEngineOutput, PreprocessedRequest}, + }, + worker_type::WorkerType, +}; + +struct RequestKvEngine { + sources: Vec>>, + calls: AtomicUsize, +} + +#[async_trait] +impl AsyncEngine, ManyOut>, Error> + for RequestKvEngine +{ + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + let (request, context) = request.transfer(()); + let rank = request + .routing + .as_ref() + .and_then(|routing| routing.dp_rank) + .unwrap(); + self.calls.fetch_add(1, Ordering::SeqCst); + // Each first request emits exactly one native KV event. Re-sending it + // would conceal a subscription that became ready too early. + let payload = rmp_serde::to_vec_named(&json!([ + 0.0, [{"type":"BlockStored", "block_hashes":[rank as u64 + 10], + "parent_block_hash":null, "token_ids":request.token_ids, + "block_size":4, "lora_id":null}], rank + ]))?; + self.sources[rank as usize] + .lock() + .await + .as_mut() + .expect("ready rank must have an engine publisher") + .send(vec![Vec::new(), 1_u64.to_be_bytes().to_vec(), payload]) + .await?; + let output = LLMEngineOutput { + token_ids: vec![42], + engine_data: Some(json!({"sglang_response":{ + "text":"first-token", "output_ids":[42], + "meta_info":{"finish_reason":{"type":"length","length":1}} + }})), + ..Default::default() + }; + Ok(ResponseStream::new( + Box::pin(futures::stream::iter([Annotated::from_data(output)])), + context.context(), + )) + } +} + +async fn kv_relay_readiness(dp_size: u32) { + let runtime = Runtime::from_current().unwrap(); + let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) + .await + .unwrap(); + let namespace = format!("kv-relay-readiness-{dp_size}"); + let endpoint = drt + .namespace(&namespace) + .unwrap() + .component("backend") + .unwrap() + .endpoint("generate"); + let worker_id = drt.discovery().instance_id(); + let mut card = ModelDeploymentCard::with_name_only("kv-relay-model"); + card.model_input = ModelInput::Tokens; + card.model_type = ModelType::Chat; + card.worker_type = Some(WorkerType::Aggregated); + card.kv_cache_block_size = 4; + card.runtime_config.data_parallel_size = dp_size; + card.runtime_config + .set_engine_specific(SGLANG_GENERATE_CAPABILITY, true) + .unwrap(); + card.runtime_config + .set_engine_specific("require_kv_event_source_readiness", true) + .unwrap(); + + // This is the leader registration that telemetry followers discover. It + // must exist before the final rank's relay can begin starting. + endpoint.register_endpoint_instance().await.unwrap(); + register_model_card(&endpoint, &card).await.unwrap(); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let service = HttpService::builder() + .port(port) + .discovery(Some(drt.discovery())) + .enable_engine_apis(true) + .build() + .unwrap(); + let manager = service.state().manager_clone(); + manager.require_model_readiness(); + let router = manager + .kv_chooser_for_with_worker_role( + &endpoint, + 4, + Some(KvRouterConfig { + skip_initial_worker_wait: true, + router_event_threads: 1, + router_track_active_blocks: false, + ..Default::default() + }), + None, + Some(WorkerType::Aggregated), + "decode", + Some(card.name().to_owned()), + false, + ) + .await + .unwrap(); + let mut configs = manager + .get_or_create_runtime_config_watcher(&endpoint) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while !configs.borrow().contains_key(&worker_id) { + configs.changed().await.unwrap(); + } + }) + .await + .expect("the leader must be discoverable without follower relays"); + + let zmq_context = tmq::Context::new(); + let mut sources = Vec::new(); + let mut addresses = Vec::new(); + for rank in 0..dp_size { + if rank == dp_size - 1 { + let reserved = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + addresses.push(format!("tcp://{}", reserved.local_addr().unwrap())); + sources.push(Mutex::new(None)); + // Keep the final engine publisher absent until after its relay + // object exists, distinguishing construction from connection. + drop(reserved); + continue; + } + let source = tmq::publish::publish(&zmq_context) + .set_linger(0) + .bind("tcp://127.0.0.1:*") + .unwrap(); + addresses.push(source.get_socket().get_last_endpoint().unwrap().unwrap()); + sources.push(Mutex::new(Some(source))); + } + let engine = Arc::new(RequestKvEngine { + sources, + calls: AtomicUsize::new(0), + }); + let mut worker_set = WorkerSet::new(namespace.clone(), card.mdcsum().to_owned(), card); + worker_set.generate_engine = Some(engine.clone()); + worker_set.kv_event_readiness = router.kv_event_readiness(); + worker_set.set_instance_watcher(router.client().instance_avail_watcher()); + assert!(manager.add_worker_set("kv-relay-model", &namespace, worker_set)); + let stop = CancellationToken::new(); + let server_stop = stop.clone(); + let http_task = + tokio::spawn(async move { service.run_with_listener(server_stop, listener).await }); + let client = reqwest::Client::new(); + let base = format!("http://127.0.0.1:{port}"); + let mut relays = Vec::new(); + + for rank in 0..dp_size { + // For DP=2 this deliberately holds back the follower after the leader + // is registered and its local feed is connected. DP=1 must obey the + // same rule for its only feed. + assert_eq!( + client + .get(format!("{base}/health")) + .send() + .await + .unwrap() + .status(), + reqwest::StatusCode::SERVICE_UNAVAILABLE + ); + assert!(!router.kv_event_sources_ready()); + assert!(configs.borrow().contains_key(&worker_id)); + assert_eq!( + drt.discovery() + .list(DiscoveryQuery::Endpoint { + namespace: namespace.clone(), + component: "backend".into(), + endpoint: "generate".into(), + }) + .await + .unwrap() + .len(), + 1 + ); + assert_eq!(engine.calls.load(Ordering::SeqCst), 0); + let rejected = client.post(format!("{base}/generate")) + .header("x-data-parallel-rank", rank.to_string()) + .json(&json!({"input_ids":[11,12,13,14], "sampling_params":{"max_new_tokens":1}, "stream":true})) + .send().await.unwrap(); + assert_eq!(rejected.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(engine.calls.load(Ordering::SeqCst), 0); + + let relay = KvEventPublisher::new_with_local_indexer_and_worker_id_at( + endpoint.clone(), + endpoint.id(), + Some(worker_id), + 4, + Some(KvEventSourceConfig::Zmq { + endpoint: addresses[rank as usize].clone(), + topic: String::new(), + image_token_id: None, + video_token_id: None, + }), + false, + rank, + None, + ) + .unwrap(); + if rank == dp_size - 1 { + assert!( + tokio::time::timeout(Duration::from_millis(100), relay.wait_ready()) + .await + .is_err(), + "creating a relay without its engine feed must not mark it ready" + ); + assert_eq!( + drt.discovery() + .list(DiscoveryQuery::EventSources( + EventSourceQuery::endpoint_topic(endpoint.id(), KV_EVENT_SUBJECT) + )) + .await + .unwrap() + .len(), + rank as usize, + "an unconnected feed must not advertise a ready source" + ); + assert!(configs.borrow().contains_key(&worker_id)); + assert_eq!( + drt.discovery() + .list(DiscoveryQuery::Endpoint { + namespace: namespace.clone(), + component: "backend".into(), + endpoint: "generate".into(), + }) + .await + .unwrap() + .len(), + 1 + ); + assert_eq!( + client + .get(format!("{base}/health")) + .send() + .await + .unwrap() + .status(), + reqwest::StatusCode::SERVICE_UNAVAILABLE + ); + let source = tmq::publish::publish(&zmq_context) + .set_linger(0) + .bind(&addresses[rank as usize]) + .unwrap(); + *engine.sources[rank as usize].lock().await = Some(source); + } + tokio::time::timeout(Duration::from_secs(5), relay.wait_ready()) + .await + .unwrap() + .unwrap(); + relays.push(relay); + } + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if client + .get(format!("{base}/health")) + .send() + .await + .unwrap() + .status() + .is_success() + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("all connected rank feeds must make the frontend ready"); + assert!(router.kv_event_sources_ready()); + + // No warmup events, fixed sleeps, or retrying inference requests are used. + // The first admitted request on every rank must reach the router indexer. + for rank in 0..dp_size { + let tokens = vec![11 + rank * 10, 12, 13, 14]; + let response = client + .post(format!("{base}/generate")) + .header("x-data-parallel-rank", rank.to_string()) + .json( + &json!({"input_ids":tokens, "sampling_params":{"max_new_tokens":1}, "stream":true}), + ) + .send() + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::OK); + assert!(response.text().await.unwrap().contains("first-token")); + tokio::time::timeout(Duration::from_secs(5), async { + while router + .get_overlap_blocks( + &tokens, + None, + WorkerWithDpRank::new(worker_id, rank), + None, + None, + ) + .await + .unwrap() + != 1 + { + tokio::task::yield_now().await; + } + }) + .await + .expect("the first request's single KV event must reach the router"); + } + assert_eq!(engine.calls.load(Ordering::SeqCst), dp_size as usize); + stop.cancel(); + http_task.await.unwrap().unwrap(); + drop(relays); + drt.shutdown(); +} + +#[tokio::test] +async fn frontend_waits_for_delayed_follower_and_receives_first_request_events() { + // Single-node and leader/follower startup must obey the same contract. + for dp_size in [1, 2] { + kv_relay_readiness(dp_size).await; + } +} diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index ff9c69d5e321..e0bfb886d3e3 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -812,6 +812,8 @@ where endpoint.clone(), indexer.clone(), membership_watch, + workers_with_configs.clone(), + client.instance_avail_watcher(), block_size, model_name.clone().unwrap_or_else(|| "unknown".to_string()), worker_role, @@ -902,6 +904,19 @@ where &self.indexer } + pub(crate) fn kv_event_readiness(&self) -> Option> { + self.kv_event_subscription + .as_ref() + .map(|subscription| subscription.readiness()) + } + + /// Whether every opted-in worker's required KV feeds are connected to this router. + pub fn kv_event_sources_ready(&self) -> bool { + self.kv_event_subscription + .as_ref() + .is_none_or(|subscription| subscription.is_ready()) + } + pub fn kv_router_config(&self) -> &KvRouterConfig { &self.kv_router_config } diff --git a/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs b/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs index 528e895ac0f3..fe0d64a350bf 100644 --- a/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs +++ b/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs @@ -271,6 +271,7 @@ pub(super) async fn run_broker_zmq_supervisor( let mut retry_delay = INITIAL_BACKOFF; loop { + client.clear_transport_readiness(); let view = membership_watch.borrow_and_update().clone(); update_mismatch_metric( &status_metrics, @@ -282,7 +283,7 @@ pub(super) async fn run_broker_zmq_supervisor( ); let subscriber = if let Some(kv_state_endpoint) = view.resolved_kv_state_endpoint() { - match EventSubscriber::for_endpoint_id_with_transport( + match EventSubscriber::for_endpoint_id_with_transport_ready( component.drt(), kv_state_endpoint, KV_EVENT_SUBJECT, @@ -365,6 +366,7 @@ pub(super) async fn run_broker_zmq_supervisor( ScopeExit::Rebind => retry_delay = INITIAL_BACKOFF, ScopeExit::Retry => { ingress_metrics.increment_lifecycle("reconnect"); + client.clear_transport_readiness(); let view = client.sync_membership().await; update_subscription_failure_metric( &status_metrics, @@ -423,6 +425,7 @@ async fn consume_scope( membership_tx, membership_cancel.clone(), )); + client.publish_ready_sources(); let exit = loop { tokio::select! { _ = cancellation_token.cancelled() => break ScopeExit::Stop, @@ -442,6 +445,7 @@ async fn consume_scope( ); active_publishers = self::active_publishers(&view); lanes.reconcile(&active_publishers); + client.publish_ready_sources(); } Some(MembershipUpdate::Rebind) => break ScopeExit::Rebind, Some(MembershipUpdate::Stop) => break ScopeExit::Stop, @@ -468,6 +472,7 @@ async fn consume_scope( } }; membership_cancel.cancel(); + client.clear_transport_readiness(); drop(membership_rx); if let Err(error) = membership_handle.await && !error.is_cancelled() diff --git a/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs b/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs index 9e41ccbef394..62ca627c9069 100644 --- a/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs +++ b/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs @@ -732,7 +732,7 @@ async fn run_source( let connection = tokio::select! { biased; _ = cancel.cancelled() => return, - connection = group_pool.connect(publisher_id, &endpoint, task_generation) => connection, + connection = group_pool.connect_ready(publisher_id, &endpoint, task_generation) => connection, }; let mut connection = match connection { Ok(connection) => connection, diff --git a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs b/lib/llm/src/kv_router/indexer/recovery/subscriber.rs index 9dad1493a5d5..94caea4b2878 100644 --- a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs +++ b/lib/llm/src/kv_router/indexer/recovery/subscriber.rs @@ -1,10 +1,14 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::{sync::Arc, time::Duration}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Duration, +}; use anyhow::Result; -use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent}; +use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent, WorkerId, WorkerWithDpRank}; use dynamo_runtime::{ component::{Component, Endpoint}, discovery::EventTransportKind, @@ -12,7 +16,7 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, transports::event_plane::{EventSubscriber, TypedEventSubscriber, uses_direct_zmq}, }; -use tokio::sync::{Semaphore, oneshot}; +use tokio::sync::{Semaphore, oneshot, watch}; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; @@ -22,8 +26,12 @@ use super::{ worker_query::WorkerQueryClient, }; use crate::{ - discovery::{KvSourceMembershipView, KvSourceMembershipWatch, KvSourceStatus}, + discovery::{ + KvSourceId, KvSourceMembershipView, KvSourceMembershipWatch, KvSourceStatus, + RuntimeConfigWatch, + }, kv_router::{Indexer, KvEventSourceRequirement, metrics::RouterWorkerStatusMetrics}, + local_model::runtime_config::ModelRuntimeConfig, worker_type::WorkerType, }; @@ -59,6 +67,7 @@ async fn run_subscription_supervisor( let mut retry_delay = SUBSCRIPTION_INITIAL_BACKOFF; loop { + client.clear_transport_readiness(); let view = membership_watch.borrow_and_update().clone(); update_mismatch_metric( &metrics, @@ -125,6 +134,9 @@ async fn run_subscription_supervisor( continue; } client.sync_membership().await; + if subscriber.is_some() { + client.publish_ready_sources(); + } // Subscriber construction establishes buffering before membership activation starts // initial recovery. Re-reading the watch above rejects a stale endpoint binding. @@ -163,6 +175,7 @@ async fn run_subscription_supervisor( continue; } ScopeExit::Retry => { + client.clear_transport_readiness(); let view = client.sync_membership().await; update_subscription_failure_metric( &metrics, @@ -226,6 +239,7 @@ async fn consume_scope( if view.resolved_kv_state_endpoint() != Some(kv_state_endpoint) { return ScopeExit::Rebind; } + client.publish_ready_sources(); } result = subscriber.next() => { let Some(result) = result else { @@ -369,9 +383,18 @@ pub(crate) struct KvEventSubscriptionHandle { completions: Vec>, runtime: tokio::runtime::Handle, task_guard: Option, + readiness: watch::Receiver, } impl KvEventSubscriptionHandle { + pub(crate) fn readiness(&self) -> watch::Receiver { + self.readiness.clone() + } + + pub(crate) fn is_ready(&self) -> bool { + *self.readiness.borrow() + } + pub(crate) fn set_task_guard( &mut self, task_guard: dynamo_runtime::engine::EngineContextGuard, @@ -408,6 +431,8 @@ pub async fn start_subscriber( endpoint: Endpoint, indexer: Indexer, membership_watch: KvSourceMembershipWatch, + runtime_configs: RuntimeConfigWatch, + admitted_workers: watch::Receiver>, block_size: u32, model: String, worker_role: Option, @@ -451,6 +476,13 @@ pub async fn start_subscriber( cancel.child_token(), ) .await?; + let (readiness, readiness_completion) = start_readiness_watch( + runtime_configs, + admitted_workers, + membership_watch.fork_receiver(), + client.ready_sources(), + cancel.child_token(), + ); let health_completion = source_health::spawn( membership_watch.fork_receiver(), model.clone(), @@ -487,9 +519,15 @@ pub async fn start_subscriber( let cancel = cancellation_guard.disarm(); return Ok(KvEventSubscriptionHandle { cancel, - completions: vec![completion_rx, health_completion, state_completion], + completions: vec![ + completion_rx, + health_completion, + state_completion, + readiness_completion, + ], runtime, task_guard: None, + readiness, }); } @@ -523,9 +561,15 @@ pub async fn start_subscriber( let cancel = cancellation_guard.disarm(); return Ok(KvEventSubscriptionHandle { cancel, - completions: vec![completion_rx, health_completion, state_completion], + completions: vec![ + completion_rx, + health_completion, + state_completion, + readiness_completion, + ], runtime, task_guard: None, + readiness, }); } @@ -558,12 +602,89 @@ pub async fn start_subscriber( let cancel = cancellation_guard.disarm(); Ok(KvEventSubscriptionHandle { cancel, - completions: vec![completion_rx, health_completion, state_completion], + completions: vec![ + completion_rx, + health_completion, + state_completion, + readiness_completion, + ], runtime, task_guard: None, + readiness, + }) +} + +fn required_sources_ready( + configs: &HashMap, + membership: &KvSourceMembershipView, + connected: &HashSet, +) -> bool { + // The serving card may arrive before its runtime-config watch has caught up. + if configs.is_empty() { + return false; + } + configs.iter().all(|(&worker_id, config)| { + if !config.requires_kv_event_source_readiness() { + return true; + } + let expected_endpoint = config + .kv_state_endpoint + .as_ref() + .unwrap_or(&membership.serving_endpoint); + if membership.resolved_kv_state_endpoint() != Some(expected_endpoint) { + return false; + } + let Ok(ranks) = config.data_parallel_rank_range() else { + return false; + }; + ranks.into_iter().all(|dp_rank| { + membership + .status(&WorkerWithDpRank::new(worker_id, dp_rank)) + .and_then(KvSourceStatus::active_source) + .is_some_and(|source| connected.contains(&source.source_id())) + }) }) } +fn start_readiness_watch( + mut configs: RuntimeConfigWatch, + mut admitted_workers: watch::Receiver>, + mut membership: KvSourceMembershipWatch, + mut connected: watch::Receiver>, + cancel: CancellationToken, +) -> (watch::Receiver, oneshot::Receiver<()>) { + let (sender, readiness) = watch::channel(false); + let (completion_tx, completion) = oneshot::channel(); + tokio::spawn(async move { + loop { + let admitted = admitted_workers.borrow_and_update().clone(); + let mut current_configs = configs.borrow_and_update().clone(); + current_configs.retain(|worker, _| admitted.contains(worker)); + let ready = current_configs.len() == admitted.len() + && required_sources_ready( + ¤t_configs, + &membership.borrow_and_update(), + &connected.borrow_and_update(), + ); + sender.send_if_modified(|current| { + let changed = *current != ready; + *current = ready; + changed + }); + tokio::select! { + _ = cancel.cancelled() => break, + changed = configs.changed() => if changed.is_err() { break; }, + changed = admitted_workers.changed() => if changed.is_err() { break; }, + changed = membership.changed() => if changed.is_err() { break; }, + changed = connected.changed() => if changed.is_err() { break; }, + } + } + sender.send_replace(false); + let _ = completion_tx.send(()); + }); + (readiness, completion) +} + pub(crate) struct RecoverySupervisor { client: Arc>, cancel: CancellationToken, @@ -641,6 +762,76 @@ mod tests { use crate::discovery::{KvEventSource, KvStateEndpointResolution}; + #[test] + fn readiness_requires_connected_sources_for_every_opted_in_rank() { + let mut config = ModelRuntimeConfig { + kv_event_publishing_enabled: Some(true), + ..Default::default() + }; + config.runtime_data.insert( + "require_kv_event_source_readiness".to_string(), + serde_json::json!(true), + ); + for size in [1, 2] { + config.data_parallel_size = size; + let configs = HashMap::from([(7, config.clone())]); + let mut view = metric_view(KvSourceStatus::Missing, Some(true)); + for rank in 0..size { + view.sources + .insert(WorkerWithDpRank::new(7, rank), KvSourceStatus::Missing); + } + let mut connected = HashSet::new(); + assert!(!required_sources_ready(&configs, &view, &connected)); + for rank in 0..size { + let worker = WorkerWithDpRank::new(7, rank); + let source = KvEventSource { + kv_state_endpoint: view.serving_endpoint.clone(), + worker, + publisher_id: 11 + u64::from(rank), + recovery_target: None, + }; + view.sources + .insert(worker, KvSourceStatus::ActiveLiveOnly(source.clone())); + // Advertising an object alone does not establish its subscription. + assert!(!required_sources_ready(&configs, &view, &connected)); + connected.insert(source.source_id()); + assert_eq!( + required_sources_ready(&configs, &view, &connected), + rank + 1 == size + ); + } + let follower = WorkerWithDpRank::new(7, size - 1); + view.sources.insert(follower, KvSourceStatus::Missing); + assert!( + !required_sources_ready(&configs, &view, &connected), + "relay removal withdraws readiness" + ); + } + } + + #[test] + fn readiness_preserves_legacy_and_disabled_event_workers() { + let mut config = ModelRuntimeConfig::default(); + let view = metric_view(KvSourceStatus::Missing, None); + let connected = HashSet::new(); + assert!(!required_sources_ready(&HashMap::new(), &view, &connected)); + assert!(required_sources_ready( + &HashMap::from([(7, config.clone())]), + &view, + &connected + )); + config.runtime_data.insert( + "require_kv_event_source_readiness".to_string(), + serde_json::json!(true), + ); + config.kv_event_publishing_enabled = Some(false); + assert!(required_sources_ready( + &HashMap::from([(7, config)]), + &view, + &connected + )); + } + fn metric_view(status: KvSourceStatus, capability: Option) -> KvSourceMembershipView { let serving_endpoint = EndpointId { namespace: "ns".to_string(), @@ -725,6 +916,7 @@ mod tests { completions, runtime: tokio::runtime::Handle::current(), task_guard: None, + readiness: watch::channel(false).1, }; tokio::time::timeout(Duration::from_secs(1), handle.shutdown()) @@ -748,6 +940,7 @@ mod tests { completions: vec![completion_rx], runtime: tokio::runtime::Handle::current(), task_guard: None, + readiness: watch::channel(false).1, }; handle.set_task_guard(teardown.clone()); drop(teardown); diff --git a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs index f9da8b554656..4d9c383a1126 100644 --- a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs +++ b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs @@ -108,6 +108,8 @@ pub(crate) struct WorkerQueryClient { slots: DashMap>>, /// Immutable publisher binding and rank slot lookup performed once per event envelope. publisher_bindings: DashMap, + /// Sources activated behind an established event subscription. + ready_sources: watch::Sender>, recovery_lane: RecoveryLane, recovery_attempt_timeout: Duration, cancellation_token: CancellationToken, @@ -149,6 +151,7 @@ impl WorkerQueryClient { membership_sync: Mutex::new(()), slots: DashMap::new(), publisher_bindings: DashMap::new(), + ready_sources: watch::channel(HashSet::new()).0, recovery_lane: RecoveryLane::with_semaphore(recovery_semaphore), recovery_attempt_timeout, cancellation_token, @@ -188,6 +191,7 @@ impl WorkerQueryClient { membership_sync: Mutex::new(()), slots: DashMap::new(), publisher_bindings: DashMap::new(), + ready_sources: watch::channel(HashSet::new()).0, recovery_lane: RecoveryLane::with_semaphore(recovery_semaphore), recovery_attempt_timeout, cancellation_token: CancellationToken::new(), @@ -204,6 +208,23 @@ impl WorkerQueryClient { view } + pub(crate) fn ready_sources(&self) -> watch::Receiver> { + self.ready_sources.subscribe() + } + + pub(crate) fn clear_transport_readiness(&self) { + self.ready_sources.send_replace(HashSet::new()); + } + + pub(crate) fn publish_ready_sources(&self) { + self.ready_sources.send_replace( + self.publisher_bindings + .iter() + .map(|binding| binding.binding.source_id.clone()) + .collect(), + ); + } + /// Apply membership only for source incarnations whose direct transport is preconnected. /// /// The semantic membership watch remains authoritative. Transport readiness can only @@ -225,6 +246,7 @@ impl WorkerQueryClient { } } self.reconcile_view(effective).await; + self.publish_ready_sources(); view } @@ -406,6 +428,7 @@ impl WorkerQueryClient { } pub(crate) async fn shutdown(self: &Arc) { + self.clear_transport_readiness(); self.cancellation_token.cancel(); self.deactivate_all().await; } @@ -523,6 +546,9 @@ impl WorkerQueryClient { /// Unlike a protocol rejection, the same source may reactivate after a replacement socket is /// preconnected. The reset barrier makes any already-enqueued old event visible first. pub(crate) async fn fence_transport(self: &Arc, publisher_id: PublisherId) -> bool { + self.ready_sources.send_modify(|sources| { + sources.retain(|source| source.publisher_id != publisher_id); + }); let _sync = self.membership_sync.lock().await; let Some(active) = self .publisher_bindings diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index 07f051f69239..59af6136d674 100644 --- a/lib/llm/src/kv_router/prefill_router/activation.rs +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -683,6 +683,39 @@ mod tests { card } + #[test] + fn pending_prefill_requires_feeds_only_for_opted_in_kv_routing() { + use super::super::prefill_requires_kv_event_readiness; + let mut selected = card(None); + assert!(!prefill_requires_kv_event_readiness( + &selected, + RouterMode::KV + )); + selected + .runtime_config + .set_engine_specific("require_kv_event_source_readiness", true) + .unwrap(); + assert!(prefill_requires_kv_event_readiness( + &selected, + RouterMode::KV + )); + assert!(!prefill_requires_kv_event_readiness( + &selected, + RouterMode::RoundRobin + )); + selected.router_config = Some(RouterConfig::new( + RouterMode::KV, + KvRouterConfig { + use_kv_events: false, + ..Default::default() + }, + )); + assert!(!prefill_requires_kv_event_readiness( + &selected, + RouterMode::KV + )); + } + #[test] fn inherits_decode_mode_when_card_advertises_nothing() { // The pre-override behavior: a prefill worker that says nothing is diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index 53b2cdceee5b..a3518109ecb2 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -285,6 +285,7 @@ where pub(crate) trait PrefillRouterLifecycle: Send + Sync { fn set_target(&self, target: Option); + fn kv_event_sources_ready(&self) -> bool; } impl PrefillRouterLifecycle for PrefillRouter @@ -294,6 +295,41 @@ where fn set_target(&self, target: Option) { self.set_target(target); } + + fn kv_event_sources_ready(&self) -> bool { + if self.lifecycle_state() == PrefillLifecycleState::Unavailable { + return true; + } + if let Some(binding) = self.binding.load().as_ref() { + return binding + .router + .kv_router_if_enabled() + .is_none_or(|router| router.kv_event_sources_ready()); + } + // An opted-in prefill target can be discovered before its chooser has + // finished connecting. Keep the decode WorkerSet unready during that gap. + self.target_tx + .as_ref() + .is_none_or(|targets| match targets.borrow().as_ref() { + Some(WorkerSetTarget::Committed(target)) => { + !prefill_requires_kv_event_readiness(&target.card, self.decode_router_mode) + } + _ => true, + }) + } +} + +fn prefill_requires_kv_event_readiness( + card: &crate::model_card::ModelDeploymentCard, + fallback_mode: RouterMode, +) -> bool { + card.runtime_config.requires_kv_event_source_readiness() + && card + .router_config + .as_ref() + .map_or(fallback_mode.is_kv_routing(), |config| { + config.router_mode.is_kv_routing() && config.kv_router_config.use_kv_events + }) } impl Drop for PrefillRouter diff --git a/lib/llm/src/kv_router/publisher/mod.rs b/lib/llm/src/kv_router/publisher/mod.rs index eeb9a60d98f7..39abbd66b87c 100644 --- a/lib/llm/src/kv_router/publisher/mod.rs +++ b/lib/llm/src/kv_router/publisher/mod.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use anyhow::Result; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, oneshot, watch}; use tokio_util::sync::CancellationToken; use dynamo_kv_router::indexer::{KvIndexerMetrics, LocalKvIndexer}; @@ -120,6 +120,7 @@ async fn supervise_zmq_listener( } impl KvEventSource { + #[allow(clippy::too_many_arguments)] fn start( component: Component, worker_id: WorkerId, @@ -128,6 +129,7 @@ impl KvEventSource { cancellation_token: CancellationToken, tx: mpsc::UnboundedSender>, next_event_id: Arc, + ready: oneshot::Sender<()>, ) -> Result { match source_config { KvEventSourceConfig::Zmq { @@ -151,6 +153,7 @@ impl KvEventSource { next_event_id, image_token_id, video_token_id, + Some(ready), )); let listener_abort_handle = listener_handle.abort_handle(); let supervisor_handle = @@ -207,6 +210,8 @@ pub struct KvEventPublisher { tx: mpsc::UnboundedSender>, /// Internal monotonic event ID counter. Shared with the ZMQ listener if present. next_event_id: Arc, + /// Local ingress and outbound publication have been established. + ready: watch::Receiver, } impl KvEventPublisher { @@ -333,8 +338,12 @@ impl KvEventPublisher { let next_event_id = Arc::new(AtomicU64::new(0)); + let (ready_tx, ready) = watch::channel(false); + let mut input_ready = None; let mut source = None; if let Some(config) = source_config { + let (input_ready_tx, input_ready_rx) = oneshot::channel(); + input_ready = Some(input_ready_rx); source = Some(KvEventSource::start( component.clone(), worker_id, @@ -343,6 +352,7 @@ impl KvEventPublisher { cancellation_token.clone(), tx.clone(), next_event_id.clone(), + input_ready_tx, )?); } @@ -380,6 +390,31 @@ impl KvEventPublisher { }; let publisher_id = event_publisher.publisher_id(); + tokio::select! { + _ = cancellation_token_clone.cancelled() => return, + result = event_publisher.wait_ready() => { + if let Err(error) = result { + tracing::error!(%error, worker_id, dp_rank, "KV event publication failed to connect"); + return; + } + } + } + + // A source advertisement is the relay's readiness contract. Do not + // advertise while the engine SUB socket is still connecting. The + // serving worker registers independently so followers can find it. + if let Some(input_ready) = input_ready { + tokio::select! { + _ = cancellation_token_clone.cancelled() => return, + result = input_ready => { + if result.is_err() { + tracing::error!(worker_id, dp_rank, "KV source closed before connecting"); + return; + } + } + } + } + let recovery_endpoint = if let Some(local_indexer) = local_indexer_clone.as_ref() { match start_worker_kv_query_endpoint( component.clone(), @@ -449,6 +484,8 @@ impl KvEventPublisher { } }; + ready_tx.send_replace(true); + start_event_processor( EventPlanePublisher(event_publisher), worker_id, @@ -459,6 +496,8 @@ impl KvEventPublisher { ) .await; + ready_tx.send_replace(false); + if let Err(error) = component .drt() .discovery() @@ -481,9 +520,27 @@ impl KvEventPublisher { worker_id, tx, next_event_id, + ready, }) } + /// Wait for local ingress and outbound publication to be established. + /// + /// This does not wait for frontend subscribers. Frontend KV readiness also + /// requires its subscriptions to every advertised DP-rank source. + pub async fn wait_ready(&self) -> Result<()> { + let mut ready = self.ready.clone(); + tokio::select! { + biased; + _ = self.cancellation_token.cancelled() => { + anyhow::bail!("KV event publisher stopped before becoming ready") + } + result = ready.wait_for(|ready| *ready) => { + result.map(|_| ()).map_err(|_| anyhow::anyhow!("KV event publisher failed to become ready")) + } + } + } + pub fn publish(&self, event: KvCacheEvent) -> Result<(), mpsc::error::SendError> { self.send_singleton(PlacementEvent::local_gpu(self.worker_id, event)) } diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 40784899b80b..0ca23b036636 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -30,6 +30,7 @@ mod test_event_processing { worker_id: 7, tx, next_event_id: Arc::new(AtomicU64::new(0)), + ready: watch::channel(true).1, }; publisher.publish_batch(Vec::new()).unwrap(); @@ -70,6 +71,7 @@ mod test_event_processing { worker_id: 7, tx, next_event_id: Arc::new(AtomicU64::new(0)), + ready: watch::channel(true).1, }; drop(rx); @@ -105,6 +107,7 @@ mod test_event_processing { worker_id: 7, tx, next_event_id: Arc::new(AtomicU64::new(0)), + ready: watch::channel(true).1, }; publisher @@ -1301,6 +1304,7 @@ mod tests_startup_helpers { next_event_id, None, None, + None, ) }); @@ -1432,6 +1436,7 @@ mod tests_startup_helpers { Arc::new(AtomicU64::new(0)), None, None, + None, ) }); @@ -1531,6 +1536,7 @@ mod tests_startup_helpers { Arc::new(AtomicU64::new(0)), None, None, + None, ) }); @@ -1613,7 +1619,18 @@ mod tests_startup_helpers { let listener_handle = tokio::spawn({ let token = token.clone(); let endpoint = endpoint.clone(); - start_zmq_listener(endpoint, topic, 1, tx, token, 4, next_event_id, None, None) + start_zmq_listener( + endpoint, + topic, + 1, + tx, + token, + 4, + next_event_id, + None, + None, + None, + ) }); tokio::time::sleep(tokio::time::Duration::from_millis(150)).await; diff --git a/lib/llm/src/kv_router/publisher/zmq_listener.rs b/lib/llm/src/kv_router/publisher/zmq_listener.rs index 9084655f0710..78d857d4ed7c 100644 --- a/lib/llm/src/kv_router/publisher/zmq_listener.rs +++ b/lib/llm/src/kv_router/publisher/zmq_listener.rs @@ -6,14 +6,16 @@ use std::sync::atomic::{AtomicU64, Ordering}; use anyhow::{Context, Result}; use futures::StreamExt; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, oneshot}; use tokio_util::sync::CancellationToken; use dynamo_kv_router::protocols::*; use dynamo_kv_router::zmq_wire::*; use crate::kv_router::metrics::kv_publisher_metrics; -use crate::utils::zmq::{connect_sub_socket, multipart_message}; +use dynamo_runtime::transports::event_plane::zmq_transport::connect_subscriber_with_readiness; + +use crate::utils::zmq::multipart_message; pub(super) struct DecodedZmqKvBatch { pub(super) source_cursor: u64, @@ -55,6 +57,7 @@ pub(super) async fn start_zmq_listener( next_event_id: Arc, image_token_id: Option, video_token_id: Option, + ready: Option>, ) { tracing::debug!( "KVEventPublisher connecting to ZMQ endpoint {} (topic '{}')", @@ -65,19 +68,35 @@ pub(super) async fn start_zmq_listener( let mut normalizer = ZmqEventNormalizer::new(kv_block_size) .with_image_token_id(image_token_id) .with_video_token_id(video_token_id); - let socket = match connect_sub_socket(&zmq_endpoint, Some(&zmq_topic)).await { - Ok(socket) => socket, + // Preserve the native-ingress socket's default receive high-water mark. + let (mut socket, connection) = match connect_subscriber_with_readiness( + &zmq_endpoint, + &zmq_topic, + 1_000, + ) { + Ok(connection) => connection, Err(error) => { tracing::error!(endpoint = %zmq_endpoint, topic = %zmq_topic, error = %error, "ZMQ listener failed to connect"); return; } }; - let mut socket = socket; + tokio::select! { + _ = cancellation_token.cancelled() => return, + result = connection.wait_connected(&zmq_endpoint) => { + if let Err(error) = result { + tracing::error!(endpoint = %zmq_endpoint, %error, "ZMQ listener connection failed"); + return; + } + } + } let metrics = kv_publisher_metrics(); if cancellation_token.is_cancelled() { return; } + if let Some(ready) = ready { + let _ = ready.send(()); + } let mut messages_processed = 0u64; diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 306369e3a138..5279d27dac96 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -406,6 +406,17 @@ impl Default for ModelRuntimeConfig { } impl ModelRuntimeConfig { + /// Upgraded relays opt into frontend readiness for every declared DP rank. + /// Absence preserves the discovery contract of older workers. + pub fn requires_kv_event_source_readiness(&self) -> bool { + self.kv_event_publishing_enabled != Some(false) + && self + .runtime_data + .get("require_kv_event_source_readiness") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + } + /// Check whether a runtime boolean is explicitly enabled. /// /// Rust callers commonly store booleans, while compatibility cards may diff --git a/lib/runtime/Cargo.toml b/lib/runtime/Cargo.toml index 5bf28152d53e..9e75158b3dd1 100644 --- a/lib/runtime/Cargo.toml +++ b/lib/runtime/Cargo.toml @@ -37,6 +37,7 @@ async-nats = { workspace = true } async-stream = { workspace = true } async-trait = { workspace = true } tmq = { workspace = true } +zmq = "0.10" lru = { version = "0.16" } axum = { workspace = true } blake3 = { workspace = true } diff --git a/lib/runtime/src/distributed.rs b/lib/runtime/src/distributed.rs index 45e9bed1d07a..4a70f2f941f6 100644 --- a/lib/runtime/src/distributed.rs +++ b/lib/runtime/src/distributed.rs @@ -666,7 +666,11 @@ impl DistributedRuntime { let Some(nats_client) = self.nats_client.as_ref() else { anyhow::bail!("KV router's EventSubscriber requires NATS"); }; - Ok(nats_client.client().subscribe(subject).await?) + let subscriber = nats_client.client().subscribe(subject).await?; + // Establish the subscription at the server before frontend readiness can + // release the first request (and its first KV event). + nats_client.client().flush().await?; + Ok(subscriber) } /// TODO (karenc): This is a temporary KV router measure for worker query requests. diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index 2a8126f054f1..8b72b79dbca9 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -315,6 +315,12 @@ pub struct EventPublisher { } impl EventPublisher { + /// Wait for the outbound transport to establish its forwarding connection. + /// Direct publishers already have a bound socket; broker publishers await ZMTP. + pub async fn wait_ready(&self) -> Result<()> { + self.tx.wait_ready().await + } + /// Create a publisher for an endpoint-scoped topic. pub async fn for_endpoint(endpoint: &Endpoint, topic: impl Into) -> Result { let transport_kind = endpoint.drt().default_event_transport_kind(); @@ -663,6 +669,26 @@ pub struct EventSubscriber { } impl EventSubscriber { + /// Establish broker connections before returning the subscriber. Direct ZMQ + /// sources have their own per-source readiness in the KV router. + pub async fn for_endpoint_id_with_transport_ready( + drt: &DistributedRuntime, + endpoint: &EndpointId, + topic: impl Into, + transport_kind: EventTransportKind, + ) -> Result { + Self::new_internal_with_readiness( + drt, + EventScope::Endpoint { + endpoint: endpoint.clone(), + }, + topic.into(), + transport_kind, + true, + ) + .await + } + /// Create a subscriber for an endpoint-scoped topic. pub async fn for_endpoint(endpoint: &Endpoint, topic: impl Into) -> Result { let transport_kind = endpoint.drt().default_event_transport_kind(); @@ -762,6 +788,16 @@ impl EventSubscriber { scope: EventScope, topic: String, transport_kind: EventTransportKind, + ) -> Result { + Self::new_internal_with_readiness(drt, scope, topic, transport_kind, false).await + } + + async fn new_internal_with_readiness( + drt: &DistributedRuntime, + scope: EventScope, + topic: String, + transport_kind: EventTransportKind, + wait_for_connections: bool, ) -> Result { let discovery = drt.discovery(); let routing_key = scope.subject(&topic); @@ -780,7 +816,21 @@ impl EventSubscriber { // BROKER MODE: Connect to broker's XPUB (single or multiple endpoints) let codec = Arc::new(Codec::Msgpack(MsgpackCodec)); - let stream: WireStream = if broker.xpub_endpoints.len() == 1 { + let stream: WireStream = if wait_for_connections { + let stream = + zmq_transport::ZmqSubTransport::connect_single_consumer_multiple_ready( + &broker.xpub_endpoints, + &routing_key, + ) + .await?; + let stream: WireStream = + Box::pin(stream.map(|result| result.map(|message| message.payload))); + if broker.xpub_endpoints.len() > 1 { + Box::pin(DeduplicatingStream::new(stream, codec.clone(), 100_000)) + } else { + stream + } + } else if broker.xpub_endpoints.len() == 1 { // One EventSubscriber has one consumer. Poll the ZMQ socket // directly instead of forwarding through a lossy broadcast channel. let stream = zmq_transport::ZmqSubTransport::connect_single_consumer( diff --git a/lib/runtime/src/transports/event_plane/transport.rs b/lib/runtime/src/transports/event_plane/transport.rs index 980f60fec8af..3bc2fa0e784d 100644 --- a/lib/runtime/src/transports/event_plane/transport.rs +++ b/lib/runtime/src/transports/event_plane/transport.rs @@ -21,6 +21,11 @@ pub type WireStream = Pin> + Send>>; /// Transport-agnostic interface for publishing events. #[async_trait] pub trait EventTransportTx: Send + Sync { + /// Wait for any asynchronous transport startup before advertising a ready source. + async fn wait_ready(&self) -> Result<()> { + Ok(()) + } + /// Publish raw envelope bytes to a subject/topic. async fn publish(&self, subject: &str, envelope_bytes: Bytes) -> Result<()>; diff --git a/lib/runtime/src/transports/event_plane/zmq_transport.rs b/lib/runtime/src/transports/event_plane/zmq_transport.rs index e3c8ff53a3bf..3a9bc494615e 100644 --- a/lib/runtime/src/transports/event_plane/zmq_transport.rs +++ b/lib/runtime/src/transports/event_plane/zmq_transport.rs @@ -18,17 +18,20 @@ use anyhow::{Context as _, Result, anyhow}; use async_stream::stream; use async_trait::async_trait; use bytes::Bytes; -use futures::{SinkExt, StreamExt}; +use futures::{SinkExt, Stream, StreamExt}; use once_cell::sync::OnceCell; use std::ffi::OsStr; +use std::ops::{Deref, DerefMut}; +use std::pin::Pin; use std::sync::Arc; +use std::task::{Context as TaskContext, Poll}; use thiserror::Error; use tmq::{ AsZmqSocket, Context, Message, Multipart, SocketBuilder, publish::{Publish, publish}, subscribe::{Subscribe, subscribe}, }; -use tokio::sync::{Mutex, broadcast}; +use tokio::sync::{Mutex, broadcast, watch}; use tokio_util::task::AbortOnDropHandle; /// Returns the process-wide shared ZMQ context. @@ -73,6 +76,218 @@ const ZMQ_RCVHWM: i32 = 100_000; // Receive buffer: 100K messages const ZMQ_SNDTIMEOUT_MS: i32 = 0; // Send timeout: fail fast under pressure const ZMQ_RCVTIMEOUT_MS: i32 = 100; // Receive timeout: 100ms (avoids blocking forever) +/// Keeps a socket's monitor alive until event production has been disabled. +/// +/// libzmq sends monitor events synchronously. Dropping its PAIR receiver while +/// monitoring remains enabled can block an I/O thread on a later disconnect. +pub struct MonitoredZmqSocket { + socket: S, + readiness: Option, +} + +impl MonitoredZmqSocket { + fn unmonitored(socket: S) -> Self { + Self { + socket, + readiness: None, + } + } + + fn monitored(socket: S, context: &Context) -> Result { + let readiness = monitor_connection(socket.get_socket(), context)?; + Ok(Self { + socket, + readiness: Some(readiness), + }) + } + + fn readiness(&self) -> ZmqConnectionReadiness { + self.readiness + .as_ref() + .expect("socket is monitored") + .clone() + } +} + +impl AsZmqSocket for MonitoredZmqSocket { + fn get_socket(&self) -> &zmq::Socket { + self.socket.get_socket() + } +} + +impl Deref for MonitoredZmqSocket { + type Target = S; + fn deref(&self) -> &S { + &self.socket + } +} + +impl DerefMut for MonitoredZmqSocket { + fn deref_mut(&mut self) -> &mut S { + &mut self.socket + } +} + +impl Stream for MonitoredZmqSocket { + type Item = S::Item; + fn poll_next(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { + Pin::new(&mut self.get_mut().socket).poll_next(cx) + } +} + +impl Drop for MonitoredZmqSocket { + fn drop(&mut self) { + if self.readiness.is_some() { + // rust-zmq cannot pass the null endpoint used to stop monitoring. + // Replacing the monitor with a zero-event monitor disables event + // production synchronously, before the old receiver is aborted. + // A fresh endpoint avoids racing the old monitor's asynchronous close. + let endpoint = format!("inproc://dynamo-zmq-disabled-{}", uuid::Uuid::new_v4()); + if let Err(error) = self.socket.get_socket().monitor(&endpoint, 0) { + tracing::debug!(%error, "Unable to disable closing ZMQ socket monitor"); + } + } + } +} + +/// Tracks completed ZMTP handshakes independently of the data receive loop. +/// +/// The subscription is configured before connecting. A completed handshake means +/// the transport can forward events; PUB/SUB does not acknowledge the remote +/// publisher's processing of SUBSCRIBE and this is not a delivery guarantee. +#[derive(Clone)] +pub struct ZmqConnectionReadiness { + connected: watch::Receiver>, + endpoints: Arc>>, + // Keep the receiver connected even if Tokio cancels its reader first during + // shutdown. The owning data socket can then disable monitoring safely. + _monitor_socket: Arc>, + _monitor_task: Arc>, +} + +impl ZmqConnectionReadiness { + /// Wait for the endpoint's transport handshake. Callers own the startup + /// deadline and cancellation policy; dropping this future is safe. + pub async fn wait_connected(&self, endpoint: &str) -> Result<()> { + let mut aliases = std::collections::HashSet::from([endpoint.to_string()]); + if let Some(recorded) = self.endpoints.lock().get(endpoint) { + aliases.insert(recorded.clone()); + } + if let Some(address) = endpoint.strip_prefix("tcp://") { + if let Ok(address) = address.parse::() { + aliases.insert(format!("tcp://{address}")); + } else { + // libzmq may report the hostname on its initial connection and + // the resolved IP after reconnecting. LAST_ENDPOINT alone does + // not cover both spellings. Resolve only startup hostnames; + // numeric TCP and IPC endpoints do not perform a DNS lookup. + let addresses = tokio::net::lookup_host(address) + .await + .with_context(|| format!("Failed to resolve ZMQ endpoint {endpoint}"))?; + aliases.extend(addresses.map(|address| format!("tcp://{address}"))); + } + } + let mut connected = self.connected.clone(); + loop { + if !connected.borrow_and_update().is_disjoint(&aliases) { + return Ok(()); + } + connected + .changed() + .await + .context("ZMQ connection monitor stopped")?; + } + } + + fn record_endpoint(&self, socket: &zmq::Socket, endpoint: &str) -> Result<()> { + // Preserve libzmq's initial spelling; reconnects may use a resolved IP. + let resolved = socket + .get_last_endpoint()? + .map_err(|_| anyhow!("ZMQ endpoint is not valid UTF-8"))?; + self.endpoints.lock().insert(endpoint.to_string(), resolved); + Ok(()) + } +} + +/// Start a SUB socket with an observable connection handshake. +/// +/// Socket construction remains nonblocking. Readiness is observed separately so +/// adding an unavailable peer does not block an existing grouped receive loop. +pub fn connect_subscriber_with_readiness( + endpoint: &str, + topic: &str, + rcvhwm: i32, +) -> Result<(MonitoredZmqSocket, ZmqConnectionReadiness)> { + anyhow::ensure!(rcvhwm > 0, "ZMQ receive HWM must be greater than zero"); + anyhow::ensure!( + endpoint.starts_with("tcp://") || endpoint.starts_with("ipc://"), + "ZMQ connection readiness requires a TCP or IPC endpoint" + ); + let context = shared_zmq_context()?; + let socket = context + .socket(zmq::SUB) + .map_err(|error| map_socket_creation_error(error.into()))?; + socket.set_rcvhwm(rcvhwm)?; + socket.set_rcvtimeo(ZMQ_RCVTIMEOUT_MS)?; + socket.set_linger(0)?; + socket.set_reconnect_ivl(100)?; + socket.set_reconnect_ivl_max(5000)?; + socket.set_tcp_keepalive(1)?; + socket.set_subscribe(topic.as_bytes())?; + + let socket = + >::from_zmq_socket(socket)? + .subscribe(topic.as_bytes())?; + let socket = MonitoredZmqSocket::monitored(socket, &context)?; + let readiness = socket.readiness(); + socket.get_socket().connect(endpoint)?; + readiness.record_endpoint(socket.get_socket(), endpoint)?; + Ok((socket, readiness)) +} + +fn monitor_connection(socket: &zmq::Socket, context: &Context) -> Result { + let monitor_endpoint = format!("inproc://dynamo-zmq-ready-{}", uuid::Uuid::new_v4()); + let events = zmq::SocketEvent::HANDSHAKE_SUCCEEDED.to_raw() + | zmq::SocketEvent::DISCONNECTED.to_raw() + | zmq::SocketEvent::MONITOR_STOPPED.to_raw(); + let monitor = tmq::pair(context) + .set_linger(0) + .connect(&monitor_endpoint)?; + socket.monitor(&monitor_endpoint, i32::from(events))?; + + let monitor = Arc::new(Mutex::new(monitor)); + let reader = monitor.clone(); + let (connected_tx, connected) = watch::channel(std::collections::HashSet::new()); + let monitor_task = tokio::spawn(async move { + let mut monitor = reader.lock().await; + while let Some(Ok(frames)) = monitor.next().await { + if frames.len() != 2 || frames[0].len() < 2 { + continue; + } + let event = u16::from_ne_bytes([frames[0][0], frames[0][1]]); + if event == zmq::SocketEvent::MONITOR_STOPPED.to_raw() { + break; + } + let Ok(endpoint) = std::str::from_utf8(&frames[1]) else { + continue; + }; + connected_tx.send_modify(|connected| { + if event == zmq::SocketEvent::HANDSHAKE_SUCCEEDED.to_raw() { + connected.insert(endpoint.to_string()); + } else if event == zmq::SocketEvent::DISCONNECTED.to_raw() { + connected.remove(endpoint); + } + }); + } + }); + Ok(ZmqConnectionReadiness { + connected, + endpoints: Arc::default(), + _monitor_socket: monitor, + _monitor_task: Arc::new(AbortOnDropHandle::new(monitor_task)), + }) +} + const ZMQ_SOCKET_LIMIT_GUIDANCE: &str = "ZMQ could not create another socket. The process may have reached libzmq's ZMQ_MAX_SOCKETS limit or its file-descriptor limit. Reduce direct-ZMQ peers or raise the limit with `ulimit -n`"; const PROCESS_FD_LIMIT_GUIDANCE: &str = "The process reached its file-descriptor limit. Reduce open file descriptors or raise the limit with `ulimit -n`"; @@ -164,8 +379,9 @@ impl AsRef<[u8]> for ZmqMessageOwner { /// ZMQ PUB transport for publishing events. pub struct ZmqPubTransport { - socket: Arc>, + socket: Arc>>, topic: String, + readiness: Option<(ZmqConnectionReadiness, Vec)>, } impl ZmqPubTransport { @@ -199,8 +415,9 @@ impl ZmqPubTransport { Ok(( Self { - socket: Arc::new(Mutex::new(socket)), + socket: Arc::new(Mutex::new(MonitoredZmqSocket::unmonitored(socket))), topic: topic.to_string(), + readiness: None, }, actual_endpoint, )) @@ -212,20 +429,7 @@ impl ZmqPubTransport { /// Connect to single broker XSUB endpoint (broker mode) pub async fn connect(xsub_endpoint: &str, topic: &str) -> Result { - let ctx = shared_zmq_context()?; - let socket = connect_tmq_socket(configure_publish_builder(publish(&ctx)), xsub_endpoint)?; - - tracing::info!( - endpoint = %xsub_endpoint, - topic = %topic, - sndhwm = ZMQ_SNDHWM, - "ZMQ PUB transport connected to broker XSUB" - ); - - Ok(Self { - socket: Arc::new(Mutex::new(socket)), - topic: topic.to_string(), - }) + Self::connect_multiple(&[xsub_endpoint.to_string()], topic).await } /// Connect to multiple broker XSUB endpoints (HA mode) @@ -236,10 +440,22 @@ impl ZmqPubTransport { }; let ctx = shared_zmq_context()?; - let socket = connect_tmq_socket(configure_publish_builder(publish(&ctx)), first_endpoint)?; + let socket = ctx + .socket(zmq::PUB) + .map_err(|error| map_socket_creation_error(error.into()))?; + socket.set_sndhwm(ZMQ_SNDHWM)?; + socket.set_sndtimeo(ZMQ_SNDTIMEOUT_MS)?; + let socket = MonitoredZmqSocket::monitored( + >::from_zmq_socket(socket)?, + &ctx, + )?; + let readiness = socket.readiness(); + socket.get_socket().connect(first_endpoint)?; + readiness.record_endpoint(socket.get_socket(), first_endpoint)?; for endpoint in endpoints { socket.get_socket().connect(endpoint)?; + readiness.record_endpoint(socket.get_socket(), endpoint)?; tracing::debug!(endpoint = %endpoint, "ZMQ PUB connected to broker XSUB"); } @@ -253,12 +469,22 @@ impl ZmqPubTransport { Ok(Self { socket: Arc::new(Mutex::new(socket)), topic: topic.to_string(), + readiness: Some((readiness, xsub_endpoints.to_vec())), }) } } #[async_trait] impl EventTransportTx for ZmqPubTransport { + async fn wait_ready(&self) -> Result<()> { + if let Some((readiness, endpoints)) = &self.readiness { + for endpoint in endpoints { + readiness.wait_connected(endpoint).await?; + } + } + Ok(()) + } + async fn publish(&self, _subject: &str, envelope_bytes: Bytes) -> Result<()> { let codec = MsgpackCodec; let (publisher_id, sequence) = codec.decode_envelope_identity(&envelope_bytes)?; @@ -307,24 +533,46 @@ pub type ZmqWireStream = /// /// The caller must keep this value in one task. ZMQ SUB sockets are not thread-safe. pub struct DynamicZmqSubSocket { - socket: Subscribe, + socket: MonitoredZmqSocket, expected_topic: Vec, + readiness: Option, } impl DynamicZmqSubSocket { + /// Start a dynamically managed socket and expose its connection handshakes. + pub fn connect_with_readiness( + endpoint: &str, + topic: &str, + rcvhwm: i32, + ) -> Result<(Self, ZmqConnectionReadiness)> { + let (socket, readiness) = connect_subscriber_with_readiness(endpoint, topic, rcvhwm)?; + Ok(( + Self { + socket, + expected_topic: topic.as_bytes().to_vec(), + readiness: Some(readiness.clone()), + }, + readiness, + )) + } + /// Connect a new SUB socket with an explicit receive high-water mark. pub fn connect_with_rcvhwm(endpoint: &str, topic: &str, rcvhwm: i32) -> Result { let socket = ZmqSubTransport::connect_socket_with_rcvhwm(endpoint, topic, rcvhwm)?; tracing::info!(endpoint, topic, rcvhwm, "Dynamic ZMQ SUB socket connected"); Ok(Self { - socket, + socket: MonitoredZmqSocket::unmonitored(socket), expected_topic: topic.as_bytes().to_vec(), + readiness: None, }) } /// Connect this SUB socket to one more publisher endpoint. pub fn add_endpoint(&mut self, endpoint: &str) -> Result<()> { self.socket.get_socket().connect(endpoint)?; + if let Some(readiness) = &self.readiness { + readiness.record_endpoint(self.socket.get_socket(), endpoint)?; + } Ok(()) } @@ -390,6 +638,21 @@ pub struct ValidatedZmqSource { } impl ValidatedZmqSource { + /// Connect after the subscribed socket has completed its transport handshake. + pub async fn connect_ready( + endpoint: &str, + topic: &str, + expected_publisher_id: u64, + rcvhwm: i32, + ) -> Result { + Ok(Self { + stream: ZmqSubTransport::connect_single_consumer_ready(endpoint, topic, rcvhwm).await?, + expected_topic: topic.to_string(), + expected_publisher_id, + codec: Codec::default(), + }) + } + pub async fn connect_default( endpoint: &str, topic: &str, @@ -453,6 +716,36 @@ impl ValidatedZmqSource { } impl ZmqSubTransport { + /// Connect one consumer after the subscribed socket's transport handshake. + pub async fn connect_single_consumer_ready( + endpoint: &str, + topic: &str, + rcvhwm: i32, + ) -> Result { + let (socket, readiness) = connect_subscriber_with_readiness(endpoint, topic, rcvhwm)?; + readiness.wait_connected(endpoint).await?; + Self::single_consumer_stream(socket, topic) + } + + /// Connect a consumer after every broker transport handshake completes. + pub async fn connect_single_consumer_multiple_ready( + endpoints: &[String], + topic: &str, + ) -> Result { + let (first, rest) = endpoints + .split_first() + .context("Cannot connect to zero endpoints")?; + let (socket, readiness) = connect_subscriber_with_readiness(first, topic, ZMQ_RCVHWM)?; + for endpoint in rest { + socket.get_socket().connect(endpoint)?; + readiness.record_endpoint(socket.get_socket(), endpoint)?; + } + for endpoint in endpoints { + readiness.wait_connected(endpoint).await?; + } + Self::single_consumer_stream(socket, topic) + } + fn connect_socket(endpoint: &str, topic: &str) -> Result { Self::connect_socket_with_rcvhwm(endpoint, topic, ZMQ_RCVHWM) } @@ -541,7 +834,10 @@ impl ZmqSubTransport { Self::single_consumer_stream(socket, topic) } - fn single_consumer_stream(mut socket: Subscribe, topic: &str) -> Result { + fn single_consumer_stream(mut socket: S, topic: &str) -> Result + where + S: Stream> + Send + Unpin + 'static, + { let expected_topic = topic.as_bytes().to_vec(); let stream = stream! { @@ -744,6 +1040,102 @@ mod tests { use std::collections::HashSet; use tokio::time::{Duration, timeout}; + #[tokio::test] + async fn subscriber_readiness_waits_for_handshake_with_hostname() { + // Reserve the port with a plain TCP listener: a TCP connection alone is + // insufficient, and no ZMTP publisher exists until the explicit bind. + let reservation = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let port = reservation.local_addr().unwrap().port(); + let endpoint = format!("tcp://localhost:{port}"); + let (_socket, readiness) = + connect_subscriber_with_readiness(&endpoint, "kv-events", 128).unwrap(); + assert!( + timeout( + Duration::from_millis(25), + readiness.wait_connected(&endpoint) + ) + .await + .is_err(), + "socket construction or a plain TCP connection must not report ready" + ); + + drop(reservation); + let (_publisher, _) = + ZmqPubTransport::bind(&format!("tcp://127.0.0.1:{port}"), "kv-events") + .await + .unwrap(); + timeout(Duration::from_secs(5), readiness.wait_connected(&endpoint)) + .await + .unwrap_or_else(|error| { + panic!( + "publisher handshake did not arrive: {error}; aliases={:?}, connected={:?}", + readiness.endpoints.lock(), + readiness.connected.borrow() + ) + }) + .unwrap(); + } + + #[tokio::test] + async fn subscriber_readiness_teardown_preserves_shared_context() { + // Closing a monitored socket must not leave a libzmq I/O thread blocked + // trying to send to a monitor receiver that has already been dropped. + // Reuse the same context across enough connections to exercise every I/O + // thread, and require the first publication after each connection. + let context = shared_zmq_context().unwrap(); + let topic = "monitor-lifetime"; + for sequence in 0..16_u64 { + let publisher = context.socket(zmq::XPUB).unwrap(); + publisher.set_linger(0).unwrap(); + publisher.set_rcvtimeo(2_000).unwrap(); + publisher.bind("tcp://127.0.0.1:*").unwrap(); + let endpoint = publisher.get_last_endpoint().unwrap().unwrap(); + let mut stream = timeout(Duration::from_secs(2), async { + if sequence % 2 == 0 { + ZmqSubTransport::connect_single_consumer_ready(&endpoint, topic, 128).await + } else { + ZmqSubTransport::connect_single_consumer_multiple_ready(&[endpoint], topic) + .await + } + }) + .await + .expect("previous monitor teardown must not block another connection") + .unwrap(); + let payload = encoded_event(topic, 23, sequence); + let published = payload.clone(); + let send = tokio::task::spawn_blocking(move || { + // XPUB observes the subscription so this teardown regression + // does not depend on PUB/SUB's unrelated slow-joiner timing. + let subscription = publisher.recv_bytes(0).unwrap(); + assert_eq!(&subscription[1..], topic.as_bytes()); + assert_eq!(subscription[0], 1); + publisher + .send_multipart( + [ + topic.as_bytes().to_vec(), + 23_u64.to_be_bytes().to_vec(), + sequence.to_be_bytes().to_vec(), + Frame::new(published).encode().to_vec(), + ], + 0, + ) + .unwrap(); + publisher + }); + let message = timeout(Duration::from_secs(2), stream.next()) + .await + .expect("the first event must arrive after each connection") + .unwrap() + .unwrap(); + assert_eq!(message.publisher_id, 23); + assert_eq!(message.sequence, sequence); + assert_eq!(message.payload, payload); + let publisher = send.await.unwrap(); + drop(stream); + drop(publisher); + } + } + #[test] fn emfile_errno_selects_source_specific_guidance() { assert_eq!( diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index 6cd13fbb4902..236f769c69e0 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -69,6 +69,17 @@ local KV events to Dynamo without accepting inference requests. A TP-only follower that holds part of a replica but has no local KV publisher does not need a sidecar. +For KV routing, the frontend's `/health` stays unready until every registered +DP rank has an advertised relay and the router has connected its event +subscription. The leader registers before this gate so followers can discover +it. Each relay advertises only after its engine subscription has established a +connection and outbound publishing is available; creating a publisher object +alone does not mark it ready. This contract applies to single-node sidecars as +well. Launchers can wait for frontend readiness before sending requests. + +Readiness describes connected feeds, not durable event delivery: ZeroMQ PUB/SUB +does not acknowledge individual subscriptions or KV events. + Launch each sidecar separately with `--grpc-endpoint` pointing to its local SGLang server. This requires an SGLang build exposing node-local `kv_event_sources` through `GetServerInfo` on leaders and followers started diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index b5783176ed2a..4c3b51504e61 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -955,6 +955,20 @@ fn build_engine_config( "grpc_service".to_string(), Value::String("sglang.runtime.v1.SglangService".to_string()), ); + // The leader declares readiness for the whole worker, even when all KV + // publishers are remote. Source descriptors are validated before registration. + if discovery + .server_info + .get("kv_events") + .is_some_and(|value| !value.is_null()) + || discovery + .server_info + .get("kv_event_sources") + .and_then(Value::as_array) + .is_some_and(|sources| !sources.is_empty()) + { + runtime_data.insert("require_kv_event_source_readiness".into(), true.into()); + } if let Some(total_tokens) = hicache_native_offloading_capacity(&discovery.server_info, &discovery.model_info) { @@ -1278,6 +1292,10 @@ mod tests { let sources = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap(); + assert_eq!( + config.runtime_data["require_kv_event_source_readiness"], + true + ); assert_eq!( sources, [ @@ -1335,6 +1353,12 @@ mod tests { // An explicit empty source list must not fall back to all global ranks. discovery.server_info["kv_event_sources"] = json!([]); + config = + build_engine_config(&discovery, DisaggregationMode::Aggregated, None, None).unwrap(); + assert_eq!( + config.runtime_data["require_kv_event_source_readiness"], + true + ); assert!( discover_kv_event_sources(&discovery, &mut config, &endpoint) .unwrap() @@ -1346,6 +1370,16 @@ mod tests { ); assert_eq!(config.llm.as_ref().unwrap().data_parallel_size, Some(8)); + let mut disabled = discovery.clone(); + disabled.server_info["kv_events"] = serde_json::Value::Null; + let disabled_config = + build_engine_config(&disabled, DisaggregationMode::Aggregated, None, None).unwrap(); + assert!( + !disabled_config + .runtime_data + .contains_key("require_kv_event_source_readiness") + ); + // Engines without the new metadata field retain the old multinode guard. discovery .server_info diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index 32dc28fe9f9c..78c9d7d5dd10 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -135,7 +135,12 @@ impl HeadlessSidecar { wait_for_leader(&mut configs, group_id, &shutdown)).await .context("timed out waiting for SGLang leader; check namespace, component, endpoint, and dist_init_addr")??; let leader = validate_leader(metadata, &config)?; - let _publishers = start_publishers(&endpoint, metadata, worker_id, &config, &leader)?; + let _publishers = tokio::time::timeout( + self.transport.startup_deadline, + start_publishers(&endpoint, metadata, worker_id, &config, &leader), + ) + .await + .context("timed out connecting local SGLang KV feeds")??; tracing::info!(node_rank = metadata.node_rank, worker_id, group = group_id, local_ranks = ?metadata.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), "SGLang headless sidecar publishing local KV events"); @@ -280,7 +285,7 @@ fn validate_leader(context: &NodeMetadata, config: &ModelRuntimeConfig) -> Resul Ok(leader) } -fn start_publishers( +async fn start_publishers( endpoint: &Endpoint, context: &NodeMetadata, worker_id: u64, @@ -309,8 +314,12 @@ fn start_publishers( ) }) .collect::>>()?; - // There is no serving worker to mark this process ready. Publish readiness - // only after all local publishers have been created. + for publisher in &publishers { + publisher.wait_ready().await?; + } + // There is no serving worker to mark this process ready. Local health + // requires connected ingress and publication; frontend readiness separately + // requires subscriptions to every DP rank in the leader's configuration. endpoint .drt() .system_health() @@ -584,9 +593,13 @@ mod tests { #[tokio::test] async fn local_zmq_events_keep_global_rank_and_leader_identity_without_serving() { use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent}; - use dynamo_runtime::discovery::{DiscoveryQuery, EventSourceQuery}; - use dynamo_runtime::transports::event_plane::EventSubscriber; - use futures::SinkExt; + use dynamo_runtime::discovery::{ + DiscoveryInstance, DiscoveryQuery, EventChannelQuery, EventSourceQuery, EventTransport, + }; + use dynamo_runtime::transports::event_plane::{ + Codec, Frame, zmq_transport::connect_subscriber_with_readiness, + }; + use futures::{SinkExt, StreamExt}; let runtime = Runtime::from_current().unwrap(); let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) @@ -610,36 +623,53 @@ mod tests { let (context, _) = follower_metadata(&raw).unwrap(); let config = leader("A"); let metadata = validate_leader(&context, &config).unwrap(); - let mut subscriber = EventSubscriber::for_endpoint(&endpoint, KV_EVENT_SUBJECT) - .await - .unwrap() - .typed::>(); assert!(!drt.system_health().lock().get_health_status().0); - let publishers = start_publishers(&endpoint, &context, 42, &config, &metadata).unwrap(); + let publishers = start_publishers(&endpoint, &context, 42, &config, &metadata) + .await + .unwrap(); assert!(drt.system_health().lock().get_health_status().0); - // Repeat until the real ZMQ subscriptions are connected; no fixed sleep. - let received = tokio::time::timeout(Duration::from_secs(5), async { - let mut ticks = tokio::time::interval(Duration::from_millis(50)); - let mut seq = 0_u64; - loop { - tokio::select! { - batch = subscriber.next() => { - let (envelope, events) = batch.unwrap().unwrap(); - if let Some(event) = events.first() { - break (envelope, event.clone()); - } - } - _ = ticks.tick() => { - seq += 1; - let payload = rmp_serde::to_vec_named(&json!([ - 0.0, [{"type":"BlockRemoved", "block_hashes":[42]}], 4 - ])).unwrap(); - source.send(vec![Vec::new(), seq.to_be_bytes().to_vec(), payload]).await.unwrap(); - } - } - } - }).await.expect("headless relay should forward local events"); + let channels = drt + .discovery() + .list(DiscoveryQuery::EventChannels( + EventChannelQuery::endpoint_topic(endpoint.id(), KV_EVENT_SUBJECT), + )) + .await + .unwrap(); + assert_eq!(channels.len(), 1); + let DiscoveryInstance::EventChannel { + transport: EventTransport::Zmq { endpoint: address }, + .. + } = &channels[0] + else { + panic!("expected the relay's direct ZMQ channel"); + }; + let (mut subscriber, connection) = + connect_subscriber_with_readiness(address, KV_EVENT_SUBJECT, 1000).unwrap(); + tokio::time::timeout(Duration::from_secs(5), connection.wait_connected(address)) + .await + .unwrap() + .unwrap(); + let payload = rmp_serde::to_vec_named(&json!([ + 0.0, [{"type":"BlockRemoved", "block_hashes":[42]}], 4 + ])) + .unwrap(); + source + .send(vec![Vec::new(), 1_u64.to_be_bytes().to_vec(), payload]) + .await + .unwrap(); + let frames = tokio::time::timeout(Duration::from_secs(5), subscriber.next()) + .await + .expect("headless relay should forward its first local event") + .unwrap() + .unwrap(); + assert_eq!(frames.len(), 4); + let frame = Frame::decode(frames[3].as_ref()).unwrap(); + let envelope = Codec::default().decode_envelope(&frame.payload).unwrap(); + let mut events: Vec = + Codec::default().decode_payload(&envelope.payload).unwrap(); + assert_eq!(events.len(), 1); + let received = (envelope, events.remove(0)); assert_eq!(received.1.worker_id, 42); assert_eq!(received.1.event.dp_rank, 4); assert_ne!(received.0.publisher_id, 42); From 88666972d0fa0ef1907360ffb55de8429da149f2 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Thu, 17 Sep 2026 13:31:14 -0700 Subject: [PATCH 10/17] fix(sglang): address multinode launcher review feedback Signed-off-by: jain-ria --- lib/sidecar/sglang/README.md | 21 ++++++++++-------- ...uter.sh => multinode_kv_router_sidecar.sh} | 7 +++--- lib/sidecar/sglang/src/headless.rs | 2 -- lib/sidecar/sglang/tests/executable.rs | 22 +++++++++---------- 4 files changed, 26 insertions(+), 26 deletions(-) rename lib/sidecar/sglang/launch/{multinode_kv_router.sh => multinode_kv_router_sidecar.sh} (95%) diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index 236f769c69e0..7218cd20faca 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -85,15 +85,18 @@ SGLang server. This requires an SGLang build exposing node-local `kv_event_sources` through `GetServerInfo` on leaders and followers started with `--grpc-port`. -The `launch/multinode_kv_router.sh` example also requires SGLang's explicit -`"bind": true` KV-event option (included in -[SGLang #39659](https://github.com/sgl-project/sglang/pull/39659)). It binds each -rank's publisher to `127.0.0.1`, so the engine and its sidecar must share a -network namespace. KV events can include token IDs and request metadata; -loopback prevents remote access but does not isolate other processes in that -namespace. If overriding this to a non-loopback bind, restrict every publisher -port (`SGLANG_KV_EVENT_PORT + dp_rank`) to authorized consumers using a firewall -or network policy. +For `launch/multinode_kv_router_sidecar.sh`, use an SGLang build including commit +[`c50b251`](https://github.com/sgl-project/sglang/commit/c50b251da8e455dfbb594f8c3b7ccae6eca72129), +which provides node-local KV source discovery and explicit publisher binding. +The stock v0.5.19 image guidance below applies to the single-node examples. + +The multinode example also requires SGLang's explicit +`"bind": true` KV-event option. It binds each rank's publisher to `127.0.0.1`, so +the engine and its sidecar must share a network namespace. KV events can include +token IDs and request metadata; loopback prevents remote access but does not +isolate other processes in that namespace. If overriding this to a non-loopback +bind, restrict every publisher port (`SGLANG_KV_EVENT_PORT + dp_rank`) to authorized +consumers using a firewall or network policy. ## Deploy on Kubernetes (quick start) diff --git a/lib/sidecar/sglang/launch/multinode_kv_router.sh b/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh similarity index 95% rename from lib/sidecar/sglang/launch/multinode_kv_router.sh rename to lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh index 6c7787b6699b..637021650d13 100755 --- a/lib/sidecar/sglang/launch/multinode_kv_router.sh +++ b/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh @@ -3,7 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 # # One node of an externally managed SGLang attention-DP group with local KV events. -# Requires SGLang #39659 and this Dynamo sidecar on every node. Start the frontend separately. +# Requires node-local KV source discovery and this Dynamo sidecar on every node. +# Start the frontend separately. # Defaults: two nodes, one GPU per node, global TP=2 / attention DP=2. set -e @@ -21,7 +22,7 @@ if [[ "${1:-}" == --help || "${1:-}" == -h ]]; then echo "Required: NODE_RANK, DIST_INIT_ADDR" echo "Topology: NNODES=2, TP_SIZE=NNODES, DP_SIZE=NNODES (global sizes; PP=1)" echo " DP_SIZE must be divisible by NNODES so every node owns publishers." - echo "Model: MODEL=Qwen/Qwen3-0.6B, MAX_MODEL_LEN=4096, MAX_CONCURRENT_SEQS=32" + echo "Model: MODEL=Qwen/Qwen3-32B, MAX_MODEL_LEN=4096, MAX_CONCURRENT_SEQS=32" echo "Role: ROLE=aggregated|prefill|decode (default: aggregated)" echo "PD: SGLANG_BOOTSTRAP_HOST= on prefill node 0" echo "Ports: SGLANG_HTTP_PORT=30000, SGLANG_GRPC_PORT=30001, SGLANG_KV_EVENT_PORT=5557" @@ -52,7 +53,7 @@ if (( NNODES < 2 || TP_SIZE % DP_SIZE != 0 || DP_SIZE % NNODES != 0 )); then exit 1 fi -MODEL="${MODEL:-Qwen/Qwen3-0.6B}" +MODEL="${MODEL:-Qwen/Qwen3-32B}" ROLE="${ROLE:-aggregated}" SGLANG_PYTHON="${SGLANG_PYTHON:-python3}" SGLANG_HOST="${SGLANG_HOST:-0.0.0.0}" diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index 78c9d7d5dd10..7688395de8be 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -154,8 +154,6 @@ impl HeadlessSidecar { } } } - // Publisher Drop cancels local subscriptions; the follower never owns - // the leader's serving registration. } } diff --git a/lib/sidecar/sglang/tests/executable.rs b/lib/sidecar/sglang/tests/executable.rs index a00353c84162..a176830ebfe5 100644 --- a/lib/sidecar/sglang/tests/executable.rs +++ b/lib/sidecar/sglang/tests/executable.rs @@ -35,16 +35,14 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { } #[test] -fn both_modes_require_a_local_grpc_endpoint() { - for mode in [None, Some("--telemetry-only")] { - let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) - .args(mode) - .env_remove("DYN_SIDECAR_GRPC_ENDPOINT") - .output() - .unwrap(); - assert_eq!(output.status.code(), Some(2)); - let stderr = String::from_utf8_lossy(&output.stderr); - assert!(stderr.contains("the following required arguments were not provided")); - assert!(stderr.contains("--grpc-endpoint ")); - } +fn telemetry_mode_requires_a_local_grpc_endpoint() { + let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) + .arg("--telemetry-only") + .env_remove("DYN_SIDECAR_GRPC_ENDPOINT") + .output() + .unwrap(); + assert_eq!(output.status.code(), Some(2)); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("the following required arguments were not provided")); + assert!(stderr.contains("--grpc-endpoint ")); } From d088cdab36298479df5000fb08a74d7e5d31a99b Mon Sep 17 00:00:00 2001 From: jain-ria Date: Thu, 17 Sep 2026 13:57:53 -0700 Subject: [PATCH 11/17] revert(sglang): remove KV relay readiness expansion This reverts commit 50cb6a1d0e8195bd75d18ea1efa0694dff4ea78b. Keep the launcher rename, Qwen3-32B default, README clarification, and small review cleanups from 88666972d0fa0ef1907360ffb55de8429da149f2. Signed-off-by: jain-ria --- Cargo.lock | 1 - lib/llm/src/direct_zmq_sub_pool.rs | 47 +- lib/llm/src/discovery/model.rs | 27 +- lib/llm/src/discovery/model_manager.rs | 22 - lib/llm/src/discovery/watcher.rs | 11 - lib/llm/src/discovery/worker_set.rs | 15 - lib/llm/src/http/service.rs | 3 - lib/llm/src/http/service/health.rs | 2 +- .../http/service/kv_relay_readiness_tests.rs | 370 --------------- lib/llm/src/kv_router.rs | 15 - .../kv_router/indexer/recovery/broker_zmq.rs | 7 +- .../kv_router/indexer/recovery/direct_zmq.rs | 2 +- .../kv_router/indexer/recovery/subscriber.rs | 207 +-------- .../indexer/recovery/worker_query.rs | 26 -- .../kv_router/prefill_router/activation.rs | 33 -- lib/llm/src/kv_router/prefill_router/mod.rs | 36 -- lib/llm/src/kv_router/publisher/mod.rs | 59 +-- lib/llm/src/kv_router/publisher/tests.rs | 19 +- .../src/kv_router/publisher/zmq_listener.rs | 29 +- lib/llm/src/local_model/runtime_config.rs | 11 - lib/runtime/Cargo.toml | 1 - lib/runtime/src/distributed.rs | 6 +- lib/runtime/src/transports/event_plane/mod.rs | 52 +-- .../src/transports/event_plane/transport.rs | 5 - .../transports/event_plane/zmq_transport.rs | 436 +----------------- lib/sidecar/sglang/README.md | 11 - lib/sidecar/sglang/src/engine.rs | 34 -- lib/sidecar/sglang/src/headless.rs | 98 ++-- 28 files changed, 81 insertions(+), 1504 deletions(-) delete mode 100644 lib/llm/src/http/service/kv_relay_readiness_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 5459bb870d40..cc2c91f1120c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3024,7 +3024,6 @@ dependencies = [ "uuid", "validator", "xxhash-rust", - "zmq", ] [[package]] diff --git a/lib/llm/src/direct_zmq_sub_pool.rs b/lib/llm/src/direct_zmq_sub_pool.rs index 60d3932f58f6..640ad700471d 100644 --- a/lib/llm/src/direct_zmq_sub_pool.rs +++ b/lib/llm/src/direct_zmq_sub_pool.rs @@ -11,7 +11,6 @@ use std::{ }; use anyhow::Result; -use dynamo_runtime::transports::event_plane::zmq_transport::ZmqConnectionReadiness; use dynamo_runtime::transports::event_plane::{ Codec, DynamicZmqSubSocket, ValidatedEnvelope, ValidatedZmqSource, ZmqWireMessage, }; @@ -98,7 +97,6 @@ struct SocketGroup { command_tx: mpsc::UnboundedSender, cancel: CancellationToken, handle: JoinHandle<()>, - readiness: ZmqConnectionReadiness, } struct PoolInner { @@ -124,7 +122,6 @@ pub(crate) struct DirectZmqSubRegistration { publisher_id: u64, generation: u64, armed: bool, - readiness: ZmqConnectionReadiness, } pub(crate) enum DirectZmqSubConnection { @@ -190,35 +187,6 @@ impl Drop for DirectZmqSubRegistration { } impl DirectZmqSubPool { - /// Unlike socket construction, readiness requires a completed transport handshake. - pub(crate) async fn connect_ready( - &self, - publisher_id: u64, - endpoint: &str, - generation: u64, - ) -> Result { - if self.endpoints_per_sub == 1 { - anyhow::ensure!( - !self.inner.lock().closed, - "direct-ZMQ socket pool is closed" - ); - let source = - ValidatedZmqSource::connect_ready(endpoint, &self.topic, publisher_id, self.rcvhwm) - .await?; - anyhow::ensure!( - !self.inner.lock().closed, - "direct-ZMQ socket pool is closed" - ); - return Ok(DirectZmqSubConnection::Dedicated(source)); - } - let registration = self - .register_grouped(publisher_id, endpoint, generation) - .await?; - // The group task keeps forwarding existing peers while this peer connects. - registration.readiness.wait_connected(endpoint).await?; - Ok(DirectZmqSubConnection::Grouped(registration)) - } - pub(crate) fn new( topic: impl Into>, endpoints_per_sub: usize, @@ -274,7 +242,7 @@ impl DirectZmqSubPool { ) -> Result { let (sender, receiver) = mpsc::channel(self.rcvhwm as usize); let disconnected = CancellationToken::new(); - let (group_id, completion, readiness) = { + let (group_id, completion) = { let mut inner = self.inner.lock(); Self::reap_failed_groups(&mut inner); anyhow::ensure!(!inner.closed, "direct-ZMQ socket pool is closed"); @@ -318,15 +286,12 @@ impl DirectZmqSubPool { group.assignments.remove(&publisher_id); anyhow::bail!("direct-ZMQ socket group stopped"); } - (group_id, Some(completion), group.readiness.clone()) + (group_id, Some(completion)) } else { // libzmq connects asynchronously. Keep group creation serialized so // concurrent registrations cannot create excess transient sockets. - let (socket, readiness) = DynamicZmqSubSocket::connect_with_readiness( - endpoint, - &self.topic, - self.rcvhwm, - )?; + let socket = + DynamicZmqSubSocket::connect_with_rcvhwm(endpoint, &self.topic, self.rcvhwm)?; let group_id = inner.next_group_id; inner.next_group_id = inner.next_group_id.wrapping_add(1); let (command_tx, command_rx) = mpsc::unbounded_channel(); @@ -355,10 +320,9 @@ impl DirectZmqSubPool { command_tx, cancel, handle, - readiness: readiness.clone(), }, ); - (group_id, None, readiness) + (group_id, None) } }; let registration = DirectZmqSubRegistration { @@ -369,7 +333,6 @@ impl DirectZmqSubPool { receiver, disconnected, armed: true, - readiness, }; let Some(completion) = completion else { diff --git a/lib/llm/src/discovery/model.rs b/lib/llm/src/discovery/model.rs index 6b98cb07cb1a..bfc37dbcc178 100644 --- a/lib/llm/src/discovery/model.rs +++ b/lib/llm/src/discovery/model.rs @@ -385,7 +385,7 @@ impl Model { warn_legacy_readiness_once(&self.name, wsets[0].namespace()); } NamespaceReadinessEval { - ready: eval.ready && wsets.iter().all(|ws| ws.kv_event_sources_ready()), + ready: eval.ready, has_legacy: eval.has_legacy, legacy_live_workers: eval.legacy_live_workers, present: eval.present, @@ -408,12 +408,6 @@ impl Model { self.first_ready_workers().is_some() } - pub(crate) fn kv_event_sources_ready(&self) -> bool { - self.worker_sets - .iter() - .all(|entry| entry.value().kv_event_sources_ready()) - } - /// Structured per-namespace worker readiness for this model — the data /// behind the `GET /v1/models/{model}/ready` observability endpoint. /// @@ -477,8 +471,6 @@ impl Model { } else { None } - } else if wsets.iter().any(|ws| !ws.kv_event_sources_ready()) { - Some("waiting for required KV event feeds".to_string()) } else if eval.has_legacy { Some("legacy worker(s) present but no live worker".to_string()) } else if !eval.ambiguous.is_empty() { @@ -1337,23 +1329,6 @@ mod tests { assert!(!model.is_workers_ready("dynamo")); } - #[test] - fn kv_relay_readiness_does_not_withdraw_serving_membership() { - let model = Model::new("llama".to_string()); - let (mut worker_set, _workers) = - ws_with_type("dynamo", "mdc", WorkerType::Aggregated, vec![], vec![7]); - let (ready_tx, ready_rx) = watch::channel(false); - Arc::get_mut(&mut worker_set).unwrap().kv_event_readiness = Some(ready_rx); - model.add_worker_set("dynamo".to_string(), worker_set.clone()); - assert_eq!(worker_set.worker_count(), 1); - assert!(!model.is_workers_ready("dynamo")); - ready_tx.send(true).unwrap(); - assert!(model.is_workers_ready("dynamo")); - ready_tx.send(false).unwrap(); - assert!(!model.is_workers_ready("dynamo")); - assert_eq!(worker_set.worker_count(), 1); - } - #[test] fn readiness_pd_pair_ready() { let model = Model::new("llama".to_string()); diff --git a/lib/llm/src/discovery/model_manager.rs b/lib/llm/src/discovery/model_manager.rs index 84fde0ed9a34..f3aefdea39ad 100644 --- a/lib/llm/src/discovery/model_manager.rs +++ b/lib/llm/src/discovery/model_manager.rs @@ -157,8 +157,6 @@ pub(crate) struct RemovedDiscoveryGroup { /// /// Note: Don't implement Clone for this, put it in an Arc instead. pub struct ModelManager { - /// KV frontends must not report ready before any model and its feeds are usable. - require_ready_model: AtomicBool, /// Model name → Model (which contains WorkerSets with engines) models: DashMap>, @@ -217,7 +215,6 @@ impl Default for ModelManager { impl ModelManager { pub fn new() -> Self { Self { - require_ready_model: AtomicBool::new(false), models: DashMap::new(), catalog: ArcSwap::from_pointee(CommittedCatalog::default()), cards: DashMap::new(), @@ -1098,25 +1095,6 @@ impl ModelManager { .any(|model| model.is_ready_to_serve()) } - pub(crate) fn require_model_readiness(&self) { - self.require_ready_model.store(true, Ordering::Relaxed); - } - - pub(crate) fn frontend_routing_ready(&self) -> bool { - if !self.require_ready_model.load(Ordering::Relaxed) { - return true; - } - let catalog = self.catalog.load(); - catalog - .models - .values() - .any(|model| model.is_ready_to_serve()) - && catalog - .models - .values() - .all(|model| model.kv_event_sources_ready()) - } - pub fn model_display_names(&self) -> HashSet { self.catalog .load() diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index ae3563446c8e..6f7859e72260 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -317,11 +317,6 @@ where require_typed_worker_role: bool, worker_selector_factory: WorkerSelectorFactory, ) -> Self { - if router_config.router_mode == RouterMode::KV - && router_config.kv_router_config.use_kv_events - { - model_manager.require_model_readiness(); - } Self { manager: model_manager, drt: runtime, @@ -443,11 +438,6 @@ where // are process-local, so preserve them when the MDC supplies the base config. let router_config = effective_router_config(card.router_config.as_ref(), &self.router_config); - if router_config.router_mode == RouterMode::KV - && router_config.kv_router_config.use_kv_events - { - self.manager.require_model_readiness(); - } let component = self .drt @@ -621,7 +611,6 @@ where Arc::get_mut(&mut chooser) .expect("new KV chooser must have one owner") .set_teardown_task_guard(allocator_trim.clone()); - worker_set.kv_event_readiness = chooser.kv_event_readiness(); Some(chooser) } else { None diff --git a/lib/llm/src/discovery/worker_set.rs b/lib/llm/src/discovery/worker_set.rs index b1004382ded8..352a3cc708a2 100644 --- a/lib/llm/src/discovery/worker_set.rs +++ b/lib/llm/src/discovery/worker_set.rs @@ -234,9 +234,6 @@ pub struct WorkerSet { /// None for in-process models (http/grpc) which don't have a discovery client. instance_count_rx: Option>>, - /// Frontend-local KV subscription readiness, independent of serving registration. - pub(crate) kv_event_readiness: Option>, - /// Cancels background work created while materializing this WorkerSet. lifecycle_cancellation: Option, @@ -269,7 +266,6 @@ impl WorkerSet { prefill_router: None, encoder_router: None, instance_count_rx: None, - kv_event_readiness: None, lifecycle_cancellation: None, allocator_trim: None, allocator_trim_wrapped: false, @@ -437,16 +433,6 @@ impl WorkerSet { } } - pub(crate) fn kv_event_sources_ready(&self) -> bool { - self.kv_event_readiness - .as_ref() - .is_none_or(|ready| *ready.borrow()) - && self - .prefill_router - .as_ref() - .is_none_or(|router| router.kv_event_sources_ready()) - } - /// Store the instance watcher from the Client's discovery system. /// Must be called before the WorkerSet is wrapped in Arc. pub fn set_instance_watcher(&mut self, rx: watch::Receiver>) { @@ -524,7 +510,6 @@ impl WorkerSet { prefill_router: self.prefill_router.clone(), encoder_router: self.encoder_router.clone(), instance_count_rx: self.instance_count_rx.clone(), - kv_event_readiness: self.kv_event_readiness.clone(), lifecycle_cancellation: None, allocator_trim: None, allocator_trim_wrapped: false, diff --git a/lib/llm/src/http/service.rs b/lib/llm/src/http/service.rs index 4d77d16cb1ef..f4d2a5fbe838 100644 --- a/lib/llm/src/http/service.rs +++ b/lib/llm/src/http/service.rs @@ -34,9 +34,6 @@ pub mod realtime; pub mod service_v2; pub mod sglang_generate; -#[cfg(test)] -mod kv_relay_readiness_tests; - pub use axum; pub use frontend_extension::{ FrontendExtensionContext, FrontendRouteExtension, FrontendRouteSet, diff --git a/lib/llm/src/http/service/health.rs b/lib/llm/src/http/service/health.rs index 81441cbbf388..1b6c1da8ee86 100644 --- a/lib/llm/src/http/service/health.rs +++ b/lib/llm/src/http/service/health.rs @@ -63,7 +63,7 @@ async fn live_handler( async fn health_handler( axum::extract::State(state): axum::extract::State>, ) -> impl IntoResponse { - if !state.is_ready() || !state.manager().frontend_routing_ready() { + if !state.is_ready() { return ( StatusCode::SERVICE_UNAVAILABLE, Json(json!({ diff --git a/lib/llm/src/http/service/kv_relay_readiness_tests.rs b/lib/llm/src/http/service/kv_relay_readiness_tests.rs deleted file mode 100644 index 439e35ee0c51..000000000000 --- a/lib/llm/src/http/service/kv_relay_readiness_tests.rs +++ /dev/null @@ -1,370 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Exercise the HTTP readiness gate with real engine and router ZMQ feeds. - -use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, -}; -use std::time::Duration; - -use dynamo_kv_router::{ - config::KvRouterConfig, - protocols::{KV_EVENT_SUBJECT, WorkerWithDpRank}, -}; -use dynamo_runtime::{ - CancellationToken, Runtime, - discovery::{DiscoveryQuery, EventSourceQuery}, - distributed::{DistributedConfig, DistributedRuntime}, - pipeline::{ - AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, ResponseStream, SingleIn, - async_trait, - }, -}; -use futures::SinkExt; -use serde_json::json; -use tmq::AsZmqSocket; -use tokio::sync::Mutex; - -use super::service_v2::HttpService; -use crate::{ - discovery::WorkerSet, - kv_router::publisher::{KvEventPublisher, KvEventSourceConfig}, - local_model::{register_model_card, runtime_config::SGLANG_GENERATE_CAPABILITY}, - model_card::ModelDeploymentCard, - model_type::{ModelInput, ModelType}, - protocols::{ - Annotated, - common::llm_backend::{LLMEngineOutput, PreprocessedRequest}, - }, - worker_type::WorkerType, -}; - -struct RequestKvEngine { - sources: Vec>>, - calls: AtomicUsize, -} - -#[async_trait] -impl AsyncEngine, ManyOut>, Error> - for RequestKvEngine -{ - async fn generate( - &self, - request: SingleIn, - ) -> Result>, Error> { - let (request, context) = request.transfer(()); - let rank = request - .routing - .as_ref() - .and_then(|routing| routing.dp_rank) - .unwrap(); - self.calls.fetch_add(1, Ordering::SeqCst); - // Each first request emits exactly one native KV event. Re-sending it - // would conceal a subscription that became ready too early. - let payload = rmp_serde::to_vec_named(&json!([ - 0.0, [{"type":"BlockStored", "block_hashes":[rank as u64 + 10], - "parent_block_hash":null, "token_ids":request.token_ids, - "block_size":4, "lora_id":null}], rank - ]))?; - self.sources[rank as usize] - .lock() - .await - .as_mut() - .expect("ready rank must have an engine publisher") - .send(vec![Vec::new(), 1_u64.to_be_bytes().to_vec(), payload]) - .await?; - let output = LLMEngineOutput { - token_ids: vec![42], - engine_data: Some(json!({"sglang_response":{ - "text":"first-token", "output_ids":[42], - "meta_info":{"finish_reason":{"type":"length","length":1}} - }})), - ..Default::default() - }; - Ok(ResponseStream::new( - Box::pin(futures::stream::iter([Annotated::from_data(output)])), - context.context(), - )) - } -} - -async fn kv_relay_readiness(dp_size: u32) { - let runtime = Runtime::from_current().unwrap(); - let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) - .await - .unwrap(); - let namespace = format!("kv-relay-readiness-{dp_size}"); - let endpoint = drt - .namespace(&namespace) - .unwrap() - .component("backend") - .unwrap() - .endpoint("generate"); - let worker_id = drt.discovery().instance_id(); - let mut card = ModelDeploymentCard::with_name_only("kv-relay-model"); - card.model_input = ModelInput::Tokens; - card.model_type = ModelType::Chat; - card.worker_type = Some(WorkerType::Aggregated); - card.kv_cache_block_size = 4; - card.runtime_config.data_parallel_size = dp_size; - card.runtime_config - .set_engine_specific(SGLANG_GENERATE_CAPABILITY, true) - .unwrap(); - card.runtime_config - .set_engine_specific("require_kv_event_source_readiness", true) - .unwrap(); - - // This is the leader registration that telemetry followers discover. It - // must exist before the final rank's relay can begin starting. - endpoint.register_endpoint_instance().await.unwrap(); - register_model_card(&endpoint, &card).await.unwrap(); - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let service = HttpService::builder() - .port(port) - .discovery(Some(drt.discovery())) - .enable_engine_apis(true) - .build() - .unwrap(); - let manager = service.state().manager_clone(); - manager.require_model_readiness(); - let router = manager - .kv_chooser_for_with_worker_role( - &endpoint, - 4, - Some(KvRouterConfig { - skip_initial_worker_wait: true, - router_event_threads: 1, - router_track_active_blocks: false, - ..Default::default() - }), - None, - Some(WorkerType::Aggregated), - "decode", - Some(card.name().to_owned()), - false, - ) - .await - .unwrap(); - let mut configs = manager - .get_or_create_runtime_config_watcher(&endpoint) - .await - .unwrap(); - tokio::time::timeout(Duration::from_secs(5), async { - while !configs.borrow().contains_key(&worker_id) { - configs.changed().await.unwrap(); - } - }) - .await - .expect("the leader must be discoverable without follower relays"); - - let zmq_context = tmq::Context::new(); - let mut sources = Vec::new(); - let mut addresses = Vec::new(); - for rank in 0..dp_size { - if rank == dp_size - 1 { - let reserved = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - addresses.push(format!("tcp://{}", reserved.local_addr().unwrap())); - sources.push(Mutex::new(None)); - // Keep the final engine publisher absent until after its relay - // object exists, distinguishing construction from connection. - drop(reserved); - continue; - } - let source = tmq::publish::publish(&zmq_context) - .set_linger(0) - .bind("tcp://127.0.0.1:*") - .unwrap(); - addresses.push(source.get_socket().get_last_endpoint().unwrap().unwrap()); - sources.push(Mutex::new(Some(source))); - } - let engine = Arc::new(RequestKvEngine { - sources, - calls: AtomicUsize::new(0), - }); - let mut worker_set = WorkerSet::new(namespace.clone(), card.mdcsum().to_owned(), card); - worker_set.generate_engine = Some(engine.clone()); - worker_set.kv_event_readiness = router.kv_event_readiness(); - worker_set.set_instance_watcher(router.client().instance_avail_watcher()); - assert!(manager.add_worker_set("kv-relay-model", &namespace, worker_set)); - let stop = CancellationToken::new(); - let server_stop = stop.clone(); - let http_task = - tokio::spawn(async move { service.run_with_listener(server_stop, listener).await }); - let client = reqwest::Client::new(); - let base = format!("http://127.0.0.1:{port}"); - let mut relays = Vec::new(); - - for rank in 0..dp_size { - // For DP=2 this deliberately holds back the follower after the leader - // is registered and its local feed is connected. DP=1 must obey the - // same rule for its only feed. - assert_eq!( - client - .get(format!("{base}/health")) - .send() - .await - .unwrap() - .status(), - reqwest::StatusCode::SERVICE_UNAVAILABLE - ); - assert!(!router.kv_event_sources_ready()); - assert!(configs.borrow().contains_key(&worker_id)); - assert_eq!( - drt.discovery() - .list(DiscoveryQuery::Endpoint { - namespace: namespace.clone(), - component: "backend".into(), - endpoint: "generate".into(), - }) - .await - .unwrap() - .len(), - 1 - ); - assert_eq!(engine.calls.load(Ordering::SeqCst), 0); - let rejected = client.post(format!("{base}/generate")) - .header("x-data-parallel-rank", rank.to_string()) - .json(&json!({"input_ids":[11,12,13,14], "sampling_params":{"max_new_tokens":1}, "stream":true})) - .send().await.unwrap(); - assert_eq!(rejected.status(), reqwest::StatusCode::SERVICE_UNAVAILABLE); - assert_eq!(engine.calls.load(Ordering::SeqCst), 0); - - let relay = KvEventPublisher::new_with_local_indexer_and_worker_id_at( - endpoint.clone(), - endpoint.id(), - Some(worker_id), - 4, - Some(KvEventSourceConfig::Zmq { - endpoint: addresses[rank as usize].clone(), - topic: String::new(), - image_token_id: None, - video_token_id: None, - }), - false, - rank, - None, - ) - .unwrap(); - if rank == dp_size - 1 { - assert!( - tokio::time::timeout(Duration::from_millis(100), relay.wait_ready()) - .await - .is_err(), - "creating a relay without its engine feed must not mark it ready" - ); - assert_eq!( - drt.discovery() - .list(DiscoveryQuery::EventSources( - EventSourceQuery::endpoint_topic(endpoint.id(), KV_EVENT_SUBJECT) - )) - .await - .unwrap() - .len(), - rank as usize, - "an unconnected feed must not advertise a ready source" - ); - assert!(configs.borrow().contains_key(&worker_id)); - assert_eq!( - drt.discovery() - .list(DiscoveryQuery::Endpoint { - namespace: namespace.clone(), - component: "backend".into(), - endpoint: "generate".into(), - }) - .await - .unwrap() - .len(), - 1 - ); - assert_eq!( - client - .get(format!("{base}/health")) - .send() - .await - .unwrap() - .status(), - reqwest::StatusCode::SERVICE_UNAVAILABLE - ); - let source = tmq::publish::publish(&zmq_context) - .set_linger(0) - .bind(&addresses[rank as usize]) - .unwrap(); - *engine.sources[rank as usize].lock().await = Some(source); - } - tokio::time::timeout(Duration::from_secs(5), relay.wait_ready()) - .await - .unwrap() - .unwrap(); - relays.push(relay); - } - tokio::time::timeout(Duration::from_secs(5), async { - loop { - if client - .get(format!("{base}/health")) - .send() - .await - .unwrap() - .status() - .is_success() - { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("all connected rank feeds must make the frontend ready"); - assert!(router.kv_event_sources_ready()); - - // No warmup events, fixed sleeps, or retrying inference requests are used. - // The first admitted request on every rank must reach the router indexer. - for rank in 0..dp_size { - let tokens = vec![11 + rank * 10, 12, 13, 14]; - let response = client - .post(format!("{base}/generate")) - .header("x-data-parallel-rank", rank.to_string()) - .json( - &json!({"input_ids":tokens, "sampling_params":{"max_new_tokens":1}, "stream":true}), - ) - .send() - .await - .unwrap(); - assert_eq!(response.status(), reqwest::StatusCode::OK); - assert!(response.text().await.unwrap().contains("first-token")); - tokio::time::timeout(Duration::from_secs(5), async { - while router - .get_overlap_blocks( - &tokens, - None, - WorkerWithDpRank::new(worker_id, rank), - None, - None, - ) - .await - .unwrap() - != 1 - { - tokio::task::yield_now().await; - } - }) - .await - .expect("the first request's single KV event must reach the router"); - } - assert_eq!(engine.calls.load(Ordering::SeqCst), dp_size as usize); - stop.cancel(); - http_task.await.unwrap().unwrap(); - drop(relays); - drt.shutdown(); -} - -#[tokio::test] -async fn frontend_waits_for_delayed_follower_and_receives_first_request_events() { - // Single-node and leader/follower startup must obey the same contract. - for dp_size in [1, 2] { - kv_relay_readiness(dp_size).await; - } -} diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index e0bfb886d3e3..ff9c69d5e321 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -812,8 +812,6 @@ where endpoint.clone(), indexer.clone(), membership_watch, - workers_with_configs.clone(), - client.instance_avail_watcher(), block_size, model_name.clone().unwrap_or_else(|| "unknown".to_string()), worker_role, @@ -904,19 +902,6 @@ where &self.indexer } - pub(crate) fn kv_event_readiness(&self) -> Option> { - self.kv_event_subscription - .as_ref() - .map(|subscription| subscription.readiness()) - } - - /// Whether every opted-in worker's required KV feeds are connected to this router. - pub fn kv_event_sources_ready(&self) -> bool { - self.kv_event_subscription - .as_ref() - .is_none_or(|subscription| subscription.is_ready()) - } - pub fn kv_router_config(&self) -> &KvRouterConfig { &self.kv_router_config } diff --git a/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs b/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs index fe0d64a350bf..528e895ac0f3 100644 --- a/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs +++ b/lib/llm/src/kv_router/indexer/recovery/broker_zmq.rs @@ -271,7 +271,6 @@ pub(super) async fn run_broker_zmq_supervisor( let mut retry_delay = INITIAL_BACKOFF; loop { - client.clear_transport_readiness(); let view = membership_watch.borrow_and_update().clone(); update_mismatch_metric( &status_metrics, @@ -283,7 +282,7 @@ pub(super) async fn run_broker_zmq_supervisor( ); let subscriber = if let Some(kv_state_endpoint) = view.resolved_kv_state_endpoint() { - match EventSubscriber::for_endpoint_id_with_transport_ready( + match EventSubscriber::for_endpoint_id_with_transport( component.drt(), kv_state_endpoint, KV_EVENT_SUBJECT, @@ -366,7 +365,6 @@ pub(super) async fn run_broker_zmq_supervisor( ScopeExit::Rebind => retry_delay = INITIAL_BACKOFF, ScopeExit::Retry => { ingress_metrics.increment_lifecycle("reconnect"); - client.clear_transport_readiness(); let view = client.sync_membership().await; update_subscription_failure_metric( &status_metrics, @@ -425,7 +423,6 @@ async fn consume_scope( membership_tx, membership_cancel.clone(), )); - client.publish_ready_sources(); let exit = loop { tokio::select! { _ = cancellation_token.cancelled() => break ScopeExit::Stop, @@ -445,7 +442,6 @@ async fn consume_scope( ); active_publishers = self::active_publishers(&view); lanes.reconcile(&active_publishers); - client.publish_ready_sources(); } Some(MembershipUpdate::Rebind) => break ScopeExit::Rebind, Some(MembershipUpdate::Stop) => break ScopeExit::Stop, @@ -472,7 +468,6 @@ async fn consume_scope( } }; membership_cancel.cancel(); - client.clear_transport_readiness(); drop(membership_rx); if let Err(error) = membership_handle.await && !error.is_cancelled() diff --git a/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs b/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs index 62ca627c9069..9e41ccbef394 100644 --- a/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs +++ b/lib/llm/src/kv_router/indexer/recovery/direct_zmq.rs @@ -732,7 +732,7 @@ async fn run_source( let connection = tokio::select! { biased; _ = cancel.cancelled() => return, - connection = group_pool.connect_ready(publisher_id, &endpoint, task_generation) => connection, + connection = group_pool.connect(publisher_id, &endpoint, task_generation) => connection, }; let mut connection = match connection { Ok(connection) => connection, diff --git a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs b/lib/llm/src/kv_router/indexer/recovery/subscriber.rs index 94caea4b2878..9dad1493a5d5 100644 --- a/lib/llm/src/kv_router/indexer/recovery/subscriber.rs +++ b/lib/llm/src/kv_router/indexer/recovery/subscriber.rs @@ -1,14 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::{ - collections::{HashMap, HashSet}, - sync::Arc, - time::Duration, -}; +use std::{sync::Arc, time::Duration}; use anyhow::Result; -use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent, WorkerId, WorkerWithDpRank}; +use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent}; use dynamo_runtime::{ component::{Component, Endpoint}, discovery::EventTransportKind, @@ -16,7 +12,7 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, transports::event_plane::{EventSubscriber, TypedEventSubscriber, uses_direct_zmq}, }; -use tokio::sync::{Semaphore, oneshot, watch}; +use tokio::sync::{Semaphore, oneshot}; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; @@ -26,12 +22,8 @@ use super::{ worker_query::WorkerQueryClient, }; use crate::{ - discovery::{ - KvSourceId, KvSourceMembershipView, KvSourceMembershipWatch, KvSourceStatus, - RuntimeConfigWatch, - }, + discovery::{KvSourceMembershipView, KvSourceMembershipWatch, KvSourceStatus}, kv_router::{Indexer, KvEventSourceRequirement, metrics::RouterWorkerStatusMetrics}, - local_model::runtime_config::ModelRuntimeConfig, worker_type::WorkerType, }; @@ -67,7 +59,6 @@ async fn run_subscription_supervisor( let mut retry_delay = SUBSCRIPTION_INITIAL_BACKOFF; loop { - client.clear_transport_readiness(); let view = membership_watch.borrow_and_update().clone(); update_mismatch_metric( &metrics, @@ -134,9 +125,6 @@ async fn run_subscription_supervisor( continue; } client.sync_membership().await; - if subscriber.is_some() { - client.publish_ready_sources(); - } // Subscriber construction establishes buffering before membership activation starts // initial recovery. Re-reading the watch above rejects a stale endpoint binding. @@ -175,7 +163,6 @@ async fn run_subscription_supervisor( continue; } ScopeExit::Retry => { - client.clear_transport_readiness(); let view = client.sync_membership().await; update_subscription_failure_metric( &metrics, @@ -239,7 +226,6 @@ async fn consume_scope( if view.resolved_kv_state_endpoint() != Some(kv_state_endpoint) { return ScopeExit::Rebind; } - client.publish_ready_sources(); } result = subscriber.next() => { let Some(result) = result else { @@ -383,18 +369,9 @@ pub(crate) struct KvEventSubscriptionHandle { completions: Vec>, runtime: tokio::runtime::Handle, task_guard: Option, - readiness: watch::Receiver, } impl KvEventSubscriptionHandle { - pub(crate) fn readiness(&self) -> watch::Receiver { - self.readiness.clone() - } - - pub(crate) fn is_ready(&self) -> bool { - *self.readiness.borrow() - } - pub(crate) fn set_task_guard( &mut self, task_guard: dynamo_runtime::engine::EngineContextGuard, @@ -431,8 +408,6 @@ pub async fn start_subscriber( endpoint: Endpoint, indexer: Indexer, membership_watch: KvSourceMembershipWatch, - runtime_configs: RuntimeConfigWatch, - admitted_workers: watch::Receiver>, block_size: u32, model: String, worker_role: Option, @@ -476,13 +451,6 @@ pub async fn start_subscriber( cancel.child_token(), ) .await?; - let (readiness, readiness_completion) = start_readiness_watch( - runtime_configs, - admitted_workers, - membership_watch.fork_receiver(), - client.ready_sources(), - cancel.child_token(), - ); let health_completion = source_health::spawn( membership_watch.fork_receiver(), model.clone(), @@ -519,15 +487,9 @@ pub async fn start_subscriber( let cancel = cancellation_guard.disarm(); return Ok(KvEventSubscriptionHandle { cancel, - completions: vec![ - completion_rx, - health_completion, - state_completion, - readiness_completion, - ], + completions: vec![completion_rx, health_completion, state_completion], runtime, task_guard: None, - readiness, }); } @@ -561,15 +523,9 @@ pub async fn start_subscriber( let cancel = cancellation_guard.disarm(); return Ok(KvEventSubscriptionHandle { cancel, - completions: vec![ - completion_rx, - health_completion, - state_completion, - readiness_completion, - ], + completions: vec![completion_rx, health_completion, state_completion], runtime, task_guard: None, - readiness, }); } @@ -602,89 +558,12 @@ pub async fn start_subscriber( let cancel = cancellation_guard.disarm(); Ok(KvEventSubscriptionHandle { cancel, - completions: vec![ - completion_rx, - health_completion, - state_completion, - readiness_completion, - ], + completions: vec![completion_rx, health_completion, state_completion], runtime, task_guard: None, - readiness, - }) -} - -fn required_sources_ready( - configs: &HashMap, - membership: &KvSourceMembershipView, - connected: &HashSet, -) -> bool { - // The serving card may arrive before its runtime-config watch has caught up. - if configs.is_empty() { - return false; - } - configs.iter().all(|(&worker_id, config)| { - if !config.requires_kv_event_source_readiness() { - return true; - } - let expected_endpoint = config - .kv_state_endpoint - .as_ref() - .unwrap_or(&membership.serving_endpoint); - if membership.resolved_kv_state_endpoint() != Some(expected_endpoint) { - return false; - } - let Ok(ranks) = config.data_parallel_rank_range() else { - return false; - }; - ranks.into_iter().all(|dp_rank| { - membership - .status(&WorkerWithDpRank::new(worker_id, dp_rank)) - .and_then(KvSourceStatus::active_source) - .is_some_and(|source| connected.contains(&source.source_id())) - }) }) } -fn start_readiness_watch( - mut configs: RuntimeConfigWatch, - mut admitted_workers: watch::Receiver>, - mut membership: KvSourceMembershipWatch, - mut connected: watch::Receiver>, - cancel: CancellationToken, -) -> (watch::Receiver, oneshot::Receiver<()>) { - let (sender, readiness) = watch::channel(false); - let (completion_tx, completion) = oneshot::channel(); - tokio::spawn(async move { - loop { - let admitted = admitted_workers.borrow_and_update().clone(); - let mut current_configs = configs.borrow_and_update().clone(); - current_configs.retain(|worker, _| admitted.contains(worker)); - let ready = current_configs.len() == admitted.len() - && required_sources_ready( - ¤t_configs, - &membership.borrow_and_update(), - &connected.borrow_and_update(), - ); - sender.send_if_modified(|current| { - let changed = *current != ready; - *current = ready; - changed - }); - tokio::select! { - _ = cancel.cancelled() => break, - changed = configs.changed() => if changed.is_err() { break; }, - changed = admitted_workers.changed() => if changed.is_err() { break; }, - changed = membership.changed() => if changed.is_err() { break; }, - changed = connected.changed() => if changed.is_err() { break; }, - } - } - sender.send_replace(false); - let _ = completion_tx.send(()); - }); - (readiness, completion) -} - pub(crate) struct RecoverySupervisor { client: Arc>, cancel: CancellationToken, @@ -762,76 +641,6 @@ mod tests { use crate::discovery::{KvEventSource, KvStateEndpointResolution}; - #[test] - fn readiness_requires_connected_sources_for_every_opted_in_rank() { - let mut config = ModelRuntimeConfig { - kv_event_publishing_enabled: Some(true), - ..Default::default() - }; - config.runtime_data.insert( - "require_kv_event_source_readiness".to_string(), - serde_json::json!(true), - ); - for size in [1, 2] { - config.data_parallel_size = size; - let configs = HashMap::from([(7, config.clone())]); - let mut view = metric_view(KvSourceStatus::Missing, Some(true)); - for rank in 0..size { - view.sources - .insert(WorkerWithDpRank::new(7, rank), KvSourceStatus::Missing); - } - let mut connected = HashSet::new(); - assert!(!required_sources_ready(&configs, &view, &connected)); - for rank in 0..size { - let worker = WorkerWithDpRank::new(7, rank); - let source = KvEventSource { - kv_state_endpoint: view.serving_endpoint.clone(), - worker, - publisher_id: 11 + u64::from(rank), - recovery_target: None, - }; - view.sources - .insert(worker, KvSourceStatus::ActiveLiveOnly(source.clone())); - // Advertising an object alone does not establish its subscription. - assert!(!required_sources_ready(&configs, &view, &connected)); - connected.insert(source.source_id()); - assert_eq!( - required_sources_ready(&configs, &view, &connected), - rank + 1 == size - ); - } - let follower = WorkerWithDpRank::new(7, size - 1); - view.sources.insert(follower, KvSourceStatus::Missing); - assert!( - !required_sources_ready(&configs, &view, &connected), - "relay removal withdraws readiness" - ); - } - } - - #[test] - fn readiness_preserves_legacy_and_disabled_event_workers() { - let mut config = ModelRuntimeConfig::default(); - let view = metric_view(KvSourceStatus::Missing, None); - let connected = HashSet::new(); - assert!(!required_sources_ready(&HashMap::new(), &view, &connected)); - assert!(required_sources_ready( - &HashMap::from([(7, config.clone())]), - &view, - &connected - )); - config.runtime_data.insert( - "require_kv_event_source_readiness".to_string(), - serde_json::json!(true), - ); - config.kv_event_publishing_enabled = Some(false); - assert!(required_sources_ready( - &HashMap::from([(7, config)]), - &view, - &connected - )); - } - fn metric_view(status: KvSourceStatus, capability: Option) -> KvSourceMembershipView { let serving_endpoint = EndpointId { namespace: "ns".to_string(), @@ -916,7 +725,6 @@ mod tests { completions, runtime: tokio::runtime::Handle::current(), task_guard: None, - readiness: watch::channel(false).1, }; tokio::time::timeout(Duration::from_secs(1), handle.shutdown()) @@ -940,7 +748,6 @@ mod tests { completions: vec![completion_rx], runtime: tokio::runtime::Handle::current(), task_guard: None, - readiness: watch::channel(false).1, }; handle.set_task_guard(teardown.clone()); drop(teardown); diff --git a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs index 4d9c383a1126..f9da8b554656 100644 --- a/lib/llm/src/kv_router/indexer/recovery/worker_query.rs +++ b/lib/llm/src/kv_router/indexer/recovery/worker_query.rs @@ -108,8 +108,6 @@ pub(crate) struct WorkerQueryClient { slots: DashMap>>, /// Immutable publisher binding and rank slot lookup performed once per event envelope. publisher_bindings: DashMap, - /// Sources activated behind an established event subscription. - ready_sources: watch::Sender>, recovery_lane: RecoveryLane, recovery_attempt_timeout: Duration, cancellation_token: CancellationToken, @@ -151,7 +149,6 @@ impl WorkerQueryClient { membership_sync: Mutex::new(()), slots: DashMap::new(), publisher_bindings: DashMap::new(), - ready_sources: watch::channel(HashSet::new()).0, recovery_lane: RecoveryLane::with_semaphore(recovery_semaphore), recovery_attempt_timeout, cancellation_token, @@ -191,7 +188,6 @@ impl WorkerQueryClient { membership_sync: Mutex::new(()), slots: DashMap::new(), publisher_bindings: DashMap::new(), - ready_sources: watch::channel(HashSet::new()).0, recovery_lane: RecoveryLane::with_semaphore(recovery_semaphore), recovery_attempt_timeout, cancellation_token: CancellationToken::new(), @@ -208,23 +204,6 @@ impl WorkerQueryClient { view } - pub(crate) fn ready_sources(&self) -> watch::Receiver> { - self.ready_sources.subscribe() - } - - pub(crate) fn clear_transport_readiness(&self) { - self.ready_sources.send_replace(HashSet::new()); - } - - pub(crate) fn publish_ready_sources(&self) { - self.ready_sources.send_replace( - self.publisher_bindings - .iter() - .map(|binding| binding.binding.source_id.clone()) - .collect(), - ); - } - /// Apply membership only for source incarnations whose direct transport is preconnected. /// /// The semantic membership watch remains authoritative. Transport readiness can only @@ -246,7 +225,6 @@ impl WorkerQueryClient { } } self.reconcile_view(effective).await; - self.publish_ready_sources(); view } @@ -428,7 +406,6 @@ impl WorkerQueryClient { } pub(crate) async fn shutdown(self: &Arc) { - self.clear_transport_readiness(); self.cancellation_token.cancel(); self.deactivate_all().await; } @@ -546,9 +523,6 @@ impl WorkerQueryClient { /// Unlike a protocol rejection, the same source may reactivate after a replacement socket is /// preconnected. The reset barrier makes any already-enqueued old event visible first. pub(crate) async fn fence_transport(self: &Arc, publisher_id: PublisherId) -> bool { - self.ready_sources.send_modify(|sources| { - sources.retain(|source| source.publisher_id != publisher_id); - }); let _sync = self.membership_sync.lock().await; let Some(active) = self .publisher_bindings diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs index 59af6136d674..07f051f69239 100644 --- a/lib/llm/src/kv_router/prefill_router/activation.rs +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -683,39 +683,6 @@ mod tests { card } - #[test] - fn pending_prefill_requires_feeds_only_for_opted_in_kv_routing() { - use super::super::prefill_requires_kv_event_readiness; - let mut selected = card(None); - assert!(!prefill_requires_kv_event_readiness( - &selected, - RouterMode::KV - )); - selected - .runtime_config - .set_engine_specific("require_kv_event_source_readiness", true) - .unwrap(); - assert!(prefill_requires_kv_event_readiness( - &selected, - RouterMode::KV - )); - assert!(!prefill_requires_kv_event_readiness( - &selected, - RouterMode::RoundRobin - )); - selected.router_config = Some(RouterConfig::new( - RouterMode::KV, - KvRouterConfig { - use_kv_events: false, - ..Default::default() - }, - )); - assert!(!prefill_requires_kv_event_readiness( - &selected, - RouterMode::KV - )); - } - #[test] fn inherits_decode_mode_when_card_advertises_nothing() { // The pre-override behavior: a prefill worker that says nothing is diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index a3518109ecb2..53b2cdceee5b 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -285,7 +285,6 @@ where pub(crate) trait PrefillRouterLifecycle: Send + Sync { fn set_target(&self, target: Option); - fn kv_event_sources_ready(&self) -> bool; } impl PrefillRouterLifecycle for PrefillRouter @@ -295,41 +294,6 @@ where fn set_target(&self, target: Option) { self.set_target(target); } - - fn kv_event_sources_ready(&self) -> bool { - if self.lifecycle_state() == PrefillLifecycleState::Unavailable { - return true; - } - if let Some(binding) = self.binding.load().as_ref() { - return binding - .router - .kv_router_if_enabled() - .is_none_or(|router| router.kv_event_sources_ready()); - } - // An opted-in prefill target can be discovered before its chooser has - // finished connecting. Keep the decode WorkerSet unready during that gap. - self.target_tx - .as_ref() - .is_none_or(|targets| match targets.borrow().as_ref() { - Some(WorkerSetTarget::Committed(target)) => { - !prefill_requires_kv_event_readiness(&target.card, self.decode_router_mode) - } - _ => true, - }) - } -} - -fn prefill_requires_kv_event_readiness( - card: &crate::model_card::ModelDeploymentCard, - fallback_mode: RouterMode, -) -> bool { - card.runtime_config.requires_kv_event_source_readiness() - && card - .router_config - .as_ref() - .map_or(fallback_mode.is_kv_routing(), |config| { - config.router_mode.is_kv_routing() && config.kv_router_config.use_kv_events - }) } impl Drop for PrefillRouter diff --git a/lib/llm/src/kv_router/publisher/mod.rs b/lib/llm/src/kv_router/publisher/mod.rs index 39abbd66b87c..eeb9a60d98f7 100644 --- a/lib/llm/src/kv_router/publisher/mod.rs +++ b/lib/llm/src/kv_router/publisher/mod.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use anyhow::Result; -use tokio::sync::{mpsc, oneshot, watch}; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use dynamo_kv_router::indexer::{KvIndexerMetrics, LocalKvIndexer}; @@ -120,7 +120,6 @@ async fn supervise_zmq_listener( } impl KvEventSource { - #[allow(clippy::too_many_arguments)] fn start( component: Component, worker_id: WorkerId, @@ -129,7 +128,6 @@ impl KvEventSource { cancellation_token: CancellationToken, tx: mpsc::UnboundedSender>, next_event_id: Arc, - ready: oneshot::Sender<()>, ) -> Result { match source_config { KvEventSourceConfig::Zmq { @@ -153,7 +151,6 @@ impl KvEventSource { next_event_id, image_token_id, video_token_id, - Some(ready), )); let listener_abort_handle = listener_handle.abort_handle(); let supervisor_handle = @@ -210,8 +207,6 @@ pub struct KvEventPublisher { tx: mpsc::UnboundedSender>, /// Internal monotonic event ID counter. Shared with the ZMQ listener if present. next_event_id: Arc, - /// Local ingress and outbound publication have been established. - ready: watch::Receiver, } impl KvEventPublisher { @@ -338,12 +333,8 @@ impl KvEventPublisher { let next_event_id = Arc::new(AtomicU64::new(0)); - let (ready_tx, ready) = watch::channel(false); - let mut input_ready = None; let mut source = None; if let Some(config) = source_config { - let (input_ready_tx, input_ready_rx) = oneshot::channel(); - input_ready = Some(input_ready_rx); source = Some(KvEventSource::start( component.clone(), worker_id, @@ -352,7 +343,6 @@ impl KvEventPublisher { cancellation_token.clone(), tx.clone(), next_event_id.clone(), - input_ready_tx, )?); } @@ -390,31 +380,6 @@ impl KvEventPublisher { }; let publisher_id = event_publisher.publisher_id(); - tokio::select! { - _ = cancellation_token_clone.cancelled() => return, - result = event_publisher.wait_ready() => { - if let Err(error) = result { - tracing::error!(%error, worker_id, dp_rank, "KV event publication failed to connect"); - return; - } - } - } - - // A source advertisement is the relay's readiness contract. Do not - // advertise while the engine SUB socket is still connecting. The - // serving worker registers independently so followers can find it. - if let Some(input_ready) = input_ready { - tokio::select! { - _ = cancellation_token_clone.cancelled() => return, - result = input_ready => { - if result.is_err() { - tracing::error!(worker_id, dp_rank, "KV source closed before connecting"); - return; - } - } - } - } - let recovery_endpoint = if let Some(local_indexer) = local_indexer_clone.as_ref() { match start_worker_kv_query_endpoint( component.clone(), @@ -484,8 +449,6 @@ impl KvEventPublisher { } }; - ready_tx.send_replace(true); - start_event_processor( EventPlanePublisher(event_publisher), worker_id, @@ -496,8 +459,6 @@ impl KvEventPublisher { ) .await; - ready_tx.send_replace(false); - if let Err(error) = component .drt() .discovery() @@ -520,27 +481,9 @@ impl KvEventPublisher { worker_id, tx, next_event_id, - ready, }) } - /// Wait for local ingress and outbound publication to be established. - /// - /// This does not wait for frontend subscribers. Frontend KV readiness also - /// requires its subscriptions to every advertised DP-rank source. - pub async fn wait_ready(&self) -> Result<()> { - let mut ready = self.ready.clone(); - tokio::select! { - biased; - _ = self.cancellation_token.cancelled() => { - anyhow::bail!("KV event publisher stopped before becoming ready") - } - result = ready.wait_for(|ready| *ready) => { - result.map(|_| ()).map_err(|_| anyhow::anyhow!("KV event publisher failed to become ready")) - } - } - } - pub fn publish(&self, event: KvCacheEvent) -> Result<(), mpsc::error::SendError> { self.send_singleton(PlacementEvent::local_gpu(self.worker_id, event)) } diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 0ca23b036636..40784899b80b 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -30,7 +30,6 @@ mod test_event_processing { worker_id: 7, tx, next_event_id: Arc::new(AtomicU64::new(0)), - ready: watch::channel(true).1, }; publisher.publish_batch(Vec::new()).unwrap(); @@ -71,7 +70,6 @@ mod test_event_processing { worker_id: 7, tx, next_event_id: Arc::new(AtomicU64::new(0)), - ready: watch::channel(true).1, }; drop(rx); @@ -107,7 +105,6 @@ mod test_event_processing { worker_id: 7, tx, next_event_id: Arc::new(AtomicU64::new(0)), - ready: watch::channel(true).1, }; publisher @@ -1304,7 +1301,6 @@ mod tests_startup_helpers { next_event_id, None, None, - None, ) }); @@ -1436,7 +1432,6 @@ mod tests_startup_helpers { Arc::new(AtomicU64::new(0)), None, None, - None, ) }); @@ -1536,7 +1531,6 @@ mod tests_startup_helpers { Arc::new(AtomicU64::new(0)), None, None, - None, ) }); @@ -1619,18 +1613,7 @@ mod tests_startup_helpers { let listener_handle = tokio::spawn({ let token = token.clone(); let endpoint = endpoint.clone(); - start_zmq_listener( - endpoint, - topic, - 1, - tx, - token, - 4, - next_event_id, - None, - None, - None, - ) + start_zmq_listener(endpoint, topic, 1, tx, token, 4, next_event_id, None, None) }); tokio::time::sleep(tokio::time::Duration::from_millis(150)).await; diff --git a/lib/llm/src/kv_router/publisher/zmq_listener.rs b/lib/llm/src/kv_router/publisher/zmq_listener.rs index 78d857d4ed7c..9084655f0710 100644 --- a/lib/llm/src/kv_router/publisher/zmq_listener.rs +++ b/lib/llm/src/kv_router/publisher/zmq_listener.rs @@ -6,16 +6,14 @@ use std::sync::atomic::{AtomicU64, Ordering}; use anyhow::{Context, Result}; use futures::StreamExt; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; use dynamo_kv_router::protocols::*; use dynamo_kv_router::zmq_wire::*; use crate::kv_router::metrics::kv_publisher_metrics; -use dynamo_runtime::transports::event_plane::zmq_transport::connect_subscriber_with_readiness; - -use crate::utils::zmq::multipart_message; +use crate::utils::zmq::{connect_sub_socket, multipart_message}; pub(super) struct DecodedZmqKvBatch { pub(super) source_cursor: u64, @@ -57,7 +55,6 @@ pub(super) async fn start_zmq_listener( next_event_id: Arc, image_token_id: Option, video_token_id: Option, - ready: Option>, ) { tracing::debug!( "KVEventPublisher connecting to ZMQ endpoint {} (topic '{}')", @@ -68,35 +65,19 @@ pub(super) async fn start_zmq_listener( let mut normalizer = ZmqEventNormalizer::new(kv_block_size) .with_image_token_id(image_token_id) .with_video_token_id(video_token_id); - // Preserve the native-ingress socket's default receive high-water mark. - let (mut socket, connection) = match connect_subscriber_with_readiness( - &zmq_endpoint, - &zmq_topic, - 1_000, - ) { - Ok(connection) => connection, + let socket = match connect_sub_socket(&zmq_endpoint, Some(&zmq_topic)).await { + Ok(socket) => socket, Err(error) => { tracing::error!(endpoint = %zmq_endpoint, topic = %zmq_topic, error = %error, "ZMQ listener failed to connect"); return; } }; - tokio::select! { - _ = cancellation_token.cancelled() => return, - result = connection.wait_connected(&zmq_endpoint) => { - if let Err(error) = result { - tracing::error!(endpoint = %zmq_endpoint, %error, "ZMQ listener connection failed"); - return; - } - } - } + let mut socket = socket; let metrics = kv_publisher_metrics(); if cancellation_token.is_cancelled() { return; } - if let Some(ready) = ready { - let _ = ready.send(()); - } let mut messages_processed = 0u64; diff --git a/lib/llm/src/local_model/runtime_config.rs b/lib/llm/src/local_model/runtime_config.rs index 5279d27dac96..306369e3a138 100644 --- a/lib/llm/src/local_model/runtime_config.rs +++ b/lib/llm/src/local_model/runtime_config.rs @@ -406,17 +406,6 @@ impl Default for ModelRuntimeConfig { } impl ModelRuntimeConfig { - /// Upgraded relays opt into frontend readiness for every declared DP rank. - /// Absence preserves the discovery contract of older workers. - pub fn requires_kv_event_source_readiness(&self) -> bool { - self.kv_event_publishing_enabled != Some(false) - && self - .runtime_data - .get("require_kv_event_source_readiness") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false) - } - /// Check whether a runtime boolean is explicitly enabled. /// /// Rust callers commonly store booleans, while compatibility cards may diff --git a/lib/runtime/Cargo.toml b/lib/runtime/Cargo.toml index 9e75158b3dd1..5bf28152d53e 100644 --- a/lib/runtime/Cargo.toml +++ b/lib/runtime/Cargo.toml @@ -37,7 +37,6 @@ async-nats = { workspace = true } async-stream = { workspace = true } async-trait = { workspace = true } tmq = { workspace = true } -zmq = "0.10" lru = { version = "0.16" } axum = { workspace = true } blake3 = { workspace = true } diff --git a/lib/runtime/src/distributed.rs b/lib/runtime/src/distributed.rs index 4a70f2f941f6..45e9bed1d07a 100644 --- a/lib/runtime/src/distributed.rs +++ b/lib/runtime/src/distributed.rs @@ -666,11 +666,7 @@ impl DistributedRuntime { let Some(nats_client) = self.nats_client.as_ref() else { anyhow::bail!("KV router's EventSubscriber requires NATS"); }; - let subscriber = nats_client.client().subscribe(subject).await?; - // Establish the subscription at the server before frontend readiness can - // release the first request (and its first KV event). - nats_client.client().flush().await?; - Ok(subscriber) + Ok(nats_client.client().subscribe(subject).await?) } /// TODO (karenc): This is a temporary KV router measure for worker query requests. diff --git a/lib/runtime/src/transports/event_plane/mod.rs b/lib/runtime/src/transports/event_plane/mod.rs index 8b72b79dbca9..2a8126f054f1 100644 --- a/lib/runtime/src/transports/event_plane/mod.rs +++ b/lib/runtime/src/transports/event_plane/mod.rs @@ -315,12 +315,6 @@ pub struct EventPublisher { } impl EventPublisher { - /// Wait for the outbound transport to establish its forwarding connection. - /// Direct publishers already have a bound socket; broker publishers await ZMTP. - pub async fn wait_ready(&self) -> Result<()> { - self.tx.wait_ready().await - } - /// Create a publisher for an endpoint-scoped topic. pub async fn for_endpoint(endpoint: &Endpoint, topic: impl Into) -> Result { let transport_kind = endpoint.drt().default_event_transport_kind(); @@ -669,26 +663,6 @@ pub struct EventSubscriber { } impl EventSubscriber { - /// Establish broker connections before returning the subscriber. Direct ZMQ - /// sources have their own per-source readiness in the KV router. - pub async fn for_endpoint_id_with_transport_ready( - drt: &DistributedRuntime, - endpoint: &EndpointId, - topic: impl Into, - transport_kind: EventTransportKind, - ) -> Result { - Self::new_internal_with_readiness( - drt, - EventScope::Endpoint { - endpoint: endpoint.clone(), - }, - topic.into(), - transport_kind, - true, - ) - .await - } - /// Create a subscriber for an endpoint-scoped topic. pub async fn for_endpoint(endpoint: &Endpoint, topic: impl Into) -> Result { let transport_kind = endpoint.drt().default_event_transport_kind(); @@ -788,16 +762,6 @@ impl EventSubscriber { scope: EventScope, topic: String, transport_kind: EventTransportKind, - ) -> Result { - Self::new_internal_with_readiness(drt, scope, topic, transport_kind, false).await - } - - async fn new_internal_with_readiness( - drt: &DistributedRuntime, - scope: EventScope, - topic: String, - transport_kind: EventTransportKind, - wait_for_connections: bool, ) -> Result { let discovery = drt.discovery(); let routing_key = scope.subject(&topic); @@ -816,21 +780,7 @@ impl EventSubscriber { // BROKER MODE: Connect to broker's XPUB (single or multiple endpoints) let codec = Arc::new(Codec::Msgpack(MsgpackCodec)); - let stream: WireStream = if wait_for_connections { - let stream = - zmq_transport::ZmqSubTransport::connect_single_consumer_multiple_ready( - &broker.xpub_endpoints, - &routing_key, - ) - .await?; - let stream: WireStream = - Box::pin(stream.map(|result| result.map(|message| message.payload))); - if broker.xpub_endpoints.len() > 1 { - Box::pin(DeduplicatingStream::new(stream, codec.clone(), 100_000)) - } else { - stream - } - } else if broker.xpub_endpoints.len() == 1 { + let stream: WireStream = if broker.xpub_endpoints.len() == 1 { // One EventSubscriber has one consumer. Poll the ZMQ socket // directly instead of forwarding through a lossy broadcast channel. let stream = zmq_transport::ZmqSubTransport::connect_single_consumer( diff --git a/lib/runtime/src/transports/event_plane/transport.rs b/lib/runtime/src/transports/event_plane/transport.rs index 3bc2fa0e784d..980f60fec8af 100644 --- a/lib/runtime/src/transports/event_plane/transport.rs +++ b/lib/runtime/src/transports/event_plane/transport.rs @@ -21,11 +21,6 @@ pub type WireStream = Pin> + Send>>; /// Transport-agnostic interface for publishing events. #[async_trait] pub trait EventTransportTx: Send + Sync { - /// Wait for any asynchronous transport startup before advertising a ready source. - async fn wait_ready(&self) -> Result<()> { - Ok(()) - } - /// Publish raw envelope bytes to a subject/topic. async fn publish(&self, subject: &str, envelope_bytes: Bytes) -> Result<()>; diff --git a/lib/runtime/src/transports/event_plane/zmq_transport.rs b/lib/runtime/src/transports/event_plane/zmq_transport.rs index 3a9bc494615e..e3c8ff53a3bf 100644 --- a/lib/runtime/src/transports/event_plane/zmq_transport.rs +++ b/lib/runtime/src/transports/event_plane/zmq_transport.rs @@ -18,20 +18,17 @@ use anyhow::{Context as _, Result, anyhow}; use async_stream::stream; use async_trait::async_trait; use bytes::Bytes; -use futures::{SinkExt, Stream, StreamExt}; +use futures::{SinkExt, StreamExt}; use once_cell::sync::OnceCell; use std::ffi::OsStr; -use std::ops::{Deref, DerefMut}; -use std::pin::Pin; use std::sync::Arc; -use std::task::{Context as TaskContext, Poll}; use thiserror::Error; use tmq::{ AsZmqSocket, Context, Message, Multipart, SocketBuilder, publish::{Publish, publish}, subscribe::{Subscribe, subscribe}, }; -use tokio::sync::{Mutex, broadcast, watch}; +use tokio::sync::{Mutex, broadcast}; use tokio_util::task::AbortOnDropHandle; /// Returns the process-wide shared ZMQ context. @@ -76,218 +73,6 @@ const ZMQ_RCVHWM: i32 = 100_000; // Receive buffer: 100K messages const ZMQ_SNDTIMEOUT_MS: i32 = 0; // Send timeout: fail fast under pressure const ZMQ_RCVTIMEOUT_MS: i32 = 100; // Receive timeout: 100ms (avoids blocking forever) -/// Keeps a socket's monitor alive until event production has been disabled. -/// -/// libzmq sends monitor events synchronously. Dropping its PAIR receiver while -/// monitoring remains enabled can block an I/O thread on a later disconnect. -pub struct MonitoredZmqSocket { - socket: S, - readiness: Option, -} - -impl MonitoredZmqSocket { - fn unmonitored(socket: S) -> Self { - Self { - socket, - readiness: None, - } - } - - fn monitored(socket: S, context: &Context) -> Result { - let readiness = monitor_connection(socket.get_socket(), context)?; - Ok(Self { - socket, - readiness: Some(readiness), - }) - } - - fn readiness(&self) -> ZmqConnectionReadiness { - self.readiness - .as_ref() - .expect("socket is monitored") - .clone() - } -} - -impl AsZmqSocket for MonitoredZmqSocket { - fn get_socket(&self) -> &zmq::Socket { - self.socket.get_socket() - } -} - -impl Deref for MonitoredZmqSocket { - type Target = S; - fn deref(&self) -> &S { - &self.socket - } -} - -impl DerefMut for MonitoredZmqSocket { - fn deref_mut(&mut self) -> &mut S { - &mut self.socket - } -} - -impl Stream for MonitoredZmqSocket { - type Item = S::Item; - fn poll_next(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { - Pin::new(&mut self.get_mut().socket).poll_next(cx) - } -} - -impl Drop for MonitoredZmqSocket { - fn drop(&mut self) { - if self.readiness.is_some() { - // rust-zmq cannot pass the null endpoint used to stop monitoring. - // Replacing the monitor with a zero-event monitor disables event - // production synchronously, before the old receiver is aborted. - // A fresh endpoint avoids racing the old monitor's asynchronous close. - let endpoint = format!("inproc://dynamo-zmq-disabled-{}", uuid::Uuid::new_v4()); - if let Err(error) = self.socket.get_socket().monitor(&endpoint, 0) { - tracing::debug!(%error, "Unable to disable closing ZMQ socket monitor"); - } - } - } -} - -/// Tracks completed ZMTP handshakes independently of the data receive loop. -/// -/// The subscription is configured before connecting. A completed handshake means -/// the transport can forward events; PUB/SUB does not acknowledge the remote -/// publisher's processing of SUBSCRIBE and this is not a delivery guarantee. -#[derive(Clone)] -pub struct ZmqConnectionReadiness { - connected: watch::Receiver>, - endpoints: Arc>>, - // Keep the receiver connected even if Tokio cancels its reader first during - // shutdown. The owning data socket can then disable monitoring safely. - _monitor_socket: Arc>, - _monitor_task: Arc>, -} - -impl ZmqConnectionReadiness { - /// Wait for the endpoint's transport handshake. Callers own the startup - /// deadline and cancellation policy; dropping this future is safe. - pub async fn wait_connected(&self, endpoint: &str) -> Result<()> { - let mut aliases = std::collections::HashSet::from([endpoint.to_string()]); - if let Some(recorded) = self.endpoints.lock().get(endpoint) { - aliases.insert(recorded.clone()); - } - if let Some(address) = endpoint.strip_prefix("tcp://") { - if let Ok(address) = address.parse::() { - aliases.insert(format!("tcp://{address}")); - } else { - // libzmq may report the hostname on its initial connection and - // the resolved IP after reconnecting. LAST_ENDPOINT alone does - // not cover both spellings. Resolve only startup hostnames; - // numeric TCP and IPC endpoints do not perform a DNS lookup. - let addresses = tokio::net::lookup_host(address) - .await - .with_context(|| format!("Failed to resolve ZMQ endpoint {endpoint}"))?; - aliases.extend(addresses.map(|address| format!("tcp://{address}"))); - } - } - let mut connected = self.connected.clone(); - loop { - if !connected.borrow_and_update().is_disjoint(&aliases) { - return Ok(()); - } - connected - .changed() - .await - .context("ZMQ connection monitor stopped")?; - } - } - - fn record_endpoint(&self, socket: &zmq::Socket, endpoint: &str) -> Result<()> { - // Preserve libzmq's initial spelling; reconnects may use a resolved IP. - let resolved = socket - .get_last_endpoint()? - .map_err(|_| anyhow!("ZMQ endpoint is not valid UTF-8"))?; - self.endpoints.lock().insert(endpoint.to_string(), resolved); - Ok(()) - } -} - -/// Start a SUB socket with an observable connection handshake. -/// -/// Socket construction remains nonblocking. Readiness is observed separately so -/// adding an unavailable peer does not block an existing grouped receive loop. -pub fn connect_subscriber_with_readiness( - endpoint: &str, - topic: &str, - rcvhwm: i32, -) -> Result<(MonitoredZmqSocket, ZmqConnectionReadiness)> { - anyhow::ensure!(rcvhwm > 0, "ZMQ receive HWM must be greater than zero"); - anyhow::ensure!( - endpoint.starts_with("tcp://") || endpoint.starts_with("ipc://"), - "ZMQ connection readiness requires a TCP or IPC endpoint" - ); - let context = shared_zmq_context()?; - let socket = context - .socket(zmq::SUB) - .map_err(|error| map_socket_creation_error(error.into()))?; - socket.set_rcvhwm(rcvhwm)?; - socket.set_rcvtimeo(ZMQ_RCVTIMEOUT_MS)?; - socket.set_linger(0)?; - socket.set_reconnect_ivl(100)?; - socket.set_reconnect_ivl_max(5000)?; - socket.set_tcp_keepalive(1)?; - socket.set_subscribe(topic.as_bytes())?; - - let socket = - >::from_zmq_socket(socket)? - .subscribe(topic.as_bytes())?; - let socket = MonitoredZmqSocket::monitored(socket, &context)?; - let readiness = socket.readiness(); - socket.get_socket().connect(endpoint)?; - readiness.record_endpoint(socket.get_socket(), endpoint)?; - Ok((socket, readiness)) -} - -fn monitor_connection(socket: &zmq::Socket, context: &Context) -> Result { - let monitor_endpoint = format!("inproc://dynamo-zmq-ready-{}", uuid::Uuid::new_v4()); - let events = zmq::SocketEvent::HANDSHAKE_SUCCEEDED.to_raw() - | zmq::SocketEvent::DISCONNECTED.to_raw() - | zmq::SocketEvent::MONITOR_STOPPED.to_raw(); - let monitor = tmq::pair(context) - .set_linger(0) - .connect(&monitor_endpoint)?; - socket.monitor(&monitor_endpoint, i32::from(events))?; - - let monitor = Arc::new(Mutex::new(monitor)); - let reader = monitor.clone(); - let (connected_tx, connected) = watch::channel(std::collections::HashSet::new()); - let monitor_task = tokio::spawn(async move { - let mut monitor = reader.lock().await; - while let Some(Ok(frames)) = monitor.next().await { - if frames.len() != 2 || frames[0].len() < 2 { - continue; - } - let event = u16::from_ne_bytes([frames[0][0], frames[0][1]]); - if event == zmq::SocketEvent::MONITOR_STOPPED.to_raw() { - break; - } - let Ok(endpoint) = std::str::from_utf8(&frames[1]) else { - continue; - }; - connected_tx.send_modify(|connected| { - if event == zmq::SocketEvent::HANDSHAKE_SUCCEEDED.to_raw() { - connected.insert(endpoint.to_string()); - } else if event == zmq::SocketEvent::DISCONNECTED.to_raw() { - connected.remove(endpoint); - } - }); - } - }); - Ok(ZmqConnectionReadiness { - connected, - endpoints: Arc::default(), - _monitor_socket: monitor, - _monitor_task: Arc::new(AbortOnDropHandle::new(monitor_task)), - }) -} - const ZMQ_SOCKET_LIMIT_GUIDANCE: &str = "ZMQ could not create another socket. The process may have reached libzmq's ZMQ_MAX_SOCKETS limit or its file-descriptor limit. Reduce direct-ZMQ peers or raise the limit with `ulimit -n`"; const PROCESS_FD_LIMIT_GUIDANCE: &str = "The process reached its file-descriptor limit. Reduce open file descriptors or raise the limit with `ulimit -n`"; @@ -379,9 +164,8 @@ impl AsRef<[u8]> for ZmqMessageOwner { /// ZMQ PUB transport for publishing events. pub struct ZmqPubTransport { - socket: Arc>>, + socket: Arc>, topic: String, - readiness: Option<(ZmqConnectionReadiness, Vec)>, } impl ZmqPubTransport { @@ -415,9 +199,8 @@ impl ZmqPubTransport { Ok(( Self { - socket: Arc::new(Mutex::new(MonitoredZmqSocket::unmonitored(socket))), + socket: Arc::new(Mutex::new(socket)), topic: topic.to_string(), - readiness: None, }, actual_endpoint, )) @@ -429,7 +212,20 @@ impl ZmqPubTransport { /// Connect to single broker XSUB endpoint (broker mode) pub async fn connect(xsub_endpoint: &str, topic: &str) -> Result { - Self::connect_multiple(&[xsub_endpoint.to_string()], topic).await + let ctx = shared_zmq_context()?; + let socket = connect_tmq_socket(configure_publish_builder(publish(&ctx)), xsub_endpoint)?; + + tracing::info!( + endpoint = %xsub_endpoint, + topic = %topic, + sndhwm = ZMQ_SNDHWM, + "ZMQ PUB transport connected to broker XSUB" + ); + + Ok(Self { + socket: Arc::new(Mutex::new(socket)), + topic: topic.to_string(), + }) } /// Connect to multiple broker XSUB endpoints (HA mode) @@ -440,22 +236,10 @@ impl ZmqPubTransport { }; let ctx = shared_zmq_context()?; - let socket = ctx - .socket(zmq::PUB) - .map_err(|error| map_socket_creation_error(error.into()))?; - socket.set_sndhwm(ZMQ_SNDHWM)?; - socket.set_sndtimeo(ZMQ_SNDTIMEOUT_MS)?; - let socket = MonitoredZmqSocket::monitored( - >::from_zmq_socket(socket)?, - &ctx, - )?; - let readiness = socket.readiness(); - socket.get_socket().connect(first_endpoint)?; - readiness.record_endpoint(socket.get_socket(), first_endpoint)?; + let socket = connect_tmq_socket(configure_publish_builder(publish(&ctx)), first_endpoint)?; for endpoint in endpoints { socket.get_socket().connect(endpoint)?; - readiness.record_endpoint(socket.get_socket(), endpoint)?; tracing::debug!(endpoint = %endpoint, "ZMQ PUB connected to broker XSUB"); } @@ -469,22 +253,12 @@ impl ZmqPubTransport { Ok(Self { socket: Arc::new(Mutex::new(socket)), topic: topic.to_string(), - readiness: Some((readiness, xsub_endpoints.to_vec())), }) } } #[async_trait] impl EventTransportTx for ZmqPubTransport { - async fn wait_ready(&self) -> Result<()> { - if let Some((readiness, endpoints)) = &self.readiness { - for endpoint in endpoints { - readiness.wait_connected(endpoint).await?; - } - } - Ok(()) - } - async fn publish(&self, _subject: &str, envelope_bytes: Bytes) -> Result<()> { let codec = MsgpackCodec; let (publisher_id, sequence) = codec.decode_envelope_identity(&envelope_bytes)?; @@ -533,46 +307,24 @@ pub type ZmqWireStream = /// /// The caller must keep this value in one task. ZMQ SUB sockets are not thread-safe. pub struct DynamicZmqSubSocket { - socket: MonitoredZmqSocket, + socket: Subscribe, expected_topic: Vec, - readiness: Option, } impl DynamicZmqSubSocket { - /// Start a dynamically managed socket and expose its connection handshakes. - pub fn connect_with_readiness( - endpoint: &str, - topic: &str, - rcvhwm: i32, - ) -> Result<(Self, ZmqConnectionReadiness)> { - let (socket, readiness) = connect_subscriber_with_readiness(endpoint, topic, rcvhwm)?; - Ok(( - Self { - socket, - expected_topic: topic.as_bytes().to_vec(), - readiness: Some(readiness.clone()), - }, - readiness, - )) - } - /// Connect a new SUB socket with an explicit receive high-water mark. pub fn connect_with_rcvhwm(endpoint: &str, topic: &str, rcvhwm: i32) -> Result { let socket = ZmqSubTransport::connect_socket_with_rcvhwm(endpoint, topic, rcvhwm)?; tracing::info!(endpoint, topic, rcvhwm, "Dynamic ZMQ SUB socket connected"); Ok(Self { - socket: MonitoredZmqSocket::unmonitored(socket), + socket, expected_topic: topic.as_bytes().to_vec(), - readiness: None, }) } /// Connect this SUB socket to one more publisher endpoint. pub fn add_endpoint(&mut self, endpoint: &str) -> Result<()> { self.socket.get_socket().connect(endpoint)?; - if let Some(readiness) = &self.readiness { - readiness.record_endpoint(self.socket.get_socket(), endpoint)?; - } Ok(()) } @@ -638,21 +390,6 @@ pub struct ValidatedZmqSource { } impl ValidatedZmqSource { - /// Connect after the subscribed socket has completed its transport handshake. - pub async fn connect_ready( - endpoint: &str, - topic: &str, - expected_publisher_id: u64, - rcvhwm: i32, - ) -> Result { - Ok(Self { - stream: ZmqSubTransport::connect_single_consumer_ready(endpoint, topic, rcvhwm).await?, - expected_topic: topic.to_string(), - expected_publisher_id, - codec: Codec::default(), - }) - } - pub async fn connect_default( endpoint: &str, topic: &str, @@ -716,36 +453,6 @@ impl ValidatedZmqSource { } impl ZmqSubTransport { - /// Connect one consumer after the subscribed socket's transport handshake. - pub async fn connect_single_consumer_ready( - endpoint: &str, - topic: &str, - rcvhwm: i32, - ) -> Result { - let (socket, readiness) = connect_subscriber_with_readiness(endpoint, topic, rcvhwm)?; - readiness.wait_connected(endpoint).await?; - Self::single_consumer_stream(socket, topic) - } - - /// Connect a consumer after every broker transport handshake completes. - pub async fn connect_single_consumer_multiple_ready( - endpoints: &[String], - topic: &str, - ) -> Result { - let (first, rest) = endpoints - .split_first() - .context("Cannot connect to zero endpoints")?; - let (socket, readiness) = connect_subscriber_with_readiness(first, topic, ZMQ_RCVHWM)?; - for endpoint in rest { - socket.get_socket().connect(endpoint)?; - readiness.record_endpoint(socket.get_socket(), endpoint)?; - } - for endpoint in endpoints { - readiness.wait_connected(endpoint).await?; - } - Self::single_consumer_stream(socket, topic) - } - fn connect_socket(endpoint: &str, topic: &str) -> Result { Self::connect_socket_with_rcvhwm(endpoint, topic, ZMQ_RCVHWM) } @@ -834,10 +541,7 @@ impl ZmqSubTransport { Self::single_consumer_stream(socket, topic) } - fn single_consumer_stream(mut socket: S, topic: &str) -> Result - where - S: Stream> + Send + Unpin + 'static, - { + fn single_consumer_stream(mut socket: Subscribe, topic: &str) -> Result { let expected_topic = topic.as_bytes().to_vec(); let stream = stream! { @@ -1040,102 +744,6 @@ mod tests { use std::collections::HashSet; use tokio::time::{Duration, timeout}; - #[tokio::test] - async fn subscriber_readiness_waits_for_handshake_with_hostname() { - // Reserve the port with a plain TCP listener: a TCP connection alone is - // insufficient, and no ZMTP publisher exists until the explicit bind. - let reservation = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - let port = reservation.local_addr().unwrap().port(); - let endpoint = format!("tcp://localhost:{port}"); - let (_socket, readiness) = - connect_subscriber_with_readiness(&endpoint, "kv-events", 128).unwrap(); - assert!( - timeout( - Duration::from_millis(25), - readiness.wait_connected(&endpoint) - ) - .await - .is_err(), - "socket construction or a plain TCP connection must not report ready" - ); - - drop(reservation); - let (_publisher, _) = - ZmqPubTransport::bind(&format!("tcp://127.0.0.1:{port}"), "kv-events") - .await - .unwrap(); - timeout(Duration::from_secs(5), readiness.wait_connected(&endpoint)) - .await - .unwrap_or_else(|error| { - panic!( - "publisher handshake did not arrive: {error}; aliases={:?}, connected={:?}", - readiness.endpoints.lock(), - readiness.connected.borrow() - ) - }) - .unwrap(); - } - - #[tokio::test] - async fn subscriber_readiness_teardown_preserves_shared_context() { - // Closing a monitored socket must not leave a libzmq I/O thread blocked - // trying to send to a monitor receiver that has already been dropped. - // Reuse the same context across enough connections to exercise every I/O - // thread, and require the first publication after each connection. - let context = shared_zmq_context().unwrap(); - let topic = "monitor-lifetime"; - for sequence in 0..16_u64 { - let publisher = context.socket(zmq::XPUB).unwrap(); - publisher.set_linger(0).unwrap(); - publisher.set_rcvtimeo(2_000).unwrap(); - publisher.bind("tcp://127.0.0.1:*").unwrap(); - let endpoint = publisher.get_last_endpoint().unwrap().unwrap(); - let mut stream = timeout(Duration::from_secs(2), async { - if sequence % 2 == 0 { - ZmqSubTransport::connect_single_consumer_ready(&endpoint, topic, 128).await - } else { - ZmqSubTransport::connect_single_consumer_multiple_ready(&[endpoint], topic) - .await - } - }) - .await - .expect("previous monitor teardown must not block another connection") - .unwrap(); - let payload = encoded_event(topic, 23, sequence); - let published = payload.clone(); - let send = tokio::task::spawn_blocking(move || { - // XPUB observes the subscription so this teardown regression - // does not depend on PUB/SUB's unrelated slow-joiner timing. - let subscription = publisher.recv_bytes(0).unwrap(); - assert_eq!(&subscription[1..], topic.as_bytes()); - assert_eq!(subscription[0], 1); - publisher - .send_multipart( - [ - topic.as_bytes().to_vec(), - 23_u64.to_be_bytes().to_vec(), - sequence.to_be_bytes().to_vec(), - Frame::new(published).encode().to_vec(), - ], - 0, - ) - .unwrap(); - publisher - }); - let message = timeout(Duration::from_secs(2), stream.next()) - .await - .expect("the first event must arrive after each connection") - .unwrap() - .unwrap(); - assert_eq!(message.publisher_id, 23); - assert_eq!(message.sequence, sequence); - assert_eq!(message.payload, payload); - let publisher = send.await.unwrap(); - drop(stream); - drop(publisher); - } - } - #[test] fn emfile_errno_selects_source_specific_guidance() { assert_eq!( diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index 7218cd20faca..ec1a87ec7294 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -69,17 +69,6 @@ local KV events to Dynamo without accepting inference requests. A TP-only follower that holds part of a replica but has no local KV publisher does not need a sidecar. -For KV routing, the frontend's `/health` stays unready until every registered -DP rank has an advertised relay and the router has connected its event -subscription. The leader registers before this gate so followers can discover -it. Each relay advertises only after its engine subscription has established a -connection and outbound publishing is available; creating a publisher object -alone does not mark it ready. This contract applies to single-node sidecars as -well. Launchers can wait for frontend readiness before sending requests. - -Readiness describes connected feeds, not durable event delivery: ZeroMQ PUB/SUB -does not acknowledge individual subscriptions or KV events. - Launch each sidecar separately with `--grpc-endpoint` pointing to its local SGLang server. This requires an SGLang build exposing node-local `kv_event_sources` through `GetServerInfo` on leaders and followers started diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index 4c3b51504e61..b5783176ed2a 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -955,20 +955,6 @@ fn build_engine_config( "grpc_service".to_string(), Value::String("sglang.runtime.v1.SglangService".to_string()), ); - // The leader declares readiness for the whole worker, even when all KV - // publishers are remote. Source descriptors are validated before registration. - if discovery - .server_info - .get("kv_events") - .is_some_and(|value| !value.is_null()) - || discovery - .server_info - .get("kv_event_sources") - .and_then(Value::as_array) - .is_some_and(|sources| !sources.is_empty()) - { - runtime_data.insert("require_kv_event_source_readiness".into(), true.into()); - } if let Some(total_tokens) = hicache_native_offloading_capacity(&discovery.server_info, &discovery.model_info) { @@ -1292,10 +1278,6 @@ mod tests { let sources = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap(); - assert_eq!( - config.runtime_data["require_kv_event_source_readiness"], - true - ); assert_eq!( sources, [ @@ -1353,12 +1335,6 @@ mod tests { // An explicit empty source list must not fall back to all global ranks. discovery.server_info["kv_event_sources"] = json!([]); - config = - build_engine_config(&discovery, DisaggregationMode::Aggregated, None, None).unwrap(); - assert_eq!( - config.runtime_data["require_kv_event_source_readiness"], - true - ); assert!( discover_kv_event_sources(&discovery, &mut config, &endpoint) .unwrap() @@ -1370,16 +1346,6 @@ mod tests { ); assert_eq!(config.llm.as_ref().unwrap().data_parallel_size, Some(8)); - let mut disabled = discovery.clone(); - disabled.server_info["kv_events"] = serde_json::Value::Null; - let disabled_config = - build_engine_config(&disabled, DisaggregationMode::Aggregated, None, None).unwrap(); - assert!( - !disabled_config - .runtime_data - .contains_key("require_kv_event_source_readiness") - ); - // Engines without the new metadata field retain the old multinode guard. discovery .server_info diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index 7688395de8be..459873d38c74 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -135,12 +135,7 @@ impl HeadlessSidecar { wait_for_leader(&mut configs, group_id, &shutdown)).await .context("timed out waiting for SGLang leader; check namespace, component, endpoint, and dist_init_addr")??; let leader = validate_leader(metadata, &config)?; - let _publishers = tokio::time::timeout( - self.transport.startup_deadline, - start_publishers(&endpoint, metadata, worker_id, &config, &leader), - ) - .await - .context("timed out connecting local SGLang KV feeds")??; + let _publishers = start_publishers(&endpoint, metadata, worker_id, &config, &leader)?; tracing::info!(node_rank = metadata.node_rank, worker_id, group = group_id, local_ranks = ?metadata.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), "SGLang headless sidecar publishing local KV events"); @@ -283,7 +278,7 @@ fn validate_leader(context: &NodeMetadata, config: &ModelRuntimeConfig) -> Resul Ok(leader) } -async fn start_publishers( +fn start_publishers( endpoint: &Endpoint, context: &NodeMetadata, worker_id: u64, @@ -312,12 +307,8 @@ async fn start_publishers( ) }) .collect::>>()?; - for publisher in &publishers { - publisher.wait_ready().await?; - } - // There is no serving worker to mark this process ready. Local health - // requires connected ingress and publication; frontend readiness separately - // requires subscriptions to every DP rank in the leader's configuration. + // There is no serving worker to mark this process ready. Publish readiness + // only after all local publishers have been created. endpoint .drt() .system_health() @@ -591,13 +582,9 @@ mod tests { #[tokio::test] async fn local_zmq_events_keep_global_rank_and_leader_identity_without_serving() { use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent}; - use dynamo_runtime::discovery::{ - DiscoveryInstance, DiscoveryQuery, EventChannelQuery, EventSourceQuery, EventTransport, - }; - use dynamo_runtime::transports::event_plane::{ - Codec, Frame, zmq_transport::connect_subscriber_with_readiness, - }; - use futures::{SinkExt, StreamExt}; + use dynamo_runtime::discovery::{DiscoveryQuery, EventSourceQuery}; + use dynamo_runtime::transports::event_plane::EventSubscriber; + use futures::SinkExt; let runtime = Runtime::from_current().unwrap(); let drt = DistributedRuntime::new(runtime, DistributedConfig::process_local()) @@ -621,53 +608,36 @@ mod tests { let (context, _) = follower_metadata(&raw).unwrap(); let config = leader("A"); let metadata = validate_leader(&context, &config).unwrap(); - assert!(!drt.system_health().lock().get_health_status().0); - let publishers = start_publishers(&endpoint, &context, 42, &config, &metadata) + let mut subscriber = EventSubscriber::for_endpoint(&endpoint, KV_EVENT_SUBJECT) .await - .unwrap(); + .unwrap() + .typed::>(); + assert!(!drt.system_health().lock().get_health_status().0); + let publishers = start_publishers(&endpoint, &context, 42, &config, &metadata).unwrap(); assert!(drt.system_health().lock().get_health_status().0); - let channels = drt - .discovery() - .list(DiscoveryQuery::EventChannels( - EventChannelQuery::endpoint_topic(endpoint.id(), KV_EVENT_SUBJECT), - )) - .await - .unwrap(); - assert_eq!(channels.len(), 1); - let DiscoveryInstance::EventChannel { - transport: EventTransport::Zmq { endpoint: address }, - .. - } = &channels[0] - else { - panic!("expected the relay's direct ZMQ channel"); - }; - let (mut subscriber, connection) = - connect_subscriber_with_readiness(address, KV_EVENT_SUBJECT, 1000).unwrap(); - tokio::time::timeout(Duration::from_secs(5), connection.wait_connected(address)) - .await - .unwrap() - .unwrap(); - let payload = rmp_serde::to_vec_named(&json!([ - 0.0, [{"type":"BlockRemoved", "block_hashes":[42]}], 4 - ])) - .unwrap(); - source - .send(vec![Vec::new(), 1_u64.to_be_bytes().to_vec(), payload]) - .await - .unwrap(); - let frames = tokio::time::timeout(Duration::from_secs(5), subscriber.next()) - .await - .expect("headless relay should forward its first local event") - .unwrap() - .unwrap(); - assert_eq!(frames.len(), 4); - let frame = Frame::decode(frames[3].as_ref()).unwrap(); - let envelope = Codec::default().decode_envelope(&frame.payload).unwrap(); - let mut events: Vec = - Codec::default().decode_payload(&envelope.payload).unwrap(); - assert_eq!(events.len(), 1); - let received = (envelope, events.remove(0)); + // Repeat until the real ZMQ subscriptions are connected; no fixed sleep. + let received = tokio::time::timeout(Duration::from_secs(5), async { + let mut ticks = tokio::time::interval(Duration::from_millis(50)); + let mut seq = 0_u64; + loop { + tokio::select! { + batch = subscriber.next() => { + let (envelope, events) = batch.unwrap().unwrap(); + if let Some(event) = events.first() { + break (envelope, event.clone()); + } + } + _ = ticks.tick() => { + seq += 1; + let payload = rmp_serde::to_vec_named(&json!([ + 0.0, [{"type":"BlockRemoved", "block_hashes":[42]}], 4 + ])).unwrap(); + source.send(vec![Vec::new(), seq.to_be_bytes().to_vec(), payload]).await.unwrap(); + } + } + } + }).await.expect("headless relay should forward local events"); assert_eq!(received.1.worker_id, 42); assert_eq!(received.1.event.dp_rank, 4); assert_ne!(received.0.publisher_id, 42); From abcc0655d26dcf12c8c191cbd830c9179e15d1b0 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Fri, 18 Sep 2026 10:36:28 -0700 Subject: [PATCH 12/17] Update lib/sidecar/sglang/src/lib.rs Co-authored-by: Dmitry Tokarev Signed-off-by: jain-ria --- lib/sidecar/sglang/src/lib.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/lib/sidecar/sglang/src/lib.rs b/lib/sidecar/sglang/src/lib.rs index 60ef6be7347a..ec54f478a678 100644 --- a/lib/sidecar/sglang/src/lib.rs +++ b/lib/sidecar/sglang/src/lib.rs @@ -9,8 +9,9 @@ //! role, parallelism, KV block sizing, and context length are discovered from //! the engine's gRPC metadata RPCs. //! -//! The crate never depends on `sglang` or any engine crate — only -//! `dynamo-backend-common`, `tonic`/`prost`, `clap`, and tokio. +//! The crate never depends on `sglang` or any engine crate. It uses +//! `dynamo-backend-common`, `dynamo-llm`, `dynamo-runtime`, `tonic`/`prost`, +//! `clap`, and tokio. use std::sync::Arc; From 9ab50011f8bf1400759d856a40dcb0414b77a163 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Fri, 18 Sep 2026 11:47:04 -0700 Subject: [PATCH 13/17] fix(sglang): infer sidecar mode from engine metadata Signed-off-by: jain-ria --- lib/sidecar/sglang/README.md | 13 ++--- .../launch/multinode_kv_router_sidecar.sh | 4 -- lib/sidecar/sglang/src/args.rs | 4 -- lib/sidecar/sglang/src/client.rs | 40 ++++++++++++++- lib/sidecar/sglang/src/engine.rs | 33 +++++-------- lib/sidecar/sglang/src/headless.rs | 49 +++++++++++++++++-- lib/sidecar/sglang/src/lib.rs | 34 +++++-------- lib/sidecar/sglang/tests/executable.rs | 4 +- 8 files changed, 114 insertions(+), 67 deletions(-) diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index ec1a87ec7294..0d4127a6fa4c 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -62,12 +62,13 @@ same mode dispatcher as the standalone executable. ## Multinode DP and KV routing For multinode deployments with multiple data-parallel (DP) replicas and KV-aware -routing, run a full sidecar on the leader (`node_rank=0`) and a separate -`--telemetry-only` sidecar on each follower node that owns a DP scheduler/KV -publisher. Each DP rank has its own KV cache; follower sidecars publish those -local KV events to Dynamo without accepting inference requests. A TP-only -follower that holds part of a replica but has no local KV publisher does not -need a sidecar. +routing, launch the sidecar on the leader (`node_rank=0`) and on each follower +node that owns a DP scheduler/KV publisher. The sidecar automatically selects +inference or telemetry-only mode from the local engine's `GetServerInfo` metadata. +Each DP rank has its own KV cache; follower sidecars publish those local KV events +to Dynamo without accepting inference requests. +A TP-only follower that holds part of a replica but has no local KV publisher +does not need a sidecar. Launch each sidecar separately with `--grpc-endpoint` pointing to its local SGLang server. This requires an SGLang build exposing node-local diff --git a/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh b/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh index 637021650d13..a2be8aff153a 100755 --- a/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh +++ b/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh @@ -86,10 +86,6 @@ case "$ROLE" in ;; *) echo "ROLE must be aggregated, prefill or decode" >&2; exit 1 ;; esac -if (( NODE_RANK > 0 )); then - SIDECAR_ARGS+=(--telemetry-only) -fi - print_launch_banner --no-curl "SGLang multinode KV sidecar ($ROLE, node $NODE_RANK)" "$MODEL" "$HTTP_PORT" \ "Topology: NNODES=$NNODES, TP_SIZE=$TP_SIZE, DP_SIZE=$DP_SIZE (attention DP)" \ "Rendezvous: $DIST_INIT_ADDR; local gRPC: $SGLANG_GRPC_PORT" \ diff --git a/lib/sidecar/sglang/src/args.rs b/lib/sidecar/sglang/src/args.rs index 397c8be15709..ae72e22002b5 100644 --- a/lib/sidecar/sglang/src/args.rs +++ b/lib/sidecar/sglang/src/args.rs @@ -15,10 +15,6 @@ pub struct Args { #[command(flatten)] pub sidecar: SidecarArgs, - /// Relay a follower node's local KV events without registering a request endpoint. - #[arg(long)] - pub telemetry_only: bool, - /// Maximum wait for a matching leader registration in telemetry mode. #[arg(long, default_value_t = 1800, value_parser = clap::value_parser!(u64).range(1..))] pub leader_discovery_timeout_secs: u64, diff --git a/lib/sidecar/sglang/src/client.rs b/lib/sidecar/sglang/src/client.rs index c394961e9354..239813510707 100644 --- a/lib/sidecar/sglang/src/client.rs +++ b/lib/sidecar/sglang/src/client.rs @@ -201,6 +201,36 @@ pub struct Discovery { pub server_info: Value, } +#[derive(Debug)] +pub(crate) enum StartupDiscovery { + Leader(Discovery), + Follower, +} + +pub(crate) fn bootstrap_discover( + endpoint: &GrpcEndpoint, + transport: &GrpcTransportConfig, +) -> Result { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|err| engine_shutdown(format!("bootstrap runtime: {err}")))?; + runtime.block_on(async { + let deadline = Instant::now() + transport.startup_deadline; + let mut client = connect(endpoint, transport, deadline).await?; + let server_info = get_server_info(&mut client, deadline).await?; + if json_u32(&server_info, "node_rank").is_some_and(|rank| rank > 0) { + // Followers expose metadata only. Their local KV sources are + // validated by the headless startup path before relaying. + Ok(StartupDiscovery::Follower) + } else { + discover_with_server_info(&mut client, server_info, deadline) + .await + .map(StartupDiscovery::Leader) + } + }) +} + pub async fn connect( uri: &GrpcEndpoint, cfg: &GrpcTransportConfig, @@ -292,9 +322,17 @@ impl Pool { pub async fn discover(client: &mut Client, deadline: Instant) -> Result { let server_info = get_server_info(client, deadline).await?; + discover_with_server_info(client, server_info, deadline).await +} + +async fn discover_with_server_info( + client: &mut Client, + server_info: Value, + deadline: Instant, +) -> Result { if json_u32(&server_info, "node_rank").is_some_and(|rank| rank > 0) { return Err(invalid_arg( - "full sidecar requires node_rank=0; use --telemetry-only for a follower with local KV sources", + "inference discovery requires node_rank=0; followers expose only GetServerInfo", )); } let model = rpc_with_deadline( diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index b5783176ed2a..9920540be279 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -70,11 +70,18 @@ impl SglangSidecarEngine { } pub(crate) fn from_parsed(args: Args) -> Result<(Self, WorkerConfig), DynamoError> { - if args.telemetry_only { - return Err(client::invalid_arg( - "--telemetry-only requires the headless sidecar entry point", - )); - } + let discovery = + client::bootstrap_discover(&args.sidecar.grpc_endpoint, &args.sidecar.grpc.config())?; + let client::StartupDiscovery::Leader(discovery) = discovery else { + return Err(client::invalid_arg("inference engine requires node_rank=0")); + }; + Self::from_discovery(args, discovery) + } + + pub(crate) fn from_discovery( + args: Args, + discovery: Discovery, + ) -> Result<(Self, WorkerConfig), DynamoError> { if args.sidecar.common.route_to_encoder { return Err(client::invalid_arg( "route-to-encoder is not supported by the SGLang sidecar", @@ -83,7 +90,6 @@ impl SglangSidecarEngine { let endpoint = args.sidecar.grpc_endpoint; let transport = args.sidecar.grpc.config(); - let discovery = bootstrap_discover(&endpoint, &transport)?; let disaggregation_mode = client::discovery_mode(&discovery.server_info)?; let bootstrap_host = if disaggregation_mode.is_prefill() { resolve_bootstrap_host( @@ -484,21 +490,6 @@ impl LLMEngine for SglangSidecarEngine { } } -fn bootstrap_discover( - endpoint: &GrpcEndpoint, - transport: &GrpcTransportConfig, -) -> Result { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .map_err(|err| client::engine_shutdown(format!("bootstrap runtime: {err}")))?; - runtime.block_on(async { - let deadline = Instant::now() + transport.startup_deadline; - let mut grpc_client = client::connect(endpoint, transport, deadline).await?; - client::discover(&mut grpc_client, deadline).await - }) -} - fn discovery_string(value: &Value, key: &str) -> Option { value .get(key) diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index 459873d38c74..6e0ddcaa69c9 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -157,7 +157,7 @@ fn follower_metadata(info: &Value) -> Result<(NodeMetadata, DisaggregationMode)> .context("telemetry-only mode requires GetServerInfo.kv_event_sources; upgrade SGLang")?; ensure!( metadata.node_rank > 0 && metadata.nnodes > 1, - "--telemetry-only requires a follower node; run the leader sidecar without this flag" + "KV relay requires a follower node" ); ensure!( !metadata.kv_event_sources.is_empty(), @@ -445,12 +445,47 @@ mod tests { assert!(follower_metadata(&info).is_err()); } + #[tokio::test] + async fn startup_selects_mode_from_server_info() { + let server = MetadataServer::start().await; + // A follower needs only metadata; leaders (including legacy metadata) + // must continue to model discovery, which this follower server rejects. + for (info, follower) in [ + (follower_info(), true), + (json!({"node_rank": 0}), false), + (json!({}), false), + ] { + *server.service.info.lock().unwrap() = info; + server.service.calls.lock().unwrap().clear(); + let endpoint = GrpcEndpoint::parse(&server.endpoint, "test").unwrap(); + let result = tokio::task::spawn_blocking(move || { + let transport = GrpcTransportConfig { + startup_deadline: Duration::from_secs(5), + ..Default::default() + }; + client::bootstrap_discover(&endpoint, &transport) + }) + .await + .unwrap(); + let mut expected = vec!["/sglang.runtime.v1.SglangService/GetServerInfo"]; + if follower { + assert!(matches!( + result.unwrap(), + client::StartupDiscovery::Follower + )); + } else { + assert!(result.unwrap_err().to_string().contains("GetModelInfo")); + expected.push("/sglang.runtime.v1.SglangService/GetModelInfo"); + } + assert_eq!(*server.service.calls.lock().unwrap(), expected); + } + } + #[tokio::test] async fn follower_discovers_and_monitors_using_only_server_info() { let server = MetadataServer::start().await; let args = Args::try_parse_from([ "sidecar", - "--telemetry-only", "--grpc-endpoint", &server.endpoint, "--grpc-startup-deadline-secs", @@ -461,12 +496,16 @@ mod tests { let (mut client, original, mode) = sidecar.connect_local_engine().await.unwrap(); assert_eq!(mode, DisaggregationMode::Prefill); let timeout = Duration::from_secs(1); - // A misplaced full sidecar fails with an actionable mode error, without - // asking the follower for any of the unsupported inference RPCs. + // Full inference discovery still rejects a follower without asking + // for any of its unsupported model RPCs. let error = client::discover(&mut client, Instant::now() + timeout) .await .unwrap_err(); - assert!(error.to_string().contains("use --telemetry-only")); + assert!( + error + .to_string() + .contains("inference discovery requires node_rank=0") + ); check_local_engine(&mut client, &original, mode, timeout) .await .unwrap(); diff --git a/lib/sidecar/sglang/src/lib.rs b/lib/sidecar/sglang/src/lib.rs index ec54f478a678..5db14e063597 100644 --- a/lib/sidecar/sglang/src/lib.rs +++ b/lib/sidecar/sglang/src/lib.rs @@ -40,14 +40,18 @@ pub use engine::SglangSidecarEngine; /// distinguish invalid configuration from runtime failures. pub fn run(argv: Vec) -> anyhow::Result<()> { let args = Args::try_parse_from(argv).map_err(SidecarStartupError::from)?; - if args.telemetry_only { - HeadlessSidecar::from_args(args) + let discovery = + client::bootstrap_discover(&args.sidecar.grpc_endpoint, &args.sidecar.grpc.config()) + .map_err(SidecarStartupError::from)?; + match discovery { + client::StartupDiscovery::Follower => HeadlessSidecar::from_args(args) .map_err(SidecarStartupError::from)? - .run() - } else { - let (engine, config) = - SglangSidecarEngine::from_parsed(args).map_err(SidecarStartupError::from)?; - dynamo_backend_common::run(Arc::new(engine), config) + .run(), + client::StartupDiscovery::Leader(discovery) => { + let (engine, config) = SglangSidecarEngine::from_discovery(args, discovery) + .map_err(SidecarStartupError::from)?; + dynamo_backend_common::run(Arc::new(engine), config) + } } } @@ -55,22 +59,6 @@ pub fn run(argv: Vec) -> anyhow::Result<()> { mod tests { use super::*; - #[test] - fn telemetry_uses_headless_validation_before_grpc() { - let error = run(vec![ - "sidecar".into(), - "--telemetry-only".into(), - "--grpc-endpoint".into(), - "not-a-grpc-address".into(), - // Stop at headless validation, before starting a runtime. - "--route-to-encoder".into(), - ]) - .unwrap_err(); - let error = error.downcast::().unwrap(); - assert!(matches!(error, SidecarStartupError::Dynamo(ref error) - if error.to_string().contains("telemetry mode cannot register encoder or RL request routes"))); - } - #[test] fn help_retains_structured_cli_exit() { let error = run(vec!["sidecar".into(), "--help".into()]).unwrap_err(); diff --git a/lib/sidecar/sglang/tests/executable.rs b/lib/sidecar/sglang/tests/executable.rs index a176830ebfe5..26aeea945868 100644 --- a/lib/sidecar/sglang/tests/executable.rs +++ b/lib/sidecar/sglang/tests/executable.rs @@ -27,7 +27,6 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { "DYN_SIDECAR_GRPC_RETRY_INTERVAL_SECS", "--grpc-startup-deadline-secs", "DYN_SIDECAR_GRPC_STARTUP_DEADLINE_SECS", - "--telemetry-only", "--leader-discovery-timeout-secs", ] { assert!(stdout.contains(expected), "help omits {expected}"); @@ -35,9 +34,8 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { } #[test] -fn telemetry_mode_requires_a_local_grpc_endpoint() { +fn sidecar_requires_a_local_grpc_endpoint() { let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) - .arg("--telemetry-only") .env_remove("DYN_SIDECAR_GRPC_ENDPOINT") .output() .unwrap(); From e47072982c968f76189cb89854e0c6ce49cefa12 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Fri, 18 Sep 2026 13:02:37 -0700 Subject: [PATCH 14/17] docs(sglang): add multinode disaggregated sidecar launcher Signed-off-by: jain-ria --- lib/sidecar/sglang/README.md | 11 ++++++ .../multinode_disagg_kv_router_sidecar.sh | 36 +++++++++++++++++++ 2 files changed, 47 insertions(+) create mode 100755 lib/sidecar/sglang/launch/multinode_disagg_kv_router_sidecar.sh diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index 0d4127a6fa4c..806202708ce4 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -88,6 +88,17 @@ isolate other processes in that namespace. If overriding this to a non-loopback bind, restrict every publisher port (`SGLANG_KV_EVENT_PORT + dp_rank`) to authorized consumers using a firewall or network policy. +For disaggregated multinode attention DP, use +`launch/multinode_disagg_kv_router_sidecar.sh prefill` on each prefill node and +`launch/multinode_disagg_kv_router_sidecar.sh decode` on each decode node. +Each role is a separate distributed engine group: give it its own +`DIST_INIT_ADDR` and number its nodes from zero. The defaults require four +nodes in total, with `NNODES=2`, `TP_SIZE=2`, and `DP_SIZE=2` per role. +Set `SGLANG_BOOTSTRAP_HOST` to the reachable prefill leader address on prefill +node zero. All nodes must share the Dynamo namespace, discovery and event +services; start `python3 -m dynamo.frontend --router-mode kv` separately. +Run the script with `--help` for the four per-node commands. + ## Deploy on Kubernetes (quick start) `deploy/agg.yaml` runs an aggregated deployment (a frontend plus one worker pod diff --git a/lib/sidecar/sglang/launch/multinode_disagg_kv_router_sidecar.sh b/lib/sidecar/sglang/launch/multinode_disagg_kv_router_sidecar.sh new file mode 100755 index 000000000000..9dc60d5ec55d --- /dev/null +++ b/lib/sidecar/sglang/launch/multinode_disagg_kv_router_sidecar.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# One node of a multinode prefill or decode group with NIXL transfer and KV routing. +# Defaults: two nodes per role, one GPU per node, Qwen/Qwen3-32B. + +set -e + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" + +if [[ "${1:-}" == --help || "${1:-}" == -h ]]; then + echo "Usage: NODE_RANK= DIST_INIT_ADDR= bash $0 [SGLang options...]" + echo "Run once per node, with separate rendezvous addresses for prefill and decode." + echo "All nodes and the frontend must share Dynamo namespace, discovery and event services." + echo + echo "Example: run each command on its respective node (four nodes total):" + echo " NODE_RANK=0 DIST_INIT_ADDR=prefill-leader:29500 SGLANG_BOOTSTRAP_HOST=prefill-leader bash $0 prefill" + echo " NODE_RANK=1 DIST_INIT_ADDR=prefill-leader:29500 bash $0 prefill" + echo " NODE_RANK=0 DIST_INIT_ADDR=decode-leader:29500 bash $0 decode" + echo " NODE_RANK=1 DIST_INIT_ADDR=decode-leader:29500 bash $0 decode" + echo + echo "Start the frontend separately: python3 -m dynamo.frontend --router-mode kv" + echo "Use reachable leader addresses; NIXL requires connectivity between the two groups." + echo "For model, topology and port overrides: bash $SCRIPT_DIR/multinode_kv_router_sidecar.sh --help" + exit 0 +fi + +case "${1:-}" in + prefill|decode) export ROLE="$1"; shift ;; + *) echo "First argument must be prefill or decode; use --help for examples" >&2; exit 1 ;; +esac + +# Reuse per-node engine/sidecar startup, validation, GPU budgeting and cleanup. +# ROLE enables NIXL; node 0 serves requests and followers relay local KV events. +exec bash "$SCRIPT_DIR/multinode_kv_router_sidecar.sh" "$@" From 24579affda400ebf9c428e0fd9256dee0a189be7 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Fri, 18 Sep 2026 13:33:37 -0700 Subject: [PATCH 15/17] test(sglang): assert advertised KV source identity Signed-off-by: jain-ria --- lib/sidecar/sglang/src/headless.rs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index 6e0ddcaa69c9..2bdf3409181b 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -621,7 +621,8 @@ mod tests { #[tokio::test] async fn local_zmq_events_keep_global_rank_and_leader_identity_without_serving() { use dynamo_kv_router::protocols::{KV_EVENT_SUBJECT, RouterEvent}; - use dynamo_runtime::discovery::{DiscoveryQuery, EventSourceQuery}; + use dynamo_llm::discovery::KvEventSource; + use dynamo_runtime::discovery::{DiscoveryInstance, DiscoveryQuery, EventSourceQuery}; use dynamo_runtime::transports::event_plane::EventSubscriber; use futures::SinkExt; @@ -696,14 +697,14 @@ mod tests { endpoint.id(), KV_EVENT_SUBJECT, )); - assert_eq!( - drt.discovery() - .list(source_query.clone()) - .await - .unwrap() - .len(), - 1 - ); + let sources = drt.discovery().list(source_query.clone()).await.unwrap(); + assert_eq!(sources.len(), 1); + let DiscoveryInstance::EventSource { metadata, .. } = &sources[0] else { + panic!("expected a KV event source"); + }; + let advertised: KvEventSource = serde_json::from_value(metadata.clone()).unwrap(); + assert_eq!(advertised.worker.dp_rank, 4); + assert_eq!(advertised.worker.worker_id, 42); drop(publishers); tokio::time::timeout(Duration::from_secs(5), async { while !drt From 551cbf2540eeb701da019c77ea1311feb6d326b0 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Mon, 5 Oct 2026 14:18:09 -0700 Subject: [PATCH 16/17] fix(sglang): reuse startup deadline for leader discovery Signed-off-by: jain-ria --- lib/sidecar/sglang/src/args.rs | 4 ---- lib/sidecar/sglang/src/headless.rs | 6 +++--- lib/sidecar/sglang/tests/executable.rs | 1 - 3 files changed, 3 insertions(+), 8 deletions(-) diff --git a/lib/sidecar/sglang/src/args.rs b/lib/sidecar/sglang/src/args.rs index ae72e22002b5..153bec1549a7 100644 --- a/lib/sidecar/sglang/src/args.rs +++ b/lib/sidecar/sglang/src/args.rs @@ -15,10 +15,6 @@ pub struct Args { #[command(flatten)] pub sidecar: SidecarArgs, - /// Maximum wait for a matching leader registration in telemetry mode. - #[arg(long, default_value_t = 1800, value_parser = clap::value_parser!(u64).range(1..))] - pub leader_discovery_timeout_secs: u64, - /// Reachable host that decode workers use to connect to a prefill worker's /// SGLang disaggregation bootstrap port. By default this is derived from /// SGLang's concrete `host`, then `dist_init_addr`, then a routable local diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index 20b9b831f931..a26acd89aae9 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -31,7 +31,6 @@ pub(crate) struct HeadlessSidecar { grpc_endpoint: GrpcEndpoint, transport: GrpcTransportConfig, common: CommonArgs, - discovery_timeout: Duration, } #[derive(Clone, Debug, Deserialize)] @@ -51,7 +50,6 @@ impl HeadlessSidecar { grpc_endpoint: args.sidecar.grpc_endpoint, transport: args.sidecar.grpc.config(), common: args.sidecar.common, - discovery_timeout: Duration::from_secs(args.leader_discovery_timeout_secs), }) } @@ -108,7 +106,9 @@ impl HeadlessSidecar { tracing::info!(node_rank = metadata.node_rank, group = group_id, endpoint = %endpoint.id(), "Waiting for SGLang leader for local KV publishing"); let mut configs = runtime_config_watch(&endpoint, shutdown.clone()).await?; - let (worker_id, config) = tokio::time::timeout(self.discovery_timeout, + // Leader registration follows engine readiness, so use the same startup + // allowance as local engine discovery rather than a separate setting. + let (worker_id, config) = tokio::time::timeout(self.transport.startup_deadline, wait_for_leader(&mut configs, group_id, &shutdown)).await .context("timed out waiting for SGLang leader; check namespace, component, endpoint, and dist_init_addr")??; let leader = validate_leader(metadata, &config)?; diff --git a/lib/sidecar/sglang/tests/executable.rs b/lib/sidecar/sglang/tests/executable.rs index 7e88a27978f8..3adbba59d684 100644 --- a/lib/sidecar/sglang/tests/executable.rs +++ b/lib/sidecar/sglang/tests/executable.rs @@ -27,7 +27,6 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { "DYN_SIDECAR_GRPC_RETRY_INTERVAL_SECS", "--grpc-startup-deadline-secs", "DYN_SIDECAR_GRPC_STARTUP_DEADLINE_SECS", - "--leader-discovery-timeout-secs", ] { assert!(stdout.contains(expected), "help omits {expected}"); } From a9d386e3286c37d88fa5838786a9667163378374 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Mon, 5 Oct 2026 15:17:14 -0700 Subject: [PATCH 17/17] fix(sglang): bound rendezvous DNS without blocking startup cancellation Signed-off-by: jain-ria --- lib/sidecar/sglang/src/client.rs | 151 ++++++++++++++++++++++--- lib/sidecar/sglang/src/engine.rs | 173 ++++++++++++++++++----------- lib/sidecar/sglang/src/headless.rs | 3 +- 3 files changed, 242 insertions(+), 85 deletions(-) diff --git a/lib/sidecar/sglang/src/client.rs b/lib/sidecar/sglang/src/client.rs index 5cfe497fb41a..5a33fc735f62 100644 --- a/lib/sidecar/sglang/src/client.rs +++ b/lib/sidecar/sglang/src/client.rs @@ -124,7 +124,10 @@ impl NodeMetadata { Ok(()) } - pub(crate) fn worker_group_id(&self) -> anyhow::Result> { + pub(crate) async fn worker_group_id( + &self, + deadline: Instant, + ) -> anyhow::Result> { if self.nnodes == 1 { return Ok(None); } @@ -146,10 +149,38 @@ impl NodeMetadata { validate_endpoint(&url)?; // Match the in-process group key using the shared rendezvous address, // not the local source or gRPC address. - worker_group_id_from_addresses(address.socket_addrs(|| None)?).map(Some) + resolve_rendezvous(move || address.socket_addrs(|| None), deadline) + .await + .and_then(worker_group_id_from_addresses) + .map(Some) } } +// getaddrinfo cannot be cancelled. Keep it off Tokio's workers and blocking +// pool so dropping startup or shutting down the runtime never waits for DNS. +// Each startup performs one lookup; an abandoned thread exits when DNS returns. +async fn resolve_rendezvous( + resolve: impl FnOnce() -> std::io::Result> + Send + 'static, + deadline: Instant, +) -> anyhow::Result> { + anyhow::ensure!( + Instant::now() < deadline, + "rendezvous DNS startup deadline elapsed" + ); + let (tx, rx) = tokio::sync::oneshot::channel(); + std::thread::Builder::new() + .name("sglang-rendezvous-dns".into()) + .spawn(move || { + let _ = tx.send(resolve()); + }) + .context("failed to start rendezvous DNS resolver")?; + timeout_at(deadline, rx) + .await + .context("timed out resolving SGLang rendezvous address")? + .context("rendezvous DNS resolver stopped")? + .context("failed to resolve SGLang rendezvous address") +} + fn worker_group_id_from_addresses(mut resolved: Vec) -> anyhow::Result { // Nodes may receive the same DNS answers in different orders. resolved.sort_unstable(); @@ -667,8 +698,8 @@ mod tests { }) } - #[test] - fn parses_node_local_sources_and_ignores_unrelated_server_fields() { + #[tokio::test] + async fn parses_node_local_sources_and_ignores_unrelated_server_fields() { let mut raw = node_metadata_json(); raw["model_path"] = json!("model-repo"); // Live-only relaying does not interpret the engine's optional replay field. @@ -676,7 +707,11 @@ mod tests { let metadata = NodeMetadata::from_server_info(&raw).unwrap().unwrap(); assert_eq!(metadata.kv_event_sources[0].dp_rank, 4); assert_eq!( - metadata.worker_group_id().unwrap().as_deref(), + metadata + .worker_group_id(Instant::now() + std::time::Duration::from_secs(1)) + .await + .unwrap() + .as_deref(), Some("dist_init:tcp://127.0.0.1:2345") ); metadata.validate_registration(8, Some(64)).unwrap(); @@ -726,18 +761,82 @@ mod tests { assert!(NodeMetadata::from_server_info(&raw).is_err()); } - #[test] - fn normalizes_ipv6_group_id_and_accepts_bound_ipc_sources() { + #[tokio::test] + async fn normalizes_ipv6_group_id_and_accepts_bound_ipc_sources() { let mut raw = node_metadata_json(); raw["dist_init_addr"] = json!("tcp://[::1]:2345"); raw["kv_event_sources"][0]["endpoint"] = json!("ipc:///engine/kv-events"); let metadata = NodeMetadata::from_server_info(&raw).unwrap().unwrap(); assert_eq!( - metadata.worker_group_id().unwrap().as_deref(), + metadata + .worker_group_id(Instant::now() + std::time::Duration::from_secs(1)) + .await + .unwrap() + .as_deref(), Some("dist_init:tcp://[::1]:2345") ); } + #[tokio::test] + async fn rendezvous_dns_deadline_bounds_a_blocked_resolver() { + let (release, blocked) = std::sync::mpsc::channel::<()>(); + let error = super::resolve_rendezvous( + move || { + let _ = blocked.recv_timeout(std::time::Duration::from_secs(3)); + Ok(vec![]) + }, + Instant::now() + std::time::Duration::from_millis(25), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("timed out resolving")); + drop(release); + } + + #[test] + fn cancelling_dns_does_not_hold_up_runtime_shutdown() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release, blocked) = std::sync::mpsc::channel::<()>(); + runtime.block_on(async { + let lookup = super::resolve_rendezvous( + move || { + let _ = started_tx.send(()); + // Bound the injected delay even if runtime shutdown regresses. + let _ = blocked.recv_timeout(std::time::Duration::from_secs(3)); + Ok(vec![]) + }, + Instant::now() + std::time::Duration::from_secs(10), + ); + tokio::select! { + result = lookup => panic!("resolver returned before cancellation: {result:?}"), + _ = started_rx => {}, + } + }); + let shutdown_started = std::time::Instant::now(); + drop(runtime); + let elapsed = shutdown_started.elapsed(); + drop(release); + assert!( + elapsed < std::time::Duration::from_secs(1), + "shutdown waited for DNS: {elapsed:?}" + ); + } + + #[tokio::test] + async fn rendezvous_dns_propagates_resolution_errors() { + let error = super::resolve_rendezvous( + || Err(std::io::Error::other("injected resolver failure")), + Instant::now() + std::time::Duration::from_secs(1), + ) + .await + .unwrap_err(); + assert!(format!("{error:#}").contains("injected resolver failure")); + } + #[test] fn worker_group_id_is_independent_of_dns_answer_order() { let addresses = vec![ @@ -758,22 +857,38 @@ mod tests { } #[test] - fn discovery_mode_reads_only_server_metadata() { + fn discovery_roles_accept_native_aliases_and_reject_unknown_strings() { use dynamo_backend_common::DisaggregationMode; + for (value, expected) in [ + (json!(null), DisaggregationMode::Aggregated), + (json!("null"), DisaggregationMode::Aggregated), + (json!("agg"), DisaggregationMode::Aggregated), + (json!("aggregated"), DisaggregationMode::Aggregated), + (json!("prefill"), DisaggregationMode::Prefill), + (json!("decode"), DisaggregationMode::Decode), + ] { + assert_eq!( + discovery_mode(&json!({"disaggregation_mode": value})).unwrap(), + expected + ); + } assert_eq!( discovery_mode(&json!({})).unwrap(), DisaggregationMode::Aggregated ); - assert_eq!( - discovery_mode(&json!({"disaggregation_mode": "prefill"})).unwrap(), - DisaggregationMode::Prefill - ); - assert_eq!( - discovery_mode(&json!({"disaggregation_mode": "decode"})).unwrap(), - DisaggregationMode::Decode - ); - assert!(discovery_mode(&json!({"disaggregation_mode": "unknown"})).is_err()); + for mode in ["encode", "unknown", ""] { + let error = discovery_mode(&json!({"disaggregation_mode": mode})).unwrap_err(); + assert_eq!( + error.error_type(), + ErrorType::Backend(BackendError::Unknown) + ); + assert!( + error + .to_string() + .contains("unsupported SGLang disaggregation_mode") + ); + } } #[test] diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index 30ae23235adb..c456763365e6 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -294,7 +294,8 @@ impl LLMEngine for SglangSidecarEngine { .runtime_data .insert("sglang_generate".into(), true.into()); } - let kv_event_sources = discover_kv_event_sources(&discovery, &mut config, &self.endpoint)?; + let kv_event_sources = + discover_kv_event_sources(&discovery, &mut config, &self.endpoint, deadline).await?; let connection_count = pool.len(); let kv_event_source_count = kv_event_sources.len(); self.state @@ -686,9 +687,10 @@ fn is_routable_host(host: &str) -> bool { .unwrap_or(true) } -fn local_kv_event_sources( +async fn local_kv_event_sources( metadata: &NodeMetadata, config: &mut EngineConfig, + deadline: Instant, ) -> Result, DynamoError> { let llm = config .llm @@ -698,7 +700,8 @@ fn local_kv_event_sources( .validate_registration(llm.data_parallel_size.unwrap_or(1), llm.kv_cache_block_size) .map_err(|error| client::invalid_arg(error.to_string()))?; if let Some(group_id) = metadata - .worker_group_id() + .worker_group_id(deadline) + .await .map_err(|error| client::invalid_arg(error.to_string()))? { config @@ -720,17 +723,18 @@ fn local_kv_event_sources( .collect()) } -fn discover_kv_event_sources( +async fn discover_kv_event_sources( discovery: &Discovery, engine_config: &mut EngineConfig, grpc_endpoint: &GrpcEndpoint, + deadline: Instant, ) -> Result, DynamoError> { // An explicit list is authoritative, including an empty list. Only older // engines without this field may use the legacy single-node descriptor. if let Some(metadata) = NodeMetadata::from_server_info(&discovery.server_info) .map_err(|error| client::protocol_error(error.to_string()))? { - return local_kv_event_sources(&metadata, engine_config); + return local_kv_event_sources(&metadata, engine_config, deadline).await; } let Some(descriptor) = discovery.server_info.get("kv_events") else { return Ok(Vec::new()); @@ -1071,6 +1075,9 @@ fn build_engine_config( #[cfg(test)] mod tests { + use std::time::Duration; + use tokio::time::Instant; + use clap::Parser; use dynamo_backend_common::{BackendError, ErrorType, ModelInput}; use dynamo_sidecar_common::GrpcEndpoint; @@ -1222,40 +1229,6 @@ mod tests { ); } - #[test] - fn discovery_roles_accept_native_aliases_and_reject_unknown_strings() { - for (value, expected) in [ - (json!(null), DisaggregationMode::Aggregated), - (json!("null"), DisaggregationMode::Aggregated), - (json!("agg"), DisaggregationMode::Aggregated), - (json!("aggregated"), DisaggregationMode::Aggregated), - (json!("prefill"), DisaggregationMode::Prefill), - (json!("decode"), DisaggregationMode::Decode), - ] { - assert_eq!( - crate::client::discovery_mode(&json!({"disaggregation_mode": value})).unwrap(), - expected - ); - } - assert_eq!( - crate::client::discovery_mode(&json!({})).unwrap(), - DisaggregationMode::Aggregated - ); - for mode in ["encode", "unknown", ""] { - let error = - crate::client::discovery_mode(&json!({"disaggregation_mode": mode})).unwrap_err(); - assert_eq!( - error.error_type(), - ErrorType::Backend(BackendError::Unknown) - ); - assert!( - error - .to_string() - .contains("unsupported SGLang disaggregation_mode") - ); - } - } - #[test] fn bootstrap_port_requires_a_nonzero_u16() { for (value, expected) in [(json!(1), 1), (json!("65535"), u16::MAX)] { @@ -1420,12 +1393,18 @@ mod tests { (discovery, config, endpoint) } - #[test] - fn kv_event_descriptor_is_optional_but_validated_when_present() { + #[tokio::test] + async fn kv_event_descriptor_is_optional_but_validated_when_present() { let (mut discovery, mut config, endpoint) = kv_discovery(); for value in [json!(null), json!([]), json!("zmq")] { discovery.server_info["kv_events"] = value.clone(); - let result = discover_kv_event_sources(&discovery, &mut config, &endpoint); + let result = discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1), + ) + .await; if value.is_null() { assert!(result.unwrap().is_empty()); } else { @@ -1439,9 +1418,15 @@ mod tests { } discovery.server_info = json!({}); assert!( - discover_kv_event_sources(&discovery, &mut config, &endpoint) - .unwrap() - .is_empty() + discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1) + ) + .await + .unwrap() + .is_empty() ); for (field, value) in [ @@ -1458,7 +1443,14 @@ mod tests { ] { let (mut discovery, mut config, endpoint) = kv_discovery(); discovery.server_info["kv_events"][field] = value; - let error = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap_err(); + let error = discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1), + ) + .await + .unwrap_err(); assert_eq!( error.error_type(), ErrorType::Backend(BackendError::Unknown) @@ -1467,8 +1459,8 @@ mod tests { } } - #[test] - fn kv_events_require_matching_registration_and_rank_coverage() { + #[tokio::test] + async fn kv_events_require_matching_registration_and_rank_coverage() { for (block_size, dp_size, dp_start, message) in [ (64, 2, 0, "block size"), (128, 1, 0, "reports 2 DP ranks"), @@ -1480,7 +1472,14 @@ mod tests { llm.kv_cache_block_size = Some(block_size); llm.data_parallel_size = Some(dp_size); llm.data_parallel_start_rank = Some(dp_start); - let error = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap_err(); + let error = discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1), + ) + .await + .unwrap_err(); assert_eq!( error.error_type(), ErrorType::Backend(BackendError::Unknown) @@ -1490,15 +1489,21 @@ mod tests { let (discovery, mut config, endpoint) = kv_discovery(); config.llm = None; assert!( - discover_kv_event_sources(&discovery, &mut config, &endpoint) - .unwrap_err() - .to_string() - .contains("require an LLM engine registration") + discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1) + ) + .await + .unwrap_err() + .to_string() + .contains("require an LLM engine registration") ); } - #[test] - fn kv_events_reject_unmapped_multi_node_dp_and_port_overflow() { + #[tokio::test] + async fn kv_events_reject_unmapped_multi_node_dp_and_port_overflow() { for (nnodes, base_port, message) in [ (2, 5557, "multi-node DP are unsupported"), (1, 65535, "port overflows 65535"), @@ -1506,7 +1511,14 @@ mod tests { let (mut discovery, mut config, endpoint) = kv_discovery(); discovery.server_info["nnodes"] = json!(nnodes); discovery.server_info["kv_events"]["endpoint_port_base"] = json!(base_port); - let error = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap_err(); + let error = discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1), + ) + .await + .unwrap_err(); assert_eq!( error.error_type(), ErrorType::Backend(BackendError::Unknown) @@ -1799,8 +1811,8 @@ mod tests { ); } - #[test] - fn discovers_ranked_kv_event_sources_from_server_info() { + #[tokio::test] + async fn discovers_ranked_kv_event_sources_from_server_info() { let discovery = discovery(json!({ "page_size": 64, "dcp_size": 2, @@ -1819,7 +1831,14 @@ mod tests { build_engine_config(&discovery, DisaggregationMode::Aggregated, None, None).unwrap(); let endpoint = GrpcEndpoint::parse("http://worker.example:30001", "test").unwrap(); - let sources = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap(); + let sources = discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1), + ) + .await + .unwrap(); assert_eq!( sources, @@ -1838,8 +1857,8 @@ mod tests { ); } - #[test] - fn local_metadata_overrides_global_sources_without_shrinking_registration() { + #[tokio::test] + async fn local_metadata_overrides_global_sources_without_shrinking_registration() { use crate::client::{KV_CONFIG_KEY, WORKER_GROUP_KEY}; let mut discovery = discovery(json!({ @@ -1860,7 +1879,14 @@ mod tests { let mut config = build_engine_config(&discovery, DisaggregationMode::Aggregated, None, None).unwrap(); let endpoint = GrpcEndpoint::parse("http://127.0.0.1:30001", "test").unwrap(); - let sources = discover_kv_event_sources(&discovery, &mut config, &endpoint).unwrap(); + let sources = discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1), + ) + .await + .unwrap(); assert_eq!( sources .iter() @@ -1879,9 +1905,15 @@ mod tests { // An explicit empty source list must not fall back to all global ranks. discovery.server_info["kv_event_sources"] = json!([]); assert!( - discover_kv_event_sources(&discovery, &mut config, &endpoint) - .unwrap() - .is_empty() + discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1) + ) + .await + .unwrap() + .is_empty() ); assert_eq!( config.runtime_data[KV_CONFIG_KEY]["local_dp_ranks"], @@ -1895,6 +1927,15 @@ mod tests { .as_object_mut() .unwrap() .remove("kv_event_sources"); - assert!(discover_kv_event_sources(&discovery, &mut config, &endpoint).is_err()); + assert!( + discover_kv_event_sources( + &discovery, + &mut config, + &endpoint, + Instant::now() + Duration::from_secs(1) + ) + .await + .is_err() + ); } } diff --git a/lib/sidecar/sglang/src/headless.rs b/lib/sidecar/sglang/src/headless.rs index a26acd89aae9..20cc8068da8a 100644 --- a/lib/sidecar/sglang/src/headless.rs +++ b/lib/sidecar/sglang/src/headless.rs @@ -62,7 +62,8 @@ impl HeadlessSidecar { // model discovery, HealthCheck, connection pool, or native HTTP client. let (mut client, metadata, mode) = self.connect_local_engine().await?; let group_id = metadata - .worker_group_id()? + .worker_group_id(Instant::now() + self.transport.startup_deadline) + .await? .context("telemetry-only mode requires a multinode group")?; // Supervise the local engine while waiting for the leader too. Engines