diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 4d9f7ec08..076aa7c2d 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2205,6 +2205,10 @@ def _execute_server_warmup(server_args: ServerArgs): url = server_args.url() if get_serving().api_key: headers["Authorization"] = f"Bearer {get_serving().api_key}" + if envs.SGLANG_RUST_SERVER.get(): + # The Rust listener binds before this request so /model_info is + # available, but health stays 503 until this marked request succeeds. + headers["x-sglang-startup-warmup"] = "1" ssl_verify = ssl_verify_of(server_args) diff --git a/python/sglang/srt/rust_server/config.py b/python/sglang/srt/rust_server/config.py index 06f622d7e..7eccfd8a2 100644 --- a/python/sglang/srt/rust_server/config.py +++ b/python/sglang/srt/rust_server/config.py @@ -62,6 +62,7 @@ def _build_server_args(scheduler: Scheduler) -> ServerArgs: tokenizer_worker_num=get_serving().tokenizer_worker_num, detokenizer_worker_num=get_serving().detokenizer_worker_num, skip_tokenizer_init=get_serving().skip_tokenizer_init, + skip_server_warmup=get_serving().skip_server_warmup, incremental_streaming_output=get_serving().incremental_streaming_output, disaggregation_mode=disaggregation_mode, model_config=ext.ModelConfig( diff --git a/rust/sglang-server/src/api_server/app.rs b/rust/sglang-server/src/api_server/app.rs index 13ec61143..77190940e 100644 --- a/rust/sglang-server/src/api_server/app.rs +++ b/rust/sglang-server/src/api_server/app.rs @@ -2,9 +2,17 @@ //! registers its routes here, and [`serve`] runs the assembled app on the //! pre-bound listener until shutdown. -use std::sync::Arc; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; -use axum::Router; +use axum::{ + Router, + extract::{Request, State}, + middleware::Next, + response::Response, +}; use super::disaggregation::bootstrap as pd_bootstrap; use super::{common, log, native_api, openai}; @@ -26,6 +34,57 @@ pub(super) struct AppState { pub(super) chat_formatter: Option, /// Response heartbeat (bumped per drained ring frame). pub(super) response_activity: ActivityCounter, + /// Whether the main process's startup warmup has completed. The listener + /// binds before warmup so `/model_info` is available to construct that + /// request, but health endpoints must not advertise readiness yet. + pub(super) startup_readiness: StartupReadiness, +} + +pub(super) struct StartupReadiness(AtomicBool); + +impl StartupReadiness { + fn new(skip_server_warmup: bool) -> Self { + Self(AtomicBool::new(skip_server_warmup)) + } + + pub(super) fn is_ready(&self) -> bool { + self.0.load(Ordering::Acquire) + } + + fn record_warmup_status(&self, status: axum::http::StatusCode) { + if status.is_success() { + self.0.store(true, Ordering::Release); + } + } +} + +impl Default for StartupReadiness { + fn default() -> Self { + Self::new(false) + } +} + +/// Private marker attached by the main process to its startup warmup request. +/// The middleware flips readiness only after that request returns successfully. +const STARTUP_WARMUP_HEADER: &str = "x-sglang-startup-warmup"; + +async fn mark_startup_ready( + State(state): State>, + req: Request, + next: Next, +) -> Response { + let is_startup_warmup = req.headers().contains_key(STARTUP_WARMUP_HEADER) + && matches!( + req.uri().path(), + "/generate" | "/encode" | "/v1/chat/completions" + ); + let response = next.run(req).await; + if is_startup_warmup && response.status().is_success() { + state + .startup_readiness + .record_warmup_status(response.status()); + } + response } pub async fn serve( @@ -47,6 +106,7 @@ pub async fn serve( server_args: server_args.clone(), chat_formatter, response_activity, + startup_readiness: StartupReadiness::new(server_args.skip_server_warmup), }); // Each endpoint module registers its own routes and merges here. let router = Router::new() @@ -61,6 +121,10 @@ pub async fn serve( // No body limit, matching the Python server. let mut app = router .layer(axum::extract::DefaultBodyLimit::disable()) + .layer(axum::middleware::from_fn_with_state( + state.clone(), + mark_startup_ready, + )) .with_state(state); // Prefill-only KV bootstrap registry. Merged AFTER `with_state` — its @@ -102,3 +166,23 @@ pub async fn serve( } } } + +#[cfg(test)] +mod tests { + use super::*; + use axum::http::StatusCode; + + #[test] + fn startup_readiness_requires_successful_warmup_unless_skipped() { + let readiness = StartupReadiness::new(false); + assert!(!readiness.is_ready()); + + readiness.record_warmup_status(StatusCode::INTERNAL_SERVER_ERROR); + assert!(!readiness.is_ready()); + + readiness.record_warmup_status(StatusCode::OK); + assert!(readiness.is_ready()); + + assert!(StartupReadiness::new(true).is_ready()); + } +} diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs index 71b791f55..ff7abf123 100644 --- a/rust/sglang-server/src/api_server/native_api.rs +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -96,8 +96,8 @@ pub(super) fn native_error(code: StatusCode, message: &str, stream: bool) -> Res /// restart. The deep-probe handler is built once with /// `SGLANG_HEALTH_CHECK_TIMEOUT` frozen in and serves `/health_generate` /// always; `SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION` (default true, mirroring -/// Python) decides whether `/health` shares it or is a plain 200 (routing the -/// request already proves the frontend is up). +/// Python) decides whether `/health` shares it or, after startup warmup, is a +/// plain 200 (routing the request proves the frontend is up). fn health_routes() -> Router> { let timeout = std::time::Duration::from_secs( environ::env_i64("SGLANG_HEALTH_CHECK_TIMEOUT", 20).max(0) as u64, @@ -106,13 +106,21 @@ fn health_routes() -> Router> { let health = if environ::env_bool("SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION", true) { probe.clone() } else { - get(|| async { StatusCode::OK.into_response() }) + get(health_without_generation) }; Router::new() .route("/health", health) .route("/health_generate", probe) } +async fn health_without_generation(State(state): State>) -> Response { + if state.startup_readiness.is_ready() { + StatusCode::OK.into_response() + } else { + StatusCode::SERVICE_UNAVAILABLE.into_response() + } +} + /// Sentinel host that makes the KV connector no-op. Parity with /// `sglang.srt.disaggregation.utils.FAKE_BOOTSTRAP_HOST`. const FAKE_BOOTSTRAP_HOST: &str = "2.2.2.2"; @@ -120,7 +128,8 @@ const FAKE_BOOTSTRAP_HOST: &str = "2.2.2.2"; /// `GET /health_generate` — deep health: confirm the scheduler → detok path is /// producing output. 200 if the response heartbeat advances within `timeout` /// (from `SGLANG_HEALTH_CHECK_TIMEOUT`, frozen at router build), else 503. -/// (`/health` uses the same handler when its env gate is on.) +/// It also returns 503 until startup warmup completes. (`/health` uses the same +/// handler when its env gate is on.) /// /// Fires a pre-tokenized 1-token probe (`input_ids = [0]`, skips the tokenizer) so /// an idle pipeline produces a frame, then watches the *global* @@ -131,6 +140,10 @@ async fn health_generate( State(state): State>, timeout: std::time::Duration, ) -> Response { + if !state.startup_readiness.is_ready() { + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + let baseline = state .response_activity .load(std::sync::atomic::Ordering::Relaxed); @@ -609,6 +622,21 @@ mod tests { ) } + #[tokio::test] + async fn health_is_unavailable_before_startup_warmup_finishes() { + let state = Arc::new(AppState { + senders: senders(), + response_buf: 8, + server_args: Arc::new(crate::message::config::ServerArgs::default()), + chat_formatter: None, + response_activity: Default::default(), + startup_readiness: Default::default(), + }); + + let response = health_generate(State(state), Duration::ZERO).await; + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + } + #[test] fn request_timing_records_ttft_once_and_e2e_on_finish() { let mut timing = RequestTiming { diff --git a/rust/sglang-server/src/api_server/openai/test_utils.rs b/rust/sglang-server/src/api_server/openai/test_utils.rs index d93303e04..aaf3d00df 100644 --- a/rust/sglang-server/src/api_server/openai/test_utils.rs +++ b/rust/sglang-server/src/api_server/openai/test_utils.rs @@ -104,6 +104,7 @@ pub(super) fn app_state(senders: Senders) -> Arc { server_args: server_args(), chat_formatter: None, response_activity: Default::default(), + startup_readiness: Default::default(), }) } diff --git a/rust/sglang-server/src/message/config.rs b/rust/sglang-server/src/message/config.rs index b22ef9218..8ecdf4814 100644 --- a/rust/sglang-server/src/message/config.rs +++ b/rust/sglang-server/src/message/config.rs @@ -119,6 +119,9 @@ pub struct ServerArgs { /// Token-ids-in / token-ids-out mode: no tokenizer load, raw `output_ids` /// frames. pub skip_tokenizer_init: bool, + /// Start accepting health checks immediately instead of waiting for the + /// main process's startup warmup request to finish. + pub skip_server_warmup: bool, /// Streamed `/generate` frames carry per-step deltas instead of cumulative /// text. Matches the Python `TokenizerManager`. pub incremental_streaming_output: bool, @@ -171,6 +174,7 @@ impl ServerArgs { tokenizer_worker_num, detokenizer_worker_num, skip_tokenizer_init, + skip_server_warmup, incremental_streaming_output, disaggregation_mode, model_config, @@ -202,6 +206,7 @@ impl ServerArgs { tokenizer_worker_num: usize, detokenizer_worker_num: usize, skip_tokenizer_init: bool, + skip_server_warmup: bool, incremental_streaming_output: bool, disaggregation_mode: DisaggregationMode, model_config: ModelConfig, @@ -231,6 +236,7 @@ impl ServerArgs { tokenizer_worker_num, detokenizer_worker_num, skip_tokenizer_init, + skip_server_warmup, incremental_streaming_output, disaggregation_mode, model_config, @@ -268,6 +274,7 @@ impl Default for ServerArgs { tokenizer_worker_num: 1, detokenizer_worker_num: 1, skip_tokenizer_init: false, + skip_server_warmup: false, incremental_streaming_output: false, disaggregation_mode: DisaggregationMode::Null, model_config: ModelConfig::default(),