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

[sgl-router] Stream outcome observability for 2xx SSE streams (#38737)

Browse Source
This commit is contained in:
Kan Wu
2026-09-13 19:46:44 -07:00
committed by GitHub
parent bf9773e1da
commit f539c1fc65
6 changed files with 498 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experimental/sgl-router/monitoring/README.md
+1
View File
@@ -25,6 +25,7 @@ The dashboard graphs every family the router emits:
| `sgl_router_worker_requests_total` | Counter | Per-worker **dispatches** by `worker_url`, `model_id`, `mode`, `outcome` (recorded after dispatch; blind to pre-dispatch drops) | | `sgl_router_worker_requests_total` | Counter | Per-worker **dispatches** by `worker_url`, `model_id`, `mode`, `outcome` (recorded after dispatch; blind to pre-dispatch drops) |
| `sgl_router_request_duration_seconds` | Histogram | End-to-end request latency by `model_id` | | `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_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_active_load` | Gauge | Per-worker prefill-token / decode-block load | | `sgl_router_active_load` | Gauge | Per-worker prefill-token / decode-block load |
| `sgl_router_workers` | Gauge | Registered worker count by `mode` | | `sgl_router_workers` | Gauge | Registered worker count by `mode` |
| `sgl_router_worker_health` | Gauge | Per-worker health (1=breaker admits, 0=open) | | `sgl_router_worker_health` | Gauge | Per-worker health (1=breaker admits, 0=open) |
experimental/sgl-router/src/proxy/mod.rs
+15 -7
View File
@@ -164,7 +164,7 @@ impl Proxy {
/// for the full streaming lifetime — without which a long-running SSE /// for the full streaming lifetime — without which a long-running SSE
/// response would under-report load. /// response would under-report load.
// Each parameter is a distinct, required input to a single upstream // Each parameter is a distinct, required input to a single upstream
// forward (target, breaker, path, headers, body, plus the two // forward (target, breaker, path, headers, body, plus the
// streaming-lifetime callbacks). Bundling them into a struct purely to // streaming-lifetime callbacks). Bundling them into a struct purely to
// satisfy the arg-count heuristic would add indirection without clarity. // satisfy the arg-count heuristic would add indirection without clarity.
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
@@ -177,6 +177,7 @@ impl Proxy {
body: Bytes, body: Bytes,
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>>,
) -> Result<Response<Body>, ApiError> { ) -> Result<Response<Body>, ApiError> {
if !breaker.allow() { if !breaker.allow() {
return Err(ApiError::BreakerOpen { return Err(ApiError::BreakerOpen {
@@ -217,23 +218,30 @@ impl Proxy {
// is recorded as a failure. For 5xx headers we record_failure // is recorded as a failure. For 5xx headers we record_failure
// up front and skip the pump hook (the body we surface is the // up front and skip the pump hook (the body we surface is the
// error response — its stream completing is not a worker win). // error response — its stream completing is not a worker win).
let on_complete: Option<Box<dyn FnOnce(bool) + Send + 'static>> = let caller_end_hook = if status.is_success() {
on_stream_end
} else {
None
};
let on_complete: Option<Box<dyn FnOnce(sse::StreamEnd) + Send + 'static>> =
if status.is_server_error() { if status.is_server_error() {
breaker.record_failure(); breaker.record_failure();
None None
} else { } else {
let breaker_for_hook = Arc::clone(breaker); let breaker_for_hook = Arc::clone(breaker);
Some(Box::new(move |ok| { Some(Box::new(move |end| {
if ok { if end.transport_ok {
breaker_for_hook.record_success(); breaker_for_hook.record_success();
} else { } else {
breaker_for_hook.record_failure(); breaker_for_hook.record_failure();
} }
if let Some(hook) = caller_end_hook {
hook(end);
}
})) }))
}; };
// Only record TTFT for successful streams — a 4xx/5xx error body // Only record TTFT for successful streams; error-body chunks are not
// streaming back is not a generated token, so drop the hook for // generated tokens.
// non-2xx responses.
let first_byte_hook = if status.is_success() { let first_byte_hook = if status.is_success() {
on_first_byte on_first_byte
} else { } else {
experimental/sgl-router/src/proxy/sse.rs
+208 -23
View File
@@ -11,6 +11,54 @@ use bytes::Bytes;
use futures::{FutureExt, StreamExt}; use futures::{FutureExt, StreamExt};
use tokio_stream::wrappers::ReceiverStream; use tokio_stream::wrappers::ReceiverStream;
/// How the SSE pump ended, reported to the `on_complete` hook.
#[derive(Debug, Clone, Copy)]
pub struct StreamEnd {
/// No upstream stream error and no pump panic.
pub transport_ok: bool,
/// An SSE error event (`data: {"error"...}`) rode the stream.
pub saw_error_event: bool,
/// The client dropped the response body before upstream finished.
pub client_disconnect: bool,
}
/// A `data:` line whose payload's first JSON key is `error` — tolerant of
/// SSE-legal framing variants (no space after `data:`, whitespace after `{`),
/// so the match is anchored to the spec rather than one serializer's bytes.
fn is_error_event_line(line: &[u8]) -> bool {
line.strip_prefix(b"data:")
.map(|p| p.trim_ascii_start())
.and_then(|p| p.strip_prefix(b"{"))
.map(|p| p.trim_ascii_start())
.is_some_and(|p| p.starts_with(b"\"error\""))
}
/// Line-start bytes that suffice to decide `is_error_event_line`.
const LINE_PROBE: usize = 32;
/// Finds error events emitted after an SSE response commits a 200.
/// Line-anchored, so lookalike text inside event payloads cannot match.
#[derive(Default)]
struct ErrorEventScanner {
line_start: Vec<u8>,
}
impl ErrorEventScanner {
fn feed(&mut self, chunk: &[u8]) -> bool {
let mut hit = false;
for (i, segment) in chunk.split(|&b| b == b'\n').enumerate() {
if i > 0 {
hit |= is_error_event_line(&self.line_start);
self.line_start.clear();
}
let room = LINE_PROBE - self.line_start.len();
self.line_start
.extend_from_slice(&segment[..segment.len().min(room)]);
}
hit
}
}
/// Bridge a byte stream into an axum Body that streams chunks unchanged. /// Bridge a byte stream into an axum Body that streams chunks unchanged.
/// ///
/// Spawns one tokio task per stream so the handler can return immediately. /// Spawns one tokio task per stream so the handler can return immediately.
@@ -47,14 +95,8 @@ use tokio_stream::wrappers::ReceiverStream;
/// guard scope). /// guard scope).
/// ///
/// # Completion hook /// # Completion hook
/// When `on_complete` is `Some`, the closure runs exactly once when the /// When `on_complete` is `Some`, it runs exactly once when the pump task
/// pump task finishes. The bool argument is `true` on clean stream end /// finishes with the transport, SSE error-event, and client-disconnect state.
/// (including a clean client disconnect after at least the headers
/// landed cleanly), `false` on upstream stream error or pump panic.
/// `forward_streaming_to` passes a closure that records the worker's
/// circuit-breaker outcome — without this hook, a worker that returns
/// 2xx headers and then drops the stream mid-flight would stay credited
/// as healthy.
/// ///
/// # First-byte hook /// # First-byte hook
/// When `on_first_byte` is `Some`, the closure runs exactly once, the moment /// When `on_first_byte` is `Some`, the closure runs exactly once, the moment
@@ -65,7 +107,7 @@ use tokio_stream::wrappers::ReceiverStream;
pub fn bytes_stream_to_body<S, E>( pub fn bytes_stream_to_body<S, E>(
stream: S, stream: S,
stream_guards: Option<Box<dyn Send + 'static>>, stream_guards: Option<Box<dyn Send + 'static>>,
on_complete: Option<Box<dyn FnOnce(bool) + Send + 'static>>, on_complete: Option<Box<dyn FnOnce(StreamEnd) + Send + 'static>>,
on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>, on_first_byte: Option<Box<dyn FnOnce() + Send + 'static>>,
) -> Body ) -> Body
where where
@@ -75,11 +117,14 @@ where
let (tx, rx) = tokio::sync::mpsc::channel(64); let (tx, rx) = tokio::sync::mpsc::channel(64);
tokio::spawn(async move { tokio::spawn(async move {
let tx_for_panic = tx.clone(); let tx_for_panic = tx.clone();
// Capture the pump's outcome so we can report it through `on_complete` let outcome = Arc::new(parking_lot::Mutex::new(StreamEnd {
// AFTER `pump.catch_unwind()` settles. The closure inside owns transport_ok: true,
// `outcome_setter`; the outer scope reads `outcome_holder` once. saw_error_event: false,
let outcome_holder = Arc::new(parking_lot::Mutex::new(true)); client_disconnect: false,
let outcome_setter = Arc::clone(&outcome_holder); }));
let outcome_setter = Arc::clone(&outcome);
// `None` once an error event is found — the scan is done for good.
let mut scanner = Some(ErrorEventScanner::default());
let pump = AssertUnwindSafe(async move { let pump = AssertUnwindSafe(async move {
// Hold the guards for the task's lifetime — dropped when this // Hold the guards for the task's lifetime — dropped when this
// block exits (stream done or client disconnect). Leading // block exits (stream done or client disconnect). Leading
@@ -95,17 +140,19 @@ where
std::io::Error::other(msg) std::io::Error::other(msg)
}); });
let is_err_chunk = item.is_err(); let is_err_chunk = item.is_err();
// Fire the time-to-first-token hook on the first successful match &item {
// chunk from upstream. `take()` makes it fire at most once; Ok(bytes) => {
// an error-first stream never produced a token, so it's left // TTFT hook: at most once (`take()`); an error-first
// unfired (and dropped on task end). // stream never produced a token, so it stays unfired.
if !is_err_chunk {
if let Some(hook) = on_first_byte.take() { if let Some(hook) = on_first_byte.take() {
hook(); hook();
} }
if scanner.as_mut().is_some_and(|scanner| scanner.feed(bytes)) {
outcome_setter.lock().saw_error_event = true;
scanner = None;
} }
if is_err_chunk { }
*outcome_setter.lock() = false; Err(_) => outcome_setter.lock().transport_ok = false,
} }
if tx.send(item).await.is_err() { if tx.send(item).await.is_err() {
// Receiver dropped. If we were about to ship an upstream // Receiver dropped. If we were about to ship an upstream
@@ -114,6 +161,7 @@ where
// not a router-side fault. // not a router-side fault.
if !is_err_chunk { if !is_err_chunk {
tracing::debug!("SSE client disconnected mid-stream"); tracing::debug!("SSE client disconnected mid-stream");
outcome_setter.lock().client_disconnect = true;
} }
break; break;
} }
@@ -139,8 +187,9 @@ where
.await; .await;
} }
if let Some(hook) = on_complete { if let Some(hook) = on_complete {
let ok = !panicked && *outcome_holder.lock(); let mut end = *outcome.lock();
hook(ok); end.transport_ok &= !panicked;
hook(end);
} }
}); });
Body::from_stream(ReceiverStream::new(rx)) Body::from_stream(ReceiverStream::new(rx))
@@ -394,4 +443,140 @@ mod tests {
"pump drained the entire upstream after client disconnect ({final_polls} polls); the break-on-tx.send-err path is dead" "pump drained the entire upstream after client disconnect ({final_polls} polls); the break-on-tx.send-err path is dead"
); );
} }
#[test]
fn error_event_scanner_detects_engine_error_event() {
let mut scanner = ErrorEventScanner::default();
assert!(!scanner.feed(b"data: {\"choices\": [{\"delta\": {\"content\": \"hi\"}}]}\n\n"));
assert!(
scanner.feed(b"data: {\"error\": {\"message\": \"queue is full\", \"code\": 503}}\n\n")
);
}
#[test]
fn error_event_scanner_detects_error_split_across_chunks() {
let mut scanner = ErrorEventScanner::default();
assert!(!scanner.feed(b"data: {\"err"));
assert!(scanner.feed(b"or\": {\"code\": 503}}\n\n"));
}
#[test]
fn error_event_scanner_ignores_error_text_inside_content() {
let mut scanner = ErrorEventScanner::default();
assert!(!scanner.feed(
b"data: {\"choices\": [{\"delta\": {\"content\": \"data: {\\\"error\\\" is how it looks\"}}]}\n\n",
));
assert!(!scanner.feed(b"data: [DONE]\n\n"));
}
#[test]
fn error_event_scanner_accepts_sse_framing_variants() {
for event in [
&b"data:{\"error\": {\"code\": 503}}\n\n"[..],
b"data: { \"error\": {\"code\": 503}}\n\n",
b"data: {\"error\": \"queue full\"}\n\n",
] {
assert!(
ErrorEventScanner::default().feed(event),
"missed variant: {}",
String::from_utf8_lossy(event)
);
}
}
#[test]
fn error_event_scanner_bounds_line_buffer() {
let mut scanner = ErrorEventScanner::default();
let big = vec![b'x'; 1 << 20];
assert!(!scanner.feed(&big));
assert_eq!(scanner.line_start.len(), LINE_PROBE);
}
fn body_with_completion(
chunks: Vec<Result<Bytes, std::io::Error>>,
) -> (Body, tokio::sync::oneshot::Receiver<StreamEnd>) {
let (tx, rx) = tokio::sync::oneshot::channel();
let body = bytes_stream_to_body(
stream::iter(chunks),
None,
Some(Box::new(move |end| {
let _ = tx.send(end);
})),
None,
);
(body, rx)
}
async fn stream_end(rx: tokio::sync::oneshot::Receiver<StreamEnd>) -> StreamEnd {
rx.await.expect("completion hook dropped")
}
#[tokio::test]
async fn completion_reports_error_event() {
let chunks = vec![
Ok::<Bytes, std::io::Error>(Bytes::from_static(b"data: {\"err")),
Ok(Bytes::from_static(b"or\": {\"code\": 503}}\n\n")),
];
let (body, completion) = body_with_completion(chunks);
let _ = body.collect().await.unwrap();
let end = stream_end(completion).await;
assert!(end.transport_ok);
assert!(end.saw_error_event);
assert!(!end.client_disconnect);
}
#[tokio::test]
async fn completion_reports_error_event_then_transport_error() {
let chunks = vec![
Ok(Bytes::from_static(
b"data: {\"error\": {\"code\": 503}}\n\n",
)),
Err(std::io::Error::other("connection reset")),
];
let (body, completion) = body_with_completion(chunks);
let _ = body.collect().await;
let end = stream_end(completion).await;
assert!(!end.transport_ok);
assert!(end.saw_error_event);
}
#[tokio::test]
async fn completion_reports_upstream_error() {
let chunks = vec![Err(std::io::Error::other("upstream failed"))];
let (body, completion) = body_with_completion(chunks);
let _ = body.collect().await;
let end = stream_end(completion).await;
assert!(!end.transport_ok);
assert!(!end.saw_error_event);
assert!(!end.client_disconnect);
}
#[tokio::test]
async fn completion_reports_clean_end() {
let chunks = vec![Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"data: [DONE]\n\n",
))];
let (body, completion) = body_with_completion(chunks);
let _ = body.collect().await.unwrap();
let end = stream_end(completion).await;
assert!(end.transport_ok);
assert!(!end.saw_error_event);
assert!(!end.client_disconnect);
}
#[tokio::test]
async fn completion_reports_client_disconnect() {
let chunks = std::iter::repeat_with(|| {
Ok::<Bytes, std::io::Error>(Bytes::from_static(b"data: x\n\n"))
})
.take(1000)
.collect();
let (body, completion) = body_with_completion(chunks);
let mut stream = body.into_data_stream();
let _ = stream.next().await;
drop(stream);
let end = stream_end(completion).await;
assert!(end.transport_ok);
assert!(end.client_disconnect);
}
} }
experimental/sgl-router/src/server/metrics.rs
+140
View File
@@ -23,6 +23,7 @@
//! | `sgl_router_worker_requests_total` | Counter | `worker_url`, `model_id`, `mode`, `outcome` | //! | `sgl_router_worker_requests_total` | Counter | `worker_url`, `model_id`, `mode`, `outcome` |
//! | `sgl_router_request_duration_seconds` | Histogram | `model_id` | //! | `sgl_router_request_duration_seconds` | Histogram | `model_id` |
//! | `sgl_router_ttft_seconds` | Histogram | `model_id` | //! | `sgl_router_ttft_seconds` | Histogram | `model_id` |
//! | `sgl_router_stream_outcome_total` | Counter | `worker_url`, `model_id`, `outcome` |
//! | `sgl_router_active_load` | Gauge | `worker_url`, `kind` | //! | `sgl_router_active_load` | Gauge | `worker_url`, `kind` |
//! | `sgl_router_workers` | Gauge | `mode` | //! | `sgl_router_workers` | Gauge | `mode` |
//! | `sgl_router_worker_health` | Gauge | `worker_url` | //! | `sgl_router_worker_health` | Gauge | `worker_url` |
@@ -49,6 +50,7 @@
//! The exposition is text/plain; version=0.0.4 per the Prometheus spec. //! The exposition is text/plain; version=0.0.4 per the Prometheus spec.
use crate::config::PolicyKind; use crate::config::PolicyKind;
use crate::proxy::sse::StreamEnd;
use parking_lot::Mutex; use parking_lot::Mutex;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
@@ -104,6 +106,39 @@ impl RequestOutcome {
} }
} }
/// Final outcome of a 2xx SSE stream.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamOutcome {
/// Stream ended without errors.
Ok,
/// The engine sent a `data: {"error"...}` SSE event.
StreamErrorEvent,
/// The upstream byte stream failed.
UpstreamError,
/// The client disconnected before the stream finished.
ClientDisconnect,
}
pub(crate) fn classify_stream_end(end: StreamEnd) -> StreamOutcome {
match (end.transport_ok, end.saw_error_event, end.client_disconnect) {
(false, _, _) => StreamOutcome::UpstreamError,
(_, true, _) => StreamOutcome::StreamErrorEvent,
(_, _, true) => StreamOutcome::ClientDisconnect,
_ => StreamOutcome::Ok,
}
}
impl StreamOutcome {
fn as_str(self) -> &'static str {
match self {
Self::Ok => "ok",
Self::StreamErrorEvent => "stream_error_event",
Self::UpstreamError => "upstream_error",
Self::ClientDisconnect => "client_disconnect",
}
}
}
/// Worker dispatch mode label — narrowed to the three modes the policy /// Worker dispatch mode label — narrowed to the three modes the policy
/// resolver distinguishes. The `Plain` variant covers the non-PD case. /// resolver distinguishes. The `Plain` variant covers the non-PD case.
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
@@ -233,6 +268,7 @@ pub struct MetricsRegistry {
// on `worker_requests_total` / the worker gauges instead. // on `worker_requests_total` / the worker gauges instead.
request_duration: Mutex<HashMap<String, Histogram>>, request_duration: Mutex<HashMap<String, Histogram>>,
ttft_seconds: Mutex<HashMap<String, Histogram>>, ttft_seconds: Mutex<HashMap<String, Histogram>>,
stream_outcome_total: Mutex<HashMap<StreamOutcomeKey, Arc<AtomicU64>>>,
active_load: Mutex<HashMap<ActiveLoadKey, Arc<AtomicI64>>>, active_load: Mutex<HashMap<ActiveLoadKey, Arc<AtomicI64>>>,
stale_requests_total: Mutex<HashMap<&'static str, Arc<AtomicU64>>>, stale_requests_total: Mutex<HashMap<&'static str, Arc<AtomicU64>>>,
decode_affinity_total: Mutex<HashMap<&'static str, Arc<AtomicU64>>>, decode_affinity_total: Mutex<HashMap<&'static str, Arc<AtomicU64>>>,
@@ -271,6 +307,14 @@ struct EdgeResponseKey {
status_code: u16, status_code: u16,
} }
/// Labels for `sgl_router_stream_outcome_total`.
#[derive(Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Clone)]
struct StreamOutcomeKey {
worker_url: String,
model_id: String,
outcome: &'static str,
}
/// Per-worker state sampled from the [`crate::workers::WorkerRegistry`] at /// Per-worker state sampled from the [`crate::workers::WorkerRegistry`] at
/// scrape time and rendered as the `sgl_router_workers` / /// scrape time and rendered as the `sgl_router_workers` /
/// `sgl_router_worker_*` gauge families. Built by the `/metrics` route from /// `sgl_router_worker_*` gauge families. Built by the `/metrics` route from
@@ -429,6 +473,22 @@ impl MetricsRegistry {
hist.observe(seconds); hist.observe(seconds);
} }
/// Record the final outcome of a 2xx stream.
pub fn record_stream_outcome(&self, worker_url: &str, model_id: &str, outcome: StreamOutcome) {
let key = StreamOutcomeKey {
worker_url: worker_url.to_owned(),
model_id: model_id.to_owned(),
outcome: outcome.as_str(),
};
let counter = self
.stream_outcome_total
.lock()
.entry(key)
.or_default()
.clone();
counter.fetch_add(1, Ordering::Relaxed);
}
/// Bump the edge counter `responses_total{route,method,status_code}`. Called /// Bump the edge counter `responses_total{route,method,status_code}`. Called
/// at the middleware, so it captures every outcome — incl. early-exit /// at the middleware, so it captures every outcome — incl. early-exit
/// 400/413/503 that the old per-handler site skipped. /// 400/413/503 that the old per-handler site skipped.
@@ -684,6 +744,26 @@ impl MetricsRegistry {
} }
drop(guard); drop(guard);
// Final outcomes observed after a 2xx stream's headers are committed.
out.push_str("# HELP sgl_router_stream_outcome_total Final outcome of a 2xx stream.\n");
out.push_str("# TYPE sgl_router_stream_outcome_total counter\n");
let guard = self.stream_outcome_total.lock();
let mut entries: Vec<(&StreamOutcomeKey, u64)> = guard
.iter()
.map(|(k, v)| (k, v.load(Ordering::Relaxed)))
.collect();
entries.sort();
for (key, value) in entries {
out.push_str(&format!(
"sgl_router_stream_outcome_total{{worker_url=\"{}\",model_id=\"{}\",outcome=\"{}\"}} {}\n",
escape_label(&key.worker_url),
escape_label(&key.model_id),
key.outcome,
value,
));
}
drop(guard);
// responses_total — edge, by route/method/status (incl. early-exit 400/413/503) // responses_total — edge, by route/method/status (incl. early-exit 400/413/503)
out.push_str( out.push_str(
"# HELP sgl_router_responses_total Responses returned at the router HTTP edge, by route, method and HTTP status code.\n", "# HELP sgl_router_responses_total Responses returned at the router HTTP edge, by route, method and HTTP status code.\n",
@@ -1011,6 +1091,13 @@ fn escape_label(s: &str) -> String {
mod tests { mod tests {
use super::*; use super::*;
fn assert_metric_line(output: &str, expected: &str) {
assert!(
output.lines().any(|line| line == expected),
"missing metric line `{expected}`; got:\n{output}"
);
}
#[test] #[test]
fn empty_registry_renders_only_help_lines() { fn empty_registry_renders_only_help_lines() {
let reg = MetricsRegistry::new(); let reg = MetricsRegistry::new();
@@ -1151,6 +1238,59 @@ mod tests {
} }
} }
#[test]
fn stream_outcome_precedence() {
use StreamOutcome::*;
for (transport_ok, saw_error_event, client_disconnect, expected) in [
(false, false, false, UpstreamError),
(false, false, true, UpstreamError),
(false, true, false, UpstreamError),
(false, true, true, UpstreamError),
(true, false, false, Ok),
(true, false, true, ClientDisconnect),
(true, true, false, StreamErrorEvent),
(true, true, true, StreamErrorEvent),
] {
let end = StreamEnd {
transport_ok,
saw_error_event,
client_disconnect,
};
assert_eq!(classify_stream_end(end), expected, "{end:?}");
}
}
#[test]
fn record_stream_outcome_emits_labelled_counter_lines() {
let reg = MetricsRegistry::new();
reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::Ok);
reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::Ok);
reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::StreamErrorEvent);
reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::UpstreamError);
reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::ClientDisconnect);
let out = reg.render();
for expected in [
r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="ok"} 2"#,
r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="stream_error_event"} 1"#,
r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="upstream_error"} 1"#,
r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="client_disconnect"} 1"#,
] {
assert_metric_line(&out, expected);
}
}
#[test]
fn stream_outcome_absent_until_recorded() {
let reg = MetricsRegistry::new();
let out = reg.render();
assert!(out.contains("# TYPE sgl_router_stream_outcome_total counter"));
assert!(
!out.contains("sgl_router_stream_outcome_total{"),
"no series until an outcome is recorded; got:\n{out}",
);
}
#[test] #[test]
fn record_response_counts_by_route_method_status_code() { fn record_response_counts_by_route_method_status_code() {
let reg = MetricsRegistry::new(); let reg = MetricsRegistry::new();
experimental/sgl-router/src/server/routes/chat.rs
+15 -2
View File
@@ -17,11 +17,12 @@ use crate::policies::{
request_tokens_for, ExternalPrefixSignal, PrefillProposal, ProposalKind, RequestTokens, request_tokens_for, ExternalPrefixSignal, PrefillProposal, ProposalKind, RequestTokens,
SelectionContext, SelectionContext,
}; };
use crate::proxy::sse::StreamEnd;
use crate::server::app_context::AppContext; use crate::server::app_context::AppContext;
use crate::server::error::ApiError; use crate::server::error::ApiError;
use crate::server::metrics::{ use crate::server::metrics::{
MetricsRegistry, PolicySelectionFailureReason, RequestOutcome, StaleRequestOutcome, classify_stream_end, MetricsRegistry, PolicySelectionFailureReason, RequestOutcome,
WorkerModeLabel, StaleRequestOutcome, WorkerModeLabel,
}; };
use crate::workers::{LoadGuard, Worker}; use crate::workers::{LoadGuard, Worker};
use axum::body::Body; use axum::body::Body;
@@ -764,6 +765,16 @@ pub async fn chat_completions(
start, start,
}; };
// Classifies a 2xx stream after its headers are committed. Takes the
// streaming worker's URL (Final D in PD mode).
let make_stream_end_hook = |worker_url: String| -> Box<dyn FnOnce(StreamEnd) + Send + 'static> {
let metrics = Arc::clone(&ctx.metrics);
let model = metrics_model.clone();
Box::new(move |end| {
metrics.record_stream_outcome(&worker_url, &model, classify_stream_end(end));
})
};
// Forward the router-computed tokens to the engine as `input_ids` so it // Forward the router-computed tokens to the engine as `input_ids` so it
// skips re-tokenizing the same prompt — but only when they are // skips re-tokenizing the same prompt — but only when they are
// engine-equivalent (chat-encoder path) AND the request contains nothing // engine-equivalent (chat-encoder path) AND the request contains nothing
@@ -904,6 +915,7 @@ pub async fn chat_completions(
outgoing_body, outgoing_body,
Some(stream_guards), Some(stream_guards),
Some(make_ttft_hook()), Some(make_ttft_hook()),
Some(make_stream_end_hook(decode_worker.url.clone())),
); );
tokio::select! { tokio::select! {
biased; biased;
@@ -939,6 +951,7 @@ pub async fn chat_completions(
outgoing_body, outgoing_body,
Some(stream_guards), Some(stream_guards),
Some(make_ttft_hook()), Some(make_ttft_hook()),
Some(make_stream_end_hook(worker.url.clone())),
); );
// Bias `fetch` over the cancellation branch: a successful // Bias `fetch` over the cancellation branch: a successful
// response that completes in the same poll as the token firing // response that completes in the same poll as the token firing
experimental/sgl-router/tests/proxy/chat_routing.rs
+116 -22
View File
@@ -950,6 +950,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() {
body, body,
None, None,
None, None,
None,
) )
.await; .await;
@@ -1319,35 +1320,20 @@ async fn streaming_active_load_drops_on_client_disconnect() {
Duration::from_millis(100), Duration::from_millis(100),
) )
.await; .await;
let ctx = build_ctx_with_worker(&worker.url); let (ctx, body) = stream_chat(&worker.url).await;
let active_load = Arc::clone(&ctx.active_load); let active_load = Arc::clone(&ctx.active_load);
let app = build_router(ctx);
let req = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(
serde_json::to_vec(&serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
"stream": true,
}))
.unwrap(),
))
.unwrap();
let res = app.oneshot(req).await.unwrap();
// Read one chunk to confirm the stream is live, then drop the body. // Read one chunk to confirm the stream is live, then drop the body.
use futures::StreamExt; use futures::StreamExt;
let mut data_stream = res.into_body().into_data_stream(); let mut data_stream = body.into_data_stream();
let _first = data_stream.next().await; let _first = data_stream.next().await;
drop(data_stream); drop(data_stream);
// Wait long enough for the SSE pump to notice the receiver-drop and let expected = format!(
// exit (per `bytes_stream_to_body_breaks_on_client_disconnect` test r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="client_disconnect"}} 1"#,
// in sse.rs, that takes well under 200 ms). worker.url,
tokio::time::sleep(Duration::from_millis(300)).await; );
wait_for_metric(&ctx, &expected).await;
assert_eq!( assert_eq!(
active_load.inflight_count(), active_load.inflight_count(),
@@ -1477,3 +1463,111 @@ async fn non_streaming_error_path_drops_active_load_guard() {
"error path must drop the active-load guard", "error path must drop the active-load guard",
); );
} }
fn has_metric_line(metrics: &str, expected: &str) -> bool {
metrics.lines().any(|line| line == expected)
}
/// Send a streaming request and wait for the expected metric.
async fn stream_chat_and_render(
worker_url: &str,
expected_metric: &str,
) -> (Arc<AppContext>, String) {
let (ctx, body) = stream_chat(worker_url).await;
body.collect().await.unwrap();
let metrics = wait_for_metric(&ctx, expected_metric).await;
(ctx, metrics)
}
async fn stream_chat(worker_url: &str) -> (Arc<AppContext>, Body) {
let ctx = build_ctx_with_worker(worker_url);
let app = build_router(ctx.clone());
let req = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(Body::from(
serde_json::json!({
"model": "tiny",
"messages": [{"role": "user", "content": "hi"}],
"stream": true
})
.to_string(),
))
.unwrap();
let res = app.oneshot(req).await.unwrap();
assert_eq!(res.status(), StatusCode::OK);
(ctx, res.into_body())
}
async fn wait_for_metric(ctx: &AppContext, expected_metric: &str) -> String {
let deadline = std::time::Instant::now() + Duration::from_secs(2);
loop {
let metrics = ctx.metrics.render();
if has_metric_line(&metrics, expected_metric) {
return metrics;
}
assert!(
std::time::Instant::now() < deadline,
"timed out waiting for `{expected_metric}`; got:\n{metrics}"
);
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
/// A post-200 SSE error event is classified without affecting routing health.
#[tokio::test]
async fn streaming_error_event_records_outcome_without_tripping_breaker() {
let worker = crate::common::mock_worker::MockWorker::start(vec![
"data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n",
"data: {\"error\": {\"message\": \"The request queue is full.\", \"code\": 503}}\n\n",
"data: [DONE]\n\n",
])
.await;
let expected = format!(
r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="stream_error_event"}} 1"#,
worker.url,
);
let (ctx, metrics) = stream_chat_and_render(&worker.url, &expected).await;
assert!(has_metric_line(
&metrics,
r#"sgl_router_responses_total{route="/v1/chat/completions",method="POST",status_code="200"} 1"#
));
assert!(
ctx.registry
.all()
.iter()
.all(|worker| worker.breaker.would_allow()),
"SSE error event must not trip the circuit breaker",
);
}
#[tokio::test]
async fn streaming_clean_completion_records_ok() {
let worker = crate::common::mock_worker::MockWorker::start(vec![
"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
"data: [DONE]\n\n",
])
.await;
let expected = format!(
r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="ok"}} 1"#,
worker.url,
);
stream_chat_and_render(&worker.url, &expected).await;
}
#[tokio::test]
async fn streaming_error_event_then_transport_failure_records_upstream_error() {
let worker = crate::common::mock_worker::MockWorker::start_returning_partial_body(
StatusCode::OK,
b"data: {\"error\": {\"code\": 503}}\n\n",
)
.await;
let (ctx, body) = stream_chat(&worker.url).await;
assert!(body.collect().await.is_err());
let expected = format!(
r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="upstream_error"}} 1"#,
worker.url,
);
wait_for_metric(&ctx, &expected).await;
}