diff --git a/Cargo.lock b/Cargo.lock index 831fc057dc0b..a4e7a77759f3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3085,13 +3085,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 63d1be2d3bbd..30e7451cf252 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -2569,12 +2569,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 9d68aa449fa6..5e00cd1fc8c9 100644 --- a/lib/bindings/python/rust/backend.rs +++ b/lib/bindings/python/rust/backend.rs @@ -108,10 +108,7 @@ 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 bootstrap = dynamo_sglang_sidecar::SglangSidecarEngine::try_from_args_async(cli_argv) - .map_err(sidecar_startup_to_pyerr)?; - - py.allow_threads(move || dynamo_sidecar_common::run(bootstrap)) + py.allow_threads(move || dynamo_sglang_sidecar::run(cli_argv)) .map_err(sidecar_run_to_pyerr) } diff --git a/lib/sidecar/common/src/lib.rs b/lib/sidecar/common/src/lib.rs index 34f117a427a4..83efd9aebf2f 100644 --- a/lib/sidecar/common/src/lib.rs +++ b/lib/sidecar/common/src/lib.rs @@ -33,7 +33,7 @@ pub use transport::{ DEFAULT_MAX_GRPC_MESSAGE_SIZE, GrpcChannelPool, format_error_chain, startup_deadline, }; -pub use run::run; +pub use run::{run, run_task}; #[cfg(test)] #[path = "transport/tests.rs"] diff --git a/lib/sidecar/common/src/run.rs b/lib/sidecar/common/src/run.rs index 0d39b03b9a5c..d3f0b8de74d8 100644 --- a/lib/sidecar/common/src/run.rs +++ b/lib/sidecar/common/src/run.rs @@ -16,6 +16,18 @@ use crate::SidecarStartupError; pub fn run( bootstrap: impl Future>, ) -> anyhow::Result<()> { + run_task( + |runtime, shutdown| async move { run_until_shutdown(bootstrap, &runtime, shutdown).await }, + ) +} + +/// Run a sidecar task with the shared runtime and process-signal lifecycle. +/// The task owns discovery and may serve inference or relay telemetry only. +pub fn run_task(task: F) -> anyhow::Result<()> +where + F: FnOnce(Runtime, CancellationToken) -> Fut, + Fut: Future>, +{ logging::init(); let runtime = Runtime::from_settings()?; let secondary = runtime.secondary(); @@ -35,7 +47,7 @@ pub fn run( signal_runtime.mark_shutting_down(); }); - let result = run_until_shutdown(bootstrap, &runtime, shutdown.clone()).await; + let result = task(runtime.clone(), shutdown.clone()).await; shutdown.cancel(); signal_handle.abort(); let _ = signal_handle.await; diff --git a/lib/sidecar/sglang/Cargo.toml b/lib/sidecar/sglang/Cargo.toml index 4f0b9a9ddbb4..10264a1d40e6 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,4 +47,8 @@ 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 } tracing-subscriber = { workspace = true } diff --git a/lib/sidecar/sglang/README.md b/lib/sidecar/sglang/README.md index 6342e51546a9..d4e8b5096ff8 100644 --- a/lib/sidecar/sglang/README.md +++ b/lib/sidecar/sglang/README.md @@ -36,7 +36,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. @@ -55,7 +55,47 @@ 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. + +## Multinode DP and KV routing + +For multinode deployments with multiple data-parallel (DP) replicas and KV-aware +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 +`kv_event_sources` through `GetServerInfo` on leaders and followers started +with `--grpc-port`. + +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. + +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) 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" "$@" diff --git a/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh b/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh new file mode 100755 index 000000000000..a2be8aff153a --- /dev/null +++ b/lib/sidecar/sglang/launch/multinode_kv_router_sidecar.sh @@ -0,0 +1,114 @@ +#!/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 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 + +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-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" + 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-32B}" +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 +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. +# 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" \ + --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}\",\"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 "$@" & + +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/client.rs b/lib/sidecar/sglang/src/client.rs index 594e44ca2dd8..5a33fc735f62 100644 --- a/lib/sidecar/sglang/src/client.rs +++ b/lib/sidecar/sglang/src/client.rs @@ -3,14 +3,18 @@ //! Thin sidecar client for SGLang's native `sglang.runtime.v1.SglangService`. +use std::collections::HashSet; use std::future::Future; +use std::net::SocketAddr; 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, format_error_chain, }; +use serde::Deserialize; use serde_json::Value; use tokio::time::{Instant, timeout_at}; use tonic::transport::{Channel, Endpoint}; @@ -20,6 +24,204 @@ 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) async fn worker_group_id( + &self, + deadline: Instant, + ) -> 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. + 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(); + 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!( + (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(()) +} const RETRY_LOG_INTERVAL: Duration = Duration::from_secs(30); /// Metadata exposed by SGLang's model/server discovery RPCs. @@ -33,6 +235,31 @@ pub struct Discovery { pub server_info: Value, } +#[derive(Debug)] +pub(crate) enum StartupDiscovery { + Leader(Box), + Follower, +} + +pub(crate) async fn bootstrap_discover( + endpoint: &GrpcEndpoint, + transport: &GrpcTransportConfig, + bootstrap: bool, +) -> Result { + let deadline = Instant::now() + transport.startup_deadline; + let mut client = connect(endpoint, transport, deadline, bootstrap).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(|d| StartupDiscovery::Leader(Box::new(d))) + } +} + /// `bootstrap`: true for synchronous constructors before logging setup; /// false for deferred launcher discovery and `LLMEngine::start`. pub async fn connect( @@ -156,6 +383,20 @@ 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( + "inference discovery requires node_rank=0; followers expose only GetServerInfo", + )); + } let model = rpc_with_deadline( "GetModelInfo", deadline, @@ -163,13 +404,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 { @@ -217,11 +467,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)?; // Generate responses are forwarded as token deltas. Accepting cumulative // output here would duplicate tokens and inflate completion usage. if server_info @@ -287,6 +536,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}")))?; @@ -376,8 +640,8 @@ mod tests { use tonic::transport::Endpoint; use super::{ - client_from_channel, discover, json_u32, json_u64, parse_discovery, rpc_with_deadline, - status_to_dynamo, + NodeMetadata, client_from_channel, discover, discovery_mode, json_u32, json_u64, + parse_discovery, rpc_with_deadline, status_to_dynamo, worker_group_id_from_addresses, }; use crate::proto as pb; @@ -415,9 +679,7 @@ mod tests { model_path: "model-repo".to_string(), json_info: json!({"tokenizer_path": "tokenizer-repo"}).to_string(), }, - pb::GetServerInfoResponse { - json_info: json!({"incremental_streaming_output": true}).to_string(), - }, + json!({"incremental_streaming_output": true}), Vec::new(), ) .unwrap(); @@ -425,6 +687,210 @@ 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 + }] + }) + } + + #[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. + 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(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(); + 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()); + } + + #[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(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![ + "[::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_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 + ); + 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] fn discovery_requires_incremental_streaming() { for info in [ @@ -438,9 +904,7 @@ mod tests { model_path: "model-repo".to_string(), json_info: "{}".to_string(), }, - pb::GetServerInfoResponse { - json_info: info.to_string(), - }, + info, Vec::new(), ) .unwrap_err(); @@ -459,26 +923,7 @@ mod tests { fn discovery_rejects_malformed_json_and_non_object_metadata() { for label in ["GetModelInfo.json_info", "GetServerInfo.json_info"] { for raw in ["{", "null", "[]", "1", "\"metadata\""] { - let error = parse_discovery( - pb::GetModelInfoResponse { - model_path: "model".into(), - json_info: if label.starts_with("GetModelInfo") { - raw - } else { - "{}" - } - .into(), - }, - pb::GetServerInfoResponse { - json_info: if label.starts_with("GetServerInfo") { - raw.into() - } else { - json!({"incremental_streaming_output": true}).to_string() - }, - }, - vec![], - ) - .unwrap_err(); + let error = super::parse_json_object(label, raw).unwrap_err(); assert_eq!( error.error_type(), ErrorType::Backend(BackendError::Unknown) @@ -502,9 +947,7 @@ mod tests { json_info: json!({"model_path": json_path, "tokenizer_path": tokenizer}) .to_string(), }, - pb::GetServerInfoResponse { - json_info: json!({"incremental_streaming_output": true}).to_string(), - }, + json!({"incremental_streaming_output": true}), vec![], ); if let Some(expected) = expected { @@ -543,9 +986,7 @@ mod tests { model_path: "model".into(), json_info: model_info.to_string(), }, - pb::GetServerInfoResponse { - json_info: server_info.to_string(), - }, + server_info.clone(), vec![ pb::ModelCard { id: "unrelated".into(), @@ -590,15 +1031,12 @@ mod tests { model_path: "model".into(), json_info: "{}".into(), }, - pb::GetServerInfoResponse { - json_info: json!({ - "incremental_streaming_output": true, - "served_model_name": "", - "context_length": context, - "max_req_input_len": input_limit, - }) - .to_string(), - }, + json!({ + "incremental_streaming_output": true, + "served_model_name": "", + "context_length": context, + "max_req_input_len": input_limit, + }), vec![pb::ModelCard { id: "first-alias".into(), max_model_len: card_limit, diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index 7bf44cfe26c1..c456763365e6 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -22,7 +22,7 @@ use tokio::time::Instant; use tokio_util::sync::CancellationToken; use crate::args::Args; -use crate::client::{self, Client, Discovery, Pool}; +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::{ @@ -104,7 +104,7 @@ impl SglangSidecarEngine { runtime.block_on(Self::from_parsed_async(args, true)) } - fn validate_args(args: &Args) -> Result<(), DynamoError> { + pub(crate) fn validate_args(args: &Args) -> Result<(), DynamoError> { if args.sidecar.common.route_to_encoder { return Err(client::invalid_arg( "route-to-encoder is not supported by the SGLang sidecar", @@ -126,13 +126,14 @@ impl SglangSidecarEngine { Self::from_discovered(args, discovery) } - fn from_discovered( + pub(crate) fn from_discovered( args: Args, discovery: Discovery, ) -> Result<(Self, WorkerConfig), DynamoError> { + Self::validate_args(&args)?; let endpoint = args.sidecar.grpc_endpoint; let transport = args.sidecar.grpc.config(); - 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(), @@ -252,7 +253,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 {:?}", @@ -293,7 +294,8 @@ 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 = + 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 @@ -574,24 +576,11 @@ async fn bootstrap_discover( transport: &GrpcTransportConfig, bootstrap: bool, ) -> Result { - let deadline = Instant::now() + transport.startup_deadline; - let mut grpc_client = client::connect(endpoint, transport, deadline, bootstrap).await?; - client::discover(&mut grpc_client, deadline).await -} - -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}`" - ))), + match client::bootstrap_discover(endpoint, transport, bootstrap).await? { + client::StartupDiscovery::Leader(discovery) => Ok(*discovery), + client::StartupDiscovery::Follower => { + Err(client::invalid_arg("inference engine requires node_rank=0")) + } } } @@ -698,11 +687,55 @@ fn is_routable_host(host: &str) -> bool { .unwrap_or(true) } -fn discover_kv_event_sources( +async fn local_kv_event_sources( + metadata: &NodeMetadata, + config: &mut EngineConfig, + deadline: Instant, +) -> Result, DynamoError> { + let llm = config + .llm + .as_ref() + .ok_or_else(|| client::invalid_arg("KV sources require LLM registration"))?; + 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) = metadata + .worker_group_id(deadline) + .await + .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": metadata.kv_event_sources.iter().map(|source| source.dp_rank).collect::>(), + })); + Ok(metadata + .kv_event_sources + .iter() + .map(|source| DiscoveredKvEventSource { + endpoint: source.endpoint.clone(), + topic: source.topic.clone(), + dp_rank: source.dp_rank, + }) + .collect()) +} + +async fn discover_kv_event_sources( discovery: &Discovery, - engine_config: &EngineConfig, + 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, deadline).await; + } let Some(descriptor) = discovery.server_info.get("kv_events") else { return Ok(Vec::new()); }; @@ -1042,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; @@ -1050,8 +1086,8 @@ mod tests { use super::{ Args, DisaggregationMode, DiscoveredKvEventSource, Discovery, EngineConfig, SglangSidecarEngine, build_engine_config, discover_kv_event_sources, - discovery_bootstrap_port, discovery_mode, hicache_native_offloading_capacity, - kv_event_connect_host, resolve_bootstrap_host_with_local, sglang_eagle_enabled, + discovery_bootstrap_port, hicache_native_offloading_capacity, kv_event_connect_host, + resolve_bootstrap_host_with_local, sglang_eagle_enabled, }; fn discovery(server_info: serde_json::Value) -> Discovery { @@ -1193,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!( - discovery_mode(&discovery(json!({"disaggregation_mode": value}))).unwrap(), - expected - ); - } - assert_eq!( - discovery_mode(&discovery(json!({}))).unwrap(), - DisaggregationMode::Aggregated - ); - for mode in ["encode", "unknown", ""] { - let error = - discovery_mode(&discovery(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)] { @@ -1391,12 +1393,18 @@ mod tests { (discovery, config, endpoint) } - #[test] - fn kv_event_descriptor_is_optional_but_validated_when_present() { - let (mut discovery, config, endpoint) = kv_discovery(); + #[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, &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 { @@ -1410,9 +1418,15 @@ mod tests { } discovery.server_info = json!({}); assert!( - discover_kv_event_sources(&discovery, &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 [ @@ -1427,9 +1441,16 @@ mod tests { ("block_size", json!(u64::from(u32::MAX) + 1)), ("dp_size", json!(0)), ] { - let (mut discovery, config, endpoint) = kv_discovery(); + let (mut discovery, mut config, endpoint) = kv_discovery(); discovery.server_info["kv_events"][field] = value; - let error = discover_kv_event_sources(&discovery, &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) @@ -1438,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"), @@ -1451,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, &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) @@ -1461,23 +1489,36 @@ mod tests { let (discovery, mut config, endpoint) = kv_discovery(); config.llm = None; assert!( - discover_kv_event_sources(&discovery, &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"), ] { - let (mut discovery, config, endpoint) = kv_discovery(); + 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, &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) @@ -1770,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, @@ -1786,11 +1827,18 @@ 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, + Instant::now() + Duration::from_secs(1), + ) + .await + .unwrap(); assert_eq!( sources, @@ -1808,4 +1856,86 @@ mod tests { ] ); } + + #[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!({ + "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, + "dp_size": 8 + } + })); + 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, + Instant::now() + Duration::from_secs(1), + ) + .await + .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], + "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. + discovery.server_info["kv_event_sources"] = json!([]); + assert!( + 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"], + json!([]) + ); + assert_eq!(config.llm.as_ref().unwrap().data_parallel_size, Some(8)); + + // 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, + 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 new file mode 100644 index 000000000000..20cc8068da8a --- /dev/null +++ b/lib/sidecar/sglang/src/headless.rs @@ -0,0 +1,699 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! 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, 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::config::HealthStatus; +use dynamo_runtime::distributed::DistributedRuntime; +use dynamo_runtime::prelude::DistributedRuntimeProvider; +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::{self, KV_CONFIG_KEY, NodeMetadata, WORKER_GROUP_KEY}; + +const METADATA_POLL_INTERVAL: Duration = Duration::from_secs(5); + +pub(crate) struct HeadlessSidecar { + grpc_endpoint: GrpcEndpoint, + transport: GrpcTransportConfig, + common: CommonArgs, +} + +#[derive(Clone, Debug, Deserialize)] +struct LeaderKvConfig { + block_size: u32, + local_dp_ranks: Vec, +} + +impl HeadlessSidecar { + pub(crate) fn from_args(args: Args) -> Result { + 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 { + grpc_endpoint: args.sidecar.grpc_endpoint, + transport: args.sidecar.grpc.config(), + common: args.sidecar.common, + }) + } + + pub(crate) async fn run_inner( + &self, + drt: DistributedRuntime, + 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(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 + // 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(drt, 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, false).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, + drt: DistributedRuntime, + shutdown: CancellationToken, + metadata: &NodeMetadata, + mode: DisaggregationMode, + group_id: &str, + ) -> Result<()> { + let component = if mode == DisaggregationMode::Aggregated { + &self.common.component + } else { + mode.discovery_component() + }; + let endpoint = drt + .namespace(&self.common.namespace)? + .component(component)? + .endpoint(&self.common.endpoint); + 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?; + // 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)?; + 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 { + tokio::select! { + _ = shutdown.cancelled() => return Ok(()), + result = configs.changed() => { + result.context("leader discovery watch closed")?; + check_leader_unchanged(&configs.borrow_and_update(), group_id, worker_id, &config)?; + } + } + } + } +} + +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, + "KV relay requires a follower node" + ); + 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, + 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: &NodeMetadata, 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: &NodeMetadata, + worker_id: u64, + config: &ModelRuntimeConfig, + leader: &LeaderKvConfig, +) -> Result> { + let publishers = 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::>>()?; + // 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)] +mod tests { + use super::*; + use dynamo_runtime::{Runtime, distributed::DistributedConfig}; + 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 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 transport = GrpcTransportConfig { + startup_deadline: Duration::from_secs(5), + ..Default::default() + }; + let result = client::bootstrap_discover(&endpoint, &transport, false).await; + 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", + "--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); + // 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("inference discovery requires node_rank=0") + ); + 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 { + 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, _) = follower_metadata(&follower_info()).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_llm::discovery::KvEventSource; + use dynamo_runtime::discovery::{DiscoveryInstance, 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 = follower_info(); + raw["kv_event_sources"][0]["endpoint"] = json!(source_address); + 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(); + 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"); + 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, + )); + 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 + .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..8292219111d1 100644 --- a/lib/sidecar/sglang/src/lib.rs +++ b/lib/sidecar/sglang/src/lib.rs @@ -9,12 +9,22 @@ //! 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; + +use clap::Parser; +use dynamo_sidecar_common::SidecarStartupError; + +use args::Args; +use headless::HeadlessSidecar; pub mod args; pub mod client; pub mod engine; +mod headless; mod native_http; /// Generated SGLang gRPC types, temporarily exposed for the Mocker server @@ -24,3 +34,79 @@ pub mod proto; mod protocol; pub use engine::SglangSidecarEngine; + +/// 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)?; + SglangSidecarEngine::validate_args(&args).map_err(SidecarStartupError::from)?; + dynamo_sidecar_common::run_task(|runtime, shutdown| async move { + use dynamo_runtime::system_status_server::SystemProbePolicy; + use dynamo_runtime::{DistributedRuntime, distributed::DistributedConfig}; + let startup = async { + let drt = DistributedRuntime::new_with_probe_policy( + runtime.clone(), + DistributedConfig::try_from_settings()?, + SystemProbePolicy::RuntimeOnly, + ) + .await?; + tracing::info!("Sidecar runtime connected; discovering engine metadata"); + let discovery = client::bootstrap_discover( + &args.sidecar.grpc_endpoint, + &args.sidecar.grpc.config(), + false, + ) + .await + .map_err(SidecarStartupError::from)?; + Ok::<_, anyhow::Error>((drt, discovery)) + }; + let runtime_shutdown = runtime.shutdown_started_token(); + let result = tokio::select! { + biased; + _ = shutdown.cancelled() => return Ok(()), + _ = runtime_shutdown.cancelled() => anyhow::bail!("runtime shut down during sidecar initialization"), + result = startup => result, + }; + if runtime.is_shutting_down() { + if shutdown.is_cancelled() { + return Ok(()); + } + anyhow::bail!("runtime shut down during sidecar initialization"); + } + let (drt, discovery) = result?; + match discovery { + client::StartupDiscovery::Follower => { + let follower = + HeadlessSidecar::from_args(args).map_err(SidecarStartupError::from)?; + tokio::select! { + biased; + _ = shutdown.cancelled() => Ok(()), + _ = runtime_shutdown.cancelled() => anyhow::bail!("runtime shut down during follower relay"), + result = follower.run_inner(drt, shutdown.clone()) => result, + } + } + client::StartupDiscovery::Leader(discovery) => { + let (engine, config) = SglangSidecarEngine::from_discovered(args, *discovery) + .map_err(SidecarStartupError::from)?; + dynamo_backend_common::Worker::new(Arc::new(engine), config) + .run_with_drt(drt, shutdown) + .await + .map_err(Into::into) + } + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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 b6aad57e6ee0..c0efb9a0c9b7 100644 --- a/lib/sidecar/sglang/src/main.rs +++ b/lib/sidecar/sglang/src/main.rs @@ -2,5 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 fn main() -> anyhow::Result<()> { - dynamo_sidecar_common::run(dynamo_sglang_sidecar::SglangSidecarEngine::from_cli()?) + 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/tests/executable.rs b/lib/sidecar/sglang/tests/executable.rs index e6fb4c9c5ad3..3adbba59d684 100644 --- a/lib/sidecar/sglang/tests/executable.rs +++ b/lib/sidecar/sglang/tests/executable.rs @@ -32,6 +32,18 @@ fn executable_exposes_sglang_and_shared_sidecar_contracts() { } } +#[test] +fn sidecar_requires_a_local_grpc_endpoint() { + let output = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar")) + .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 ")); +} + #[test] fn invalid_arguments_fail_before_runtime_configuration() { let mut command = Command::new(env!("CARGO_BIN_EXE_dynamo-sglang-sidecar"));