[sgl-router] Fix readiness, IPv6 discovery, logging, and model validation (#40604)

This commit is contained in:
Kan Wu
2026-09-22 11:13:06 +08:00
committed by GitHub
parent 1d025491f3
commit 27f796ca6c
5 changed files with 79 additions and 32 deletions
+15 -1
View File
@@ -118,7 +118,12 @@ fn extract_workers(es: &EndpointSlice, mode: WorkerMode) -> Vec<WorkerSpec> {
}
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);
+1 -15
View File
@@ -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")?;
@@ -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())?;
@@ -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<Arc<AppContext>>) -> 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<AppContext> {
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");
@@ -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