[Rust] Gate health on startup warmup completion (#37994)

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<openai::ChatFormatter>,
     /// 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<Arc<AppState>>,
+    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());
+    }
+}

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<Arc<AppState>> {
     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<Arc<AppState>> {
     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<Arc<AppState>>) -> 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<Arc<AppState>>,
     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 {

rust/sglang-server/src/api_server/openai/test_utils.rs

@@ -104,6 +104,7 @@ pub(super) fn app_state(senders: Senders) -> Arc<super::AppState> {
         server_args: server_args(),
         chat_formatter: None,
         response_activity: Default::default(),
+        startup_readiness: Default::default(),
     })
 }
 

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(),