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

[Router] Abort the engine when a client disconnects mid-request (#39461)

Browse Source 
Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Kan Wu <wukanustc@gmail.com>
This commit is contained in:
Kangyan-Zhou
2026-09-22 12:35:56 +08:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 5 Shangming Cai Kan Wu
parent a9f02b0fa4
commit 2032f3a071
15 changed files with 629 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/src/proxy/abort.rs
+47
View File
@@ -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<RequestBuilder>);
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");
}
});
}
}
}
experimental/sgl-router/src/proxy/mod.rs
+24
View File
@@ -3,8 +3,11 @@
//! HTTP proxy — forwards requests to the upstream SGLang worker. //! HTTP proxy — forwards requests to the upstream SGLang worker.
mod abort;
pub mod sse; pub mod sse;
use abort::AbortOnDrop;
use crate::health::circuit_breaker::CircuitBreaker; use crate::health::circuit_breaker::CircuitBreaker;
use crate::server::error::ApiError; use crate::server::error::ApiError;
use crate::server::header_utils::should_forward_request_header; 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 /// path concatenation (no double-slash) and pass a typed URL to the
/// split error variants (`UpstreamUnreachable` / `UpstreamTimeout` / /// split error variants (`UpstreamUnreachable` / `UpstreamTimeout` /
/// `UpstreamStatus`). /// `UpstreamStatus`).
#[allow(clippy::too_many_arguments)]
pub async fn forward_json_to( pub async fn forward_json_to(
&self, &self,
worker_url: &str, worker_url: &str,
@@ -211,6 +215,7 @@ impl Proxy {
path: &str, path: &str,
headers: &HeaderMap, headers: &HeaderMap,
body: Bytes, body: Bytes,
abort_rid: Option<&str>,
) -> Result<Response<Body>, ApiError> { ) -> Result<Response<Body>, ApiError> {
let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen { let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen {
worker: worker_url.to_string(), worker: worker_url.to_string(),
@@ -228,6 +233,8 @@ impl Proxy {
req = req req = req
.header("content-type", "application/json") .header("content-type", "application/json")
.timeout(self.request_timeout); .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| { let resp = req.send().await.map_err(|e| {
breaker.record_failure(); breaker.record_failure();
Self::classify_reqwest_error_for(worker_url.clone(), e, path) Self::classify_reqwest_error_for(worker_url.clone(), e, path)
@@ -257,6 +264,7 @@ impl Proxy {
return Err(ApiError::UpstreamStatus { status }); return Err(ApiError::UpstreamStatus { status });
} }
}; };
abort.disarm();
match breaker_outcome(status) { match breaker_outcome(status) {
BreakerOutcome::Failure => breaker.record_failure(), BreakerOutcome::Failure => breaker.record_failure(),
BreakerOutcome::Success => breaker.record_success(), BreakerOutcome::Success => breaker.record_success(),
@@ -301,6 +309,7 @@ impl Proxy {
path: &str, path: &str,
headers: &HeaderMap, headers: &HeaderMap,
body: Bytes, body: Bytes,
abort_rid: Option<&str>,
stream_guards: Option<Box<dyn Send + 'static>>, stream_guards: Option<Box<dyn Send + 'static>>,
on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>, on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>,
on_stream_end: Option<Box<dyn FnOnce(sse::StreamEnd) + Send + 'static>>, on_stream_end: Option<Box<dyn FnOnce(sse::StreamEnd) + Send + 'static>>,
@@ -322,11 +331,16 @@ impl Proxy {
req = req req = req
.header("content-type", "application/json") .header("content-type", "application/json")
.header("accept", "text/event-stream"); .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| { let resp = req.send().await.map_err(|e| {
breaker.record_failure(); breaker.record_failure();
Self::classify_reqwest_error_for(worker_url.clone(), e, path) Self::classify_reqwest_error_for(worker_url.clone(), e, path)
})?; })?;
let status = resp.status(); let status = resp.status();
if !status.is_success() {
abort.disarm();
}
let upstream_ct = resp let upstream_ct = resp
.headers() .headers()
.get(reqwest::header::CONTENT_TYPE) .get(reqwest::header::CONTENT_TYPE)
@@ -366,6 +380,9 @@ impl Proxy {
BreakerOutcome::Success => { BreakerOutcome::Success => {
let breaker_for_hook = Arc::clone(breaker); let breaker_for_hook = Arc::clone(breaker);
Some(Box::new(move |end| { Some(Box::new(move |end| {
if end.reason == sse::StreamEndReason::Completed {
abort.disarm();
}
match stream_breaker_outcome(end) { match stream_breaker_outcome(end) {
BreakerOutcome::Success => breaker_for_hook.record_success(), BreakerOutcome::Success => breaker_for_hook.record_success(),
BreakerOutcome::Failure => breaker_for_hook.record_failure(), BreakerOutcome::Failure => breaker_for_hook.record_failure(),
@@ -465,6 +482,7 @@ mod tests {
"/chat", "/chat",
&headers, &headers,
Bytes::new(), Bytes::new(),
None,
) )
.now_or_never() .now_or_never()
.is_none()); .is_none());
@@ -481,6 +499,7 @@ mod tests {
None, None,
None, None,
None, None,
None,
) )
.now_or_never() .now_or_never()
.is_none()); .is_none());
@@ -583,6 +602,7 @@ mod tests {
None, None,
None, None,
None, None,
None,
expiration, expiration,
) )
.await .await
@@ -694,6 +714,7 @@ mod tests {
"/v1/chat/completions", "/v1/chat/completions",
&headers, &headers,
Bytes::from_static(b"{}"), Bytes::from_static(b"{}"),
None,
) )
.await .await
.expect("dispatch should reach the worker (breaker must stay closed)"); .expect("dispatch should reach the worker (breaker must stay closed)");
@@ -737,6 +758,7 @@ mod tests {
"/v1/chat/completions", "/v1/chat/completions",
&headers, &headers,
Bytes::from_static(b"{}"), Bytes::from_static(b"{}"),
None,
) )
.await; .await;
} }
@@ -780,6 +802,7 @@ mod tests {
"/v1/chat/completions", "/v1/chat/completions",
&headers, &headers,
Bytes::from_static(b"{}"), Bytes::from_static(b"{}"),
None,
) )
.await .await
.expect("the half-open probe must be admitted and reach the worker"); .expect("the half-open probe must be admitted and reach the worker");
@@ -819,6 +842,7 @@ mod tests {
None, None,
None, None,
None, None,
None,
) )
.await .await
.expect("streaming dispatch should reach the worker"); .expect("streaming dispatch should reach the worker");
experimental/sgl-router/src/server/app.rs
+9
View File
@@ -150,6 +150,9 @@ async fn access_log_and_record(
.unwrap_or_else(|| outcome_from_status(status.as_u16())) .unwrap_or_else(|| outcome_from_status(status.as_u16()))
.as_str(), .as_str(),
worker = log_ctx.map(|c| c.worker_url.as_str()).unwrap_or(""), 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(""), model = log_ctx.map(|c| c.model_id.as_str()).unwrap_or(""),
stream = log_ctx.is_some_and(|c| c.streaming), stream = log_ctx.is_some_and(|c| c.streaming),
latency_ms, latency_ms,
@@ -446,6 +449,7 @@ mod tests {
model_id: "tiny".into(), model_id: "tiny".into(),
streaming: false, streaming: false,
outcome: RequestOutcome::Cancelled, outcome: RequestOutcome::Cancelled,
engine_rid: Some("1f0c2b7a4e9d4f3ab6c5d8e7f0a1b2c3".into()),
}); });
resp resp
}), }),
@@ -468,6 +472,11 @@ mod tests {
logs.contains("worker=\"http://worker-a:30000\"") && logs.contains("model=\"tiny\""), 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}", "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 — // The handler's outcome must win over the status-derived fallback —
// otherwise the log and `worker_requests_total` can disagree about a // otherwise the log and `worker_requests_total` can disagree about a
// request the handler classified itself (here, a cancellation served // request the handler classified itself (here, a cancellation served
experimental/sgl-router/src/server/metrics.rs
+2
View File
@@ -223,6 +223,8 @@ pub struct RequestLogContext {
/// line and `worker_requests_total` cannot disagree — the middleware can /// line and `worker_requests_total` cannot disagree — the middleware can
/// only see the status, which cannot express a router-side cancellation. /// only see the status, which cannot express a router-side cancellation.
pub outcome: RequestOutcome, pub outcome: RequestOutcome,
/// Router-minted engine ID, logged beside the caller's correlation ID.
pub engine_rid: Option<String>,
} }
/// Final outcome of a 2xx SSE stream. /// Final outcome of a 2xx SSE stream.
experimental/sgl-router/src/server/routes/chat/forward.rs
+15 -2
View File
@@ -81,7 +81,12 @@ pub(super) async fn forward_chat_request(
}; };
(decode, bootstrap) (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); let prefill_load_guards = (worker_load_guard, active_request_guard);
// In PD mode, prefill runs independently and decode supplies the client response. // 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, &response_worker,
&headers, &headers,
body, body,
engine_rid.as_deref(),
response_load_guards, response_load_guards,
&metrics, &metrics,
expiration_token.clone(), expiration_token.clone(),
@@ -124,7 +130,7 @@ pub(super) async fn forward_chat_request(
model: metrics.model.clone(), 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. // Materialize dispatch errors here so the access log retains the selected worker.
let mut response = match result { let mut response = match result {
Ok(mut response) => { Ok(mut response) => {
@@ -171,6 +177,7 @@ fn spawn_prefill_request(
CHAT_PATH, CHAT_PATH,
&headers, &headers,
body, body,
None,
) )
.await .await
{ {
@@ -188,11 +195,13 @@ fn spawn_prefill_request(
}); });
} }
#[allow(clippy::too_many_arguments)]
async fn forward_to_response_worker( async fn forward_to_response_worker(
ctx: &AppContext, ctx: &AppContext,
worker: &Worker, worker: &Worker,
headers: &HeaderMap, headers: &HeaderMap,
body: Bytes, body: Bytes,
engine_rid: Option<&str>,
load_guards: LoadGuards, load_guards: LoadGuards,
metrics: &DispatchMetrics, metrics: &DispatchMetrics,
expiration: CancellationToken, expiration: CancellationToken,
@@ -209,6 +218,7 @@ async fn forward_to_response_worker(
CHAT_PATH, CHAT_PATH,
headers, headers,
body, body,
engine_rid,
Some(stream_guards), Some(stream_guards),
Some(metrics.first_byte_callback()), Some(metrics.first_byte_callback()),
Some(metrics.stream_end_callback(worker.url.clone())), Some(metrics.stream_end_callback(worker.url.clone())),
@@ -226,6 +236,7 @@ async fn forward_to_response_worker(
CHAT_PATH, CHAT_PATH,
headers, headers,
body, body,
engine_rid,
) )
.await .await
} }
@@ -295,6 +306,7 @@ impl DispatchMetrics {
fn record_dispatch_result( fn record_dispatch_result(
&self, &self,
result: &Result<Response<Body>, ApiError>, result: &Result<Response<Body>, ApiError>,
engine_rid: Option<String>,
) -> RequestLogContext { ) -> RequestLogContext {
let http_status = match result { let http_status = match result {
Ok(response) => response.status().as_u16(), Ok(response) => response.status().as_u16(),
@@ -326,6 +338,7 @@ impl DispatchMetrics {
model_id: self.model.clone(), model_id: self.model.clone(),
streaming: self.streaming, streaming: self.streaming,
outcome, outcome,
engine_rid,
} }
} }
} }
experimental/sgl-router/src/server/routes/chat/preparation.rs
+96 -12
View File
@@ -26,6 +26,8 @@ pub(super) struct PreparedChatRequest {
pub(super) tokens: Option<RequestTokens>, pub(super) tokens: Option<RequestTokens>,
/// Token count for routing/load accounting; estimated from body size when unavailable. /// Token count for routing/load accounting; estimated from body size when unavailable.
pub(super) input_token_count: usize, pub(super) input_token_count: usize,
caller_set_rid: bool,
fans_out: bool,
can_forward_input_ids: bool, can_forward_input_ids: bool,
parsed_body: Option<Value>, parsed_body: Option<Value>,
sampling_defaults: Vec<(SamplingField, Number)>, sampling_defaults: Vec<(SamplingField, Number)>,
@@ -69,16 +71,27 @@ impl PreparedChatRequest {
body, body,
tokens, tokens,
input_token_count, input_token_count,
caller_set_rid: fields.caller_set_rid,
fans_out: requests_multiple_samples(&fields, &sampling_defaults),
can_forward_input_ids, can_forward_input_ids,
parsed_body, parsed_body,
sampling_defaults, sampling_defaults,
}) })
} }
pub(super) fn engine_rid(&self, pd_mode: bool) -> Option<String> {
// 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( pub(super) fn into_outgoing_body(
self, self,
ctx: &AppContext, ctx: &AppContext,
bootstrap: Option<&BootstrapFields>, bootstrap: Option<&BootstrapFields>,
engine_rid: Option<&str>,
) -> Result<Bytes, ApiError> { ) -> Result<Bytes, ApiError> {
// Routing tokens can replace engine tokenization only for supported chat templates. // Routing tokens can replace engine tokenization only for supported chat templates.
let input_ids = match (self.tokens.as_ref(), self.parsed_body.as_ref()) { let input_ids = match (self.tokens.as_ref(), self.parsed_body.as_ref()) {
@@ -104,6 +117,7 @@ impl PreparedChatRequest {
input_ids, input_ids,
bootstrap, bootstrap,
&self.sampling_defaults, &self.sampling_defaults,
engine_rid,
) )
} }
} }
@@ -116,6 +130,8 @@ pub(super) struct RoutingFields {
max_tokens: Option<u64>, max_tokens: Option<u64>,
max_completion_tokens: Option<u64>, max_completion_tokens: Option<u64>,
sampling: [SamplingValue; SamplingField::ALL.len()], 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. /// Null is absent; unrepresentable values are rejected only under a sampling contract.
@@ -248,6 +264,7 @@ impl RoutingKey {
enum RequestKey { enum RequestKey {
Routing(RoutingKey), Routing(RoutingKey),
Sampling(SamplingField), Sampling(SamplingField),
Rid,
Other, Other,
} }
@@ -267,6 +284,7 @@ impl<'de> Deserialize<'de> for RequestKey {
"model" => RequestKey::Routing(RoutingKey::Model), "model" => RequestKey::Routing(RoutingKey::Model),
"max_tokens" => RequestKey::Routing(RoutingKey::MaxTokens), "max_tokens" => RequestKey::Routing(RoutingKey::MaxTokens),
"max_completion_tokens" => RequestKey::Routing(RoutingKey::MaxCompletionTokens), "max_completion_tokens" => RequestKey::Routing(RoutingKey::MaxCompletionTokens),
"rid" => RequestKey::Rid,
other => match SamplingField::from_wire_name(other) { other => match SamplingField::from_wire_name(other) {
Some(field) => RequestKey::Sampling(field), Some(field) => RequestKey::Sampling(field),
None => RequestKey::Other, None => RequestKey::Other,
@@ -324,6 +342,9 @@ impl<'de> serde::de::Visitor<'de> for RoutingFieldsVisitor {
SamplingValue::Unusable SamplingValue::Unusable
}; };
} }
RequestKey::Rid => {
fields.caller_set_rid = map.next_value::<Option<IgnoredAny>>()?.is_some();
}
RequestKey::Other => { RequestKey::Other => {
// Validate unrelated JSON without retaining its contents. // Validate unrelated JSON without retaining its contents.
map.next_value::<IgnoredAny>()?; map.next_value::<IgnoredAny>()?;
@@ -374,6 +395,22 @@ fn should_tokenize_request(
can_forward_input_ids || policy_needs_request_tokens || bucket_routing_enabled 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 { fn estimate_prefill_tokens(body: &Bytes) -> usize {
// Never 0: a zero-load entry is invisible to the cache-aware imbalance fast path. // Never 0: a zero-load entry is invisible to the cache-aware imbalance fast path.
(body.len() / BYTES_PER_TOKEN_ESTIMATE).max(1) (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. /// Append before the closing brace so injected values win over explicit nulls.
fn append_sampling_defaults( fn append_top_level_fields(
body: &Bytes, body: &Bytes,
sampling_defaults: &[(SamplingField, Number)], sampling_defaults: &[(SamplingField, Number)],
rid: Option<&str>,
) -> Option<Bytes> { ) -> Option<Bytes> {
use std::io::Write as _; use std::io::Write as _;
@@ -405,13 +443,23 @@ fn append_sampling_defaults(
let has_members = body[open + 1..close] let has_members = body[open + 1..close]
.iter() .iter()
.any(|b| !b.is_ascii_whitespace()); .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]); output.extend_from_slice(&body[..close]);
for (i, (field, value)) in sampling_defaults.iter().enumerate() { let mut wrote_any = has_members;
if has_members || i > 0 { for (field, value) in sampling_defaults {
if wrote_any {
output.push(b','); output.push(b',');
} }
write!(output, "\"{}\":{}", field.wire_name(), value).ok()?; 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..]); output.extend_from_slice(&body[close..]);
Some(Bytes::from(output)) Some(Bytes::from(output))
@@ -424,14 +472,15 @@ fn build_outgoing_body(
input_ids: Option<&[u32]>, input_ids: Option<&[u32]>,
bootstrap: Option<&BootstrapFields>, bootstrap: Option<&BootstrapFields>,
sampling_defaults: &[(SamplingField, Number)], sampling_defaults: &[(SamplingField, Number)],
rid: Option<&str>,
) -> Result<Bytes, ApiError> { ) -> Result<Bytes, ApiError> {
let sampling_only = input_ids.is_none() && bootstrap.is_none(); let needs_parse = input_ids.is_some() || bootstrap.is_some();
if sampling_only && sampling_defaults.is_empty() { if !needs_parse && sampling_defaults.is_empty() && rid.is_none() {
// Cloning Bytes shares the original allocation when no injection is needed. // Cloning Bytes shares the original allocation when no injection is needed.
return Ok(body.clone()); return Ok(body.clone());
} }
if sampling_only { if !needs_parse {
if let Some(spliced) = append_sampling_defaults(body, sampling_defaults) { if let Some(spliced) = append_top_level_fields(body, sampling_defaults, rid) {
return Ok(spliced); return Ok(spliced);
} }
} }
@@ -445,6 +494,9 @@ fn build_outgoing_body(
return Err(invalid_request()); 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 { for (field, default) in sampling_defaults {
body_fields.insert(field.wire_name().into(), default.clone().into()); body_fields.insert(field.wire_name().into(), default.clone().into());
} }
@@ -739,7 +791,8 @@ mod tests {
.unwrap() .unwrap()
.extend(fields.as_object().unwrap().clone()); .extend(fields.as_object().unwrap().clone());
for value in [None, Some(original.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::<Value>(&out).unwrap(), expected); assert_eq!(serde_json::from_slice::<Value>(&out).unwrap(), expected);
} }
} }
@@ -750,7 +803,7 @@ mod tests {
for raw in [r#"{"model":"x"}"#, r#"{"model":"x","messages":[]}"#] { for raw in [r#"{"model":"x"}"#, r#"{"model":"x","messages":[]}"#] {
let body = Bytes::copy_from_slice(raw.as_bytes()); let body = Bytes::copy_from_slice(raw.as_bytes());
for value in [None, Some(serde_json::from_slice(&body).unwrap())] { 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, body);
assert_eq!(out.as_ptr(), body.as_ptr()); assert_eq!(out.as_ptr(), body.as_ptr());
} }
@@ -1108,7 +1161,7 @@ mod tests {
&metrics(), &metrics(),
) )
.unwrap(); .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!( assert_eq!(
serde_json::from_slice::<Value>(&out).unwrap(), serde_json::from_slice::<Value>(&out).unwrap(),
json!({ json!({
@@ -1240,7 +1293,7 @@ mod tests {
let inject = resolve_sampling_defaults(&config, &fields_of(raw), &metrics()).unwrap(); let inject = resolve_sampling_defaults(&config, &fields_of(raw), &metrics()).unwrap();
assert_eq!(inject.len(), 2); assert_eq!(inject.len(), 2);
for value in [None, Some(serde_json::from_slice(&body).unwrap())] { 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}"); assert_eq!(std::str::from_utf8(&out).unwrap(), expected, "{raw}");
let parsed: Value = serde_json::from_slice(&out).unwrap(); let parsed: Value = serde_json::from_slice(&out).unwrap();
assert_eq!(parsed["temperature"], json!(1.0)); assert_eq!(parsed["temperature"], json!(1.0));
@@ -1464,4 +1517,35 @@ mod tests {
"a value at the cap is still read" "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);
}
}
} }
experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs
+12 -1
View File
@@ -198,7 +198,18 @@ async fn caller_input_ids_are_used_for_routing_and_preserved() {
send(Arc::clone(&ctx), request.clone()).await, send(Arc::clone(&ctx), request.clone()).await,
StatusCode::OK 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. // Bypasses are not rendering failures.
assert!(!ctx assert!(!ctx
experimental/sgl-router/tests/proxy/chat_routing.rs
+93 -47
View File
@@ -17,10 +17,14 @@ use sgl_router::workers::{WireProtocol, Worker, WorkerRegistry};
use axum::body::Body; use axum::body::Body;
use axum::http::{Request, StatusCode}; use axum::http::{Request, StatusCode};
use http_body_util::BodyExt; use http_body_util::BodyExt;
use sgl_router::state::load_monitor::router_inflight_load::{
spawn_janitor, JanitorHandle, RouterInflightLoadRegistry,
};
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use tower::ServiceExt; use tower::ServiceExt;
mod cancellation;
mod reorg; mod reorg;
const TEST_TIMEOUT: Duration = Duration::from_secs(5); const TEST_TIMEOUT: Duration = Duration::from_secs(5);
@@ -74,6 +78,36 @@ fn build_ctx_with_worker(url: &str) -> Arc<AppContext> {
Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) 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<AppContext>, 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] #[tokio::test]
async fn non_streaming_returns_200() { async fn non_streaming_returns_200() {
let worker = crate::common::mock_worker::MockWorker::start(vec![]).await; 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", "/v1/chat/completions",
&headers, &headers,
body, body,
None,
) )
.await; .await;
assert!(res.is_err(), "body drop should surface as ApiError"); 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", "/v1/chat/completions",
&headers, &headers,
bytes::Bytes::from_static(b"{}"), bytes::Bytes::from_static(b"{}"),
None,
) )
.await; .await;
assert!(res.is_ok(), "clean OK call must succeed: {res:?}"); 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,
None, None,
None,
) )
.await; .await;
@@ -1248,6 +1285,7 @@ async fn forward_json_to_records_failure_on_5xx() {
"/v1/chat/completions", "/v1/chat/completions",
&headers, &headers,
body, body,
None,
) )
.await; .await;
@@ -1282,6 +1320,7 @@ async fn forward_json_to_rejects_when_breaker_open() {
"/v1/chat/completions", "/v1/chat/completions",
&headers, &headers,
body, body,
None,
) )
.await; .await;
@@ -1320,6 +1359,7 @@ async fn forward_json_to_malformed_url_returns_worker_misconfigured_and_trips_br
"/v1/chat/completions", "/v1/chat/completions",
&headers, &headers,
body, body,
None,
) )
.await; .await;
@@ -1607,58 +1647,17 @@ async fn streaming_active_load_drops_on_client_disconnect() {
/// returns; cancellation fires; handler returns 504. /// returns; cancellation fires; handler returns 504.
#[tokio::test] #[tokio::test]
async fn janitor_expiry_returns_504_stale_request_expired() { async fn janitor_expiry_returns_504_stale_request_expired() {
use sgl_router::state::load_monitor::router_inflight_load::{ // Upstream that takes 2s to respond — longer than the helper's 50ms
spawn_janitor, RouterInflightLoadRegistry, // stale_request_timeout, so the janitor sweeps before it answers.
};
// Upstream that takes 2s to respond — longer than our 50ms
// stale_request_timeout.
let worker = let worker =
crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(2)).await; crate::common::mock_worker::MockWorker::start_hanging(Duration::from_secs(2)).await;
let (ctx, _janitor) = build_ctx_with_janitor(&worker.url);
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 app = build_router(ctx); let app = build_router(ctx);
let req = Request::builder() let res = app
.method("POST") .oneshot(cancellation::request(serde_json::json!({})))
.uri("/v1/chat/completions") .await
.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(),
))
.unwrap(); .unwrap();
let res = app.oneshot(req).await.unwrap();
assert_eq!( assert_eq!(
res.status(), res.status(),
StatusCode::GATEWAY_TIMEOUT, StatusCode::GATEWAY_TIMEOUT,
@@ -1677,6 +1676,53 @@ async fn janitor_expiry_returns_504_stale_request_expired() {
body_str.contains("\"code\":\"stale_request_expired\""), body_str.contains("\"code\":\"stale_request_expired\""),
"504 body must encode the same code in the JSON envelope: {body_str}", "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 /// Task A: a non-streaming request that errors out (upstream
experimental/sgl-router/tests/proxy/chat_routing/cancellation.rs
+205
View File
@@ -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<Event>;
async fn chat(State(events): State<Events>, Json(body): Json<Value>) -> (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<bytes::Bytes, std::io::Error>,
>())
} 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<Events>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> 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<AppContext>,
events: mpsc::UnboundedReceiver<Event>,
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> {
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;
}
experimental/sgl-router/tests/proxy/chat_routing/reorg.rs
+14 -1
View File
@@ -230,7 +230,19 @@ async fn length_selects_plain_bucket_before_engine_selection() {
let response = app.clone().oneshot(request(body("hi"))).await.unwrap(); let response = app.clone().oneshot(request(body("hi"))).await.unwrap();
assert_eq!(response.status(), StatusCode::OK); assert_eq!(response.status(), StatusCode::OK);
let _ = response.into_body().collect().await.unwrap(); 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()); assert!(long_worker.captured.lock().unwrap().last_body.is_none());
let response = app let response = app
@@ -293,6 +305,7 @@ async fn pd_picks_both_groups_from_selected_bucket_and_shares_bootstrap() {
let d: serde_json::Value = let d: serde_json::Value =
serde_json::from_slice(decode.captured.lock().unwrap().last_body.as_ref().unwrap()) serde_json::from_slice(decode.captured.lock().unwrap().last_body.as_ref().unwrap())
.unwrap(); .unwrap();
assert!(p.get("rid").is_none() && d.get("rid").is_none());
assert!(p["bootstrap_room"].is_number()); assert!(p["bootstrap_room"].is_number());
assert_eq!(p["bootstrap_room"], d["bootstrap_room"]); assert_eq!(p["bootstrap_room"], d["bootstrap_room"]);
let calls = policy.calls.lock().unwrap(); let calls = policy.calls.lock().unwrap();
experimental/sgl-router/tests/proxy/common/mock_worker.rs
+27
View File
@@ -39,9 +39,22 @@ pub struct MockWorker {
// Used in header_forwarding_test; not every test file reads captured headers. // Used in header_forwarding_test; not every test file reads captured headers.
#[allow(dead_code)] #[allow(dead_code)]
pub captured: Arc<Mutex<CapturedHeaders>>, pub captured: Arc<Mutex<CapturedHeaders>>,
#[allow(dead_code)]
pub abort_log: Arc<Mutex<Vec<Value>>>,
_shutdown: oneshot::Sender<()>, _shutdown: oneshot::Sender<()>,
} }
#[allow(dead_code)] // shared across all axum variants
fn abort_request_route<S>(log: Arc<Mutex<Vec<Value>>>) -> axum::routing::MethodRouter<S>
where
S: Clone + Send + Sync + 'static,
{
post(move |Json(body): Json<Value>| async move {
log.lock().unwrap().push(body);
StatusCode::OK
})
}
impl MockWorker { impl MockWorker {
/// Bind to a random port on 127.0.0.1 and start serving. /// 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. #[allow(dead_code)] // Only used by some test files.
pub async fn start(stream_chunks: Vec<&'static str>) -> Self { pub async fn start(stream_chunks: Vec<&'static str>) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default())); let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let state = MockWorkerState { let state = MockWorkerState {
captured: captured.clone(), captured: captured.clone(),
stream_chunks: Arc::new(stream_chunks), stream_chunks: Arc::new(stream_chunks),
@@ -60,6 +74,7 @@ impl MockWorker {
let app = axum::Router::new() let app = axum::Router::new()
.route("/v1/chat/completions", post(chat)) .route("/v1/chat/completions", post(chat))
.route("/server_info", get(serve_tiny_server_info)) .route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state); .with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -77,6 +92,7 @@ impl MockWorker {
Self { Self {
url, url,
captured, captured,
abort_log,
_shutdown: tx, _shutdown: tx,
} }
} }
@@ -88,6 +104,7 @@ impl MockWorker {
#[allow(dead_code)] #[allow(dead_code)]
pub async fn start_hanging(delay: Duration) -> Self { pub async fn start_hanging(delay: Duration) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default())); let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
#[derive(Clone)] #[derive(Clone)]
struct HangState { struct HangState {
@@ -127,6 +144,7 @@ impl MockWorker {
let app = axum::Router::new() let app = axum::Router::new()
.route("/v1/chat/completions", post(hang_handler)) .route("/v1/chat/completions", post(hang_handler))
.route("/server_info", get(serve_tiny_server_info)) .route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state); .with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -144,6 +162,7 @@ impl MockWorker {
Self { Self {
url, url,
captured, captured,
abort_log,
_shutdown: tx, _shutdown: tx,
} }
} }
@@ -154,6 +173,7 @@ impl MockWorker {
#[allow(dead_code)] #[allow(dead_code)]
pub async fn start_slow_stream(chunks: Vec<&'static str>, delay: Duration) -> Self { pub async fn start_slow_stream(chunks: Vec<&'static str>, delay: Duration) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default())); let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
#[derive(Clone)] #[derive(Clone)]
struct SlowState { struct SlowState {
@@ -207,6 +227,7 @@ impl MockWorker {
let app = axum::Router::new() let app = axum::Router::new()
.route("/v1/chat/completions", post(slow_chat)) .route("/v1/chat/completions", post(slow_chat))
.route("/server_info", get(serve_tiny_server_info)) .route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state); .with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -224,6 +245,7 @@ impl MockWorker {
Self { Self {
url, url,
captured, captured,
abort_log,
_shutdown: tx, _shutdown: tx,
} }
} }
@@ -248,6 +270,7 @@ impl MockWorker {
partial_body_bytes: &'static [u8], partial_body_bytes: &'static [u8],
) -> Self { ) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default())); let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap(); let addr: SocketAddr = listener.local_addr().unwrap();
let url = format!("http://{addr}"); let url = format!("http://{addr}");
@@ -309,6 +332,7 @@ impl MockWorker {
Self { Self {
url, url,
captured, captured,
abort_log,
_shutdown: tx, _shutdown: tx,
} }
} }
@@ -319,6 +343,7 @@ impl MockWorker {
#[allow(dead_code)] #[allow(dead_code)]
pub async fn start_returning_error(status: StatusCode, body: Value) -> Self { pub async fn start_returning_error(status: StatusCode, body: Value) -> Self {
let captured = Arc::new(Mutex::new(CapturedHeaders::default())); let captured = Arc::new(Mutex::new(CapturedHeaders::default()));
let abort_log: Arc<Mutex<Vec<Value>>> = Arc::new(Mutex::new(Vec::new()));
let body_arc = Arc::new(body.to_string()); let body_arc = Arc::new(body.to_string());
#[derive(Clone)] #[derive(Clone)]
@@ -360,6 +385,7 @@ impl MockWorker {
let app = axum::Router::new() let app = axum::Router::new()
.route("/v1/chat/completions", post(error_handler)) .route("/v1/chat/completions", post(error_handler))
.route("/server_info", get(serve_tiny_server_info)) .route("/server_info", get(serve_tiny_server_info))
.route("/abort_request", abort_request_route(abort_log.clone()))
.with_state(state); .with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -377,6 +403,7 @@ impl MockWorker {
Self { Self {
url, url,
captured, captured,
abort_log,
_shutdown: tx, _shutdown: tx,
} }
} }
experimental/sgl-router/tests/proxy/common/mod.rs
+8
View File
@@ -6,3 +6,11 @@
pub mod cache_aware_fixture; pub mod cache_aware_fixture;
pub mod mock_worker; pub mod mock_worker;
pub mod streaming; 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())
}
experimental/sgl-router/tests/proxy/h2c_forward.rs
+3
View File
@@ -72,6 +72,7 @@ async fn h2c_client_reaches_http2_only_worker() {
"/v1/chat/completions", "/v1/chat/completions",
&axum::http::HeaderMap::new(), &axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"), Bytes::from_static(b"{}"),
None,
) )
.await .await
.expect("h2c client must reach an HTTP/2-only worker"); .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", "/v1/chat/completions",
&axum::http::HeaderMap::new(), &axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"), Bytes::from_static(b"{}"),
None,
) )
.await; .await;
assert!( assert!(
@@ -164,6 +166,7 @@ async fn h2c_client_streams_sse_from_http2_only_worker() {
&axum::http::HeaderMap::new(), &axum::http::HeaderMap::new(),
Bytes::from_static(b"{}"), Bytes::from_static(b"{}"),
None, None,
None,
Some(Box::new(move || { Some(Box::new(move || {
flag.store(true, std::sync::atomic::Ordering::SeqCst); flag.store(true, std::sync::atomic::Ordering::SeqCst);
})), })),
experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs
+58
View File
@@ -376,3 +376,61 @@ async fn pd_mode_prefill_5xx_does_not_poison_decode_response() {
let pv = parse_body(&prefill_body); let pv = parse_body(&prefill_body);
assert_eq!(bootstrap_port(&pv), Some(8997)); assert_eq!(bootstrap_port(&pv), Some(8997));
} }
fn streaming_chat_request() -> Request<Body> {
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());
}
}
experimental/sgl-router/tests/proxy/roundrobin_input_ids.rs
+16 -2
View File
@@ -106,9 +106,23 @@ fn without_forwarding(mut cfg: Config, policy: PolicyKind) -> Config {
cfg 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<AppContext>, mock: &MockWorker, request: &Value) { async fn assert_forwarded_unchanged(ctx: &Arc<AppContext>, mock: &MockWorker, request: &Value) {
assert_eq!(send(Arc::clone(ctx), request.clone()).await, StatusCode::OK); 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 assert!(!ctx
.metrics .metrics
.render() .render()
@@ -428,6 +442,6 @@ async fn kimi_ids_forward_with_engine_rendering_fallback() {
} else { } else {
assert!(ids.is_none()); assert!(ids.is_none());
} }
assert_eq!(captured(&mock), request); assert_eq!(without_minted_rid(captured(&mock)), request);
} }
} }