diff --git a/sgl-model-gateway/bindings/python/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/sglang_router/router_args.py index 8ac4b5d04..9fa3125b4 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router_args.py +++ b/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", diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index 3e7a25b10..79337d466 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -179,6 +179,7 @@ struct Router { prometheus_host: Option, prometheus_duration_buckets: Option>, request_timeout_secs: u64, + shutdown_grace_period_secs: u64, request_id_headers: Option>, 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, prometheus_duration_buckets: Option>, request_timeout_secs: u64, + shutdown_grace_period_secs: u64, request_id_headers: Option>, 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())) diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index d497d3712..9ad931f9b 100644 --- a/sgl-model-gateway/src/main.rs +++ b/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, } } } diff --git a/sgl-model-gateway/src/middleware.rs b/sgl-model-gateway/src/middleware.rs index eac6c4a04..bfcb981e5 100644 --- a/sgl-model-gateway/src/middleware.rs +++ b/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) diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 65e816583..6876012e7 100644 --- a/sgl-model-gateway/src/server.rs +++ b/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, pub request_timeout_secs: u64, pub request_id_headers: Option>, + /// 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 Result<(), Box)?; } else { - let listener = TcpListener::bind(&bind_addr) - .await - .map_err(|e| format!("Failed to bind to {}: {}", bind_addr, e))?; - serve(listener, app) - .with_graceful_shutdown(shutdown_signal()) + let addr: std::net::SocketAddr = bind_addr + .parse() + .map_err(|e| format!("Invalid address: {}", e))?; + + 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)?; }