Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion experimental/sgl-router/src/discovery/k8s.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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}")),
Expand Down Expand Up @@ -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);
Expand Down
16 changes: 1 addition & 15 deletions experimental/sgl-router/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;
Expand Down Expand Up @@ -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")?;
Expand Down
9 changes: 5 additions & 4 deletions experimental/sgl-router/src/server/routes/chat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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())?;
Expand Down
55 changes: 43 additions & 12 deletions experimental/sgl-router/src/server/routes/health.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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();
Expand All @@ -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");
Expand Down
15 changes: 15 additions & 0 deletions experimental/sgl-router/tests/proxy/chat_routing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading