[Router] Drain readiness before SIGTERM shutdown so k8s deregisters the pod first (#39016)
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
co-authored by
Kangyan Zhou
Claude Opus 5
Kangyan Zhou
Claude Opus 5
parent
279339f113
commit
3e03879f68
@@ -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::<std::time::Instant>);
|
||||
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<sgl_router::server::app_context::AppContext>,
|
||||
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<tokio::sync::oneshot::Sender<()>>,
|
||||
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);
|
||||
|
||||
@@ -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<TokenizerRegistry>,
|
||||
@@ -43,9 +52,10 @@ pub struct AppContext {
|
||||
/// [`crate::policies::kv_events::KvEventIndex::metrics_source`].
|
||||
pub kv_metrics: Option<KvIndexMetrics>,
|
||||
/// 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<InflightHttp>,
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,3 +8,4 @@ pub mod header_utils;
|
||||
pub mod inflight;
|
||||
pub mod metrics;
|
||||
pub mod routes;
|
||||
pub mod shutdown;
|
||||
|
||||
@@ -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<Output = ()>,
|
||||
) {
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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")
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user