minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Router] Drain readiness before SIGTERM shutdown so k8s deregisters the pod first (#39016)

Browse Source 
Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
2026-09-16 09:59:26 -07:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 5
parent 279339f113
commit 3e03879f68
10 changed files with 1322 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/tests/proxy/graceful_shutdown.rs
+298 -2
View File
@@ -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.