diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index 759b2e14f..91ce23a73 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -62,6 +62,19 @@ fn install_bootstrap_subscriber() { .try_init(); } +/// How often to report progress while axum drains in-flight requests. That +/// phase is unbounded, so without a heartbeat a pod SIGKILLed at +/// `terminationGracePeriodSeconds` leaves no evidence of what it was waiting on. +const DRAIN_HEARTBEAT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5); + +/// How long the in-flight drain may run before the heartbeat escalates from +/// INFO to WARN. Under this, a pod finishing a long streaming completion is +/// doing exactly what the drain is for, and logging it at WARN would fire on +/// every routine rollout — training operators to filter router WARNs, which +/// are also where `further termination signal ignored` and the drain advisory +/// land. Past it the pod is at real risk of being SIGKILLed with work open. +const DRAIN_WARN_AFTER: std::time::Duration = std::time::Duration::from_secs(30); + /// Install SIGTERM and SIGINT handlers up front so a failure here surfaces /// before `axum::serve` starts. If installation fails (rare: container /// without signal capability, seccomp policy), we return an error and the @@ -85,6 +98,35 @@ async fn main() -> Result<()> { init_tracing(&cfg.observability.log_level, cfg.observability.log_format)?; + // Before anything slow — tokenizer download, discovery, bind. kubelet can + // SIGTERM a pod mid-rollout while it is still starting, and until the + // handlers exist that signal takes the default disposition: instant death, + // no readiness flip, no drain, no log. tokio's `Signal` buffers a + // notification until first polled, so installing here loses nothing and + // costs one deferred shutdown instead of a silent kill. + let (sigterm, sigint) = install_signal_handlers()?; + + // Emitted here rather than from `Config::validate`: this is startup advice + // about the deployment, not a validation failure, and keeping it out of + // `validate` leaves that function free of side effects. It also runs after + // the configured subscriber is installed. Static message, values as + // structured fields: a message that varies with the configured seconds + // cannot be grouped or deduped by a log aggregator. + if let Some(advisory) = sgl_router::config::shutdown_drain_advisory( + cfg.server.shutdown_drain_secs, + cfg.server.termination_grace_secs, + ) { + tracing::warn!( + shutdown_drain_secs = advisory.shutdown_drain_secs, + termination_grace_secs = advisory.termination_grace_secs, + grace_declared = advisory.grace_declared, + "shutdown drain leaves no room under terminationGracePeriodSeconds for the \ + in-flight drain that follows it; raise the grace period to at least the drain \ + plus in-flight request time, or lower the drain. If the grace period is already \ + higher, declare it with --termination-grace-secs", + ); + } + tracing::info!( configured_decode_policy = ?cfg.model.decode_policy, "sgl-router {} starting on {}:{}", @@ -226,18 +268,100 @@ async fn main() -> Result<()> { .with_context(|| format!("bind {bind}"))?; tracing::info!("listening on {bind}"); - let (sigterm, sigint) = install_signal_handlers()?; + // Published the moment the readiness drain finishes, i.e. when axum starts + // its in-flight drain. That phase — not the pause, and not the uptime + // before it — is what the heartbeat below reports on. + let (inflight_drain_tx, inflight_drain_rx) = + tokio::sync::watch::channel(None::); + let shutdown_ctx = ctx.clone(); + let drain = cfg.server.shutdown_drain(); + let serve = axum::serve(listener, app).with_graceful_shutdown(async move { + shutdown_signal(sigterm, sigint, shutdown_ctx, drain).await; + let _ = inflight_drain_tx.send(Some(std::time::Instant::now())); + }); - let serve = axum::serve(listener, app).with_graceful_shutdown(shutdown_signal(sigterm, sigint)); + let heartbeat_ctx = ctx.clone(); + let mut heartbeat_rx = inflight_drain_rx.clone(); + let heartbeat = tokio::spawn(async move { + // Stay silent until the in-flight drain actually begins; an `Err` here + // means the sender went away without one, so there is nothing to report. + let Ok(started) = heartbeat_rx + .wait_for(Option::is_some) + .await + .map(|at| at.expect("wait_for only resolves once the instant is published")) + else { + return; + }; + let mut ticker = tokio::time::interval(DRAIN_HEARTBEAT_INTERVAL); + // Delay, not the default Burst: the runtime stalling is exactly the + // condition this heartbeat exists to report, and Burst would answer it + // with a clump of back-dated ticks instead of one line per interval. + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + ticker.tick().await; // the first tick completes immediately + + // One message, two severities: duplicating the text across an if/else + // is how the two arms drift apart. + macro_rules! heartbeat { + ($level:ident, $elapsed:expr) => { + tracing::$level!( + elapsed_secs = $elapsed, + // What axum is actually waiting on: every open exchange, + // on every route, until its response body finishes. + inflight_http = heartbeat_ctx.inflight_http.count(), + // The proxied subset, to separate "waiting on a worker" + // from "waiting on a client that stopped reading". + inflight_proxied = heartbeat_ctx.active_load.inflight_count(), + "still draining in-flight requests; this phase is unbounded and ends at \ + SIGKILL when terminationGracePeriodSeconds expires", + ) + }; + } + loop { + ticker.tick().await; + let elapsed = started.elapsed(); + if elapsed < DRAIN_WARN_AFTER { + heartbeat!(info, elapsed.as_secs()); + } else { + heartbeat!(warn, elapsed.as_secs()); + } + } + }); let server_result = serve.await.context("axum serve"); + heartbeat.abort(); + let inflight_drain_secs = inflight_drain_rx.borrow().map(|at| at.elapsed().as_secs()); // Best-effort: cancel discovery + manager + janitor on shutdown. // The janitor handle's drop signals cancellation; we additionally // await `shutdown` so the task joins cleanly before the process - // exits — useful for tracing tail logs. + // exits — useful for tracing tail logs. `JanitorHandle::shutdown` caps its + // own join at 2 s, so it cannot hang the exit — though those 2 s are still + // charged to terminationGracePeriodSeconds. discovery_handle.abort(); manager_handle.abort(); janitor_handle.shutdown().await; + // The ERROR arms exist because otherwise the log says "shutdown complete" + // at INFO and the error leaves the process through `Termination`, never + // through `tracing` — so a severity-based alert sees nothing wrong with a + // crashed router. `None` means the server stopped without ever reaching the + // drain, which is not the same as draining instantly, so it gets its own + // message rather than `inflight_drain_secs = 0`. + match (&server_result, inflight_drain_secs) { + (Ok(()), Some(inflight_drain_secs)) => { + tracing::info!(inflight_drain_secs, "shutdown complete") + } + (Ok(()), None) => { + tracing::info!("shutdown complete; the server stopped without a termination signal") + } + (Err(e), Some(inflight_drain_secs)) => tracing::error!( + error = %e, + inflight_drain_secs, + "shutdown complete, but the server exited with an error", + ), + (Err(e), None) => tracing::error!( + error = %e, + "the server exited with an error before any termination signal", + ), + } server_result } @@ -252,14 +376,115 @@ fn prefix_index_config( } } -/// Waits for either Unix termination signal and logs the selected cause. -async fn shutdown_signal(mut sigterm: Signal, mut sigint: Signal) { - tokio::select! { - _ = sigterm.recv() => tracing::info!("got SIGTERM, shutting down"), - _ = sigint.recv() => tracing::info!("got SIGINT, shutting down"), +/// Resolve when a termination signal arrives, then hand control to axum's +/// graceful drain. On SIGTERM (k8s pod termination) first run the readiness +/// drain — flip `/readyz` to 503 and keep serving for `drain` so the endpoint +/// removal reaches kube-proxy before we stop accepting, closing the +/// rolling-update race. SIGINT (local Ctrl-C) skips the readiness drain +/// entirely — no 503 flip, no pause — so dev iteration does not pay it. +/// +/// Either way axum's own in-flight drain runs afterwards and is unbounded: a +/// long streaming completion still holds the process until it finishes or +/// `terminationGracePeriodSeconds` expires. A further termination signal cuts +/// the pause short but cannot reach that phase; it is logged instead. +async fn shutdown_signal( + mut sigterm: Signal, + mut sigint: Signal, + ctx: Arc, + drain: std::time::Duration, +) { + let sigterm_first = tokio::select! { + _ = sigterm.recv() => { + tracing::info!("got SIGTERM, shutting down"); + true + } + _ = sigint.recv() => { + tracing::info!("got SIGINT, shutting down without the readiness drain"); + false + } + }; + + let (expedite_tx, expedite_rx) = tokio::sync::oneshot::channel::<()>(); + // Only the SIGTERM path runs a pause, so only it has something to cut + // short; on the SIGINT path the first further signal goes straight to the + // warning below. + let mut expedite_tx = sigterm_first.then_some(expedite_tx); + // Hand both streams to a task that outlives this future, on EITHER branch. + // Dropping them here would make every later signal vanish: tokio never + // restores the default disposition, so the process would neither expedite + // nor die, and nothing would be logged. + tokio::spawn(async move { + loop { + let delivered = tokio::select! { + delivered = sigterm.recv() => delivered, + delivered = sigint.recv() => delivered, + }; + if delivered.is_none() { + // The signal driver is gone (runtime shutting down). Looping + // would spin without ever receiving again. + return; + } + handle_further_signal(&mut expedite_tx, sigterm_first); + } + }); + + if sigterm_first { + let expedite = async move { + let _ = expedite_rx.await; + }; + sgl_router::server::shutdown::drain_for_termination(&ctx, drain, expedite).await; } } +/// What a termination signal past the first one achieved. +#[derive(Debug, PartialEq, Eq)] +enum FurtherSignal { + /// Cut the readiness pause short. + Expedited, + /// Arrived with no pause left to cut short, and was reported as such. + Ignored, +} + +/// Handle one termination signal past the first, and say what it did. +/// +/// A failed `send` is not a lost race to shrug at: it means the pause already +/// ended on its own and dropped the receiver with it, which is the same +/// "nothing left to cut short" state as a spent `expedite_tx` and must reach +/// the same notice. Discarding that `Err` is what made the *first* signal after +/// the pause disappear, so only a second one was ever reported — the opposite +/// of this function's contract. +/// +/// Split out of the watcher task because that task owns real `Signal` streams +/// and cannot be driven from a test; this is where the decision lives, so this +/// is what a test can pin. +fn handle_further_signal( + expedite_tx: &mut Option>, + sigterm_first: bool, +) -> FurtherSignal { + // The first further signal cuts the readiness pause short, so an operator + // watching a stuck rollout is not held for a window that has stopped being + // useful. + if let Some(tx) = expedite_tx.take() { + if tx.send(()).is_ok() { + return FurtherSignal::Expedited; + } + } + // Two messages, because the two states call for different conclusions: one + // ran a pause that has since passed, the other never ran one at all. + if sigterm_first { + tracing::warn!( + "further termination signal ignored: the readiness pause is over and the \ + in-flight drain cannot be cut short; send SIGKILL to force an immediate exit", + ); + } else { + tracing::warn!( + "further termination signal ignored: SIGINT skips the readiness pause and the \ + in-flight drain cannot be cut short; send SIGKILL to force an immediate exit", + ); + } + FurtherSignal::Ignored +} + #[cfg(test)] mod tests { use super::*; @@ -284,6 +509,58 @@ mod tests { assert!(install_signal_handlers().is_ok()); } + #[test] + fn a_further_signal_expedites_a_running_pause() { + let (tx, mut rx) = tokio::sync::oneshot::channel::<()>(); + let mut expedite_tx = Some(tx); + assert_eq!( + handle_further_signal(&mut expedite_tx, true), + FurtherSignal::Expedited, + ); + assert!( + expedite_tx.is_none(), + "the sender is spent once the pause has been expedited", + ); + assert!( + rx.try_recv().is_ok(), + "the running pause must actually be notified", + ); + } + + /// The regression this function exists for. When the pause elapses on its + /// own, the receiver goes with it while the watcher still holds the sender + /// — so the very next signal hits a `send` that fails. Discarding that + /// `Err` swallowed it, and only the signal AFTER it was ever reported. + #[test] + fn the_first_signal_after_the_pause_ends_is_reported_not_swallowed() { + let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + drop(rx); // the pause elapsed and dropped its receiver + let mut expedite_tx = Some(tx); + assert_eq!( + handle_further_signal(&mut expedite_tx, true), + FurtherSignal::Ignored, + "a send into a dropped receiver expedites nothing and must say so", + ); + // ...and every later signal behaves identically, rather than the second + // one being the first to report anything. + assert_eq!( + handle_further_signal(&mut expedite_tx, true), + FurtherSignal::Ignored, + ); + } + + /// SIGINT never runs a readiness pause, so `expedite_tx` is `None` from the + /// start and every further Ctrl-C is reported rather than expediting a + /// pause that does not exist. + #[test] + fn a_further_signal_on_the_sigint_path_is_always_ignored() { + let mut expedite_tx = None; + assert_eq!( + handle_further_signal(&mut expedite_tx, false), + FurtherSignal::Ignored, + ); + } + #[test] fn init_tracing_is_idempotent() { let _ = init_tracing("info", LogFormat::Text); diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index f860acfeb..952ed1a68 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -14,9 +14,18 @@ use crate::server::inflight::InflightHttp; use crate::server::metrics::MetricsRegistry; use crate::tokenizer::TokenizerRegistry; use crate::workers::WorkerRegistry; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; +/// `/readyz` readiness as a one-way door: `NOT_READY -> READY -> DRAINING`, +/// and never backwards. One atomic rather than a pair of bools so the latch is +/// the transition itself — `mark_ready`'s compare-exchange simply cannot +/// succeed from `DRAINING` — instead of an invariant stated in a doc comment +/// and enforced by nobody. +const READINESS_NOT_READY: u8 = 0; +const READINESS_READY: u8 = 1; +const READINESS_DRAINING: u8 = 2; + pub struct AppContext { pub config: Config, pub tokenizers: Arc, @@ -43,9 +52,10 @@ pub struct AppContext { /// [`crate::policies::kv_events::KvEventIndex::metrics_source`]. pub kv_metrics: Option, /// Open HTTP exchanges, on every route. What axum's graceful shutdown - /// waits on — `active_load` sees only the proxied subset. + /// is actually waiting on during the drain — `active_load` sees only the + /// proxied subset. pub inflight_http: Arc, - ready: AtomicBool, + readiness: AtomicU8, } impl AppContext { @@ -103,18 +113,47 @@ impl AppContext { kv_metrics: None, engine_load: EngineLoadTable::new(), inflight_http: InflightHttp::new(), - ready: AtomicBool::new(false), + readiness: AtomicU8::new(READINESS_NOT_READY), } } + /// Report bootstrap as finished, unless the pod has already begun draining. + /// The `DRAINING` state is a one-way door (see [`Self::mark_not_ready`]), + /// and a compare-exchange is what enforces it: a plain store would let any + /// later caller — a re-initialization path, a discovery-recovery hook — + /// flip `/readyz` back to 200 seconds before the listener closes, re-arming + /// the rolling-update race the drain exists to close. pub fn mark_ready(&self) { // Relaxed: this flag does not synchronize other state; readers only - // care about eventual visibility, not happens-before with surrounding ops. - self.ready.store(true, Ordering::Relaxed); + // care about eventual visibility, not happens-before with surrounding + // ops. Failure means the state was already READY or is DRAINING — + // correct in both cases, so the result is deliberately discarded. + let _ = self.readiness.compare_exchange( + READINESS_NOT_READY, + READINESS_READY, + Ordering::Relaxed, + Ordering::Relaxed, + ); } + /// Flip `/readyz` to 503, permanently. Called at the start of the SIGTERM + /// drain so probes and any probe-driven load balancer see this pod as + /// not-ready while the endpoint removal (triggered by the pod's + /// `deletionTimestamp`, not by this flip) propagates. See + /// [`crate::server::shutdown::drain_for_termination`] for which mechanism + /// the pause is sized for. + /// + /// Not the inverse of [`mark_ready`](Self::mark_ready): this transition + /// cannot be undone, because the process it announces cannot be either. + pub fn mark_not_ready(&self) { + self.readiness.store(READINESS_DRAINING, Ordering::Relaxed); + } + + /// Whether bootstrap finished — only ONE term of the `/readyz` predicate, + /// which also requires a non-empty worker registry (see + /// `server::routes::health::readyz`). pub fn is_ready(&self) -> bool { - self.ready.load(Ordering::Relaxed) + self.readiness.load(Ordering::Relaxed) == READINESS_READY } #[cfg(test)] @@ -162,7 +201,57 @@ impl AppContext { kv_metrics: None, engine_load: EngineLoadTable::new(), inflight_http: InflightHttp::new(), - ready: AtomicBool::new(false), + readiness: AtomicU8::new(READINESS_NOT_READY), } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn mark_not_ready_flips_readiness_back_off() { + let ctx = AppContext::stub(); + // stub starts not-ready; mark_ready is the readiness on-switch. + ctx.mark_ready(); + assert!(ctx.is_ready(), "mark_ready must report ready"); + // The SIGTERM drain path needs the inverse so /readyz can flip to 503 + // before the server stops accepting. + ctx.mark_not_ready(); + assert!( + !ctx.is_ready(), + "mark_not_ready must flip readiness back off", + ); + } + + /// The safety-critical direction. A `mark_ready` reaching a draining pod + /// would put `/readyz` back to 200 with the listener seconds from closing, + /// silently re-arming the rolling-update race — so the latch is tested, + /// not just documented. + #[test] + fn readiness_does_not_come_back_once_draining() { + let ctx = AppContext::stub(); + ctx.mark_ready(); + ctx.mark_not_ready(); + + ctx.mark_ready(); + assert!( + !ctx.is_ready(), + "mark_ready must not re-ready a pod that has begun draining", + ); + } + + /// `mark_ready` is idempotent: the compare-exchange failing because the + /// state is already READY must not be mistaken for the draining case. + #[test] + fn mark_ready_is_idempotent() { + let ctx = AppContext::stub(); + ctx.mark_ready(); + ctx.mark_ready(); + assert!( + ctx.is_ready(), + "a second mark_ready must keep the pod ready" + ); + } +} diff --git a/experimental/sgl-router/src/server/mod.rs b/experimental/sgl-router/src/server/mod.rs index bbfc67398..1760108d3 100644 --- a/experimental/sgl-router/src/server/mod.rs +++ b/experimental/sgl-router/src/server/mod.rs @@ -8,3 +8,4 @@ pub mod header_utils; pub mod inflight; pub mod metrics; pub mod routes; +pub mod shutdown; diff --git a/experimental/sgl-router/src/server/shutdown.rs b/experimental/sgl-router/src/server/shutdown.rs new file mode 100644 index 000000000..d4f9478e4 --- /dev/null +++ b/experimental/sgl-router/src/server/shutdown.rs @@ -0,0 +1,175 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Graceful-termination helpers. +//! +//! On SIGTERM the pod is on its way out, but the data plane does not know it +//! yet: kube-proxy keeps routing here until the endpoint removal propagates to +//! it (an external load balancer deregisters on its own probe cadence +//! instead — see [`drain_for_termination`]). Requests sent in that window +//! reach a socket that is about to stop accepting and fail at the client. The +//! drain below holds the listener open for a fixed, operator-set window sized +//! to cover that propagation, with `/readyz` already reporting 503. + +use crate::server::app_context::AppContext; +use std::future::Future; +use std::time::Duration; + +/// Begin a graceful-termination drain: flip `/readyz` to 503, then keep +/// serving for `drain` before the caller stops accepting connections. A zero +/// `drain` flips readiness and returns at once. This composes with axum's +/// `with_graceful_shutdown`: once this future resolves, axum stops accepting +/// and drains the already-in-flight requests. +/// +/// The two deregistration mechanisms, and which one the pause is sized for: +/// +/// 1. **Endpoint removal.** For a pod deletion or rolling update the +/// EndpointSlice controller marks this pod's endpoint not-ready the moment +/// the `deletionTimestamp` is stamped — it does not wait on a probe. The +/// pause covers the propagation of that to kube-proxy on every node. This +/// is the mechanism the default drain is sized for. +/// 2. **The `/readyz` flip.** Probe-driven deregistration (an external load +/// balancer, or a kubelet readiness probe on a pod that is not being +/// deleted) needs `failureThreshold` consecutive failures at +/// `periodSeconds` apart before it acts. A drain shorter than that product +/// never gets observed, so operators relying on this path must raise the +/// drain to match their own probe cadence — the default does not do it for +/// them. +/// +/// `expedite` cuts the *pause* short: if it resolves first (a further +/// termination signal), the drain returns early so the process is not held for +/// a window that has stopped being useful. It does not reach the axum +/// in-flight drain that runs afterwards, which is unbounded. Pass +/// [`std::future::pending`] to never expedite. Note that two signals delivered +/// close enough together coalesce into one notification, so a drain already +/// running takes one *further* signal to expedite. +pub async fn drain_for_termination( + ctx: &AppContext, + drain: Duration, + expedite: impl Future, +) { + ctx.mark_not_ready(); + if drain.is_zero() { + // Still say so: without this line `--shutdown-drain-secs 0` produces a + // shutdown log indistinguishable from an image that predates the drain. + tracing::info!("/readyz now 503; drain pause disabled (shutdown_drain_secs=0)"); + return; + } + tracing::info!( + drain_secs = drain.as_secs(), + "draining: /readyz now 503, waiting before the server stops accepting" + ); + tokio::select! { + _ = tokio::time::sleep(drain) => { + tracing::info!( + drain_secs = drain.as_secs(), + "drain pause elapsed; the server now stops accepting", + ); + } + _ = expedite => { + tracing::info!("drain pause expedited by a further termination signal"); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + #[tokio::test] + async fn drain_zero_flips_readiness_without_pausing() { + let ctx = AppContext::stub(); + ctx.mark_ready(); + // Real time, and an elapsed-time assertion rather than a `timeout`: a + // "minimum safe drain" floor added to the production path would still + // fit inside any timeout generous enough not to be flaky, so only + // measuring the elapsed time actually pins "0 means no pause". + let started = std::time::Instant::now(); + drain_for_termination(&ctx, Duration::ZERO, std::future::pending::<()>()).await; + let elapsed = started.elapsed(); + assert!( + elapsed < Duration::from_millis(50), + "zero drain must not pause, slept {elapsed:?}", + ); + assert!(!ctx.is_ready(), "zero drain still flips readiness off"); + } + + #[tokio::test(start_paused = true)] + async fn drain_holds_for_the_delay_after_flipping_readiness() { + let ctx = AppContext::stub(); + ctx.mark_ready(); + // A 30 s drain must not have returned within a 10 ms window... + let returned = tokio::time::timeout( + Duration::from_millis(10), + drain_for_termination(&ctx, Duration::from_secs(30), std::future::pending::<()>()), + ) + .await; + assert!( + returned.is_err(), + "drain must still be sleeping out the configured delay", + ); + // ...but readiness flipped off on entry, before the sleep. + assert!( + !ctx.is_ready(), + "readiness must flip off before the drain delay elapses", + ); + } + + /// The pause must last *the configured* time — not merely "some time". + /// Two distinct values, because a hardcoded constant substituted for + /// `drain` satisfies any single-value test. + #[tokio::test(start_paused = true)] + async fn drain_holds_for_exactly_the_configured_delay() { + for secs in [5_u64, 30] { + let ctx = Arc::new(AppContext::stub()); + ctx.mark_ready(); + let drain_ctx = Arc::clone(&ctx); + let handle = tokio::spawn(async move { + drain_for_termination( + &drain_ctx, + Duration::from_secs(secs), + std::future::pending::<()>(), + ) + .await; + }); + // Let the task register its sleep before advancing: a timer that + // has not been created yet cannot be advanced past. + tokio::task::yield_now().await; + + // One millisecond shy of the deadline the drain must still be held. + tokio::time::advance(Duration::from_millis(secs * 1000 - 1)).await; + tokio::task::yield_now().await; + assert!( + !handle.is_finished(), + "a {secs} s drain returned before its configured delay elapsed", + ); + + // Past it, it must return promptly. + tokio::time::advance(Duration::from_millis(2)).await; + tokio::time::timeout(Duration::from_secs(1), handle) + .await + .unwrap_or_else(|_| panic!("a {secs} s drain must return once its delay elapses")) + .expect("drain task joined cleanly"); + } + } + + #[tokio::test(start_paused = true)] + async fn drain_is_cut_short_when_expedite_resolves() { + let ctx = AppContext::stub(); + ctx.mark_ready(); + // A further termination signal (here: an already-resolved expedite + // future) must cut the pause short so an operator re-sending SIGTERM + // is not held for the full window. + let done = tokio::time::timeout( + Duration::from_millis(10), + drain_for_termination(&ctx, Duration::from_secs(3600), std::future::ready(())), + ) + .await; + assert!( + done.is_ok(), + "an expedite signal must cut the drain short, not wait out 3600 s" + ); + assert!(!ctx.is_ready(), "readiness still flipped off"); + } +} diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py b/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py index f3c1e5171..fb87bd118 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/conftest.py @@ -109,40 +109,30 @@ def _wait_for_replacement_pod_ready( last_observed = "no pods" while time.time() < deadline: - result = _kubectl( - "get", - "pods", - "-n", - namespace, - "-l", - selector, - "-o", - "json", - check=False, - ) - if getattr(result, "returncode", 0) == 0: - pods = json.loads(result.stdout or "{}").get("items", []) - names = [pod.get("metadata", {}).get("name", "") for pod in pods] - last_observed = ", ".join(filter(None, names)) or "no pods" + pods = _pods(selector, namespace, check=False) + names = [pod.get("metadata", {}).get("name", "") for pod in pods] + # An empty list is "nothing observed" whether the pods are gone or the + # kubectl call failed; both read the same in a timeout message. + last_observed = ", ".join(filter(None, names)) or "no pods" - if old_pod not in names: - for pod in sorted( - pods, key=lambda item: item.get("metadata", {}).get("name", "") + if old_pod not in names: + for pod in sorted( + pods, key=lambda item: item.get("metadata", {}).get("name", "") + ): + metadata = pod.get("metadata", {}) + status = pod.get("status", {}) + ready = any( + condition.get("type") == "Ready" + and condition.get("status") == "True" + for condition in status.get("conditions", []) + ) + if ( + metadata.get("name") != old_pod + and _is_live(pod) + and status.get("phase") == "Running" + and ready ): - metadata = pod.get("metadata", {}) - status = pod.get("status", {}) - ready = any( - condition.get("type") == "Ready" - and condition.get("status") == "True" - for condition in status.get("conditions", []) - ) - if ( - metadata.get("name") != old_pod - and not metadata.get("deletionTimestamp") - and status.get("phase") == "Running" - and ready - ): - return metadata["name"] + return metadata["name"] time.sleep(interval) @@ -172,14 +162,20 @@ def _port_forward_start( service: str, local_port: int, remote_port: int, + resource: str = "svc", ) -> subprocess.Popen: - """Start kubectl port-forward and wait until the port is reachable.""" + """Start kubectl port-forward and wait until the port is reachable. + + `resource="pod"` binds one specific pod instead of the Service. A draining + pod is removed from the Service's ready endpoints, so a test that needs to + keep talking to it through the drain must address the pod directly. + """ cmd = [ "kubectl", "--context", KUBECTL_CONTEXT, "port-forward", - f"svc/{service}", + f"{resource}/{service}", f"{local_port}:{remote_port}", "-n", namespace, @@ -215,6 +211,66 @@ def _cleanup_port_forward(name: str, pf: subprocess.Popen) -> None: logger.debug("Port-forward %s exited cleanly (rc=%s)", name, rc) +def _pod_json(pod: str, namespace: str = NAMESPACE) -> dict: + """One pod's full object. The `or "{}"` mirrors + `_wait_for_replacement_pod_ready`: kubectl can hand back empty stdout, and a + JSONDecodeError there says nothing about what went wrong.""" + result = _kubectl("get", "pod", pod, "-n", namespace, "-o", "json") + return json.loads(result.stdout or "{}") + + +def _pods( + selector: str, + namespace: str = NAMESPACE, + check: bool = True, +) -> list[dict]: + """Pod objects matching `selector`. `check=False` yields `[]` on a failed + kubectl instead of raising, for poll loops that expect the API server to be + briefly unavailable mid-rollout. The `or "{}"` guards kubectl handing back + empty stdout, where a JSONDecodeError would say nothing about what went + wrong.""" + result = _kubectl( + "get", "pods", "-n", namespace, "-l", selector, "-o", "json", check=check + ) + if getattr(result, "returncode", 0) != 0: + return [] + return json.loads(result.stdout or "{}").get("items", []) + + +def _is_live(pod: dict) -> bool: + """Whether a pod object is not already terminating. One predicate rather + than two copies of `deletionTimestamp`, so the replacement-pod poll and + `_pod_names` cannot drift apart on what counts as gone.""" + return not pod.get("metadata", {}).get("deletionTimestamp") + + +def _pod_names(selector: str, namespace: str = NAMESPACE) -> list[str]: + """Names of pods matching `selector`, excluding any already terminating.""" + return [p["metadata"]["name"] for p in _pods(selector, namespace) if _is_live(p)] + + +def _container_restart_count( + pod: str, + container: str, + namespace: str = NAMESPACE, +) -> int: + """`restartCount` for one container — how a test observes that the process + exited and kubelet restarted it in place (no new pod, same name).""" + statuses = _pod_json(pod, namespace).get("status", {}).get("containerStatuses", []) + for status in statuses: + if status["name"] == container: + return int(status["restartCount"]) + raise AssertionError(f"container {container!r} not found on pod {pod!r}") + + +def _pod_ready_condition(pod: str, namespace: str = NAMESPACE) -> str: + """The pod's `Ready` condition as k8s currently sees it ("True"/"False").""" + for cond in _pod_json(pod, namespace).get("status", {}).get("conditions", []): + if cond["type"] == "Ready": + return cond["status"] + return "Unknown" + + def _poll_until( predicate, description: str, diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml index d04cab9ac..54204b9b5 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml +++ b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml @@ -16,6 +16,9 @@ spec: app: sgl-router-cluster spec: serviceAccountName: sgl-router-cluster + # Must exceed --shutdown-drain-secs plus the time in-flight requests need + # after the pause, or the pod is SIGKILLed mid-drain. + terminationGracePeriodSeconds: 40 containers: - name: router image: sgl-router:e2e @@ -37,6 +40,10 @@ spec: - "--service-discovery" - "--selector" - "app=sglang,cross-ns-test=true" + # Parity with router.yaml so the cross-namespace router exercises + # the same shutdown path; no test asserts on it here. + - "--shutdown-drain-secs" + - "8" ports: - containerPort: 8091 name: http diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml index 021d8bfcb..0732d7782 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml +++ b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml @@ -14,6 +14,10 @@ spec: app: sgl-router spec: serviceAccountName: sgl-router + # Must exceed --shutdown-drain-secs plus the time in-flight requests + # need after the pause, or the pod is SIGKILLed mid-drain and the drain + # has bought nothing. + terminationGracePeriodSeconds: 40 containers: - name: router image: sgl-router:e2e @@ -44,6 +48,18 @@ spec: - "sgl-router-test" - "--selector" - "app=sglang" + # On SIGTERM, keep serving with /readyz at 503 so the endpoint + # removal reaches kube-proxy before the listener closes. Sized for + # the deletionTimestamp path, which does not wait on a probe. + # Probe-driven deregistration would instead need this above the + # readinessProbe's failureThreshold * periodSeconds below. + # test_shutdown_drain.py reads this value; do not restate it there. + - "--shutdown-drain-secs" + - "8" + # Declared so the startup advisory compares the drain against this + # pod's real grace period instead of assuming the k8s default. + - "--termination-grace-secs" + - "40" ports: - containerPort: 8090 name: http diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py index a89449860..cbb36bd6a 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_conftest_helpers.py @@ -2,6 +2,7 @@ import json from types import SimpleNamespace import conftest as k8s_conftest +import pytest def _pod(name: str, phase: str, ready: bool) -> dict: @@ -52,3 +53,97 @@ def test_wait_for_replacement_pod_ignores_old_and_pending_pods(monkeypatch): assert replacement == new_pod assert len(calls) == 4 assert all("-o" in args and "json" in args for args, _ in calls) + + +def _stub_kubectl(monkeypatch, payload): + """Point every conftest helper at a canned kubectl response. `payload` is + the object kubectl would have printed.""" + monkeypatch.setattr( + k8s_conftest, + "_kubectl", + lambda *args, **kwargs: SimpleNamespace(stdout=json.dumps(payload)), + ) + + +def test_pod_names_excludes_terminating_pods(monkeypatch): + terminating = _pod("sgl-router-old", "Running", True) + terminating["metadata"]["deletionTimestamp"] = "2026-01-01T00:00:00Z" + _stub_kubectl( + monkeypatch, + {"items": [terminating, _pod("sgl-router-new", "Running", True)]}, + ) + + assert k8s_conftest._pod_names("app=sgl-router") == ["sgl-router-new"] + + +def test_pods_returns_empty_on_a_failed_kubectl(monkeypatch): + """The poll loops call this with `check=False` precisely because the API + server can be briefly unavailable mid-rollout; a non-zero return must read + as "nothing observed", not raise out of the loop.""" + monkeypatch.setattr( + k8s_conftest, + "_kubectl", + lambda *args, **kwargs: SimpleNamespace(stdout="", returncode=1), + ) + + assert k8s_conftest._pods("app=sgl-router", check=False) == [] + + +def test_pods_tolerates_empty_stdout(monkeypatch): + monkeypatch.setattr( + k8s_conftest, + "_kubectl", + lambda *args, **kwargs: SimpleNamespace(stdout=""), + ) + + assert k8s_conftest._pods("app=sgl-router") == [] + + +def test_container_restart_count_reads_the_named_container(monkeypatch): + _stub_kubectl( + monkeypatch, + { + "status": { + "containerStatuses": [ + {"name": "sidecar", "restartCount": 9}, + {"name": "router", "restartCount": 3}, + ] + } + }, + ) + + assert k8s_conftest._container_restart_count("sgl-router-0", "router") == 3 + + +def test_container_restart_count_rejects_a_missing_container(monkeypatch): + """`containerStatuses` lags during a restart, so the absent case is a real + state — and the drain test reads its whole timing floor off this number. + Returning 0 there would silently read as "never restarted".""" + _stub_kubectl(monkeypatch, {"status": {"containerStatuses": []}}) + + with pytest.raises(AssertionError, match="router"): + k8s_conftest._container_restart_count("sgl-router-0", "router") + + +def test_pod_ready_condition_reports_the_ready_status(monkeypatch): + _stub_kubectl( + monkeypatch, + { + "status": { + "conditions": [ + {"type": "Initialized", "status": "True"}, + {"type": "Ready", "status": "False"}, + ] + } + }, + ) + + assert k8s_conftest._pod_ready_condition("sgl-router-0") == "False" + + +def test_pod_ready_condition_is_unknown_before_the_condition_exists(monkeypatch): + """A pod whose Ready condition has not been written yet must read as + "Unknown" rather than crash the drain test's diagnostic logging.""" + _stub_kubectl(monkeypatch, {"status": {"conditions": []}}) + + assert k8s_conftest._pod_ready_condition("sgl-router-0") == "Unknown" diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_shutdown_drain.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_shutdown_drain.py new file mode 100644 index 000000000..723f6672c --- /dev/null +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_shutdown_drain.py @@ -0,0 +1,258 @@ +"""SIGTERM readiness-drain integration tests. + +The drain exists to produce a *Kubernetes* behaviour, and the Rust tests can +only argue it: they substitute a channel for the real `Signal` and an +in-process `AppContext` for a real pod. These run the shipped container, so +they cover `main.rs::shutdown_signal` — the signal handler, the SIGTERM/SIGINT +branch, and the `ctx` wiring — which no in-process test reaches. + +Why `kill -TERM 1` rather than `kubectl delete pod`: deleting a pod stamps a +`deletionTimestamp`, and the endpoints controller marks the endpoint not-ready +on that alone, without ever consulting `/readyz`. A delete-based test would +therefore pass identically with the drain removed — it would look like +coverage while pinning nothing. Signalling the process directly leaves the pod +undeleted, so a `/readyz` 503 can only have come from the drain calling +`AppContext::mark_not_ready`. + +The router image is `debian:bookworm-slim` with an exec-form ENTRYPOINT, so +the binary is PID 1 and `kill -TERM 1` reaches it exactly as kubelet's SIGTERM +would. +""" + +from __future__ import annotations + +import logging +import re +import time +from pathlib import Path + +import httpx +from conftest import ( + NAMESPACE, + _cleanup_port_forward, + _container_restart_count, + _kubectl, + _pod_names, + _pod_ready_condition, + _poll_until, + _port_forward_start, + _wait_for_deployment_ready, +) + +logger = logging.getLogger(__name__) + +# Distinct from the shared 8090 forward so this test's pod-scoped forward +# cannot collide with a leaked service-scoped one from another test. +DRAIN_PORT = 8094 + +_ROUTER_MANIFEST = Path(__file__).parent / "manifests" / "router.yaml" + + +def _manifest_drain_secs() -> int: + """Read `--shutdown-drain-secs` out of the manifest the pod is started + from. Read rather than restated, because every assertion below is scaled to + the drain window: a manifest edit that this file did not track would leave + the test green while measuring the wrong window.""" + args = _ROUTER_MANIFEST.read_text() + match = re.search( + r'"--shutdown-drain-secs"\s*\n\s*-\s*"(\d+)"', + args, + ) + assert match, f"--shutdown-drain-secs not found in {_ROUTER_MANIFEST}" + return int(match.group(1)) + + +CONFIGURED_DRAIN_SECS = _manifest_drain_secs() + + +def _manifest_grace_secs() -> tuple[int, int]: + """`terminationGracePeriodSeconds` from the pod spec, and the + `--termination-grace-secs` the router is told about it. Two places by + necessity — the router cannot read its own pod spec — which is exactly why + they can drift apart.""" + manifest = _ROUTER_MANIFEST.read_text() + spec = re.search(r"terminationGracePeriodSeconds:\s*(\d+)", manifest) + assert spec, f"terminationGracePeriodSeconds not found in {_ROUTER_MANIFEST}" + declared = re.search( + r'"--termination-grace-secs"\s*\n\s*-\s*"(\d+)"', + manifest, + ) + assert declared, f"--termination-grace-secs not found in {_ROUTER_MANIFEST}" + return int(spec.group(1)), int(declared.group(1)) + + +def test_declared_grace_period_matches_the_pod_spec(): + """`--termination-grace-secs` silences the startup advisory, so a value + that has drifted from the pod's real `terminationGracePeriodSeconds` is + worse than no flag at all: it silences the warning against a budget the pod + does not have. No cluster needed — this is a manifest self-consistency + check, and it is the only thing standing between the two numbers.""" + spec_secs, declared_secs = _manifest_grace_secs() + assert declared_secs == spec_secs, ( + f"--termination-grace-secs is {declared_secs} but the pod spec grants " + f"{spec_secs}s; the advisory would be checked against the wrong budget" + ) + assert CONFIGURED_DRAIN_SECS < spec_secs, ( + f"the {CONFIGURED_DRAIN_SECS}s drain leaves no room under the {spec_secs}s " + f"grace period for the in-flight drain that follows it" + ) + + +# Budget for observing the /readyz flip, deliberately a fraction of the drain: +# the assertions that follow it must still land inside the window, so the poll +# cannot be allowed to consume the whole thing. +FLIP_OBSERVATION_SECS = max(2, CONFIGURED_DRAIN_SECS // 2) + +# Floor on a mid-drain HTTP timeout. Below this the request has no realistic +# chance on a loaded kind runner, so there is no point issuing it — the window +# has effectively closed and `_mid_drain_timeout` says so instead. +MIN_HTTP_TIMEOUT_SECS = 1.0 + +# How long past the window the container restart may take to become VISIBLE. +# Kubelet's own restart latency lands in here, and it only ever makes the +# observed time longer — so this is slack on the measurement, not a second +# claim about the drain. Sized to still catch a units regression that +# LENGTHENS the pause: the `from_secs`/`from_millis` slip that +# `ServerConfig::shutdown_drain()` exists to guard cuts both ways, and 8s +# becoming 80s satisfies every lower bound in this file. +RESTART_OBSERVATION_SLACK_SECS = 60 + + +def _mid_drain_timeout(sigterm_at: float, what: str, want: float) -> float: + """An HTTP timeout for a mid-drain assertion that cannot outlast the window + the assertion claims to run inside. + + Without this the per-request timeouts sum past the drain (a 4s flip poll + plus 5s and 10s requests against an 8s window), so on a slow runner the + listener closes with a request still open and the test dies on whichever + transport error that raised — not on the assertion written to explain the + outcome. Checking the remaining budget up front puts the explanation back. + """ + remaining = CONFIGURED_DRAIN_SECS - (time.monotonic() - sigterm_at) + assert remaining > MIN_HTTP_TIMEOUT_SECS, ( + f"no drain window left for {what}: {CONFIGURED_DRAIN_SECS - remaining:.1f}s " + f"of the {CONFIGURED_DRAIN_SECS}s window already spent. If this runner is " + f"simply slow, raise --shutdown-drain-secs in {_ROUTER_MANIFEST.name}" + ) + return min(want, remaining) + + +def _router_pod() -> str: + pods = _pod_names("app=sgl-router") + assert len(pods) == 1, f"expected exactly one live router pod, got {pods}" + return pods[0] + + +class TestReadinessDrain: + """SIGTERM must flip /readyz to 503 while the pod keeps serving.""" + + def test_sigterm_flips_readyz_while_the_pod_keeps_serving(self, k8s_cluster): + _wait_for_deployment_ready("sgl-router") + pod = _router_pod() + restarts_before = _container_restart_count(pod, "router") + + # Bind the pod, not the Service: a draining pod leaves the Service's + # ready endpoints, and the point of this test is to keep talking to it + # after that happens. + pf = _port_forward_start(NAMESPACE, pod, DRAIN_PORT, 8090, resource="pod") + base = f"http://127.0.0.1:{DRAIN_PORT}" + try: + assert httpx.get(f"{base}/readyz", timeout=5.0).status_code == 200, ( + "router must be ready before SIGTERM" + ) + + # Through `sh -c`: the slim image ships no `kill` binary, and + # `kubectl exec` execs directly rather than through a shell, so the + # builtin is the only way to signal PID 1 from outside. + sigterm_at = time.monotonic() + _kubectl("exec", "-n", NAMESPACE, pod, "--", "sh", "-c", "kill -TERM 1") + + # The flip is observable from outside the pod. + _poll_until( + lambda: httpx.get(f"{base}/readyz", timeout=3.0).status_code == 503, + "/readyz returns 503 after SIGTERM", + timeout=FLIP_OBSERVATION_SECS, + interval=0.2, + ) + + # ...and the pod is still serving while it reports not-ready. + # `/healthz` staying 200 is what stops the liveness probe + # restarting a pod that is draining on purpose. + healthz_timeout = _mid_drain_timeout(sigterm_at, "the liveness probe", 5.0) + assert ( + httpx.get(f"{base}/healthz", timeout=healthz_timeout).status_code == 200 + ), "liveness must stay green while the pod drains" + + # A proxied completion still succeeding does double duty: it is the + # request k8s may still route during the window, AND it proves the + # worker registry is non-empty — so the 503 above can only be the + # readiness flip, not `/readyz`'s other term. + chat = httpx.post( + f"{base}/v1/chat/completions", + json={ + "model": "tiny", + "messages": [{"role": "user", "content": "drain"}], + }, + timeout=_mid_drain_timeout(sigterm_at, "a proxied completion", 10.0), + ) + assert chat.status_code == 200, ( + f"a proxied request must still succeed mid-drain, got {chat.status_code}" + ) + + # Everything above claims to have run *inside* the window. Say so, + # so an overrun reads as "the window closed" and not as whichever + # transport error the closed listener happened to raise next. + mid_drain_elapsed = time.monotonic() - sigterm_at + assert mid_drain_elapsed < CONFIGURED_DRAIN_SECS, ( + f"the mid-drain assertions took {mid_drain_elapsed:.1f}s, past the " + f"{CONFIGURED_DRAIN_SECS}s window they claim to observe" + ) + + # Recorded, not asserted: k8s needs failureThreshold consecutive + # failing probes, periodSeconds apart, to mark the pod not-ready — + # longer than the drain at the values in router.yaml. That is + # exactly why the default is sized for the deletionTimestamp path + # instead, and why probe-driven setups must raise it. + logger.info("pod Ready condition mid-drain: %s", _pod_ready_condition(pod)) + + # The drain's FLOOR, pinned where it is actually observable: hold + # until just inside the window and prove the process is still up. + # Timing the floor off the restart instead is satisfiable by test + # overhead alone — the restart poll does not start until everything + # above has run, so a build whose pause was 1s would still look like + # it lasted the whole window. + still_up_at = CONFIGURED_DRAIN_SECS - 1 + time.sleep(max(0.0, still_up_at - (time.monotonic() - sigterm_at))) + assert _container_restart_count(pod, "router") == restarts_before, ( + f"the router exited within {still_up_at}s of SIGTERM, short of the " + f"configured {CONFIGURED_DRAIN_SECS}s drain" + ) + + finally: + _cleanup_port_forward(f"pod/{pod}", pf) + + # The drain must END in an exit. Read off `restartCount`: the process + # exits when the drain elapses and kubelet restarts the container in + # place, same pod. Watching this rather than the listener closing is + # deliberate — the restart is fast enough that a port-forward probe can + # miss the closed window entirely and hang, whereas `restartCount` is + # monotonic and cannot be missed. The poll's own timeout is deliberately + # looser than the ceiling below, so a lengthened drain fails on the + # assertion (which explains it) rather than on a bare TimeoutError. + _poll_until( + lambda: _container_restart_count(pod, "router") > restarts_before, + "router container restarts once the drain elapses", + timeout=CONFIGURED_DRAIN_SECS + RESTART_OBSERVATION_SLACK_SECS + 30, + interval=0.5, + ) + # The drain's CEILING. Its mirror image — the pause not being cut short + # — is the still-up assertion inside the window above; together they + # bound the pause from both sides, which neither does alone. + held_open_for = time.monotonic() - sigterm_at + assert held_open_for < CONFIGURED_DRAIN_SECS + RESTART_OBSERVATION_SLACK_SECS, ( + f"the router was still up {held_open_for:.1f}s after SIGTERM, past the " + f"configured {CONFIGURED_DRAIN_SECS}s drain by more than kubelet's restart " + f"latency can explain — check the seconds-to-Duration conversion in " + f"ServerConfig::shutdown_drain()" + ) + _wait_for_deployment_ready("sgl-router") diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 10e48bff3..5cf8d7a66 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs @@ -2,10 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 //! Pins the contract that `axum::serve(...).with_graceful_shutdown(...)` — -//! exactly as wired in `src/main.rs` — drains every in-flight streaming +//! the same combinator `src/main.rs` uses — drains every in-flight streaming //! request through the **real** `build_router(ctx)` stack before the //! server future resolves. A k8s SIGTERM must not truncate streaming -//! completions. +//! completions. (`main.rs` additionally runs the readiness drain first; the +//! later tests cover that.) //! //! Why route the test through the real router (chat handler + proxy + //! SSE pump) rather than a synthetic `Router::new().route(...)`: a @@ -13,6 +14,14 @@ //! `bytes_stream_to_body` completion hook, in `chat::chat_completions`' //! guards, or in the SSE pump's `tx.send().await` race — all of which //! would be silently skipped by a synthetic-handler test. +//! +//! The later tests pin the readiness drain that runs *before* that axum +//! drain: `server::shutdown::drain_for_termination` flips `/readyz` to 503 +//! and holds the listener open for `--shutdown-drain-secs` so the endpoint +//! removal reaches kube-proxy first. They substitute a channel for the real +//! `Signal`, so `main.rs`'s `shutdown_signal` is not exercised here; the k8s +//! integration suite (`tests/e2e/k8s_integration/test_shutdown_drain.py`) +//! signals the shipped binary and covers that wiring. use futures::future::join_all; use sgl_router::config::{ @@ -235,6 +244,293 @@ async fn shutdown_with_no_inflight_returns_promptly() { ); } +/// The readiness-drain contract: on SIGTERM the drain flips `/readyz` to 503 +/// *while the server keeps accepting* (`/healthz` stays 200, a brand-new +/// connection is still served), so the endpoint removal reaches kube-proxy +/// before the listener closes. Mirrors `src/main.rs`'s SIGTERM arm by driving +/// the shutdown future as "await the signal, then `drain_for_termination`" +/// against the real `build_router(ctx)` stack. +/// +/// The drain window is ended by the `expedite` channel rather than by wall +/// clock, so the mid-drain assertions cannot lose a race with a sleeping +/// timer on a loaded runner — and the expedite path itself gets covered. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn readyz_flips_to_503_during_drain_while_still_serving() { + let worker = crate::common::mock_worker::MockWorker::start_slow_stream( + SLOW_CHUNKS.to_vec(), + Duration::from_millis(20), + ) + .await; + let ctx = build_ctx_with_worker(&worker.url); + assert!(ctx.is_ready(), "ctx starts ready"); + + let app = build_router(ctx.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + // `sigterm_tx` stands in for SIGTERM delivery; `expedite_tx` stands in for + // the further termination signal that cuts the pause short. The drain is + // an hour so only `expedite_tx` can end it. + let ctx_for_shutdown = ctx.clone(); + let (sigterm_tx, sigterm_rx) = oneshot::channel::<()>(); + let (expedite_tx, expedite_rx) = oneshot::channel::<()>(); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = sigterm_rx.await; + sgl_router::server::shutdown::drain_for_termination( + &ctx_for_shutdown, + Duration::from_secs(3600), + async { + let _ = expedite_rx.await; + }, + ) + .await; + }) + .await + .unwrap(); + }); + + // Every probe opens its own connection: a pooled client would ride the + // pre-SIGTERM connection and keep passing even if the listener had already + // closed, which is exactly the regression this test exists to catch. + let client = reqwest::Client::builder() + .pool_max_idle_per_host(0) + .build() + .unwrap(); + let readyz = format!("http://{addr}/readyz"); + let healthz = format!("http://{addr}/healthz"); + + // Before SIGTERM: ready + worker registered ⇒ /readyz 200. + let pre = client.get(&readyz).send().await.unwrap(); + assert_eq!( + pre.status(), + reqwest::StatusCode::OK, + "ready before SIGTERM" + ); + + sigterm_tx.send(()).unwrap(); + // The drain flips readiness before its first await, but the flip and this + // observation are on different tasks — wait for it rather than sleeping. + tokio::time::timeout(Duration::from_secs(5), async { + while ctx.is_ready() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the drain must flip readiness off promptly after SIGTERM"); + + let mid_ready = client.get(&readyz).send().await.unwrap(); + assert_eq!( + mid_ready.status(), + reqwest::StatusCode::SERVICE_UNAVAILABLE, + "/readyz must flip to 503 during the drain so probes and load balancers see this pod as not-ready before the listener closes", + ); + + // State the accept explicitly rather than inferring it from a 200: this is + // the half of the contract that a pooled client would silently satisfy. + tokio::net::TcpStream::connect(addr) + .await + .expect("the listener must still accept new connections during the drain"); + let mid_health = client.get(&healthz).send().await.unwrap(); + assert_eq!( + mid_health.status(), + reqwest::StatusCode::OK, + "the server must still be serving during the drain window", + ); + + // A *real proxied* request (not just the local health handlers) must still + // be accepted and served during the drain window — this is the request k8s + // may still route before the endpoint removal reaches kube-proxy. + let chat = format!("http://{addr}/v1/chat/completions"); + let body = serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + }); + let mid_chat = client.post(&chat).json(&body).send().await.unwrap(); + assert_eq!( + mid_chat.status(), + reqwest::StatusCode::OK, + "a proxied chat request must still succeed during the drain window", + ); + + // The request the drain actually exists for: it ARRIVES during the pause + // (kube-proxy has not observed the removal yet) and is still streaming when + // the pause ends. It must survive the handover into axum's in-flight drain, + // not just the window it started in. + // + // Await the response headers here rather than inside the spawned task: that + // is the point at which the request is provably in flight, so cutting the + // pause short below cannot race the client's connect on a loaded runner. + let stream_client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .unwrap(); + let late_request = serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + }); + let late_resp = stream_client + .post(&chat) + .json(&late_request) + .send() + .await + .unwrap(); + assert!( + late_resp.status().is_success(), + "a stream started during the drain must be accepted: {}", + late_resp.status(), + ); + let late = tokio::spawn(async move { late_resp.bytes().await.unwrap() }); + + // Cut the pause short while that stream is still mid-flight; the server + // resolves without waiting out the hour. + expedite_tx.send(()).unwrap(); + + let late_body = late.await.expect("late client task joined"); + assert!( + String::from_utf8_lossy(&late_body).contains("data: [DONE]"), + "a request that arrived during the drain must still complete after the pause ends", + ); + tokio::time::timeout(Duration::from_secs(5), server) + .await + .expect("an expedite signal must end the drain instead of sleeping an hour") + .expect("server task joined cleanly"); +} + +/// After the drain elapses and the server future resolves, axum must have +/// stopped accepting: a *new* connection is refused. This is the other half of +/// the contract — the drain has to actually END in a closed listener, or the +/// pause merely postpones shutdown without ever handing traffic off. (What +/// closes the rolling-update race is the pause itself, covered by +/// `readyz_flips_to_503_during_drain_while_still_serving`.) Asserted on a raw +/// TCP connect so the failure has to be `ConnectionRefused`; a `reqwest` error +/// would also cover a timeout, which is a different (and on a loaded runner, +/// plausible) outcome. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn new_connections_refused_after_drain_completes() { + let worker = crate::common::mock_worker::MockWorker::start(vec![]).await; + let ctx = build_ctx_with_worker(&worker.url); + let app = build_router(ctx.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + // Short drain so the test is fast; the point is the post-resolve state. + let drain = Duration::from_millis(100); + let ctx_for_shutdown = ctx.clone(); + let (sigterm_tx, sigterm_rx) = oneshot::channel::<()>(); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = sigterm_rx.await; + sgl_router::server::shutdown::drain_for_termination( + &ctx_for_shutdown, + drain, + std::future::pending::<()>(), + ) + .await; + }) + .await + .unwrap(); + }); + + // Server accepts before shutdown. + tokio::net::TcpStream::connect(addr) + .await + .expect("listener accepts before SIGTERM"); + + // Fire SIGTERM and wait for the drain + server future to fully resolve. + sigterm_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(5), server) + .await + .expect("server resolves after the drain elapses") + .expect("server task joined cleanly"); + + // A fresh connection must now be refused — the listener is closed. + let err = tokio::net::TcpStream::connect(addr) + .await + .expect_err("a new connection must be refused after the drain completes"); + assert_eq!( + err.kind(), + std::io::ErrorKind::ConnectionRefused, + "expected the closed listener to refuse, got {err:?}", + ); +} + +/// End-to-end composition: SIGTERM → `drain_for_termination` (flip 503, pause) +/// → axum drains the already-attached streaming request to `[DONE]`. +/// `shutdown_drains_100_inflight_streaming_chat_completions` drives a bare +/// oneshot shutdown future; this one composes the readiness drain with the axum +/// drain, so a regression that truncates in-flight streams once the drain +/// begins is caught. It does NOT assert the flip/pause ordering — +/// `readyz_flips_to_503_during_drain_while_still_serving` covers that. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn inflight_stream_completes_through_drain_for_termination() { + let worker = crate::common::mock_worker::MockWorker::start_slow_stream( + SLOW_CHUNKS.to_vec(), + Duration::from_millis(60), + ) + .await; + let ctx = build_ctx_with_worker(&worker.url); + let app = build_router(ctx.clone()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let url = format!("http://{addr}/v1/chat/completions"); + + let drain = Duration::from_millis(50); + let ctx_for_shutdown = ctx.clone(); + let (sigterm_tx, sigterm_rx) = oneshot::channel::<()>(); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = sigterm_rx.await; + sgl_router::server::shutdown::drain_for_termination( + &ctx_for_shutdown, + drain, + std::future::pending::<()>(), + ) + .await; + }) + .await + .unwrap(); + }); + + // Start one slow stream and hand back the response only once its headers + // have arrived — that is the point at which the request is provably + // in-flight, so SIGTERM below cannot race the client's connect. + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .build() + .unwrap(); + let body = serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + }); + let resp = client.post(&url).json(&body).send().await.unwrap(); + assert!( + resp.status().is_success(), + "stream started: {}", + resp.status() + ); + let inflight = tokio::spawn(async move { resp.bytes().await.unwrap() }); + + // Fire SIGTERM mid-stream: the drain must NOT truncate the in-flight stream. + sigterm_tx.send(()).unwrap(); + + let received = inflight.await.expect("client task joined"); + let body_str = String::from_utf8_lossy(&received); + assert!( + body_str.contains("data: [DONE]"), + "the in-flight stream must terminate with `data: [DONE]` through the drain path, got: {body_str}", + ); + tokio::time::timeout(Duration::from_secs(5), server) + .await + .expect("server resolves after in-flight stream drains") + .expect("server task joined cleanly"); +} + /// Poll until `inflight_http` settles on `want`, so the assertions below do not /// race the guard drop that happens on the server task after the client has /// already seen the last byte.