diff --git a/experimental/sgl-router/monitoring/README.md b/experimental/sgl-router/monitoring/README.md index 65d1a3c9f894..b4ff52055dd8 100644 --- a/experimental/sgl-router/monitoring/README.md +++ b/experimental/sgl-router/monitoring/README.md @@ -26,7 +26,7 @@ Families the router emits. The dashboard graphs all of them except the | `sgl_router_worker_requests_total` | Counter | Per-worker **dispatches** by `worker_url`, `model_id`, `mode`, `outcome` (recorded after dispatch; blind to pre-dispatch drops). See [Dispatch outcomes](#dispatch-outcomes) | | `sgl_router_request_duration_seconds` | Histogram | End-to-end request latency by `model_id` | | `sgl_router_ttft_seconds` | Histogram | Time to first token (streaming) by `model_id` | -| `sgl_router_stream_outcome_total` | Counter | Streaming outcomes by `worker_url`, `model_id`, and `outcome` (`ok`, `stream_error_event`, `upstream_error`, or `client_disconnect`). Counts committed 2xx streams only — non-2xx responses are counted by status in `responses_total` | +| `sgl_router_stream_outcome_total` | Counter | Streaming outcomes by `worker_url`, `model_id`, and `outcome` (`ok`, `stream_error_event`, `upstream_error`, `client_disconnect`, `expired`, or `aborted`). Counts committed 2xx streams only — non-2xx responses are counted by status in `responses_total` | | `sgl_router_active_load` | Gauge | Per-worker prefill-token / decode-block load | | `sgl_router_workers` | Gauge | Registered worker count by `mode` | | `sgl_router_worker_health` | Gauge | Per-worker health (1=breaker admits, 0=open) | @@ -126,7 +126,7 @@ worker error the router forwards is a successful *proxy* operation and a failed | `client_error` | 4xx except 429 | no — the caller sent something invalid | | `backpressure` | 429, 503 | no — responsive but at capacity | | `error` | 5xx except 503, plus transport failures, timeouts and incomplete bodies | **yes** | -| `cancelled` | the router's own stale-request deadline | no | +| `cancelled` | the router's own stale-request deadline, or a PD decode dispatch dropped because prefill failed first | no | `error` is the only bucket that means *this worker failed*, which is why the Error-ratio panel uses it alone. The split matters during an incident: a @@ -136,8 +136,11 @@ statuses — so the two agree, and the error ratio keeps pointing at genuine faults instead of pegging at 100% exactly when it is being read. A hung worker surfaces as `error` (the router's upstream timeout), *not* as -`cancelled`. Only the stale-request deadline produces `cancelled`; -`sgl_router_stale_requests_total{outcome="expired"}` counts the same events. +`cancelled`. Two things produce `cancelled`: the stale-request deadline, which +`sgl_router_stale_requests_total{outcome="expired"}` counts the same events of, +and a PD decode worker whose dispatch was dropped because prefill failed first — +the client's status is then prefill's, booked against the prefill worker, so +decode's abandoned dispatch is counted here rather than disappearing. ## Access log diff --git a/experimental/sgl-router/src/proxy/mod.rs b/experimental/sgl-router/src/proxy/mod.rs index daa71ca097f3..36bc2b7fd525 100644 --- a/experimental/sgl-router/src/proxy/mod.rs +++ b/experimental/sgl-router/src/proxy/mod.rs @@ -49,7 +49,7 @@ enum BreakerOutcome { /// A real fault (5xx other than backpressure) → `record_failure`: count /// toward opening. Failure, - /// Backpressure or router-side stream expiry → + /// Backpressure or a router-side stream expiry or abort → /// `record_backpressure`: never opens the breaker and, while Closed, leaves /// an in-progress failure streak intact — but still resolves a half-open /// probe so a recovered-but-busy worker isn't wedged shut. @@ -85,13 +85,13 @@ fn breaker_outcome(status: reqwest::StatusCode) -> BreakerOutcome { } } -/// Router-side expiry says nothing about worker health. Preserve the existing +/// Router-side expiry or abort says nothing about worker health. Preserve the existing /// treatment of completed streams and client disconnects; upstream faults, /// idle timeouts, and pump panics remain failures. fn stream_breaker_outcome(end: sse::StreamEnd) -> BreakerOutcome { use sse::StreamEndReason; match end.reason { - StreamEndReason::Expired => BreakerOutcome::Neutral, + StreamEndReason::Expired | StreamEndReason::Aborted => BreakerOutcome::Neutral, StreamEndReason::Completed | StreamEndReason::ClientDisconnect => BreakerOutcome::Success, StreamEndReason::UpstreamError | StreamEndReason::IdleTimeout @@ -314,6 +314,7 @@ impl Proxy { on_first_byte: Option>, on_stream_end: Option>, expiration: Option, + stream_abort: Option, ) -> Result, ApiError> { let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen { worker: worker_url.to_string(), @@ -410,6 +411,7 @@ impl Proxy { sse::StreamLimits { idle_timeout: self.stream_idle_timeout, expiration, + abort: stream_abort.filter(|_| status.is_success()), }, ); let mut out = Response::new(body); @@ -500,6 +502,7 @@ mod tests { None, None, None, + None, ) .now_or_never() .is_none()); @@ -604,6 +607,7 @@ mod tests { None, None, expiration, + None, ) .await .unwrap(); @@ -843,6 +847,7 @@ mod tests { None, None, None, + None, ) .await .expect("streaming dispatch should reach the worker"); diff --git a/experimental/sgl-router/src/proxy/sse.rs b/experimental/sgl-router/src/proxy/sse.rs index 695af12af8d2..6fc65f0b8178 100644 --- a/experimental/sgl-router/src/proxy/sse.rs +++ b/experimental/sgl-router/src/proxy/sse.rs @@ -4,6 +4,7 @@ //! SSE passthrough — bridges a reqwest `bytes_stream()` into an axum Body. use std::panic::AssertUnwindSafe; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use axum::body::Body; @@ -20,6 +21,8 @@ pub enum StreamEndReason { IdleTimeout, /// The router's stale-request deadline expired, regardless of worker health. Expired, + /// The router aborted the stream before its first chunk. + Aborted, ClientDisconnect, PumpPanicked, } @@ -69,6 +72,11 @@ impl ErrorEventScanner { } } +/// Whether a buffered SSE body carries an error event. +pub fn has_error_event(body: &[u8]) -> bool { + ErrorEventScanner::default().feed(body) +} + /// Bounds on a streaming response beyond what the upstream stream itself provides. #[derive(Debug, Clone, Default)] pub struct StreamLimits { @@ -76,6 +84,8 @@ pub struct StreamLimits { pub idle_timeout: Option, /// Fires when the stale-request janitor expires the request. pub expiration: Option, + /// Aborts the stream if it fires before the first upstream chunk. + pub abort: Option, } /// Bridge a byte stream into an axum Body that streams chunks unchanged. @@ -83,7 +93,8 @@ pub struct StreamLimits { /// One tokio task pumps upstream chunks through a bounded 64-slot channel so a /// slow client backpressures the upstream read. The pump stops as soon as the /// client disconnects, even while upstream is silent, when `limits.idle_timeout` -/// elapses between chunks, or when `limits.expiration` fires. +/// elapses between chunks, when `limits.expiration` fires, or when +/// `limits.abort` fires before the first chunk. /// /// The terminal result travels on a separate channel and is chained after the /// data, so a full queue cannot block cleanup or turn a failed stream into a @@ -111,10 +122,12 @@ where }; let mut scanner = ErrorEventScanner::default(); let idle = limits.idle_timeout.unwrap_or(Duration::MAX); - let expired = async { - match limits.expiration { - Some(token) => token.cancelled().await, - None => std::future::pending().await, + let started = AtomicBool::new(false); + let expired = cancelled(limits.expiration); + let aborted = async { + cancelled(limits.abort).await; + if started.load(Ordering::Relaxed) { + std::future::pending::<()>().await; } }; // Disconnect and expiration race the whole forwarding loop, so they @@ -130,6 +143,10 @@ where end.reason = StreamEndReason::Expired; Err(std::io::Error::other("SSE stream exceeded stale_request_timeout")) } + _ = aborted => { + end.reason = StreamEndReason::Aborted; + Err(std::io::Error::other("SSE stream aborted by the router")) + } result = async { loop { let bytes = match tokio::time::timeout(idle, stream.next()).await { @@ -144,6 +161,7 @@ where Err(std::io::Error::other("SSE upstream idle timeout")), ), }; + started.store(true, Ordering::Relaxed); if let Some(hook) = on_first_byte.take() { hook(); } @@ -191,6 +209,13 @@ where Body::from_stream(ReceiverStream::new(rx).map(Ok).chain(terminal)) } +async fn cancelled(token: Option) { + match token { + Some(token) => token.cancelled().await, + None => std::future::pending().await, + } +} + #[cfg(test)] mod tests { use super::*; @@ -249,6 +274,7 @@ mod tests { StreamLimits { idle_timeout: Some(Duration::from_secs(1)), expiration: None, + abort: None, }, ); assert!(body @@ -282,6 +308,7 @@ mod tests { StreamLimits { idle_timeout: None, expiration: Some(token.clone()), + abort: None, }, ); tokio::task::yield_now().await; @@ -300,6 +327,39 @@ mod tests { .to_string() .contains("stale_request_timeout")); } + + #[tokio::test] + async fn abort_ends_the_stream_only_before_its_first_chunk() { + for started in [false, true] { + let token = CancellationToken::new(); + let (tx, rx) = tokio::sync::mpsc::channel(1); + let (body, end) = limited_body( + ReceiverStream::new(rx), + StreamLimits { + abort: Some(token.clone()), + ..Default::default() + }, + ); + let mut data = body.into_data_stream(); + if started { + tx.send(Ok(Bytes::from_static(b"data: a\n\n"))) + .await + .unwrap(); + data.next().await.unwrap().unwrap(); + } + token.cancel(); + drop(tx); + let rest: Vec<_> = StreamExt::collect(data).await; + assert_eq!(rest.iter().any(Result::is_err), !started); + let expected = if started { + StreamEndReason::Completed + } else { + StreamEndReason::Aborted + }; + assert_eq!(end.await.unwrap().reason, expected); + } + } + #[tokio::test] async fn passes_through_a_simple_byte_stream() { let chunks = vec![ diff --git a/experimental/sgl-router/src/server/error.rs b/experimental/sgl-router/src/server/error.rs index 5079b18b8195..1ade2852fb2d 100644 --- a/experimental/sgl-router/src/server/error.rs +++ b/experimental/sgl-router/src/server/error.rs @@ -9,13 +9,8 @@ use thiserror::Error; pub const X_ROUTER_ERROR_CODE: HeaderName = HeaderName::from_static("x-router-error-code"); -/// Carries the worker's *own* HTTP status when the router had to synthesize a -/// status of its own over a worker that did respond (today: a mid-body drop, -/// where headers arrived but the body did not). Lets a gateway / operator -/// recover what the engine actually said instead of seeing only the router's -/// synthesized 502. Absent on every other response: a forwarded worker response -/// already carries the worker's status in the status line, and a router-only -/// condition (no workers, breaker open, ...) has no upstream status to report. +/// Original worker status when a failed prefill or incomplete body becomes a router 502. +/// Absent when no upstream response arrived or its status was forwarded unchanged. pub const X_ROUTER_UPSTREAM_STATUS: HeaderName = HeaderName::from_static("x-router-upstream-status"); @@ -103,6 +98,9 @@ pub enum ApiError { #[error("upstream response body incomplete after status {status}")] UpstreamStatus { status: StatusCode }, + #[error("prefill failed before KV transfer completed")] + PrefillFailed { status: Option }, + /// Wall-clock timeout exceeded while waiting for the upstream worker's /// response (per-request `request_timeout`). /// @@ -187,7 +185,9 @@ impl ApiError { ApiError::SamplingContract { .. } => ErrorClass::BadRequest, ApiError::ModelNotFound(_) => ErrorClass::NotFound, ApiError::UpstreamUnreachable { .. } => ErrorClass::Upstream, - ApiError::UpstreamStatus { .. } => ErrorClass::Upstream, + ApiError::UpstreamStatus { .. } | ApiError::PrefillFailed { .. } => { + ErrorClass::Upstream + } ApiError::UpstreamTimeout { .. } => ErrorClass::Timeout, ApiError::NoHealthyWorkers { .. } => ErrorClass::NoTarget, ApiError::NoPrefillWorkersAvailable { .. } => ErrorClass::NoTarget, @@ -218,6 +218,7 @@ impl ApiError { // didn't. We surface a 502 but echo the worker's status in // `x-router-upstream-status` (see `into_response`). ApiError::UpstreamStatus { .. } => "upstream_body_incomplete", + ApiError::PrefillFailed { .. } => "prefill_failed", ApiError::UpstreamTimeout { .. } => "upstream_timeout", ApiError::NoHealthyWorkers { .. } => "no_healthy_workers", ApiError::NoPrefillWorkersAvailable { .. } => "no_prefill_workers_available", @@ -232,13 +233,14 @@ impl ApiError { /// The worker's *own* status to echo in `x-router-upstream-status`, for the /// case where the router synthesized its own status over a worker that did - /// respond. Today only the mid-body-drop (`UpstreamStatus`) carries one. This + /// respond, including a failed prefill or an incomplete response body. This /// is an exhaustive, wildcard-free match (not an `if let` at the call site) so /// a future "synthesized over a responding worker" variant is forced to decide /// whether it echoes a status, rather than silently inheriting `None`. fn upstream_status(&self) -> Option { match self { ApiError::UpstreamStatus { status } => Some(*status), + ApiError::PrefillFailed { status } => *status, ApiError::BadRequest(_) | ApiError::SamplingContract { .. } | ApiError::ModelNotFound(_) @@ -311,6 +313,7 @@ impl IntoResponse for ApiError { ); "upstream response was incomplete".to_string() } + ApiError::PrefillFailed { .. } => self.to_string(), ApiError::UpstreamTimeout { worker } => { tracing::warn!(upstream = %worker, "upstream request timed out"); "upstream request timed out".to_string() @@ -369,9 +372,7 @@ impl IntoResponse for ApiError { .into_response(); resp.headers_mut() .insert(X_ROUTER_ERROR_CODE, HeaderValue::from_static(code)); - // Preserve the worker's real status when we synthesized our own (today, - // only the mid-body-drop case: the worker sent a status, then dropped the - // body, so we report a 502 but don't throw away what it said). + // Preserve the worker's real status when we synthesized our own. if let Some(upstream) = self.upstream_status() { resp.headers_mut().insert( X_ROUTER_UPSTREAM_STATUS, diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index 2a15af9a4035..2ad646336ca1 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -179,8 +179,9 @@ pub enum RequestOutcome { /// The worker failed to serve the request: a 5xx fault, a transport failure, /// a timeout, or a body that never completed. Error, - /// The router cancelled the request itself — today only the stale-request - /// deadline. Never derived from a status; see [`outcome_from_status`]. + /// The router cancelled the request itself: the stale-request deadline, or + /// a PD decode dispatch abandoned because prefill failed first. Never + /// derived from a status; see [`outcome_from_status`]. Cancelled, } @@ -260,11 +261,14 @@ pub enum StreamOutcome { ClientDisconnect, /// The router's stale-request deadline aborted the stream. Expired, + /// The router aborted the stream before its first chunk. + Aborted, } pub(crate) fn classify_stream_end(end: StreamEnd) -> StreamOutcome { match end.reason { StreamEndReason::Expired => StreamOutcome::Expired, + StreamEndReason::Aborted => StreamOutcome::Aborted, StreamEndReason::UpstreamError | StreamEndReason::IdleTimeout | StreamEndReason::PumpPanicked => StreamOutcome::UpstreamError, @@ -282,6 +286,7 @@ impl StreamOutcome { Self::UpstreamError => "upstream_error", Self::ClientDisconnect => "client_disconnect", Self::Expired => "expired", + Self::Aborted => "aborted", } } } @@ -1602,6 +1607,7 @@ mod tests { (StreamEndReason::IdleTimeout, UpstreamError), (StreamEndReason::PumpPanicked, UpstreamError), (StreamEndReason::Expired, Expired), + (StreamEndReason::Aborted, Aborted), ] { for saw_error_event in [false, true] { let end = StreamEnd { diff --git a/experimental/sgl-router/src/server/routes/chat/forward.rs b/experimental/sgl-router/src/server/routes/chat/forward.rs index 80325f0662fd..51e4c5d598f2 100644 --- a/experimental/sgl-router/src/server/routes/chat/forward.rs +++ b/experimental/sgl-router/src/server/routes/chat/forward.rs @@ -5,7 +5,7 @@ use super::preparation::{generate_room_id, BootstrapFields, PreparedChatRequest}; use crate::discovery::WorkerMode; -use crate::proxy::sse::{StreamEnd, StreamEndReason}; +use crate::proxy::sse::{self, StreamEnd, StreamEndReason}; use crate::server::app_context::AppContext; use crate::server::error::ApiError; use crate::server::metrics::{ @@ -18,6 +18,7 @@ use axum::body::Body; use axum::http::{HeaderMap, HeaderName, HeaderValue, Response}; use axum::response::IntoResponse; use bytes::Bytes; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Instant; use tokio_util::sync::CancellationToken; @@ -81,7 +82,7 @@ pub(super) async fn forward_chat_request( }; (decode, bootstrap) }); - let engine_rid = request.engine_rid(pd.is_some()); + let engine_rid = request.engine_rid(); let body = request.into_outgoing_body( ctx, pd.as_ref().map(|(_, bootstrap)| bootstrap), @@ -90,29 +91,33 @@ pub(super) async fn forward_chat_request( let prefill_load_guards = (worker_load_guard, active_request_guard); // In PD mode, prefill runs independently and decode supplies the client response. - let (response_worker, response_load_guards) = if let Some((decode, bootstrap)) = pd { - spawn_prefill_request( - ctx, - prefill, - headers.clone(), - body.clone(), - prefill_load_guards, - bootstrap.room, - ); - let decode_load_guards = ( - decode.load_guard(), - ctx.router_inflight_load - .register(decode.id.clone(), decode.url.clone(), 0, 1), - ); - (decode, decode_load_guards) - } else { - (prefill, prefill_load_guards) - }; + let stream_abort = CancellationToken::new(); + let (response_worker, response_load_guards, prefill_task) = + if let Some((decode, bootstrap)) = pd { + let task = spawn_prefill_request( + ctx, + &metrics, + Arc::clone(&prefill), + headers.clone(), + body.clone(), + prefill_load_guards, + bootstrap.room, + stream_abort.clone(), + ); + let decode_load_guards = ( + decode.load_guard(), + ctx.router_inflight_load + .register(decode.id.clone(), decode.url.clone(), 0, 1), + ); + (decode, decode_load_guards, Some((task, prefill))) + } else { + (prefill, prefill_load_guards, None) + }; // In PD mode, prefill can finish before decode. Watch the registration // held by the response so expiration remains live for its full lifetime. let expiration_token = response_load_guards.1.cancel_token().clone(); - let response_future = forward_to_response_worker( + let response = forward_to_response_worker( ctx, &response_worker, &headers, @@ -121,16 +126,24 @@ pub(super) async fn forward_chat_request( response_load_guards, &metrics, expiration_token.clone(), + stream_abort, ); + let dispatch = async { + match prefill_task { + Some((task, prefill)) => forward_pd(task, prefill, response).await, + None => (response.await, None), + } + }; // A ready response wins if request expiration fires in the same poll. - let result = tokio::select! { + let (result, blamed_prefill) = tokio::select! { biased; - result = response_future => result, - _ = expiration_token.cancelled() => Err(ApiError::StaleRequestExpired { - model: metrics.model.clone(), - }), + dispatch = dispatch => dispatch, + _ = expiration_token.cancelled() => { + let model = metrics.model.clone(); + (Err(ApiError::StaleRequestExpired { model }), None) + } }; - let log_context = metrics.record_dispatch_result(&result, engine_rid); + let log_context = metrics.record_dispatch_result(&result, engine_rid, blamed_prefill.as_ref()); // Materialize dispatch errors here so the access log retains the selected worker. let mut response = match result { Ok(mut response) => { @@ -157,19 +170,27 @@ fn parse_decode_url_header(decode_url: &str) -> Option { .ok() } +/// A prefill's client-visible failure; `None` means it succeeded. +type PrefillFailure = Option, ApiError>>; + +/// Runs prefill to completion even if the client disconnects. A failure also +/// aborts decode's stream until its first token, which proves KV transfer completed. +#[allow(clippy::too_many_arguments)] fn spawn_prefill_request( ctx: &AppContext, + metrics: &DispatchMetrics, prefill_worker: Arc, headers: HeaderMap, body: Bytes, load_guards: LoadGuards, bootstrap_room: u64, -) { + stream_abort: CancellationToken, +) -> tokio::task::JoinHandle { let proxy = Arc::clone(&ctx.proxy); - // Let prefill finish KV transfer after client cancellation; router shutdown still cancels it. + let (registry, model) = (Arc::clone(&metrics.registry), metrics.model.clone()); tokio::spawn(async move { let _load_guards = load_guards; - match proxy + let result = proxy .forward_json_to( &prefill_worker.url, prefill_worker.protocol(), @@ -179,20 +200,99 @@ fn spawn_prefill_request( body, None, ) - .await - { - Ok(_) => tracing::debug!( - prefill_url = %prefill_worker.url, bootstrap_room, "prefill side completed", - ), - // Prefill failures surface to the client through decode's bootstrap timeout. - Err(error) => tracing::warn!( - prefill_url = %prefill_worker.url, - bootstrap_room, - %error, - "prefill request failed; decode will time out on bootstrap_room", + .await; + let failure = prefill_failure(result).await; + let prefill_url = &prefill_worker.url; + let outcome = failure + .as_ref() + .map_or(RequestOutcome::Success, dispatch_outcome); + registry.record_worker_request(prefill_url, &model, WorkerModeLabel::Prefill, outcome); + match &failure { + None => tracing::debug!(%prefill_url, bootstrap_room, "prefill side completed"), + Some(Ok(response)) => tracing::debug!( + %prefill_url, bootstrap_room, status = %response.status(), + "prefill rejected the request", ), + Some(Err(error)) => { + tracing::warn!(%prefill_url, bootstrap_room, %error, "prefill failed") + } } - }); + if failure.is_some() { + stream_abort.cancel(); + } + failure + }) +} + +/// Client errors and backpressure pass through; only a prefill fault becomes +/// `prefill_failed`. +async fn prefill_failure(result: Result, ApiError>) -> PrefillFailure { + let status = match result { + // A streaming prefill reports a late failure as a 200 carrying an SSE error event. + Ok(response) if response.status().is_success() => { + let status = response.status(); + match axum::body::to_bytes(response.into_body(), usize::MAX).await { + Ok(body) if !sse::has_error_event(&body) => return None, + _ => status, + } + } + Ok(response) + if matches!( + outcome_from_status(response.status().as_u16()), + RequestOutcome::Error + ) => + { + response.status() + } + result => return Some(result), + }; + Some(Err(ApiError::PrefillFailed { + status: Some(status), + })) +} + +/// Prefill's failure ended the request, so the client-visible outcome belongs to +/// `prefill` rather than to the decode worker the dispatch had selected. +struct Blame { + prefill: Arc, + /// Whether decode had already reached the wire when it was dropped. False + /// only if prefill failed before decode's future was ever polled, which is + /// a pre-dispatch drop and must stay uncounted. + decode_dispatched: bool, +} + +/// Returns decode's response as soon as decode answers, unless prefill fails +/// first; then prefill's failure is returned along with the blamed prefill. +async fn forward_pd( + task: tokio::task::JoinHandle, + prefill: Arc, + decode: impl std::future::Future, ApiError>>, +) -> (Result, ApiError>, Option) { + // Set on decode's first poll, which is where `forward_to_response_worker` + // begins and the request reaches the worker. Read from the same task, so + // `Relaxed` needs no ordering beyond what the select already gives. + let dispatched = AtomicBool::new(false); + let decode = async { + dispatched.store(true, Ordering::Relaxed); + decode.await + }; + let mut decode = std::pin::pin!(decode); + let blame = |prefill| { + Some(Blame { + prefill, + decode_dispatched: dispatched.load(Ordering::Relaxed), + }) + }; + tokio::select! { + biased; + failure = task => match failure { + Ok(None) => (decode.await, None), + Ok(Some(failure)) => (failure, blame(prefill)), + Err(_) => (Err(ApiError::PrefillFailed { status: None }), blame(prefill)), + }, + // Dropping the handle leaves prefill running. + response = &mut decode => (response, None), + } } #[allow(clippy::too_many_arguments)] @@ -205,6 +305,7 @@ async fn forward_to_response_worker( load_guards: LoadGuards, metrics: &DispatchMetrics, expiration: CancellationToken, + stream_abort: CancellationToken, ) -> Result, ApiError> { if metrics.streaming { // Load and duration guards live until the SSE pump ends, not just until headers arrive. @@ -223,6 +324,7 @@ async fn forward_to_response_worker( Some(metrics.first_byte_callback()), Some(metrics.stream_end_callback(worker.url.clone())), Some(expiration), + Some(stream_abort), ) .await } else { @@ -302,30 +404,46 @@ impl DispatchMetrics { }) } - // HTTP status determines the outcome; router cancellations and dispatch failures stay distinct. + /// Logs a blamed prefill's failure against it; its task already recorded the outcome. fn record_dispatch_result( &self, result: &Result, ApiError>, engine_rid: Option, + blame: Option<&Blame>, ) -> RequestLogContext { - let http_status = match result { - Ok(response) => response.status().as_u16(), - Err(error) => error.status_code().as_u16(), - }; - let outcome = match result { - Err(ApiError::StaleRequestExpired { .. }) => { - self.registry - .record_stale_request(StaleRequestOutcome::Expired); - RequestOutcome::Cancelled + if let Err(ApiError::StaleRequestExpired { .. }) = result { + self.registry + .record_stale_request(StaleRequestOutcome::Expired); + } + let outcome = dispatch_outcome(result); + let worker_url = match blame { + // The prefill task books its own outcome, so recording `outcome` + // here would double-count it against prefill. Decode still has to + // be accounted for: its dispatch reached the worker and was then + // abandoned, and leaving it out would silently shrink decode's + // dispatch counts during exactly the prefill incident an operator + // is reading them to understand. + Some(blame) => { + if blame.decode_dispatched { + self.registry.record_worker_request( + &self.worker_url, + &self.model, + self.mode, + RequestOutcome::Cancelled, + ); + } + &blame.prefill.url } - // These 503s come from the router, not worker backpressure. - Err(ApiError::BreakerOpen { .. } | ApiError::WorkerMisconfigured { .. }) => { - RequestOutcome::Error + None => { + self.registry.record_worker_request( + &self.worker_url, + &self.model, + self.mode, + outcome, + ); + &self.worker_url } - _ => outcome_from_status(http_status), }; - self.registry - .record_worker_request(&self.worker_url, &self.model, self.mode, outcome); if !self.streaming { self.registry.observe_request_duration( &self.model, @@ -334,7 +452,7 @@ impl DispatchMetrics { } // The app middleware emits the access log and edge counters exactly once. RequestLogContext { - worker_url: self.worker_url.clone(), + worker_url: worker_url.clone(), model_id: self.model.clone(), streaming: self.streaming, outcome, @@ -343,6 +461,19 @@ impl DispatchMetrics { } } +// HTTP status determines the outcome; router cancellations and dispatch failures stay distinct. +fn dispatch_outcome(result: &Result, ApiError>) -> RequestOutcome { + match result { + Ok(response) => outcome_from_status(response.status().as_u16()), + Err(ApiError::StaleRequestExpired { .. }) => RequestOutcome::Cancelled, + // These 503s come from the router, not worker backpressure. + Err(ApiError::BreakerOpen { .. } | ApiError::WorkerMisconfigured { .. }) => { + RequestOutcome::Error + } + Err(error) => outcome_from_status(error.status_code().as_u16()), + } +} + /// Record total request duration when streaming ends or setup fails. struct StreamDurationGuard { metrics: Arc, diff --git a/experimental/sgl-router/src/server/routes/chat/preparation.rs b/experimental/sgl-router/src/server/routes/chat/preparation.rs index 2ed4d53f3697..f755214bd4f7 100644 --- a/experimental/sgl-router/src/server/routes/chat/preparation.rs +++ b/experimental/sgl-router/src/server/routes/chat/preparation.rs @@ -84,9 +84,10 @@ impl PreparedChatRequest { }) } - 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 { + pub(super) fn engine_rid(&self) -> Option { + // Caller IDs are unsafe for prefix aborts; fan-out regenerates IDs. + // PD shares the minted ID, but only decode arms abort-on-drop. + if self.caller_set_rid || self.fans_out { return None; } Some(uuid::Uuid::new_v4().simple().to_string()) diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 93a76b1bdeaa..3da133859a0f 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -1184,6 +1184,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() { None, None, None, + None, ) .await; diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs index 1255558a945a..b0c8d05fbaaf 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs @@ -305,7 +305,8 @@ 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["rid"].is_string()); + assert_eq!(p["rid"], d["rid"]); 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/h2c_forward.rs b/experimental/sgl-router/tests/proxy/h2c_forward.rs index cb7eaa0a5ff3..50e613b96795 100644 --- a/experimental/sgl-router/tests/proxy/h2c_forward.rs +++ b/experimental/sgl-router/tests/proxy/h2c_forward.rs @@ -172,6 +172,7 @@ async fn h2c_client_streams_sse_from_http2_only_worker() { })), None, None, + None, ) .await .expect("h2c client must stream from an HTTP/2-only worker"); diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs index 082f50620a7d..c310fdec83bb 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -324,61 +324,6 @@ async fn pd_mode_bootstrap_port_matches_chosen_prefill_worker() { ); } -/// Pin Pattern B's "prefill failure is invisible to the client" -/// contract: when the spawned prefill task gets a 5xx (or any other -/// upstream error), the decode response still reaches the client -/// unmodified. The router intentionally does not wire fail-fast here — -/// the decode side will eventually hang on `bootstrap_room` and time -/// out, but the chat handler itself doesn't propagate the prefill -/// error. Matches llm-d / aibrix behaviour. -#[tokio::test] -async fn pd_mode_prefill_5xx_does_not_poison_decode_response() { - let prefill = crate::common::mock_worker::MockWorker::start_returning_error( - StatusCode::INTERNAL_SERVER_ERROR, - json!({"error": "simulated prefill failure"}), - ) - .await; - let decode = crate::common::mock_worker::MockWorker::start(vec![]).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); - - // Client must see decode's 200 — the failing prefill is invisible. - let res = app.oneshot(chat_request()).await.unwrap(); - assert_eq!( - res.status(), - StatusCode::OK, - "decode response should reach the client even when prefill returned 5xx", - ); - - // Decode received its body (proves dual dispatch fired despite - // the prefill failure). - let decode_body = await_captured_body(&decode, Duration::from_secs(2), "decode").await; - let v = parse_body(&decode_body); - assert_eq!(bootstrap_port(&v), Some(8997)); - - // Prefill also received its body — it just returned 5xx. The - // bootstrap fields are present so the engine WOULD have honoured - // the bootstrap_room if the mock had succeeded. - let prefill_body = await_captured_body(&prefill, Duration::from_secs(2), "prefill").await; - let pv = parse_body(&prefill_body); - assert_eq!(bootstrap_port(&pv), Some(8997)); -} - fn streaming_chat_request() -> Request { Request::builder() .method("POST") @@ -396,7 +341,7 @@ fn streaming_chat_request() -> Request { } #[tokio::test] -async fn pd_mode_disconnect_does_not_abort_either_worker() { +async fn pd_mode_disconnect_aborts_decode_but_not_prefill() { 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"], @@ -429,10 +374,20 @@ async fn pd_mode_disconnect_does_not_abort_either_worker() { 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()); - } + tokio::time::timeout(Duration::from_secs(2), async { + while decode.abort_log.lock().unwrap().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(prefill.abort_log.lock().unwrap().is_empty()); + let body = await_captured_body(&decode, Duration::from_secs(2), "decode").await; + let rid = parse_body(&body)["rid"].clone(); + assert_eq!( + decode.abort_log.lock().unwrap()[0], + json!({"rid": rid, "abort_all": false}) + ); } + +mod reliability; diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection/reliability.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection/reliability.rs new file mode 100644 index 000000000000..855903d5a483 --- /dev/null +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection/reliability.rs @@ -0,0 +1,288 @@ +use super::*; +use crate::common::mock_worker::MockWorker; +use futures::StreamExt; +use http_body_util::BodyExt; +use sgl_router::server::app_context::ChatRouting; + +fn pd_ctx(prefill: &str, decode: &str, reorg: bool) -> Arc { + let spec = |id: &str, url: &str, mode| WorkerSpec { + id: WorkerId(id.into()), + url: url.into(), + mode, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: Some(8997), + }; + let mut ctx = build_ctx(vec![ + spec("p", prefill, WorkerMode::Prefill), + spec("d", decode, WorkerMode::Decode), + ]); + if reorg { + let mutable = Arc::get_mut(&mut ctx).unwrap(); + mutable.config.model.policy = PolicyKind::PowerOfTwo; + let state = sgl_router::state::kv_events::KvEventIndex::new(); + let (resolver, _) = sgl_router::policies_reorg::factory::build_resolver( + &mutable.config.model, + &state, + None, + ) + .unwrap(); + mutable.chat_routing = ChatRouting::Reorg([(ModelId("tiny".into()), resolver)].into()); + } + ctx +} + +fn chat(stream: bool) -> Request { + Request::post("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({"model": "tiny", "stream": stream}).to_string(), + )) + .unwrap() +} + +/// A prefill worker that answers `status` once `release` is notified. +async fn start_prefill(release: Arc, status: StatusCode) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let app = axum::Router::new().route( + "/v1/chat/completions", + axum::routing::post(move || { + let release = release.clone(); + async move { + release.notified().await; + (status, axum::Json(json!({}))) + } + }), + ); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + url +} + +async fn wait_until(mut done: impl FnMut() -> bool) { + tokio::time::timeout(Duration::from_secs(2), async { + while !done() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); +} + +/// Request outcomes recorded against the worker at `url`. +fn outcomes(ctx: &AppContext, url: &str, mode: &str) -> Vec { + let prefix = format!( + r#"sgl_router_worker_requests_total{{worker_url="{url}",model_id="tiny",mode="{mode}",outcome=""# + ); + let metrics = ctx.metrics.render(); + metrics + .lines() + .filter_map(|line| Some(line.strip_prefix(&prefix)?.split('"').next()?.to_string())) + .collect() +} + +/// Either side's failure ends the request without waiting for the other side. +#[tokio::test] +async fn pd_failure_returns_without_waiting_for_the_other_side() { + let rejected = json!({"error": "rejected"}); + let error = |status| MockWorker::start_returning_error(status, rejected.clone()); + let hanging = || MockWorker::start_hanging(Duration::from_secs(10)); + for reorg in [false, true] { + for stream in [false, true] { + // (prefill, decode, client status, router error code, whether prefill is blamed) + let cases = [ + ( + error(StatusCode::BAD_REQUEST).await, + hanging().await, + StatusCode::BAD_REQUEST, + None, + true, + ), + ( + error(StatusCode::SERVICE_UNAVAILABLE).await, + hanging().await, + StatusCode::SERVICE_UNAVAILABLE, + None, + true, + ), + ( + error(StatusCode::INTERNAL_SERVER_ERROR).await, + hanging().await, + StatusCode::BAD_GATEWAY, + Some("prefill_failed"), + true, + ), + // Transport failure: the prefill body is cut off. + ( + MockWorker::start_returning_partial_body(StatusCode::OK, b"{").await, + hanging().await, + StatusCode::BAD_GATEWAY, + Some("upstream_body_incomplete"), + true, + ), + ( + hanging().await, + error(StatusCode::TOO_MANY_REQUESTS).await, + StatusCode::TOO_MANY_REQUESTS, + None, + false, + ), + ]; + for (prefill, decode, expected, code, prefill_blamed) in cases { + let ctx = pd_ctx(&prefill.url, &decode.url, reorg); + let response = tokio::time::timeout( + Duration::from_secs(1), + build_router(ctx.clone()).oneshot(chat(stream)), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(response.status(), expected); + let router_code = response.headers().get("x-router-error-code").cloned(); + assert_eq!(router_code.as_ref().map(|c| c.to_str().unwrap()), code); + let body = response.into_body().collect().await.unwrap().to_bytes(); + if code.is_none() { + assert_eq!(serde_json::from_slice::(&body).unwrap(), rejected); + } + let (blamed, other) = match prefill_blamed { + true => ((&prefill.url, "prefill"), (&decode.url, "decode")), + false => ((&decode.url, "decode"), (&prefill.url, "prefill")), + }; + assert_eq!(outcomes(&ctx, blamed.0, blamed.1).len(), 1); + // A blamed prefill means decode was dispatched and then dropped, + // which is counted rather than dropped from the per-worker series. + // A blamed decode leaves prefill still running, with nothing to + // record yet. + let abandoned: &[&str] = if prefill_blamed { &["cancelled"] } else { &[] }; + assert_eq!(outcomes(&ctx, other.0, other.1), abandoned); + let decode_worker = ctx.registry.get(&WorkerId("d".into())).unwrap(); + assert_eq!(decode_worker.router_inflight_load(), 0); + } + } + } +} + +/// Until decode streams its first token, a prefill failure aborts decode: before +/// decode's headers the client sees prefill's error, after them the stream breaks. +#[tokio::test] +async fn prefill_failure_aborts_decode_before_its_first_token() { + for reorg in [false, true] { + for early_headers in [false, true] { + let decode = if early_headers { + MockWorker::start_slow_stream(vec!["data: late\n\n"], Duration::from_secs(10)).await + } else { + MockWorker::start_hanging(Duration::from_secs(10)).await + }; + let release = Arc::new(tokio::sync::Notify::new()); + let prefill_url = start_prefill(release.clone(), StatusCode::BAD_REQUEST).await; + let ctx = pd_ctx(&prefill_url, &decode.url, reorg); + let request = tokio::spawn(build_router(ctx.clone()).oneshot(chat(early_headers))); + wait_until(|| decode.captured.lock().unwrap().last_body.is_some()).await; + if early_headers { + let response = request.await.unwrap().unwrap(); + assert_eq!(response.status(), StatusCode::OK); + release.notify_one(); + assert!(response.into_body().collect().await.is_err()); + assert!(ctx.metrics.render().contains(&format!( + r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="aborted"}} 1"#, + decode.url + ))); + } else { + release.notify_one(); + let response = request.await.unwrap().unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + let workers = ctx.registry.workers_for(&ModelId("tiny".into())); + wait_until(|| { + !decode.abort_log.lock().unwrap().is_empty() + && workers.iter().all(|w| w.router_inflight_load() == 0) + }) + .await; + let sent = parse_body(decode.captured.lock().unwrap().last_body.as_ref().unwrap()); + assert_eq!( + decode.abort_log.lock().unwrap()[0], + json!({"rid": sent["rid"], "abort_all": false}) + ); + // Decode's dispatch is proven by the `wait_until` above, so it owes + // the per-worker series an outcome either way: `success` once it + // committed 200 headers (the broken stream is `stream_outcome`'s to + // report), `cancelled` when prefill's failure dropped it first. + let expected = if early_headers { + "success" + } else { + "cancelled" + }; + assert_eq!(outcomes(&ctx, &decode.url, "decode"), [expected]); + } + } +} + +/// Decode's headers do not wait for prefill, and once decode streams its first +/// token KV transfer is done, so a later prefill failure leaves the stream intact. +#[tokio::test] +async fn prefill_failure_after_decode_first_token_keeps_the_stream() { + for reorg in [false, true] { + let decode = MockWorker::start_slow_stream( + vec!["data: a\n\n", "data: [DONE]\n\n"], + Duration::from_millis(100), + ) + .await; + let release = Arc::new(tokio::sync::Notify::new()); + let prefill_url = start_prefill(release.clone(), StatusCode::INTERNAL_SERVER_ERROR).await; + let ctx = pd_ctx(&prefill_url, &decode.url, reorg); + let response = build_router(ctx.clone()).oneshot(chat(true)).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let mut body = response.into_body().into_data_stream(); + assert_eq!(body.next().await.unwrap().unwrap(), "data: a\n\n"); + release.notify_one(); + let prefill = ctx.registry.get(&WorkerId("p".into())).unwrap(); + wait_until(|| prefill.router_inflight_load() == 0).await; + assert_eq!(body.next().await.unwrap().unwrap(), "data: [DONE]\n\n"); + assert!(body.next().await.is_none()); + assert!(decode.abort_log.lock().unwrap().is_empty()); + assert_eq!(outcomes(&ctx, &prefill_url, "prefill"), ["error"]); + assert_eq!(outcomes(&ctx, &decode.url, "decode"), ["success"]); + } +} + +/// Prefill must finish KV transfer on its own, so decode failing first leaves it running. +#[tokio::test] +async fn decode_failure_leaves_prefill_running() { + for reorg in [false, true] { + let release = Arc::new(tokio::sync::Notify::new()); + let prefill_url = start_prefill(release.clone(), StatusCode::OK).await; + let decode = MockWorker::start_returning_error(StatusCode::BAD_REQUEST, json!({})).await; + let ctx = pd_ctx(&prefill_url, &decode.url, reorg); + let response = build_router(ctx.clone()) + .oneshot(chat(false)) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let prefill = ctx.registry.get(&WorkerId("p".into())).unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!(prefill.router_inflight_load(), 1); + release.notify_one(); + wait_until(|| prefill.router_inflight_load() == 0).await; + assert_eq!(outcomes(&ctx, &prefill_url, "prefill"), ["success"]); + } +} + +/// A streaming prefill that fails after committing a 200 is still a prefill failure. +#[tokio::test] +async fn prefill_sse_error_event_is_a_prefill_failure() { + for reorg in [false, true] { + let prefill = + MockWorker::start(vec!["data: {\"error\": {\"message\": \"boom\"}}\n\n"]).await; + let decode = MockWorker::start_hanging(Duration::from_secs(10)).await; + let ctx = pd_ctx(&prefill.url, &decode.url, reorg); + let response = tokio::time::timeout( + Duration::from_secs(1), + build_router(ctx.clone()).oneshot(chat(true)), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + assert_eq!(response.headers()["x-router-error-code"], "prefill_failed"); + assert_eq!(outcomes(&ctx, &prefill.url, "prefill"), ["error"]); + } +}