[sgl-router] Fix readiness, IPv6 discovery, logging, and model validation (#40604)
This commit is contained in:
@@ -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);
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user