[model-gateway] fix graceful shutdown for TLS/Non-TLS server (#15491)

sgl-model-gateway/bindings/python/sglang_router/router_args.py

@@ -59,6 +59,8 @@ class RouterArgs:
     request_id_headers: Optional[List[str]] = None
     # Request timeout in seconds
     request_timeout_secs: int = 1800
+    # Grace period in seconds to wait for in-flight requests during shutdown
+    shutdown_grace_period_secs: int = 180
     # Max concurrent requests for rate limiting (-1 to disable)
     max_concurrent_requests: int = -1
     # Queue size for pending requests when max concurrent limit reached
@@ -364,6 +366,12 @@ class RouterArgs:
             default=RouterArgs.request_timeout_secs,
             help="Request timeout in seconds",
         )
+        parser.add_argument(
+            f"--{prefix}shutdown-grace-period-secs",
+            type=int,
+            default=RouterArgs.shutdown_grace_period_secs,
+            help="Grace period in seconds to wait for in-flight requests during shutdown",
+        )
         # Retry configuration
         parser.add_argument(
             f"--{prefix}retry-max-retries",

sgl-model-gateway/bindings/python/src/lib.rs

@@ -179,6 +179,7 @@ struct Router {
     prometheus_host: Option<String>,
     prometheus_duration_buckets: Option<Vec<f64>>,
     request_timeout_secs: u64,
+    shutdown_grace_period_secs: u64,
     request_id_headers: Option<Vec<String>>,
     pd_disaggregation: bool,
     bucket_adjust_interval_secs: usize,
@@ -448,6 +449,7 @@ impl Router {
         prometheus_host = None,
         prometheus_duration_buckets = None,
         request_timeout_secs = 1800,
+        shutdown_grace_period_secs = 180,
         request_id_headers = None,
         pd_disaggregation = false,
         bucket_adjust_interval_secs = 5,
@@ -528,6 +530,7 @@ impl Router {
         prometheus_host: Option<String>,
         prometheus_duration_buckets: Option<Vec<f64>>,
         request_timeout_secs: u64,
+        shutdown_grace_period_secs: u64,
         request_id_headers: Option<Vec<String>>,
         pd_disaggregation: bool,
         bucket_adjust_interval_secs: usize,
@@ -621,6 +624,7 @@ impl Router {
             prometheus_host,
             prometheus_duration_buckets,
             request_timeout_secs,
+            shutdown_grace_period_secs,
             request_id_headers,
             pd_disaggregation,
             bucket_adjust_interval_secs,
@@ -729,6 +733,7 @@ impl Router {
                 prometheus_config,
                 request_timeout_secs: self.request_timeout_secs,
                 request_id_headers: self.request_id_headers.clone(),
+                shutdown_grace_period_secs: self.shutdown_grace_period_secs,
             })
             .await
             .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))

sgl-model-gateway/src/main.rs

@@ -223,6 +223,12 @@ struct CliArgs {
     #[arg(long, default_value_t = 1800)]
     request_timeout_secs: u64,
 
+    /// Grace period in seconds to wait for in-flight requests during shutdown.
+    /// When the server receives SIGTERM/SIGINT, it will stop accepting new connections
+    /// and wait up to this duration for existing streaming requests to complete.
+    #[arg(long, default_value_t = 180)]
+    shutdown_grace_period_secs: u64,
+
     #[arg(long, default_value_t = -1)]
     max_concurrent_requests: i32,
 
@@ -713,6 +719,7 @@ impl CliArgs {
             } else {
                 Some(self.request_id_headers.clone())
             },
+            shutdown_grace_period_secs: self.shutdown_grace_period_secs,
         }
     }
 }

sgl-model-gateway/src/middleware.rs

@@ -599,7 +599,7 @@ pub async fn concurrency_limit_middleware(
 // HTTP Metrics Layer (Layer 1: SMG metrics)
 // ============================================================================
 
-/// Global counter for active HTTP connections
+/// Global counter for active HTTP connections (handlers currently executing)
 static ACTIVE_HTTP_CONNECTIONS: AtomicU64 = AtomicU64::new(0);
 
 /// Tower Layer for HTTP metrics collection (SMG Layer 1 metrics)

sgl-model-gateway/src/server.rs

@@ -12,12 +12,12 @@ use axum::{
     http::StatusCode,
     response::{IntoResponse, Response},
     routing::{delete, get, post},
-    serve, Json, Router,
+    Json, Router,
 };
 use rustls::crypto::ring;
 use serde::Deserialize;
 use serde_json::{json, Value};
-use tokio::{net::TcpListener, signal, spawn};
+use tokio::{signal, spawn};
 use tracing::{error, info, warn, Level};
 
 use crate::{
@@ -636,6 +636,9 @@ pub struct ServerConfig {
     pub prometheus_config: Option<PrometheusConfig>,
     pub request_timeout_secs: u64,
     pub request_id_headers: Option<Vec<String>>,
+    /// Grace period in seconds to wait for in-flight requests during shutdown.
+    /// Default is 30 seconds.
+    pub shutdown_grace_period_secs: u64,
 }
 
 pub fn build_app(
@@ -1004,9 +1007,10 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
 
         let handle = axum_server::Handle::new();
         let handle_clone = handle.clone();
+        let grace_period = Duration::from_secs(config.shutdown_grace_period_secs);
         spawn(async move {
             shutdown_signal().await;
-            handle_clone.graceful_shutdown(None);
+            handle_clone.graceful_shutdown(Some(grace_period));
         });
 
         axum_server::bind_rustls(addr, tls_config)
@@ -1015,11 +1019,21 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
             .await
             .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
     } else {
-        let listener = TcpListener::bind(&bind_addr)
+        let addr: std::net::SocketAddr = bind_addr
-            .await
+            .parse()
-            .map_err(|e| format!("Failed to bind to {}: {}", bind_addr, e))?;
+            .map_err(|e| format!("Invalid address: {}", e))?;
-        serve(listener, app)
+
-            .with_graceful_shutdown(shutdown_signal())
+        let handle = axum_server::Handle::new();
+        let handle_clone = handle.clone();
+        let grace_period = Duration::from_secs(config.shutdown_grace_period_secs);
+        spawn(async move {
+            shutdown_signal().await;
+            handle_clone.graceful_shutdown(Some(grace_period));
+        });
+
+        axum_server::bind(addr)
+            .handle(handle)
+            .serve(app.into_make_service())
             .await
             .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
     }