[router] Add /flush_cache endpoint to experimental sgl-router (#27612)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
2026-06-08 18:53:29 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent db143e5212
commit 95090b837e
4 changed files with 409 additions and 0 deletions
@@ -53,5 +53,9 @@ pub fn build_router(ctx: Arc<AppContext>) -> Router {
.layer(DefaultBodyLimit::max(MAX_CHAT_BODY_BYTES)) .layer(DefaultBodyLimit::max(MAX_CHAT_BODY_BYTES))
.layer(middleware::from_fn(log_413)), .layer(middleware::from_fn(log_413)),
) )
.route(
"/flush_cache",
post(crate::server::routes::cache::flush_cache),
)
.with_state(ctx) .with_state(ctx)
} }
@@ -0,0 +1,379 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Cache-management admin endpoints.
use crate::server::app_context::AppContext;
use crate::workers::worker::Worker;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::Json;
use futures::stream::{self, StreamExt};
use reqwest::Client;
use serde::Serialize;
use std::sync::Arc;
use std::time::Duration;
/// Cap on concurrent in-flight `/flush_cache` requests. Bounds how many
/// flushes are issued at once when a large fleet is flushed; the rest queue
/// and run as slots free up.
const MAX_CONCURRENT_FLUSH: usize = 32;
/// One worker that failed to flush, with a human-readable reason.
#[derive(Serialize)]
pub struct FailedWorker {
pub worker: String,
pub error: String,
}
/// Per-worker breakdown of a `/flush_cache` fan-out. `total_workers` is the
/// registry size snapshotted at call time; every registered worker is
/// attempted, so `successful.len() + failed.len() == total_workers`.
/// `message` is a human/log summary — the HTTP status is authoritative.
#[derive(Serialize)]
pub struct FlushCacheResult {
pub successful: Vec<String>,
pub failed: Vec<FailedWorker>,
pub total_workers: usize,
pub message: String,
}
impl FlushCacheResult {
/// Build a result from a completed fan-out, deriving `message` from the
/// outcome counts so the count/message coherence lives in one place
/// rather than at each call site.
fn from_outcomes(
total_workers: usize,
successful: Vec<String>,
failed: Vec<FailedWorker>,
) -> Self {
let message = if total_workers == 0 {
"No workers registered; nothing to flush".to_string()
} else if failed.is_empty() {
format!(
"Successfully flushed cache on all {} workers",
total_workers
)
} else {
format!(
"Cache flush: {} succeeded, {} failed",
successful.len(),
failed.len()
)
};
Self {
successful,
failed,
total_workers,
message,
}
}
}
/// `POST /flush_cache` — fan SGLang's `/flush_cache` admin call out to every
/// registered worker and report a per-worker breakdown.
///
/// Targets the whole fleet (plain, prefill, and decode workers all hold KV
/// cache), not just one model's pool. Deliberately **bypasses the circuit
/// breaker**: an operator flushing caches wants every worker hit — including
/// ones whose breaker is open — and recording breaker success/failure for an
/// out-of-band admin call would skew the state the request router uses to
/// pick workers.
///
/// Status: `200 OK` when every worker flushed successfully (or the fleet is
/// empty); `502 BAD_GATEWAY` when at least one worker failed. The JSON body
/// always carries the full breakdown so a partial failure is actionable.
pub async fn flush_cache(State(ctx): State<Arc<AppContext>>) -> Response {
let workers = ctx.registry.all();
let total_workers = workers.len();
if workers.is_empty() {
// A flush against an empty fleet is a no-op, but it usually means a
// discovery/config problem (the router knows of no workers), so warn
// rather than stay silent.
tracing::warn!("flush_cache called but no workers are registered");
return (
StatusCode::OK,
Json(FlushCacheResult::from_outcomes(0, Vec::new(), Vec::new())),
)
.into_response();
}
let (successful, failed) =
fan_out_flush(&workers, &ctx.proxy.client, ctx.proxy.request_timeout).await;
// Partial failure is an operational event an operator needs to see at the
// common production log level — match the rest of the router, which warns
// on upstream failures.
if failed.is_empty() {
tracing::info!(total_workers, "flush_cache: all workers flushed");
} else {
tracing::warn!(
total_workers,
succeeded = successful.len(),
failed = failed.len(),
"flush_cache: some workers failed to flush",
);
}
let status = if failed.is_empty() {
StatusCode::OK
} else {
StatusCode::BAD_GATEWAY
};
(
status,
Json(FlushCacheResult::from_outcomes(
total_workers,
successful,
failed,
)),
)
.into_response()
}
/// POST `/flush_cache` to each worker concurrently (bounded by
/// [`MAX_CONCURRENT_FLUSH`]) and partition the outcomes into
/// (successful URLs, failed workers). A non-2xx status or a transport
/// error both count as failures.
async fn fan_out_flush(
workers: &[Arc<Worker>],
client: &Client,
timeout: Duration,
) -> (Vec<String>, Vec<FailedWorker>) {
// Snapshot the URLs into owned Strings up front so the per-worker stream
// does not borrow the `workers` slice across the await points.
let urls: Vec<String> = workers.iter().map(|w| w.url.clone()).collect();
let outcomes = stream::iter(urls)
.map(|url| {
let client = client.clone();
async move {
let flush_url = format!("{}/flush_cache", url.trim_end_matches('/'));
let result = client.post(&flush_url).timeout(timeout).send().await;
(url, result)
}
})
.buffer_unordered(MAX_CONCURRENT_FLUSH)
.collect::<Vec<_>>()
.await;
let mut successful = Vec::new();
let mut failed = Vec::new();
for (url, result) in outcomes {
match result {
Ok(resp) if resp.status().is_success() => successful.push(url),
Ok(resp) => failed.push(FailedWorker {
worker: url,
error: format!("HTTP {}", resp.status()),
}),
// Render the full source chain (`{:#}`), not just reqwest's outer
// message, so a connect-refused / DNS / TLS / timeout cause is
// visible in the per-worker error rather than collapsed away.
Err(e) => failed.push(FailedWorker {
worker: url,
error: format!("{:#}", anyhow::Error::new(e)),
}),
}
}
(successful, failed)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use crate::server::app_context::AppContext;
use axum::body::Body;
use axum::http::Request;
use axum::routing::post;
use axum::Router;
use http_body_util::BodyExt;
use serde_json::Value;
use tokio::net::TcpListener;
use tokio::sync::oneshot;
use tower::ServiceExt;
/// Spawn a fake worker that answers `POST /flush_cache` with `status`.
/// Returns its base URL and a shutdown handle (drop or send to stop).
async fn spawn_fake_flush_worker(status: StatusCode) -> (String, oneshot::Sender<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let app = Router::new().route("/flush_cache", post(move || async move { status }));
let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move {
let _ = axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = rx.await;
})
.await;
});
(format!("http://127.0.0.1:{port}"), tx)
}
/// Reserve a port then drop the listener so a connect attempt fails fast
/// with ConnectionRefused (no waiting on the connect timeout).
fn unused_port() -> u16 {
use std::net::TcpListener;
let l = TcpListener::bind("127.0.0.1:0").unwrap();
l.local_addr().unwrap().port()
}
fn ctx_with_workers(urls: &[&str]) -> Arc<AppContext> {
let ctx = AppContext::stub();
for (i, url) in urls.iter().enumerate() {
ctx.registry
.add(WorkerSpec {
id: WorkerId(format!("w-{i}")),
url: (*url).to_string(),
mode: WorkerMode::Plain,
model_ids: vec![ModelId("stub-model".into())],
bootstrap_port: None,
})
.expect("worker accepted");
}
Arc::new(ctx)
}
async fn post_flush(ctx: Arc<AppContext>) -> (StatusCode, Value) {
let app = crate::server::app::build_router(ctx);
let res = app
.oneshot(
Request::builder()
.method("POST")
.uri("/flush_cache")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status = res.status();
let bytes = res.into_body().collect().await.unwrap().to_bytes();
let body: Value = serde_json::from_slice(&bytes).unwrap();
(status, body)
}
#[tokio::test]
async fn all_workers_succeed_returns_200() {
let (u1, _s1) = spawn_fake_flush_worker(StatusCode::OK).await;
let (u2, _s2) = spawn_fake_flush_worker(StatusCode::OK).await;
let (status, body) = post_flush(ctx_with_workers(&[&u1, &u2])).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["total_workers"], 2);
assert_eq!(body["successful"].as_array().unwrap().len(), 2);
assert!(body["failed"].as_array().unwrap().is_empty());
}
#[tokio::test]
async fn partial_failure_returns_502_with_breakdown() {
let (ok_url, _s1) = spawn_fake_flush_worker(StatusCode::OK).await;
let (err_url, _s2) = spawn_fake_flush_worker(StatusCode::INTERNAL_SERVER_ERROR).await;
let (status, body) = post_flush(ctx_with_workers(&[&ok_url, &err_url])).await;
assert_eq!(status, StatusCode::BAD_GATEWAY);
assert_eq!(body["total_workers"], 2);
assert_eq!(
body["successful"].as_array().unwrap(),
&vec![Value::String(ok_url.clone())]
);
let failed = body["failed"].as_array().unwrap();
assert_eq!(failed.len(), 1);
assert_eq!(failed[0]["worker"], err_url);
assert!(failed[0]["error"].as_str().unwrap().contains("500"));
}
#[tokio::test]
async fn empty_registry_returns_200_with_zero_workers() {
let (status, body) = post_flush(Arc::new(AppContext::stub())).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["total_workers"], 0);
assert!(body["successful"].as_array().unwrap().is_empty());
assert!(body["failed"].as_array().unwrap().is_empty());
}
#[tokio::test]
async fn unreachable_worker_is_reported_failed() {
let url = format!("http://127.0.0.1:{}", unused_port());
let (status, body) = post_flush(ctx_with_workers(&[&url])).await;
assert_eq!(status, StatusCode::BAD_GATEWAY);
let failed = body["failed"].as_array().unwrap();
assert_eq!(failed.len(), 1);
assert_eq!(failed[0]["worker"], url);
}
/// A non-5xx, non-2xx status (e.g. 404) is still a failure and still
/// drives the top-level 502, with the status echoed in the error.
#[tokio::test]
async fn non_5xx_error_status_is_reported_failed() {
let (ok_url, _s1) = spawn_fake_flush_worker(StatusCode::OK).await;
let (nf_url, _s2) = spawn_fake_flush_worker(StatusCode::NOT_FOUND).await;
let (status, body) = post_flush(ctx_with_workers(&[&ok_url, &nf_url])).await;
assert_eq!(status, StatusCode::BAD_GATEWAY);
let failed = body["failed"].as_array().unwrap();
assert_eq!(failed.len(), 1);
assert_eq!(failed[0]["worker"], nf_url);
assert!(failed[0]["error"].as_str().unwrap().contains("404"));
}
/// A worker URL with a trailing slash must still resolve to
/// `<url>/flush_cache` (not `<url>//flush_cache`). Guards the
/// `trim_end_matches('/')` in `fan_out_flush` against a regression that
/// would 404 every slash-suffixed worker.
#[tokio::test]
async fn worker_url_with_trailing_slash_is_flushed() {
let (base, _s) = spawn_fake_flush_worker(StatusCode::OK).await;
let url = format!("{base}/");
let (status, body) = post_flush(ctx_with_workers(&[&url])).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(
body["successful"].as_array().unwrap(),
&vec![Value::String(url.clone())]
);
assert!(body["failed"].as_array().unwrap().is_empty());
}
/// The fan-out targets the whole fleet, not one model's pool: prefill
/// and decode workers (which also hold KV cache) must both be flushed.
/// Asserted through the handler — `registry::all()` returning mixed modes
/// is necessary but not sufficient if a mode filter ever slips into the
/// handler path.
#[tokio::test]
async fn flushes_prefill_and_decode_workers() {
let (p_url, _s1) = spawn_fake_flush_worker(StatusCode::OK).await;
let (d_url, _s2) = spawn_fake_flush_worker(StatusCode::OK).await;
let ctx = AppContext::stub();
ctx.registry
.add(WorkerSpec {
id: WorkerId("p".into()),
url: p_url.clone(),
mode: WorkerMode::Prefill,
model_ids: vec![ModelId("stub-model".into())],
bootstrap_port: Some(8998),
})
.expect("prefill accepted");
ctx.registry
.add(WorkerSpec {
id: WorkerId("d".into()),
url: d_url.clone(),
mode: WorkerMode::Decode,
model_ids: vec![ModelId("stub-model".into())],
bootstrap_port: None,
})
.expect("decode accepted");
let (status, body) = post_flush(Arc::new(ctx)).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["total_workers"], 2);
let mut succeeded: Vec<&str> = body["successful"]
.as_array()
.unwrap()
.iter()
.map(|v| v.as_str().unwrap())
.collect();
succeeded.sort_unstable();
let mut expected = [p_url.as_str(), d_url.as_str()];
expected.sort_unstable();
assert_eq!(succeeded, expected);
}
}
@@ -1,6 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0 // SPDX-License-Identifier: Apache-2.0
pub mod cache;
pub mod chat; pub mod chat;
pub mod health; pub mod health;
pub mod metrics; pub mod metrics;
@@ -196,6 +196,15 @@ impl WorkerRegistry {
pub fn get(&self, id: &WorkerId) -> Option<Arc<Worker>> { pub fn get(&self, id: &WorkerId) -> Option<Arc<Worker>> {
self.by_id.get(id).map(|w| Arc::clone(&w)) self.by_id.get(id).map(|w| Arc::clone(&w))
} }
/// Snapshot of every registered worker, across all models and modes.
///
/// Used by fleet-wide admin fan-out (e.g. `/flush_cache`) that targets
/// every worker the router knows about rather than one model's pool.
/// Order is unspecified (iterates the underlying `DashMap`).
pub fn all(&self) -> Vec<Arc<Worker>> {
self.by_id.iter().map(|e| Arc::clone(e.value())).collect()
}
} }
/// `true` when the two modes can't coexist for the same model — i.e. /// `true` when the two modes can't coexist for the same model — i.e.
@@ -244,6 +253,22 @@ mod tests {
assert!(m_missing.is_empty()); assert!(m_missing.is_empty());
} }
#[test]
fn all_returns_every_worker_across_models_and_modes() {
let r = WorkerRegistry::default();
let _ = r.add(spec("w1", WorkerMode::Plain, &["m1"]));
let _ = r.add(spec("p", WorkerMode::Prefill, &["m2"]));
let _ = r.add(spec("d", WorkerMode::Decode, &["m2"]));
let mut ids: Vec<String> = r.all().into_iter().map(|w| w.id.0.clone()).collect();
ids.sort();
assert_eq!(ids, vec!["d", "p", "w1"]);
}
#[test]
fn all_is_empty_for_fresh_registry() {
assert!(WorkerRegistry::default().all().is_empty());
}
#[test] #[test]
fn remove_drops_from_all_models() { fn remove_drops_from_all_models() {
let r = WorkerRegistry::default(); let r = WorkerRegistry::default();