From 95090b837e0ef0190679345330babe2248f29778 Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Mon, 8 Jun 2026 18:53:29 -0700 Subject: [PATCH] [router] Add /flush_cache endpoint to experimental sgl-router (#27612) Co-authored-by: Claude Opus 4.8 --- experimental/sgl-router/src/server/app.rs | 4 + .../sgl-router/src/server/routes/cache.rs | 379 ++++++++++++++++++ .../sgl-router/src/server/routes/mod.rs | 1 + .../sgl-router/src/workers/registry.rs | 25 ++ 4 files changed, 409 insertions(+) create mode 100644 experimental/sgl-router/src/server/routes/cache.rs diff --git a/experimental/sgl-router/src/server/app.rs b/experimental/sgl-router/src/server/app.rs index c727762c5..4412809db 100644 --- a/experimental/sgl-router/src/server/app.rs +++ b/experimental/sgl-router/src/server/app.rs @@ -53,5 +53,9 @@ pub fn build_router(ctx: Arc) -> Router { .layer(DefaultBodyLimit::max(MAX_CHAT_BODY_BYTES)) .layer(middleware::from_fn(log_413)), ) + .route( + "/flush_cache", + post(crate::server::routes::cache::flush_cache), + ) .with_state(ctx) } diff --git a/experimental/sgl-router/src/server/routes/cache.rs b/experimental/sgl-router/src/server/routes/cache.rs new file mode 100644 index 000000000..ef136533f --- /dev/null +++ b/experimental/sgl-router/src/server/routes/cache.rs @@ -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, + pub failed: Vec, + 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, + failed: Vec, + ) -> 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>) -> 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], + client: &Client, + timeout: Duration, +) -> (Vec, Vec) { + // 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 = 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::>() + .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 { + 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) -> (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 + /// `/flush_cache` (not `//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); + } +} diff --git a/experimental/sgl-router/src/server/routes/mod.rs b/experimental/sgl-router/src/server/routes/mod.rs index c5d1845f4..8ffdba687 100644 --- a/experimental/sgl-router/src/server/routes/mod.rs +++ b/experimental/sgl-router/src/server/routes/mod.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 +pub mod cache; pub mod chat; pub mod health; pub mod metrics; diff --git a/experimental/sgl-router/src/workers/registry.rs b/experimental/sgl-router/src/workers/registry.rs index 11c9d360d..aa51c89bc 100644 --- a/experimental/sgl-router/src/workers/registry.rs +++ b/experimental/sgl-router/src/workers/registry.rs @@ -196,6 +196,15 @@ impl WorkerRegistry { pub fn get(&self, id: &WorkerId) -> Option> { 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> { + 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. @@ -244,6 +253,22 @@ mod tests { 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 = 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] fn remove_drops_from_all_models() { let r = WorkerRegistry::default();