From 27f796ca6c40e9a0b944dd7f9cdc241937dcd6a0 Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Mon, 21 Sep 2026 20:13:06 -0700 Subject: [PATCH] [sgl-router] Fix readiness, IPv6 discovery, logging, and model validation (#40604) --- experimental/sgl-router/src/discovery/k8s.rs | 16 +++++- experimental/sgl-router/src/main.rs | 16 +----- .../sgl-router/src/server/routes/chat.rs | 9 +-- .../sgl-router/src/server/routes/health.rs | 55 +++++++++++++++---- .../sgl-router/tests/proxy/chat_routing.rs | 15 +++++ 5 files changed, 79 insertions(+), 32 deletions(-) diff --git a/experimental/sgl-router/src/discovery/k8s.rs b/experimental/sgl-router/src/discovery/k8s.rs index ec9598b22..561efd575 100644 --- a/experimental/sgl-router/src/discovery/k8s.rs +++ b/experimental/sgl-router/src/discovery/k8s.rs @@ -118,7 +118,12 @@ fn extract_workers(es: &EndpointSlice, mode: WorkerMode) -> Vec { } let pod_uid: Option<&str> = ep.target_ref.as_ref().and_then(|r| r.uid.as_deref()); for addr in &ep.addresses { - let url = format!("http://{addr}:{port}"); + let host = if es.address_type == "IPv6" { + format!("[{addr}]") + } else { + addr.clone() + }; + let url = format!("http://{host}:{port}"); let id = match pod_uid { Some(uid) => WorkerId(format!("{ns}/{uid}")), None => WorkerId(format!("{ns}/{slice_name}/{addr}:{port}")), @@ -541,6 +546,15 @@ mod tests { assert_eq!(ws[0].mode, WorkerMode::Decode); } + #[test] + fn brackets_ipv6_worker_addresses() { + let mut slice = make_slice(&["2001:db8::1"], 30000, true); + slice.address_type = "IPv6".into(); + let workers = extract_workers(&slice, WorkerMode::Plain); + assert_eq!(workers[0].url, "http://[2001:db8::1]:30000"); + assert!(url::Url::parse(&workers[0].url).is_ok()); + } + #[test] fn skips_not_ready_endpoints() { let s = make_slice(&["10.0.0.1"], 30000, false); diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index b41c0b4c3..c9cb8608c 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -50,14 +50,10 @@ const DRAIN_WARN_AFTER: Duration = Duration::from_secs(30); async fn main() -> Result<()> { // Resolve CLI configuration and set up startup logging. let cli = Cli::parse(); - install_bootstrap_subscriber(); + init_tracing(&cli.server.log_level, cli.server.log_format)?; let config = cli .into_config() .context("resolve configuration from CLI flags")?; - init_tracing( - &config.observability.log_level, - config.observability.log_format, - )?; // Buffer termination signals before tokenizer loading or discovery can block startup. let (sigterm, sigint) = install_signal_handlers()?; @@ -150,16 +146,6 @@ fn init_tracing(default_level: &str, format: LogFormat) -> Result<()> { Ok(()) } -// Provide startup logging before configuration resolution; later installs are no-ops. -fn install_bootstrap_subscriber() { - let filter = tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); - let _ = tracing_subscriber::fmt() - .with_env_filter(filter) - .with_target(true) - .try_init(); -} - fn install_signal_handlers() -> Result<(Signal, Signal)> { let sigterm = signal(SignalKind::terminate()).context("install SIGTERM handler")?; let sigint = signal(SignalKind::interrupt()).context("install SIGINT handler")?; diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index 008b73522..36b77c16f 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -62,15 +62,16 @@ async fn chat_completions_legacy( .ok_or_else(|| ApiError::BadRequest("missing `model` field".into()))?, ); + let policy = ctx + .policies + .get(&model) + .ok_or_else(|| ApiError::ModelNotFound(model.0.clone()))?; + // Find healthy workers: the prefill pool in PD mode, otherwise the plain pool. let resolver = PdPoolResolver::new(Arc::clone(&ctx.registry)); let candidates = resolver .prefill_candidates(&model) .map_err(|error| pool_error(error, &model))?; - let policy = ctx - .policies - .get(&model) - .ok_or_else(|| ApiError::ModelNotFound(model.0.clone()))?; let request = PreparedChatRequest::prepare(ctx, model, fields, body, policy.needs_request_tokens())?; diff --git a/experimental/sgl-router/src/server/routes/health.rs b/experimental/sgl-router/src/server/routes/health.rs index 2bcd6af8b..ed67f6563 100644 --- a/experimental/sgl-router/src/server/routes/health.rs +++ b/experimental/sgl-router/src/server/routes/health.rs @@ -1,6 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 +use crate::discovery::ModelId; +use crate::policies::registry::{PdPoolResolver, PdPools}; use crate::server::app_context::AppContext; use axum::extract::State; use axum::http::StatusCode; @@ -11,18 +13,18 @@ pub async fn healthz() -> StatusCode { StatusCode::OK } -/// Readiness probe — 200 only when the pod can actually serve traffic. -/// -/// Requires BOTH: -/// 1. `AppContext::mark_ready()` was called by main (process bootstrap -/// finished — config loaded, tokenizers built, server bound), AND -/// 2. At least one worker is registered. Without this second check, -/// `/readyz` flips green before the first `DiscoveryEvent::Added` -/// has been processed — the Service starts sending traffic to a -/// pod whose registry is empty, and every request returns 503 -/// `no_healthy_workers`. +/// Ready after startup only when the configured model has a usable plain or PD pool. pub async fn readyz(State(ctx): State>) -> StatusCode { - if ctx.is_ready() && !ctx.registry.is_empty() { + if !ctx.is_ready() { + return StatusCode::SERVICE_UNAVAILABLE; + } + let resolver = PdPoolResolver::new(Arc::clone(&ctx.registry)); + let ready = match resolver.resolve(&ModelId(ctx.config.model.id.clone())) { + Ok(PdPools::Plain { workers }) => !workers.is_empty(), + Ok(PdPools::Pd { prefill, decode }) => !prefill.is_empty() && !decode.is_empty(), + Err(_) => false, + }; + if ready { StatusCode::OK } else { StatusCode::SERVICE_UNAVAILABLE @@ -104,6 +106,35 @@ mod tests { assert_eq!(res.status(), StatusCode::OK); } + #[tokio::test] + async fn readiness_requires_resolved_model_and_both_pd_roles() { + use crate::discovery::{WorkerId, WorkerMode, WorkerSpec}; + + use WorkerMode::{Decode, Plain, Prefill}; + + for (model, modes, ready) in [ + (None, vec![Plain], false), + (Some("other"), vec![Plain], false), + (Some("stub-model"), vec![Prefill], false), + (Some("stub-model"), vec![Decode], false), + (Some("stub-model"), vec![Prefill, Decode], true), + ] { + let ctx = test_ctx(true, false); + for (i, mode) in modes.into_iter().enumerate() { + ctx.registry + .add(WorkerSpec { + id: WorkerId(i.to_string()), + url: format!("http://worker-{i}:30000"), + mode, + model_ids: model.map(|m| ModelId(m.into())).into_iter().collect(), + bootstrap_port: None, + }) + .unwrap(); + } + assert_eq!(readyz(State(ctx)).await == StatusCode::OK, ready); + } + } + fn test_ctx(ready: bool, with_worker: bool) -> Arc { use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; let ctx = AppContext::stub(); @@ -116,7 +147,7 @@ mod tests { id: WorkerId("test-w".into()), url: "http://test:30000".into(), mode: WorkerMode::Plain, - model_ids: vec![ModelId("test".into())], + model_ids: vec![ModelId(ctx.config.model.id.clone())], bootstrap_port: None, }) .expect("test worker accepted"); diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index d3218c2f8..1f92bb685 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -939,6 +939,21 @@ async fn no_healthy_workers_returns_503() { ); } +#[tokio::test] +async fn unknown_model_without_workers_returns_404() { + let ctx = build_ctx_with_worker("http://127.0.0.1:1"); + ctx.registry.remove(&WorkerId("w1".into())); + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"unknown","messages":[]}"#)) + .unwrap(); + let response = build_router(ctx).oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!(response.headers()["x-router-error-code"], "model_not_found"); +} + /// A worker is registered for a model that is NOT the configured `cfg.model` (so the /// policy registry has no entry for it). The handler returns 404 /// `model_not_found` rather than 500 — clients can recover by sending a