From 2032f3a0716f641bf1e88019adf8dc223052cee3 Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Mon, 21 Sep 2026 21:35:56 -0700 Subject: [PATCH] [Router] Abort the engine when a client disconnects mid-request (#39461) Co-authored-by: Kangyan Zhou Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Shangming Cai Co-authored-by: Kan Wu --- experimental/sgl-router/src/proxy/abort.rs | 47 ++++ experimental/sgl-router/src/proxy/mod.rs | 24 ++ experimental/sgl-router/src/server/app.rs | 9 + experimental/sgl-router/src/server/metrics.rs | 2 + .../src/server/routes/chat/forward.rs | 17 +- .../src/server/routes/chat/preparation.rs | 108 ++++++++- .../tests/proxy/cache_aware_input_ids.rs | 13 +- .../sgl-router/tests/proxy/chat_routing.rs | 140 ++++++++---- .../tests/proxy/chat_routing/cancellation.rs | 205 ++++++++++++++++++ .../tests/proxy/chat_routing/reorg.rs | 15 +- .../tests/proxy/common/mock_worker.rs | 27 +++ .../sgl-router/tests/proxy/common/mod.rs | 8 + .../sgl-router/tests/proxy/h2c_forward.rs | 3 + .../tests/proxy/pd_bootstrap_injection.rs | 58 +++++ .../tests/proxy/roundrobin_input_ids.rs | 18 +- 15 files changed, 629 insertions(+), 65 deletions(-) create mode 100644 experimental/sgl-router/src/proxy/abort.rs create mode 100644 experimental/sgl-router/tests/proxy/chat_routing/cancellation.rs diff --git a/experimental/sgl-router/src/proxy/abort.rs b/experimental/sgl-router/src/proxy/abort.rs new file mode 100644 index 000000000..351aec4ac --- /dev/null +++ b/experimental/sgl-router/src/proxy/abort.rs @@ -0,0 +1,47 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use axum::http::{header::AUTHORIZATION, HeaderMap}; +use reqwest::{Client, RequestBuilder, Url}; +use std::time::Duration; + +/// Cancels unfinished engine work without delaying request cleanup. +pub(super) struct AbortOnDrop(Option); + +impl AbortOnDrop { + // Only router-minted IDs are safe: the engine aborts by prefix. + pub(super) fn new( + client: &Client, + worker: &Url, + headers: &HeaderMap, + rid: Option<&str>, + ) -> Self { + Self(rid.filter(|rid| !rid.is_empty()).map(|rid| { + let mut request = client + .post(worker.join("/abort_request").expect("validated worker URL")) + .json(&serde_json::json!({"rid": rid, "abort_all": false})) + .timeout(Duration::from_secs(5)); + if let Some(auth) = headers.get(AUTHORIZATION) { + request = request.header(AUTHORIZATION, auth); + } + request + })) + } + + pub(super) fn disarm(&mut self) { + self.0 = None; + } +} + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + let Some(request) = self.0.take() else { return }; + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + if let Err(error) = request.send().await.and_then(|r| r.error_for_status()) { + tracing::warn!(%error, "engine abort failed"); + } + }); + } + } +} diff --git a/experimental/sgl-router/src/proxy/mod.rs b/experimental/sgl-router/src/proxy/mod.rs index e23d9063f..daa71ca09 100644 --- a/experimental/sgl-router/src/proxy/mod.rs +++ b/experimental/sgl-router/src/proxy/mod.rs @@ -3,8 +3,11 @@ //! HTTP proxy — forwards requests to the upstream SGLang worker. +mod abort; pub mod sse; +use abort::AbortOnDrop; + use crate::health::circuit_breaker::CircuitBreaker; use crate::server::error::ApiError; use crate::server::header_utils::should_forward_request_header; @@ -203,6 +206,7 @@ impl Proxy { /// path concatenation (no double-slash) and pass a typed URL to the /// split error variants (`UpstreamUnreachable` / `UpstreamTimeout` / /// `UpstreamStatus`). + #[allow(clippy::too_many_arguments)] pub async fn forward_json_to( &self, worker_url: &str, @@ -211,6 +215,7 @@ impl Proxy { path: &str, headers: &HeaderMap, body: Bytes, + abort_rid: Option<&str>, ) -> Result, ApiError> { let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen { worker: worker_url.to_string(), @@ -228,6 +233,8 @@ impl Proxy { req = req .header("content-type", "application/json") .timeout(self.request_timeout); + let mut abort = + AbortOnDrop::new(self.client_for(protocol), &worker_url, headers, abort_rid); let resp = req.send().await.map_err(|e| { breaker.record_failure(); Self::classify_reqwest_error_for(worker_url.clone(), e, path) @@ -257,6 +264,7 @@ impl Proxy { return Err(ApiError::UpstreamStatus { status }); } }; + abort.disarm(); match breaker_outcome(status) { BreakerOutcome::Failure => breaker.record_failure(), BreakerOutcome::Success => breaker.record_success(), @@ -301,6 +309,7 @@ impl Proxy { path: &str, headers: &HeaderMap, body: Bytes, + abort_rid: Option<&str>, stream_guards: Option>, on_first_byte: Option>, on_stream_end: Option>, @@ -322,11 +331,16 @@ impl Proxy { req = req .header("content-type", "application/json") .header("accept", "text/event-stream"); + let mut abort = + AbortOnDrop::new(self.client_for(protocol), &worker_url, headers, abort_rid); let resp = req.send().await.map_err(|e| { breaker.record_failure(); Self::classify_reqwest_error_for(worker_url.clone(), e, path) })?; let status = resp.status(); + if !status.is_success() { + abort.disarm(); + } let upstream_ct = resp .headers() .get(reqwest::header::CONTENT_TYPE) @@ -366,6 +380,9 @@ impl Proxy { BreakerOutcome::Success => { let breaker_for_hook = Arc::clone(breaker); Some(Box::new(move |end| { + if end.reason == sse::StreamEndReason::Completed { + abort.disarm(); + } match stream_breaker_outcome(end) { BreakerOutcome::Success => breaker_for_hook.record_success(), BreakerOutcome::Failure => breaker_for_hook.record_failure(), @@ -465,6 +482,7 @@ mod tests { "/chat", &headers, Bytes::new(), + None, ) .now_or_never() .is_none()); @@ -481,6 +499,7 @@ mod tests { None, None, None, + None, ) .now_or_never() .is_none()); @@ -583,6 +602,7 @@ mod tests { None, None, None, + None, expiration, ) .await @@ -694,6 +714,7 @@ mod tests { "/v1/chat/completions", &headers, Bytes::from_static(b"{}"), + None, ) .await .expect("dispatch should reach the worker (breaker must stay closed)"); @@ -737,6 +758,7 @@ mod tests { "/v1/chat/completions", &headers, Bytes::from_static(b"{}"), + None, ) .await; } @@ -780,6 +802,7 @@ mod tests { "/v1/chat/completions", &headers, Bytes::from_static(b"{}"), + None, ) .await .expect("the half-open probe must be admitted and reach the worker"); @@ -819,6 +842,7 @@ mod tests { None, None, None, + None, ) .await .expect("streaming dispatch should reach the worker"); diff --git a/experimental/sgl-router/src/server/app.rs b/experimental/sgl-router/src/server/app.rs index a5074bc02..428b23707 100644 --- a/experimental/sgl-router/src/server/app.rs +++ b/experimental/sgl-router/src/server/app.rs @@ -150,6 +150,9 @@ async fn access_log_and_record( .unwrap_or_else(|| outcome_from_status(status.as_u16())) .as_str(), worker = log_ctx.map(|c| c.worker_url.as_str()).unwrap_or(""), + engine_rid = log_ctx + .and_then(|c| c.engine_rid.as_deref()) + .unwrap_or(""), model = log_ctx.map(|c| c.model_id.as_str()).unwrap_or(""), stream = log_ctx.is_some_and(|c| c.streaming), latency_ms, @@ -446,6 +449,7 @@ mod tests { model_id: "tiny".into(), streaming: false, outcome: RequestOutcome::Cancelled, + engine_rid: Some("1f0c2b7a4e9d4f3ab6c5d8e7f0a1b2c3".into()), }); resp }), @@ -468,6 +472,11 @@ mod tests { logs.contains("worker=\"http://worker-a:30000\"") && logs.contains("model=\"tiny\""), "a routed request must be logged with its worker and model; captured:\n{logs}", ); + assert!( + logs.contains("engine_rid=\"1f0c2b7a4e9d4f3ab6c5d8e7f0a1b2c3\"") + && logs.contains("request_id="), + "the minted rid must be logged beside the caller's request id; captured:\n{logs}", + ); // The handler's outcome must win over the status-derived fallback — // otherwise the log and `worker_requests_total` can disagree about a // request the handler classified itself (here, a cancellation served diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index 9a7b00679..c1ed54a5b 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -223,6 +223,8 @@ pub struct RequestLogContext { /// line and `worker_requests_total` cannot disagree — the middleware can /// only see the status, which cannot express a router-side cancellation. pub outcome: RequestOutcome, + /// Router-minted engine ID, logged beside the caller's correlation ID. + pub engine_rid: Option, } /// Final outcome of a 2xx SSE stream. diff --git a/experimental/sgl-router/src/server/routes/chat/forward.rs b/experimental/sgl-router/src/server/routes/chat/forward.rs index 562310ab3..80325f066 100644 --- a/experimental/sgl-router/src/server/routes/chat/forward.rs +++ b/experimental/sgl-router/src/server/routes/chat/forward.rs @@ -81,7 +81,12 @@ pub(super) async fn forward_chat_request( }; (decode, bootstrap) }); - let body = request.into_outgoing_body(ctx, pd.as_ref().map(|(_, bootstrap)| bootstrap))?; + let engine_rid = request.engine_rid(pd.is_some()); + let body = request.into_outgoing_body( + ctx, + pd.as_ref().map(|(_, bootstrap)| bootstrap), + engine_rid.as_deref(), + )?; let prefill_load_guards = (worker_load_guard, active_request_guard); // In PD mode, prefill runs independently and decode supplies the client response. @@ -112,6 +117,7 @@ pub(super) async fn forward_chat_request( &response_worker, &headers, body, + engine_rid.as_deref(), response_load_guards, &metrics, expiration_token.clone(), @@ -124,7 +130,7 @@ pub(super) async fn forward_chat_request( model: metrics.model.clone(), }), }; - let log_context = metrics.record_dispatch_result(&result); + let log_context = metrics.record_dispatch_result(&result, engine_rid); // Materialize dispatch errors here so the access log retains the selected worker. let mut response = match result { Ok(mut response) => { @@ -171,6 +177,7 @@ fn spawn_prefill_request( CHAT_PATH, &headers, body, + None, ) .await { @@ -188,11 +195,13 @@ fn spawn_prefill_request( }); } +#[allow(clippy::too_many_arguments)] async fn forward_to_response_worker( ctx: &AppContext, worker: &Worker, headers: &HeaderMap, body: Bytes, + engine_rid: Option<&str>, load_guards: LoadGuards, metrics: &DispatchMetrics, expiration: CancellationToken, @@ -209,6 +218,7 @@ async fn forward_to_response_worker( CHAT_PATH, headers, body, + engine_rid, Some(stream_guards), Some(metrics.first_byte_callback()), Some(metrics.stream_end_callback(worker.url.clone())), @@ -226,6 +236,7 @@ async fn forward_to_response_worker( CHAT_PATH, headers, body, + engine_rid, ) .await } @@ -295,6 +306,7 @@ impl DispatchMetrics { fn record_dispatch_result( &self, result: &Result, ApiError>, + engine_rid: Option, ) -> RequestLogContext { let http_status = match result { Ok(response) => response.status().as_u16(), @@ -326,6 +338,7 @@ impl DispatchMetrics { model_id: self.model.clone(), streaming: self.streaming, outcome, + engine_rid, } } } diff --git a/experimental/sgl-router/src/server/routes/chat/preparation.rs b/experimental/sgl-router/src/server/routes/chat/preparation.rs index 8d9150087..47c4d3f3c 100644 --- a/experimental/sgl-router/src/server/routes/chat/preparation.rs +++ b/experimental/sgl-router/src/server/routes/chat/preparation.rs @@ -26,6 +26,8 @@ pub(super) struct PreparedChatRequest { pub(super) tokens: Option, /// Token count for routing/load accounting; estimated from body size when unavailable. pub(super) input_token_count: usize, + caller_set_rid: bool, + fans_out: bool, can_forward_input_ids: bool, parsed_body: Option, sampling_defaults: Vec<(SamplingField, Number)>, @@ -69,16 +71,27 @@ impl PreparedChatRequest { body, tokens, input_token_count, + caller_set_rid: fields.caller_set_rid, + fans_out: requests_multiple_samples(&fields, &sampling_defaults), can_forward_input_ids, parsed_body, sampling_defaults, }) } + pub(super) fn engine_rid(&self, pd_mode: bool) -> Option { + // Caller IDs are unsafe for prefix aborts; fan-out regenerates IDs; PD must finish KV transfer. + if self.caller_set_rid || self.fans_out || pd_mode { + return None; + } + Some(uuid::Uuid::new_v4().simple().to_string()) + } + pub(super) fn into_outgoing_body( self, ctx: &AppContext, bootstrap: Option<&BootstrapFields>, + engine_rid: Option<&str>, ) -> Result { // Routing tokens can replace engine tokenization only for supported chat templates. let input_ids = match (self.tokens.as_ref(), self.parsed_body.as_ref()) { @@ -104,6 +117,7 @@ impl PreparedChatRequest { input_ids, bootstrap, &self.sampling_defaults, + engine_rid, ) } } @@ -116,6 +130,8 @@ pub(super) struct RoutingFields { max_tokens: Option, max_completion_tokens: Option, sampling: [SamplingValue; SamplingField::ALL.len()], + // Preserve both string and list IDs without retaining their contents. + caller_set_rid: bool, } /// Null is absent; unrepresentable values are rejected only under a sampling contract. @@ -248,6 +264,7 @@ impl RoutingKey { enum RequestKey { Routing(RoutingKey), Sampling(SamplingField), + Rid, Other, } @@ -267,6 +284,7 @@ impl<'de> Deserialize<'de> for RequestKey { "model" => RequestKey::Routing(RoutingKey::Model), "max_tokens" => RequestKey::Routing(RoutingKey::MaxTokens), "max_completion_tokens" => RequestKey::Routing(RoutingKey::MaxCompletionTokens), + "rid" => RequestKey::Rid, other => match SamplingField::from_wire_name(other) { Some(field) => RequestKey::Sampling(field), None => RequestKey::Other, @@ -324,6 +342,9 @@ impl<'de> serde::de::Visitor<'de> for RoutingFieldsVisitor { SamplingValue::Unusable }; } + RequestKey::Rid => { + fields.caller_set_rid = map.next_value::>()?.is_some(); + } RequestKey::Other => { // Validate unrelated JSON without retaining its contents. map.next_value::()?; @@ -374,6 +395,22 @@ fn should_tokenize_request( can_forward_input_ids || policy_needs_request_tokens || bucket_routing_enabled } +// Use the effective n, including injected defaults; unreadable values opt out. +fn requests_multiple_samples( + fields: &RoutingFields, + sampling_defaults: &[(SamplingField, Number)], +) -> bool { + match fields.sampling_field(SamplingField::N) { + SamplingValue::Number(n) => n > 1.0, + SamplingValue::Unusable => true, + SamplingValue::Absent => sampling_defaults + .iter() + .find(|(field, _)| *field == SamplingField::N) + .and_then(|(_, value)| value.as_f64()) + .is_some_and(|n| n > 1.0), + } +} + fn estimate_prefill_tokens(body: &Bytes) -> usize { // Never 0: a zero-load entry is invisible to the cache-aware imbalance fast path. (body.len() / BYTES_PER_TOKEN_ESTIMATE).max(1) @@ -391,9 +428,10 @@ pub(super) struct BootstrapFields { } /// Append before the closing brace so injected values win over explicit nulls. -fn append_sampling_defaults( +fn append_top_level_fields( body: &Bytes, sampling_defaults: &[(SamplingField, Number)], + rid: Option<&str>, ) -> Option { use std::io::Write as _; @@ -405,13 +443,23 @@ fn append_sampling_defaults( let has_members = body[open + 1..close] .iter() .any(|b| !b.is_ascii_whitespace()); - let mut output = Vec::with_capacity(body.len() + 24 * sampling_defaults.len() + 1); + let rid_budget = rid.map_or(0, |rid| rid.len() + ",\"rid\":\"\"".len()); + let mut output = Vec::with_capacity(body.len() + 24 * sampling_defaults.len() + rid_budget + 1); output.extend_from_slice(&body[..close]); - for (i, (field, value)) in sampling_defaults.iter().enumerate() { - if has_members || i > 0 { + let mut wrote_any = has_members; + for (field, value) in sampling_defaults { + if wrote_any { output.push(b','); } write!(output, "\"{}\":{}", field.wire_name(), value).ok()?; + wrote_any = true; + } + if let Some(rid) = rid { + if wrote_any { + output.push(b','); + } + output.extend_from_slice(b"\"rid\":"); + serde_json::to_writer(&mut output, rid).ok()?; } output.extend_from_slice(&body[close..]); Some(Bytes::from(output)) @@ -424,14 +472,15 @@ fn build_outgoing_body( input_ids: Option<&[u32]>, bootstrap: Option<&BootstrapFields>, sampling_defaults: &[(SamplingField, Number)], + rid: Option<&str>, ) -> Result { - let sampling_only = input_ids.is_none() && bootstrap.is_none(); - if sampling_only && sampling_defaults.is_empty() { + let needs_parse = input_ids.is_some() || bootstrap.is_some(); + if !needs_parse && sampling_defaults.is_empty() && rid.is_none() { // Cloning Bytes shares the original allocation when no injection is needed. return Ok(body.clone()); } - if sampling_only { - if let Some(spliced) = append_sampling_defaults(body, sampling_defaults) { + if !needs_parse { + if let Some(spliced) = append_top_level_fields(body, sampling_defaults, rid) { return Ok(spliced); } } @@ -445,6 +494,9 @@ fn build_outgoing_body( return Err(invalid_request()); } }; + if let Some(rid) = rid { + body_fields.insert("rid".into(), Value::String(rid.to_owned())); + } for (field, default) in sampling_defaults { body_fields.insert(field.wire_name().into(), default.clone().into()); } @@ -739,7 +791,8 @@ mod tests { .unwrap() .extend(fields.as_object().unwrap().clone()); for value in [None, Some(original.clone())] { - let out = build_outgoing_body(&body, value, ids, bootstrap.as_ref(), &[]).unwrap(); + let out = + build_outgoing_body(&body, value, ids, bootstrap.as_ref(), &[], None).unwrap(); assert_eq!(serde_json::from_slice::(&out).unwrap(), expected); } } @@ -750,7 +803,7 @@ mod tests { for raw in [r#"{"model":"x"}"#, r#"{"model":"x","messages":[]}"#] { let body = Bytes::copy_from_slice(raw.as_bytes()); for value in [None, Some(serde_json::from_slice(&body).unwrap())] { - let out = build_outgoing_body(&body, value, None, None, &[]).unwrap(); + let out = build_outgoing_body(&body, value, None, None, &[], None).unwrap(); assert_eq!(out, body); assert_eq!(out.as_ptr(), body.as_ptr()); } @@ -1108,7 +1161,7 @@ mod tests { &metrics(), ) .unwrap(); - let out = build_outgoing_body(&body, None, Some(&[1, 2, 3]), None, &inject).unwrap(); + let out = build_outgoing_body(&body, None, Some(&[1, 2, 3]), None, &inject, None).unwrap(); assert_eq!( serde_json::from_slice::(&out).unwrap(), json!({ @@ -1240,7 +1293,7 @@ mod tests { let inject = resolve_sampling_defaults(&config, &fields_of(raw), &metrics()).unwrap(); assert_eq!(inject.len(), 2); for value in [None, Some(serde_json::from_slice(&body).unwrap())] { - let out = build_outgoing_body(&body, value, None, None, &inject).unwrap(); + let out = build_outgoing_body(&body, value, None, None, &inject, None).unwrap(); assert_eq!(std::str::from_utf8(&out).unwrap(), expected, "{raw}"); let parsed: Value = serde_json::from_slice(&out).unwrap(); assert_eq!(parsed["temperature"], json!(1.0)); @@ -1464,4 +1517,35 @@ mod tests { "a value at the cap is still read" ); } + + #[test] + fn abort_opt_outs_follow_caller_rid_and_effective_sample_count() { + for raw in [r#"{"rid":"abc"}"#, r#"{"rid":["a","b"]}"#] { + assert!(fields_of(raw).caller_set_rid); + } + assert!(!fields_of(r#"{"rid":null}"#).caller_set_rid); + for (raw, fan_out) in [ + (r#"{}"#, false), + (r#"{"n":1}"#, false), + (r#"{"n":2}"#, true), + (r#"{"n":"3"}"#, true), + (r#"{"n":[2]}"#, true), + ] { + assert_eq!( + requests_multiple_samples(&fields_of(raw), &[]), + fan_out, + "{raw}" + ); + } + for (config, fan_out) in [(r#"{"n":1}"#, false), (r#"{"n":4}"#, true)] { + let fields = fields_of("{}"); + let defaults = resolve_sampling_defaults( + &overrides_of(ConflictPolicy::Reject, config), + &fields, + &metrics(), + ) + .unwrap(); + assert_eq!(requests_multiple_samples(&fields, &defaults), fan_out); + } + } } diff --git a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs index 9e21e044b..1131c1863 100644 --- a/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs @@ -198,7 +198,18 @@ async fn caller_input_ids_are_used_for_routing_and_preserved() { send(Arc::clone(&ctx), request.clone()).await, StatusCode::OK ); - assert_eq!(captured(&mock), request, "body must be forwarded untouched"); + let mut forwarded = captured(&mock); + let rid = forwarded + .as_object_mut() + .expect("a forwarded chat body is an object") + .remove("rid"); + assert!( + rid.as_ref() + .and_then(Value::as_str) + .is_some_and(crate::common::is_engine_shaped_rid), + "plain mode must mint an abort rid; got {rid:?}", + ); + assert_eq!(forwarded, request, "body must be forwarded untouched"); } // Bypasses are not rendering failures. assert!(!ctx diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 1f92bb685..0e251603b 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -17,10 +17,14 @@ use sgl_router::workers::{WireProtocol, Worker, WorkerRegistry}; use axum::body::Body; use axum::http::{Request, StatusCode}; use http_body_util::BodyExt; +use sgl_router::state::load_monitor::router_inflight_load::{ + spawn_janitor, JanitorHandle, RouterInflightLoadRegistry, +}; use std::sync::Arc; use std::time::Duration; use tower::ServiceExt; +mod cancellation; mod reorg; const TEST_TIMEOUT: Duration = Duration::from_secs(5); @@ -74,6 +78,36 @@ fn build_ctx_with_worker(url: &str) -> Arc { Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) } +/// Expire requests after 50ms; keep the janitor handle alive during the test. +fn build_ctx_with_janitor(url: &str) -> (Arc, JanitorHandle) { + let cfg = config_for(url); + let registry = Arc::new(WorkerRegistry::default()); + let _ = registry.add(WorkerSpec { + id: WorkerId("w1".into()), + url: url.to_string(), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + }); + let policies = Arc::new(build_policy_registry(&cfg).unwrap()); + let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); + let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap()); + let router_inflight_load = RouterInflightLoadRegistry::new( + Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock), + Duration::from_millis(50), + ); + let janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20)); + let ctx = Arc::new(AppContext::with_router_inflight_load( + cfg, + tokenizers, + proxy, + registry, + policies, + router_inflight_load, + )); + (ctx, janitor) +} + #[tokio::test] async fn non_streaming_returns_200() { let worker = crate::common::mock_worker::MockWorker::start(vec![]).await; @@ -1033,6 +1067,7 @@ async fn forward_json_to_records_failure_on_body_drop() { "/v1/chat/completions", &headers, body, + None, ) .await; assert!(res.is_err(), "body drop should surface as ApiError"); @@ -1090,6 +1125,7 @@ async fn forward_json_to_records_success_only_after_body_completes() { "/v1/chat/completions", &headers, bytes::Bytes::from_static(b"{}"), + None, ) .await; assert!(res.is_ok(), "clean OK call must succeed: {res:?}"); @@ -1145,6 +1181,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() { None, None, None, + None, ) .await; @@ -1248,6 +1285,7 @@ async fn forward_json_to_records_failure_on_5xx() { "/v1/chat/completions", &headers, body, + None, ) .await; @@ -1282,6 +1320,7 @@ async fn forward_json_to_rejects_when_breaker_open() { "/v1/chat/completions", &headers, body, + None, ) .await; @@ -1320,6 +1359,7 @@ async fn forward_json_to_malformed_url_returns_worker_misconfigured_and_trips_br "/v1/chat/completions", &headers, body, + None, ) .await; @@ -1607,58 +1647,17 @@ async fn streaming_active_load_drops_on_client_disconnect() { /// returns; cancellation fires; handler returns 504. #[tokio::test] async fn janitor_expiry_returns_504_stale_request_expired() { - use sgl_router::state::load_monitor::router_inflight_load::{ - spawn_janitor, RouterInflightLoadRegistry, - }; - // Upstream that takes 2s to respond — longer than our 50ms - // stale_request_timeout. + // Upstream that takes 2s to respond — longer than the helper's 50ms + // stale_request_timeout, so the janitor sweeps before it answers. let worker = crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(2)).await; - - let cfg = config_for(&worker.url); - let registry = Arc::new(WorkerRegistry::default()); - let _ = registry.add(WorkerSpec { - id: WorkerId("w1".into()), - url: worker.url.clone(), - mode: WorkerMode::Plain, - model_ids: vec![ModelId("tiny".into())], - bootstrap_port: None, - }); - let policies = Arc::new(build_policy_registry(&cfg).unwrap()); - let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); - let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap()); - // Aggressive 50ms timeout: the janitor will sweep on the next - // tick (every 20ms) and fire the cancellation token before the - // upstream returns. - let router_inflight_load = RouterInflightLoadRegistry::new( - Arc::new(sgl_router::state::load_monitor::router_inflight_load::SystemTimeClock), - Duration::from_millis(50), - ); - let _janitor = spawn_janitor(Arc::clone(&router_inflight_load), Duration::from_millis(20)); - let ctx = Arc::new(AppContext::with_router_inflight_load( - cfg, - tokenizers, - proxy, - registry, - policies, - router_inflight_load, - )); + let (ctx, _janitor) = build_ctx_with_janitor(&worker.url); let app = build_router(ctx); - let req = Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header("content-type", "application/json") - .body(Body::from( - serde_json::to_vec(&serde_json::json!({ - "model": "tiny", - "messages": [{"role": "user", "content": "hi"}], - "stream": false, - })) - .unwrap(), - )) + let res = app + .oneshot(cancellation::request(serde_json::json!({}))) + .await .unwrap(); - let res = app.oneshot(req).await.unwrap(); assert_eq!( res.status(), StatusCode::GATEWAY_TIMEOUT, @@ -1677,6 +1676,53 @@ async fn janitor_expiry_returns_504_stale_request_expired() { body_str.contains("\"code\":\"stale_request_expired\""), "504 body must encode the same code in the JSON envelope: {body_str}", ); + + assert_engine_abort(&worker).await; +} + +#[tokio::test] +async fn janitor_expiry_aborts_before_headers_and_mid_stream() { + use crate::common::mock_worker::MockWorker; + for before_headers in [true, false] { + let worker = if before_headers { + MockWorker::start_hanging(Duration::from_secs(2)).await + } else { + MockWorker::start_slow_stream(vec!["data: a\n\n"], Duration::from_secs(2)).await + }; + let (ctx, _janitor) = build_ctx_with_janitor(&worker.url); + let response = build_router(ctx) + .oneshot(cancellation::request(serde_json::json!({"stream":true}))) + .await + .unwrap(); + assert_eq!( + response.status(), + if before_headers { + StatusCode::GATEWAY_TIMEOUT + } else { + StatusCode::OK + } + ); + let result = response.into_body().collect().await; + assert_eq!(result.is_ok(), before_headers); + assert_engine_abort(&worker).await; + } +} + +async fn assert_engine_abort(worker: &crate::common::mock_worker::MockWorker) { + tokio::time::timeout(TEST_TIMEOUT, async { + while worker.abort_log.lock().unwrap().is_empty() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + let forwarded: serde_json::Value = + serde_json::from_slice(worker.captured.lock().unwrap().last_body.as_ref().unwrap()) + .unwrap(); + assert_eq!( + *worker.abort_log.lock().unwrap(), + vec![serde_json::json!({"rid":forwarded["rid"], "abort_all":false})] + ); } /// Task A: a non-streaming request that errors out (upstream diff --git a/experimental/sgl-router/tests/proxy/chat_routing/cancellation.rs b/experimental/sgl-router/tests/proxy/chat_routing/cancellation.rs new file mode 100644 index 000000000..517cbef08 --- /dev/null +++ b/experimental/sgl-router/tests/proxy/chat_routing/cancellation.rs @@ -0,0 +1,205 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use axum::{extract::State, http::HeaderMap, routing::post, Json, Router}; +use serde_json::{json, Value}; +use tokio::sync::mpsc; + +type Event = (&'static str, Value); +type Events = mpsc::UnboundedSender; + +async fn chat(State(events): State, Json(body): Json) -> (StatusCode, Body) { + events.send(("chat", body.clone())).unwrap(); + if body["before_headers"] == true || (body["hold"] == true && body["stream"] != true) { + std::future::pending::<()>().await; + } + let response = if body["hold"] == true { + Body::from_stream(futures::stream::pending::< + Result, + >()) + } else if body["stream"] == true { + Body::from("data: [DONE]\n\n") + } else { + Body::from("{}") + }; + ( + StatusCode::from_u16(body["status"].as_u64().unwrap_or(200) as u16).unwrap(), + response, + ) +} + +async fn abort( + State(events): State, + headers: HeaderMap, + Json(body): Json, +) -> StatusCode { + assert_eq!(headers["authorization"], "Bearer test"); + events.send(("abort", body)).unwrap(); + StatusCode::INTERNAL_SERVER_ERROR // Abort failures must not affect the worker's breaker. +} + +struct Harness { + ctx: Arc, + events: mpsc::UnboundedReceiver, + server: tokio::task::JoinHandle<()>, +} + +impl Harness { + async fn new() -> Self { + let (events, rx) = mpsc::unbounded_channel(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let ctx = build_ctx_with_worker(&format!("http://{}", listener.local_addr().unwrap())); + let server = tokio::spawn(async move { + axum::serve( + listener, + Router::new() + .route("/v1/chat/completions", post(chat)) + .route("/abort_request", post(abort)) + .with_state(events), + ) + .await + .unwrap(); + }); + Self { + ctx, + events: rx, + server, + } + } + + async fn event(&mut self, expected: &str) -> Value { + let (kind, body) = tokio::time::timeout(TEST_TIMEOUT, self.events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(kind, expected); + body + } + + async fn quiet(&mut self) { + assert!( + tokio::time::timeout(Duration::from_millis(50), self.events.recv()) + .await + .is_err() + ); + } +} + +impl Drop for Harness { + fn drop(&mut self) { + self.server.abort(); + } +} + +pub(super) fn request(mut body: Value) -> Request { + body["model"] = json!("tiny"); + Request::post("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer test") + .header("x-request-id", "reused-gateway-id") + .body(Body::from(body.to_string())) + .unwrap() +} + +#[tokio::test] +async fn only_unfinished_requests_abort() { + let mut h = Harness::new().await; + for (stream, hold, before_headers) in [ + (false, false, false), + (true, false, false), + (false, true, false), + (true, true, true), + (true, true, false), + ] { + let task = tokio::spawn(build_router(h.ctx.clone()).oneshot(request(json!({ + "stream": stream, "hold": hold, "before_headers": before_headers, + })))); + let forwarded = h.event("chat").await; + assert!(crate::common::is_engine_shaped_rid( + forwarded["rid"].as_str().unwrap() + )); + let worker = h.ctx.registry.get(&WorkerId("w1".into())).unwrap(); + if hold && (!stream || before_headers) { + worker.breaker.record_failure(); + worker.breaker.record_failure(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + } else { + let response = tokio::time::timeout(TEST_TIMEOUT, task) + .await + .unwrap() + .unwrap() + .unwrap(); + if hold { + drop(response); // Silent upstream: cancellation must not wait for a token. + } else { + response.into_body().collect().await.unwrap(); + } + } + if hold { + assert_eq!( + h.event("abort").await, + json!({"rid": forwarded["rid"], "abort_all": false}) + ); + } + h.quiet().await; + assert_eq!(worker.breaker.snapshot().state_code, 0); + worker.breaker.record_success(); + } +} + +#[tokio::test] +async fn caller_ids_fan_out_and_rejected_streams_do_not_abort() { + let mut h = Harness::new().await; + for fields in [ + json!({"rid":"a"}), + json!({"rid":["a","b"]}), + json!({"n":2}), + json!({"status":400}), + json!({"status":429}), + json!({"status":500}), + json!({"status":503}), + ] { + let mut body = json!({"stream":true, "hold":true}); + body.as_object_mut() + .unwrap() + .extend(fields.as_object().unwrap().clone()); + let response = build_router(h.ctx.clone()) + .oneshot(request(body)) + .await + .unwrap(); + let forwarded = h.event("chat").await; + if fields.get("status").is_none() { + assert_eq!(forwarded.get("rid"), fields.get("rid")); + } + drop(response); + h.quiet().await; + } +} + +#[tokio::test] +async fn concurrent_requests_with_the_same_header_get_distinct_abort_ids() { + let mut h = Harness::new().await; + let mut tasks = Vec::new(); + for _ in 0..2 { + tasks.push(tokio::spawn( + build_router(h.ctx.clone()).oneshot(request(json!({"hold":true}))), + )); + } + let mut rids = Vec::new(); + for _ in 0..2 { + rids.push(h.event("chat").await["rid"].clone()); + } + assert_ne!(rids[0], rids[1]); + for task in tasks { + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + } + for _ in 0..2 { + let aborted = h.event("abort").await; + let index = rids.iter().position(|rid| rid == &aborted["rid"]).unwrap(); + rids.remove(index); + } + h.quiet().await; +} diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs index e1c4711be..2f28593cf 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs @@ -230,7 +230,19 @@ async fn length_selects_plain_bucket_before_engine_selection() { let response = app.clone().oneshot(request(body("hi"))).await.unwrap(); assert_eq!(response.status(), StatusCode::OK); let _ = response.into_body().collect().await.unwrap(); - assert!(short_worker.captured.lock().unwrap().last_body.is_some()); + let forwarded: serde_json::Value = serde_json::from_slice( + short_worker + .captured + .lock() + .unwrap() + .last_body + .as_ref() + .unwrap(), + ) + .unwrap(); + assert!(crate::common::is_engine_shaped_rid( + forwarded["rid"].as_str().unwrap() + )); assert!(long_worker.captured.lock().unwrap().last_body.is_none()); let response = app @@ -293,6 +305,7 @@ async fn pd_picks_both_groups_from_selected_bucket_and_shares_bootstrap() { let d: serde_json::Value = serde_json::from_slice(decode.captured.lock().unwrap().last_body.as_ref().unwrap()) .unwrap(); + assert!(p.get("rid").is_none() && d.get("rid").is_none()); assert!(p["bootstrap_room"].is_number()); assert_eq!(p["bootstrap_room"], d["bootstrap_room"]); let calls = policy.calls.lock().unwrap(); diff --git a/experimental/sgl-router/tests/proxy/common/mock_worker.rs b/experimental/sgl-router/tests/proxy/common/mock_worker.rs index 181937c29..00bf652d5 100644 --- a/experimental/sgl-router/tests/proxy/common/mock_worker.rs +++ b/experimental/sgl-router/tests/proxy/common/mock_worker.rs @@ -39,9 +39,22 @@ pub struct MockWorker { // Used in header_forwarding_test; not every test file reads captured headers. #[allow(dead_code)] pub captured: Arc>, + #[allow(dead_code)] + pub abort_log: Arc>>, _shutdown: oneshot::Sender<()>, } +#[allow(dead_code)] // shared across all axum variants +fn abort_request_route(log: Arc>>) -> axum::routing::MethodRouter +where + S: Clone + Send + Sync + 'static, +{ + post(move |Json(body): Json| async move { + log.lock().unwrap().push(body); + StatusCode::OK + }) +} + impl MockWorker { /// Bind to a random port on 127.0.0.1 and start serving. /// @@ -50,6 +63,7 @@ impl MockWorker { #[allow(dead_code)] // Only used by some test files. pub async fn start(stream_chunks: Vec<&'static str>) -> Self { let captured = Arc::new(Mutex::new(CapturedHeaders::default())); + let abort_log: Arc>> = Arc::new(Mutex::new(Vec::new())); let state = MockWorkerState { captured: captured.clone(), stream_chunks: Arc::new(stream_chunks), @@ -60,6 +74,7 @@ impl MockWorker { let app = axum::Router::new() .route("/v1/chat/completions", post(chat)) .route("/server_info", get(serve_tiny_server_info)) + .route("/abort_request", abort_request_route(abort_log.clone())) .with_state(state); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -77,6 +92,7 @@ impl MockWorker { Self { url, captured, + abort_log, _shutdown: tx, } } @@ -88,6 +104,7 @@ impl MockWorker { #[allow(dead_code)] pub async fn start_hanging(delay: Duration) -> Self { let captured = Arc::new(Mutex::new(CapturedHeaders::default())); + let abort_log: Arc>> = Arc::new(Mutex::new(Vec::new())); #[derive(Clone)] struct HangState { @@ -127,6 +144,7 @@ impl MockWorker { let app = axum::Router::new() .route("/v1/chat/completions", post(hang_handler)) .route("/server_info", get(serve_tiny_server_info)) + .route("/abort_request", abort_request_route(abort_log.clone())) .with_state(state); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -144,6 +162,7 @@ impl MockWorker { Self { url, captured, + abort_log, _shutdown: tx, } } @@ -154,6 +173,7 @@ impl MockWorker { #[allow(dead_code)] pub async fn start_slow_stream(chunks: Vec<&'static str>, delay: Duration) -> Self { let captured = Arc::new(Mutex::new(CapturedHeaders::default())); + let abort_log: Arc>> = Arc::new(Mutex::new(Vec::new())); #[derive(Clone)] struct SlowState { @@ -207,6 +227,7 @@ impl MockWorker { let app = axum::Router::new() .route("/v1/chat/completions", post(slow_chat)) .route("/server_info", get(serve_tiny_server_info)) + .route("/abort_request", abort_request_route(abort_log.clone())) .with_state(state); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -224,6 +245,7 @@ impl MockWorker { Self { url, captured, + abort_log, _shutdown: tx, } } @@ -248,6 +270,7 @@ impl MockWorker { partial_body_bytes: &'static [u8], ) -> Self { let captured = Arc::new(Mutex::new(CapturedHeaders::default())); + let abort_log: Arc>> = Arc::new(Mutex::new(Vec::new())); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr: SocketAddr = listener.local_addr().unwrap(); let url = format!("http://{addr}"); @@ -309,6 +332,7 @@ impl MockWorker { Self { url, captured, + abort_log, _shutdown: tx, } } @@ -319,6 +343,7 @@ impl MockWorker { #[allow(dead_code)] pub async fn start_returning_error(status: StatusCode, body: Value) -> Self { let captured = Arc::new(Mutex::new(CapturedHeaders::default())); + let abort_log: Arc>> = Arc::new(Mutex::new(Vec::new())); let body_arc = Arc::new(body.to_string()); #[derive(Clone)] @@ -360,6 +385,7 @@ impl MockWorker { let app = axum::Router::new() .route("/v1/chat/completions", post(error_handler)) .route("/server_info", get(serve_tiny_server_info)) + .route("/abort_request", abort_request_route(abort_log.clone())) .with_state(state); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -377,6 +403,7 @@ impl MockWorker { Self { url, captured, + abort_log, _shutdown: tx, } } diff --git a/experimental/sgl-router/tests/proxy/common/mod.rs b/experimental/sgl-router/tests/proxy/common/mod.rs index 8c63c5722..a7fcfe1d2 100644 --- a/experimental/sgl-router/tests/proxy/common/mod.rs +++ b/experimental/sgl-router/tests/proxy/common/mod.rs @@ -6,3 +6,11 @@ pub mod cache_aware_fixture; pub mod mock_worker; pub mod streaming; + +#[allow(dead_code)] // not every test file inspects forwarded rids +pub fn is_engine_shaped_rid(rid: &str) -> bool { + rid.len() == 32 + && rid + .bytes() + .all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase()) +} diff --git a/experimental/sgl-router/tests/proxy/h2c_forward.rs b/experimental/sgl-router/tests/proxy/h2c_forward.rs index 983fce2b4..cb7eaa0a5 100644 --- a/experimental/sgl-router/tests/proxy/h2c_forward.rs +++ b/experimental/sgl-router/tests/proxy/h2c_forward.rs @@ -72,6 +72,7 @@ async fn h2c_client_reaches_http2_only_worker() { "/v1/chat/completions", &axum::http::HeaderMap::new(), Bytes::from_static(b"{}"), + None, ) .await .expect("h2c client must reach an HTTP/2-only worker"); @@ -96,6 +97,7 @@ async fn http1_client_cannot_reach_http2_only_worker() { "/v1/chat/completions", &axum::http::HeaderMap::new(), Bytes::from_static(b"{}"), + None, ) .await; assert!( @@ -164,6 +166,7 @@ async fn h2c_client_streams_sse_from_http2_only_worker() { &axum::http::HeaderMap::new(), Bytes::from_static(b"{}"), None, + None, Some(Box::new(move || { flag.store(true, std::sync::atomic::Ordering::SeqCst); })), diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs index fcbc8738f..82eb3ab77 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -376,3 +376,61 @@ async fn pd_mode_prefill_5xx_does_not_poison_decode_response() { let pv = parse_body(&prefill_body); assert_eq!(bootstrap_port(&pv), Some(8997)); } + +fn streaming_chat_request() -> Request { + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + serde_json::to_vec(&json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + })) + .unwrap(), + )) + .unwrap() +} + +#[tokio::test] +async fn pd_mode_disconnect_does_not_abort_either_worker() { + let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await; + let decode = crate::common::mock_worker::MockWorker::start_slow_stream( + vec!["data: a\n\n", "data: b\n\n", "data: c\n\n"], + Duration::from_millis(50), + ) + .await; + let ctx = build_ctx(vec![ + WorkerSpec { + id: WorkerId("p1".into()), + url: prefill.url.clone(), + mode: WorkerMode::Prefill, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: Some(8997), + }, + WorkerSpec { + id: WorkerId("d1".into()), + url: decode.url.clone(), + mode: WorkerMode::Decode, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + }, + ]); + let app = build_router(ctx); + + let res = app.oneshot(streaming_chat_request()).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + + use futures::StreamExt; + let mut data_stream = res.into_body().into_data_stream(); + assert!(data_stream.next().await.is_some()); + drop(data_stream); + + tokio::time::sleep(Duration::from_millis(100)).await; + for worker in [&prefill, &decode] { + assert!(worker.abort_log.lock().unwrap().is_empty()); + let body = await_captured_body(worker, Duration::from_secs(2), "PD worker").await; + assert!(parse_body(&body).get("rid").is_none()); + } +} diff --git a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs index 729061421..fd66bad25 100644 --- a/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs +++ b/experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs @@ -106,9 +106,23 @@ fn without_forwarding(mut cfg: Config, policy: PolicyKind) -> Config { cfg } +fn without_minted_rid(mut body: Value) -> Value { + let rid = body + .as_object_mut() + .expect("a forwarded chat body is an object") + .remove("rid"); + assert!( + rid.as_ref() + .and_then(Value::as_str) + .is_some_and(crate::common::is_engine_shaped_rid), + "plain mode must mint an abort rid; got {rid:?}", + ); + body +} + async fn assert_forwarded_unchanged(ctx: &Arc, mock: &MockWorker, request: &Value) { assert_eq!(send(Arc::clone(ctx), request.clone()).await, StatusCode::OK); - assert_eq!(captured(mock), *request); + assert_eq!(without_minted_rid(captured(mock)), *request); assert!(!ctx .metrics .render() @@ -428,6 +442,6 @@ async fn kimi_ids_forward_with_engine_rendering_fallback() { } else { assert!(ids.is_none()); } - assert_eq!(captured(&mock), request); + assert_eq!(without_minted_rid(captured(&mock)), request); } }