[Router] Add the shutdown-drain configuration surface (#39015)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
2026-09-16 00:26:50 -07:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 5
parent 5beb2fd552
commit 7835f1de9a
23 changed files with 347 additions and 12 deletions
+94 -8
View File
@@ -11,12 +11,13 @@ use std::num::NonZeroU32;
use crate::config::sampling::{parse_sampling_overrides, ConflictPolicy, SamplingOverrides};
use crate::config::{
default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs,
resolve_mode, ActiveLoadConfig, AffinityConfig, AffinityMode, CacheAwareConfig,
CachePrefixProvider, CircuitBreakerConfig, Config, DecodePolicyKind, DiscoveryBackend,
EligibilityConfig, FilterKind, FusedTerm, K8sDiscoveryConfig, KvIndexerEndpointConfig,
LogFormat, ModelConfig, ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig,
SessionAffinityMode, StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, DEFAULT_FUSE,
default_cb_cool_down, default_host, default_port, default_proxy_request_timeout_secs,
default_shutdown_drain_secs, default_stale_request_timeout_secs, resolve_mode,
ActiveLoadConfig, AffinityConfig, AffinityMode, CacheAwareConfig, CachePrefixProvider,
CircuitBreakerConfig, Config, DecodePolicyKind, DiscoveryBackend, EligibilityConfig,
FilterKind, FusedTerm, K8sDiscoveryConfig, KvIndexerEndpointConfig, LogFormat, ModelConfig,
ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode,
StaticUrlsDiscoveryConfig, StickyConfig, StickyFallbackKind, DEFAULT_FUSE,
};
const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100;
@@ -36,11 +37,30 @@ const DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT: usize = sgl_kv_indexer::DEFAULT_QUE
pub struct Cli {
// ---- server ----
/// Address to bind the HTTP server to.
#[arg(long, default_value = "127.0.0.1")]
#[arg(long, default_value_t = default_host())]
pub host: String,
/// Port to bind the HTTP server to.
#[arg(long, default_value_t = 30000)]
#[arg(long, default_value_t = default_port())]
pub port: u16,
/// Seconds to keep serving after SIGTERM, with `/readyz` returning 503,
/// before the server stops accepting — so the endpoint removal reaches
/// kube-proxy first. Leave room under terminationGracePeriodSeconds for the
/// in-flight drain that follows. If you rely on a readiness probe (rather
/// than pod deletion) to deregister, size this above your
/// failureThreshold * periodSeconds. The default equals the k8s default
/// terminationGracePeriodSeconds, so on a pod that has not raised its grace
/// period startup warns until you do — and declare it with
/// --termination-grace-secs so the check uses the real budget.
/// 0 disables the pause.
#[arg(long, default_value_t = default_shutdown_drain_secs())]
pub shutdown_drain_secs: u64,
/// The pod's terminationGracePeriodSeconds, if you have raised it from the
/// k8s default of 30. Only used to check --shutdown-drain-secs leaves room
/// for the in-flight drain at startup: the router cannot read its own pod
/// spec, so without this it warns against the default and a deliberately
/// long drain has no way to say it is safe.
#[arg(long)]
pub termination_grace_secs: Option<u64>,
// ---- model (exactly one) ----
/// Model id this router serves (the OpenAI `model` field).
@@ -633,6 +653,8 @@ impl Cli {
server: ServerConfig {
host: self.host,
port: self.port,
shutdown_drain_secs: self.shutdown_drain_secs,
termination_grace_secs: self.termination_grace_secs,
},
observability: ObservabilityConfig {
log_level: self.log_level,
@@ -784,6 +806,70 @@ mod tests {
assert_eq!(c.model.id, "qwen3-0.6b");
assert_eq!(c.proxy.request_timeout_secs, 300);
assert_eq!(c.active_load.stale_request_timeout_secs, 600);
assert_eq!(c.server.shutdown_drain_secs, 30);
}
/// Several values, not just the default: a clamp or a rescale in the
/// mapping satisfies any single-value assertion.
#[test]
fn shutdown_drain_secs_maps_into_config() {
for secs in ["0", "17", "1800"] {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://10.0.0.1:30000",
"--shutdown-drain-secs",
secs,
]))
.unwrap();
let expected: u64 = secs.parse().unwrap();
assert_eq!(
c.server.shutdown_drain_secs, expected,
"--shutdown-drain-secs {secs} must map through unchanged",
);
assert_eq!(
c.server.shutdown_drain(),
std::time::Duration::from_secs(expected),
"the Duration accessor must agree with the configured seconds",
);
}
}
/// The ceiling is enforced on the CLI path, not only on a hand-built
/// `Config`: a drain carrying a fat-fingered extra digit must fail at
/// startup rather than turn every later termination into a SIGKILL.
#[test]
fn shutdown_drain_secs_past_the_ceiling_is_rejected() {
let error = into_config_owned(with_model(&[
"--worker-urls",
"http://10.0.0.1:30000",
"--shutdown-drain-secs",
"18000",
]))
.expect_err("a drain past the ceiling must not start")
.to_string();
assert!(
error.contains("shutdown_drain_secs"),
"the error must name the flag to fix: {error}"
);
}
/// `--termination-grace-secs` exists only to feed the startup advisory, so
/// the one thing that matters is that it reaches the config — and that
/// omitting it stays `None` (assume the k8s default) rather than
/// defaulting to a number that would silently become the compared budget.
#[test]
fn termination_grace_secs_maps_into_config_and_defaults_to_none() {
let c = into_config_owned(with_model(&["--worker-urls", "http://10.0.0.1:30000"])).unwrap();
assert_eq!(c.server.termination_grace_secs, None);
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://10.0.0.1:30000",
"--termination-grace-secs",
"120",
]))
.unwrap();
assert_eq!(c.server.termination_grace_secs, Some(120));
}
/// With `--tokenizer-path` omitted, the tokenizer source defaults to the
+174 -4
View File
@@ -7,6 +7,61 @@ pub use types::*;
use anyhow::{anyhow, Result};
/// The k8s default `terminationGracePeriodSeconds`, assumed when the operator
/// has not declared the pod's real one. A `shutdown_drain_secs` at or above the
/// grace period leaves no time for the in-flight drain, so the pod is SIGKILLed
/// before it finishes — the opposite of what the drain is for.
pub const K8S_DEFAULT_GRACE_SECS: u64 = 30;
/// Ceiling on `shutdown_drain_secs`, enforced by [`Config::validate`]. Sized
/// for the workload rather than for the k8s default grace period: a single
/// streaming completion can hold the router for many minutes, so a deployment
/// that does not want terminations cutting one off runs a
/// `terminationGracePeriodSeconds` in the tens of minutes and a drain to match.
/// Deciding whether a particular drain fits a particular grace period is
/// [`shutdown_drain_advisory`]'s job — advice, because the operator can raise
/// the budget. This constant is the separate, harder gate: it rejects a value
/// that is not a drain at all, an extra digit or seconds confused with
/// milliseconds, which no grace period could ever service.
pub const MAX_SHUTDOWN_DRAIN_SECS: u64 = 1800;
/// A `shutdown_drain_secs` that leaves no room under the grace period for the
/// in-flight drain that follows the pause. Carries the compared values as
/// fields so the caller logs a static message with structured data rather than
/// interpolating the numbers into the message text.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ShutdownDrainAdvisory {
pub shutdown_drain_secs: u64,
/// The budget the drain was compared against.
pub termination_grace_secs: u64,
/// Whether that budget came from the operator or from
/// [`K8S_DEFAULT_GRACE_SECS`]. An assumed budget makes the advisory a
/// guess; a declared one makes it a fact.
pub grace_declared: bool,
}
/// Advisory (not a hard error: the drain may well be right and the grace period
/// raised to match) for a drain that leaves no room for the in-flight drain.
/// The bound is `>=`, not `>`: a drain of exactly the grace period already
/// consumes all of it.
///
/// `termination_grace_secs` is the pod's real `terminationGracePeriodSeconds`
/// when the operator declared it. The router cannot read its own pod spec, so
/// `None` falls back to the k8s default — which is why declaring the real value
/// is the way to silence this on a deployment that raised the grace period
/// deliberately, rather than lowering a drain that was correct.
pub fn shutdown_drain_advisory(
shutdown_drain_secs: u64,
termination_grace_secs: Option<u64>,
) -> Option<ShutdownDrainAdvisory> {
let grace = termination_grace_secs.unwrap_or(K8S_DEFAULT_GRACE_SECS);
(shutdown_drain_secs >= grace).then_some(ShutdownDrainAdvisory {
shutdown_drain_secs,
termination_grace_secs: grace,
grace_declared: termination_grace_secs.is_some(),
})
}
impl Config {
/// Check invariants the type system and `clap` don't already enforce.
/// Called by [`cli::Cli::into_config`] after assembling the `Config`
@@ -21,6 +76,16 @@ impl Config {
validate_bucket_config(bucket_config)?;
}
self.model.sampling_overrides.validate()?;
if self.server.shutdown_drain_secs > MAX_SHUTDOWN_DRAIN_SECS {
return Err(anyhow!(
"shutdown_drain_secs must be at most {MAX_SHUTDOWN_DRAIN_SECS} (got {}); \
past the ceiling a value is a typo rather than a drain, and the pod \
would be SIGKILLed long before the pause elapsed. A long but deliberate \
drain is fine — declare --termination-grace-secs so startup can check \
it against the pod's real budget",
self.server.shutdown_drain_secs,
));
}
match &self.discovery {
DiscoveryBackend::StaticUrls(s) => {
if s.urls.is_empty() {
@@ -209,10 +274,7 @@ mod tests {
/// the `cli` module tests; the k8s selector grammar in `types`.
fn cfg(model_id: &str, urls: &[&str]) -> Config {
Config {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 30000,
},
server: ServerConfig::default(),
observability: ObservabilityConfig::default(),
model: ModelConfig {
id: model_id.into(),
@@ -483,4 +545,112 @@ mod tests {
.to_string();
assert!(error.contains("ttft_p95_at_capcity_ms"), "got: {error}");
}
#[test]
fn shutdown_drain_advisory_is_silent_below_the_k8s_default_grace() {
// Anything strictly under the k8s default terminationGracePeriodSeconds
// (30 s) still leaves room for the in-flight drain, so it is safe
// without operator action.
assert!(shutdown_drain_advisory(29, None).is_none());
assert!(shutdown_drain_advisory(0, None).is_none());
}
/// The default drain is exactly the assumed grace period, so out of the box
/// it leaves nothing for the in-flight drain and says so on every startup.
/// Asserted rather than left implicit because it is the one case an
/// operator meets without choosing it: a later edit to either constant that
/// silenced the warning would be changing the default deployment's
/// behaviour, and should have to say so here.
#[test]
fn the_default_drain_warns_until_the_grace_period_is_raised() {
let advisory = shutdown_drain_advisory(default_shutdown_drain_secs(), None)
.expect("the default drain must warn against the assumed k8s grace period");
assert_eq!(advisory.termination_grace_secs, K8S_DEFAULT_GRACE_SECS);
assert!(
!advisory.grace_declared,
"an assumed budget must not be reported as declared",
);
// Raising the pod's grace period past the drain is what silences it —
// the action the warning asks for has to actually work.
assert!(
shutdown_drain_advisory(
default_shutdown_drain_secs(),
Some(K8S_DEFAULT_GRACE_SECS * 2),
)
.is_none(),
"a grace period declared with room for the in-flight drain must silence it",
);
}
#[test]
fn shutdown_drain_advisory_warns_once_the_drain_consumes_the_whole_grace() {
// A drain of exactly the 30 s k8s default leaves zero seconds for the
// in-flight drain, so the pod is SIGKILLed mid-drain — the boundary
// itself must warn, not just values past it. The ceiling is in the list
// because it is startable: `validate` accepts it, so the advisory is
// the only thing left to say it does not fit the default grace period.
for drain in [K8S_DEFAULT_GRACE_SECS, 120, MAX_SHUTDOWN_DRAIN_SECS] {
let advisory = shutdown_drain_advisory(drain, None)
.unwrap_or_else(|| panic!("{drain}s must warn"));
assert_eq!(advisory.shutdown_drain_secs, drain);
assert_eq!(advisory.termination_grace_secs, K8S_DEFAULT_GRACE_SECS);
assert!(
!advisory.grace_declared,
"an undeclared grace period must be reported as assumed, not as fact",
);
}
}
/// The advisory's whole purpose is to be silenceable by declaring the real
/// budget: a 60 s drain under a 120 s grace period is a correct
/// configuration, and warning about it trains operators to ignore the line.
#[test]
fn shutdown_drain_advisory_respects_a_declared_grace_period() {
assert!(
shutdown_drain_advisory(60, Some(120)).is_none(),
"a drain with room under the declared grace period must not warn",
);
let advisory = shutdown_drain_advisory(60, Some(60))
.expect("a drain consuming the whole declared grace period must warn");
assert_eq!(advisory.termination_grace_secs, 60);
assert!(
advisory.grace_declared,
"a declared grace period must be reported as declared",
);
// ...and declaring a *shorter* budget than the k8s default must be able
// to warn about a drain the default would have waved through.
assert!(
shutdown_drain_advisory(10, Some(10)).is_some(),
"a short declared grace period must still be compared against",
);
// The configuration the ceiling was raised for: a completion streaming
// for minutes wants a drain of minutes, under a grace period declared
// to match. That is correct, not merely tolerated, so it must be silent.
assert!(
shutdown_drain_advisory(MAX_SHUTDOWN_DRAIN_SECS, Some(3600)).is_none(),
"a long drain under a grace period declared to cover it must not warn",
);
}
/// `validate` is the hard gate the advisory deliberately is not: past the
/// ceiling the value can only be a typo, and starting on it would make
/// every later termination a SIGKILL.
#[test]
fn validate_rejects_a_shutdown_drain_past_the_ceiling() {
let mut config = cfg("qwen3-0.6b", &["http://10.0.0.1:30000"]);
config.server.shutdown_drain_secs = MAX_SHUTDOWN_DRAIN_SECS;
config
.validate()
.expect("the ceiling itself must remain startable");
config.server.shutdown_drain_secs = MAX_SHUTDOWN_DRAIN_SECS + 1;
let error = config
.validate()
.expect_err("a drain past the ceiling must fail startup")
.to_string();
assert!(
error.contains("shutdown_drain_secs"),
"the error must name the flag to fix: {error}"
);
}
}
@@ -248,6 +248,65 @@ impl std::fmt::Display for StickyFallbackKind {
pub struct ServerConfig {
pub host: String,
pub port: u16,
/// Seconds to keep serving after SIGTERM — with `/readyz` flipped to 503 —
/// before the HTTP server stops accepting. The default covers both
/// deregistration paths: endpoint removal reaching kube-proxy after the
/// pod's `deletionTimestamp` is stamped, and a probe-driven load balancer,
/// which cannot act until `failureThreshold * periodSeconds` of `/readyz`
/// failures have accumulated.
///
/// Note that it equals the k8s default `terminationGracePeriodSeconds`, so
/// a pod that has not raised its grace period is left with nothing for the
/// in-flight drain that follows the pause, and
/// [`shutdown_drain_advisory`](crate::config::shutdown_drain_advisory)
/// warns at every startup. That is the intended reading rather than a
/// misconfigured default: a router whose completions stream for minutes
/// cannot terminate cleanly inside 30 s at all, and the grace period is the
/// thing to raise. 0 disables the pause.
pub shutdown_drain_secs: u64,
/// The pod's actual `terminationGracePeriodSeconds`, when the operator
/// declares it. The router cannot read its own pod spec, so without this
/// the startup advisory can only compare the drain against the k8s
/// default — and warns, wrongly, about a deployment that raised the grace
/// period on purpose. `None` means "assume the default".
pub termination_grace_secs: Option<u64>,
}
impl ServerConfig {
/// [`Self::shutdown_drain_secs`] as a `Duration`. Keeps the seconds-to-
/// `Duration` conversion in the library, where a test can pin it, rather
/// than in `main.rs` where a `from_secs`/`from_millis` slip would silently
/// shorten every drain by a factor of 1000.
pub fn shutdown_drain(&self) -> std::time::Duration {
std::time::Duration::from_secs(self.shutdown_drain_secs)
}
}
pub fn default_host() -> String {
"127.0.0.1".into()
}
pub fn default_port() -> u16 {
30000
}
pub fn default_shutdown_drain_secs() -> u64 {
30
}
/// Exists so test fixtures can spell out only the fields they care about
/// (`tests/` is a separate crate, so a `#[cfg(test)]` constructor cannot reach
/// the integration fixtures). Keep `Cli::into_config` exhaustive so adding a
/// field still forces a decision on the production path.
impl Default for ServerConfig {
fn default() -> Self {
Self {
host: default_host(),
port: default_port(),
shutdown_drain_secs: default_shutdown_drain_secs(),
termination_grace_secs: None,
}
}
}
#[derive(Debug, Clone)]
@@ -331,6 +331,7 @@ mod tests {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: Default::default(),
model: ModelConfig {
@@ -124,6 +124,7 @@ impl AppContext {
server: crate::config::ServerConfig {
host: "x".into(),
port: 0,
..Default::default()
},
observability: Default::default(),
model: crate::config::ModelConfig {
@@ -112,6 +112,7 @@ mod tests {
server: crate::config::ServerConfig {
host: "x".into(),
port: 0,
..Default::default()
},
observability: Default::default(),
model: crate::config::ModelConfig {
@@ -226,6 +226,7 @@ mod tests {
server: crate::config::ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: Default::default(),
model: crate::config::ModelConfig {
@@ -585,6 +585,7 @@ mod tests {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: Default::default(),
model: ModelConfig {
@@ -124,6 +124,7 @@ async fn static_urls_pd_role_resolved_end_to_end() {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: sgl_router::config::ModelConfig {
@@ -55,6 +55,7 @@ fn build_app_context(
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -28,6 +28,7 @@ fn config_for(_worker_url: &str) -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -20,6 +20,7 @@ pub fn config() -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -29,6 +29,7 @@ async fn failover_when_one_worker_dies() {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: Default::default(),
model: ModelConfig {
@@ -38,6 +38,7 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -25,6 +25,7 @@ async fn forwards_whitelisted_headers_strips_others() {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -40,6 +40,7 @@ fn config() -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -39,6 +39,7 @@ fn config() -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -127,6 +127,7 @@ fn config() -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -36,6 +36,7 @@ fn config() -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -143,6 +143,7 @@ fn config(policy: PolicyKind) -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -48,6 +48,7 @@ fn config() -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -36,6 +36,7 @@ fn build_sticky_ctx(header_name: &str, worker_urls: &[String]) -> Arc<AppContext
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
@@ -32,6 +32,7 @@ fn config(_worker_url: &str) -> Config {
server: ServerConfig {
host: "0".into(),
port: 0,
..Default::default()
},
observability: ObservabilityConfig::default(),
model: ModelConfig {