From 70ee6b1714dbed9079630d24e1ea531ec2fa9315 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 18 Aug 2026 02:08:07 -0700 Subject: [PATCH] [Rust Server] Add e2e latency metadata and fix Sarashina import (#35125) --- python/sglang/srt/models/sarashina2_vision.py | 3 +- .../src/api_server/native_api.rs | 279 ++++++++++++++++-- 2 files changed, 253 insertions(+), 29 deletions(-) diff --git a/python/sglang/srt/models/sarashina2_vision.py b/python/sglang/srt/models/sarashina2_vision.py index f58908b5d..15e254b6d 100644 --- a/python/sglang/srt/models/sarashina2_vision.py +++ b/python/sglang/srt/models/sarashina2_vision.py @@ -24,11 +24,10 @@ from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.pooler import Pooler, PoolingType from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.managers.mm_utils import ( - MultimodalDataItem, - MultimodalInputs, MultiModalityDataPaddingPatternMultimodalTokens, general_mm_embed_routine, ) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.llama import LlamaForCausalLM diff --git a/rust/sglang-server/src/api_server/native_api.rs b/rust/sglang-server/src/api_server/native_api.rs index f0db30337..835c3fb81 100644 --- a/rust/sglang-server/src/api_server/native_api.rs +++ b/rust/sglang-server/src/api_server/native_api.rs @@ -7,7 +7,10 @@ //! generate-request submission (`submit`); the shared `AppState` lives in the //! parent `api_server` module. -use std::convert::Infallible; +use std::{ + convert::Infallible, + time::{Duration, Instant}, +}; use axum::{ Json, Router, @@ -30,9 +33,50 @@ use super::guard::AbortGuard; use super::submit::submit; use crate::environ::env_bool; use crate::ids::Rid; -use crate::message::{EgressItem, GenerateBody, GenerateRequest, RequestKind, SamplingParams}; +use crate::message::{ + ChunkEvent, EgressItem, GenerateBody, GenerateRequest, RequestKind, SamplingParams, +}; use crate::utils::response::{error_response, error_value}; +/// API-local timing for one request. +/// +/// Python records time-to-first-token on the first output batch and end-to-end +/// latency when that request finishes. Keep both measurements here even though +/// `/generate` currently exposes only `e2e_latency`; this avoids putting +/// API-only timestamps onto scheduler messages. +#[derive(Clone, Debug)] +struct RequestTiming { + // TODO: Move request lifecycle timing into a dedicated tracing/metrics + // module and align its design with Python's APIServerReqTimeStats. + created_at: Instant, + time_to_first_token: Option, + e2e_latency: Option, +} + +impl RequestTiming { + fn new() -> Self { + Self { + created_at: Instant::now(), + time_to_first_token: None, + e2e_latency: None, + } + } + + fn observe_first_output(&mut self) { + self.time_to_first_token + .get_or_insert_with(|| self.created_at.elapsed()); + } + + fn finish(&mut self) { + self.e2e_latency + .get_or_insert_with(|| self.created_at.elapsed()); + } + + fn terminal_latencies(&self) -> Option<(Duration, Duration)> { + Some((self.time_to_first_token?, self.e2e_latency?)) + } +} + /// The routes this module owns, mounted by `api_server::serve`. pub(super) fn routes() -> Router { Router::new() @@ -168,6 +212,11 @@ async fn generate( return native_error(code, &e.to_string(), stream); } }; + // Python starts APIServerReqTimeStats after request normalization and before + // tokenization / multimodal preprocessing / scheduler dispatch. Start at the + // equivalent boundary: into_requests() has normalized the body, while prefetch + // and every downstream stage are still ahead of us. + let timing = RequestTiming::new(); // Media I/O (URL downloads, file reads) happens here, on the API runtime // — never on the MM worker pool (see `prefetch`). if let Err(e) = super::prefetch::prefetch_all(&mut payloads).await { @@ -179,9 +228,9 @@ async fn generate( .into_iter() .next() .expect("into_requests yields >=1 payload"); - generate_single(&state, payload, stream).await + generate_single(&state, payload, stream, timing).await } else { - generate_batch(&state, payloads, stream).await + generate_batch(&state, payloads, stream, timing).await } } @@ -190,7 +239,12 @@ async fn generate( /// /// A single (non-batched) `/generate`: submit one request, then either stream its /// SSE frames or fold to one unary response. -async fn generate_single(state: &AppState, req: GenerateRequest, stream: bool) -> Response { +async fn generate_single( + state: &AppState, + req: GenerateRequest, + stream: bool, + timing: RequestTiming, +) -> Response { // `return_text_in_logprobs` is decoded on the detok shard into `*_txt`, so // `frame_value` just reads them — no tokenizer needed here. let (rid_str, mut rx) = match submit(state, RequestKind::Generate(Box::new(req)), stream).await @@ -209,13 +263,13 @@ async fn generate_single(state: &AppState, req: GenerateRequest, stream: bool) - // A single request is a 1-element batch without the `index` field — reuse // the same stream so the frame/abort/truncation logic lives in one place. use futures::StreamExt; - let s = generation_event_stream(vec![(rid_str, rx)], guard, incremental, false) + let s = generation_event_stream(vec![(rid_str, rx, timing)], guard, incremental, false) .map(|data| Ok::<_, Infallible>(Event::default().data(data))); Sse::new(s).into_response() } else { // Unary: fold to the terminal, respond once. Disarm only on a real terminal // (a truncation leaves the guard armed so the scheduler work is aborted). - let (status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await; + let (status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing(), timing).await; if terminal { guard.disarm(&rid_str); } @@ -228,12 +282,18 @@ async fn generate_single(state: &AppState, req: GenerateRequest, stream: bool) - async fn drain_unary( rx: &mut mpsc::Receiver, rid_str: &str, + mut timing: RequestTiming, ) -> (StatusCode, serde_json::Value, bool) { let mut acc = OutputAccumulator::default(); while let Some(item) = rx.recv().await { match item { - EgressItem::Frame(out) => acc.fold(&out), + EgressItem::Frame(out) => { + timing.observe_first_output(); + acc.fold(&out); + } EgressItem::Done(out) => { + timing.observe_first_output(); + timing.finish(); acc.fold(&out); let final_out = acc.into_output(); // A validation abort carries its own HTTP status + diagnostic. @@ -246,9 +306,12 @@ async fn drain_unary( StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); return (status, error_value(code, message), true); } - return (StatusCode::OK, frame_value(&final_out, rid_str), true); + let mut value = frame_value(&final_out, rid_str); + add_e2e_latency(&mut value, &timing); + return (StatusCode::OK, value, true); } EgressItem::Error(e) => { + timing.finish(); let code = e.http_status(); let status = StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); @@ -267,14 +330,16 @@ async fn drain_unary( } /// Batch `/generate`: submit all sub-requests first (scheduler runs them together), -/// then either (unary) drain each in order into a JSON array, or (streaming) -/// multiplex their streams into one SSE response, each frame carrying its `index`. +/// then either (unary) drain them concurrently into a request-ordered JSON array, +/// or (streaming) multiplex their streams into one SSE response, each frame carrying +/// its `index`. /// One [`AbortGuard`] covers the batch. A failed unary item is its own /// `{ "error": … }` entry; the batch response is 200. async fn generate_batch( state: &AppState, requests: Vec, stream: bool, + timing: RequestTiming, ) -> Response { // No cross-item rid collision to worry about: `into_requests` rejected duplicate // rids within this batch, and `Rid::from_client` made each one unique against @@ -285,7 +350,7 @@ async fn generate_batch( match submit(state, RequestKind::Generate(Box::new(req)), stream).await { Ok((rid, rx)) => { guard.arm(rid.clone()); - receivers.push((rid, rx)); + receivers.push((rid, rx, timing.clone())); } Err(resp) => return resp, } @@ -300,10 +365,20 @@ async fn generate_batch( .map(|data| Ok::<_, Infallible>(Event::default().data(data))); Sse::new(s).into_response() } else { - // Unary: drain each in order (already all submitted, so they run together). - let mut results = Vec::with_capacity(receivers.len()); - for (rid_str, mut rx) in receivers { - let (_status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await; + // Unary: poll every item concurrently, as Python's gather does. `join_all` + // preserves input order for the final JSON array, while each drain observes + // its own terminal output promptly (important for per-item e2e_latency). + let drained = futures::future::join_all(receivers.into_iter().map( + |(rid_str, mut rx, request_timing)| async move { + let client_rid = rid_str.client_facing().to_owned(); + let (_status, value, terminal) = + drain_unary(&mut rx, &client_rid, request_timing).await; + (rid_str, value, terminal) + }, + )) + .await; + let mut results = Vec::with_capacity(drained.len()); + for (rid_str, value, terminal) in drained { if terminal { guard.disarm(&rid_str); } @@ -335,7 +410,7 @@ async fn recv_indexed( /// `with_index` tags each frame (batch only), `incremental` = delta vs cumulative, /// `guard` aborts unfinished on drop. fn generation_event_stream( - receivers: Vec<(Rid, mpsc::Receiver)>, + receivers: Vec<(Rid, mpsc::Receiver, RequestTiming)>, mut guard: AbortGuard, incremental: bool, with_index: bool, @@ -344,7 +419,14 @@ fn generation_event_stream( use futures::StreamExt; let n = receivers.len(); - let rid_strs: Vec = receivers.iter().map(|(rid, _)| rid.clone()).collect(); + let rid_strs: Vec = receivers + .iter() + .map(|(rid, _, _)| rid.clone()) + .collect(); + let mut timings: Vec = receivers + .iter() + .map(|(_, _, timing)| timing.clone()) + .collect(); let mut accs: Vec = (0..n).map(|_| OutputAccumulator::default()).collect(); @@ -354,7 +436,7 @@ fn generation_event_stream( // Poll all receivers concurrently; re-arm a receiver's future after each // non-terminal frame so its stream keeps flowing. let mut futs = futures::stream::FuturesUnordered::new(); - for (i, (_, rx)) in receivers.into_iter().enumerate() { + for (i, (_, rx, _)) in receivers.into_iter().enumerate() { futs.push(recv_indexed(i, rx)); } @@ -375,6 +457,7 @@ fn generation_event_stream( for item in items { match item { EgressItem::Frame(out) => { + timings[i].observe_first_output(); accs[i].fold(&out); if incremental { yield stream_frame_string(out, &accs[i], true, rid_strs[i].client_facing(), idx(i)); @@ -383,10 +466,15 @@ fn generation_event_stream( } } EgressItem::Done(out) => { + timings[i].observe_first_output(); + timings[i].finish(); accs[i].fold(&out); terminal = Some(out); } - EgressItem::Error(e) => failed = Some(e), + EgressItem::Error(e) => { + timings[i].finish(); + failed = Some(e); + } EgressItem::Control(_) | EgressItem::Data(_) => {} // never on /generate } } @@ -399,7 +487,14 @@ fn generation_event_stream( // carries the full cumulative state, so any coalesced ones are moot. yield match out.finish_reason.as_ref().and_then(|f| f.abort_status()) { Some((code, message)) => tag_value(error_value(code, message), idx(i)), - None => stream_frame_string(out, &accs[i], incremental, rid_strs[i].client_facing(), idx(i)), + None => terminal_stream_frame_string( + out, + &accs[i], + incremental, + rid_strs[i].client_facing(), + idx(i), + &timings[i], + ), }; guard.disarm(&rid_strs[i]); // terminal → not re-pushed } else { @@ -413,12 +508,41 @@ fn generation_event_stream( } } +/// Python's `e2e_latency` is `finished_time - created_time`, in seconds, and is +/// attached only when the request finishes. The Rust native API owns the same +/// lifecycle boundary, so it adds the value while handling the terminal egress +/// item rather than putting API-only timing onto every scheduler `ChunkEvent`. +fn add_e2e_latency(value: &mut serde_json::Value, timing: &RequestTiming) { + let (time_to_first_token, e2e_latency) = timing + .terminal_latencies() + .expect("a successful terminal output has complete request timing"); + debug_assert!(time_to_first_token <= e2e_latency); + value["meta_info"]["e2e_latency"] = serde_json::json!(e2e_latency.as_secs_f64()); +} + +/// Render a terminal streaming frame. Intermediate cumulative frames keep the +/// memoized fast path; the one terminal frame uses the Value path so it can carry +/// the request-local `e2e_latency`, exactly as Python does. +fn terminal_stream_frame_string( + out: ChunkEvent, + acc: &OutputAccumulator, + incremental: bool, + rid_str: &str, + index: Option, + timing: &RequestTiming, +) -> String { + let mut value = super::frame::stream_frame_value(out, acc, incremental, rid_str); + add_e2e_latency(&mut value, timing); + tag_value(value, index) +} + #[cfg(test)] mod tests { use super::*; use crate::message::ChunkEvent; use crate::tokenizer_manager::Senders; use futures::StreamExt; + use std::time::Duration; fn senders() -> Senders { Senders { tm: flume::unbounded().0, @@ -453,6 +577,101 @@ mod tests { serde_json::from_str(s).expect("frame is JSON") } + fn timed_receiver( + rid: u64, + rx: mpsc::Receiver, + ) -> (Rid, mpsc::Receiver, RequestTiming) { + ( + Rid::from(rid.to_string()), + rx, + RequestTiming { + created_at: Instant::now() - Duration::from_millis(10), + time_to_first_token: None, + e2e_latency: None, + }, + ) + } + + #[test] + fn request_timing_records_ttft_once_and_e2e_on_finish() { + let mut timing = RequestTiming { + created_at: Instant::now() - Duration::from_millis(10), + time_to_first_token: None, + e2e_latency: None, + }; + assert!(timing.terminal_latencies().is_none()); + + timing.observe_first_output(); + let time_to_first_token = timing.time_to_first_token.unwrap(); + timing.observe_first_output(); + assert_eq!( + timing.time_to_first_token, + Some(time_to_first_token), + "later output must not overwrite TTFT" + ); + + timing.finish(); + let (recorded_ttft, e2e_latency) = timing.terminal_latencies().unwrap(); + assert_eq!(recorded_ttft, time_to_first_token); + assert!(e2e_latency >= recorded_ttft); + + timing.finish(); + assert_eq!( + timing.terminal_latencies(), + Some((recorded_ttft, e2e_latency)), + "later terminal handling must not overwrite E2E latency" + ); + } + + /// The native unary response uses the same names and meanings as Python's + /// TokenizerManager metadata, and adds e2e_latency only on the terminal + /// result. The timer is seconds from normalized-request acceptance through + /// terminal-output handling. + #[tokio::test] + async fn unary_terminal_meta_info_matches_python_semantics() { + let (tx, mut rx) = mpsc::channel(2); + tx.send(EgressItem::Done(ChunkEvent { + rid: "internal-rid".into(), + text: "ok".into(), + token_ids: vec![7, 8], + prompt_tokens: 5, + completion_tokens: 2, + finish_reason: serde_json::from_value(serde_json::json!({ + "type": "length", + "length": 2 + })) + .expect("finish reason must parse"), + ..Default::default() + })) + .await + .unwrap(); + + let timing = RequestTiming { + created_at: Instant::now() - Duration::from_millis(20), + time_to_first_token: None, + e2e_latency: None, + }; + let (status, value, terminal) = drain_unary(&mut rx, "client-rid", timing).await; + assert_eq!(status, StatusCode::OK); + assert!(terminal); + assert_eq!(value["meta_info"]["id"], "client-rid"); + assert_eq!(value["meta_info"]["prompt_tokens"], 5); + assert_eq!(value["meta_info"]["completion_tokens"], 2); + assert_eq!( + value["meta_info"]["finish_reason"], + serde_json::json!({"type": "length", "length": 2}) + ); + assert!( + value["meta_info"]["e2e_latency"].as_f64().unwrap() >= 0.020, + "latency is expressed in seconds from request creation" + ); + assert!( + value["meta_info"].get("ttft").is_none() + && value["meta_info"].get("time_to_first_token").is_none(), + "TTFT is recorded internally but is not part of this PR's API" + ); + } + /// Two sub-requests' frames interleave into one stream, each tagged with its /// batch `index`; text accumulates per item; `[DONE]` comes only after both /// terminate, then the stream ends. @@ -460,7 +679,7 @@ mod tests { async fn interleaves_indexes_and_accumulates() { let (tx0, rx0) = mpsc::channel(8); let (tx1, rx1) = mpsc::channel(8); - let receivers = vec![("10".into(), rx0), ("11".into(), rx1)]; + let receivers = vec![timed_receiver(10, rx0), timed_receiver(11, rx1)]; let stream = generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true); futures::pin_mut!(stream); @@ -481,11 +700,13 @@ mod tests { assert_eq!(v["index"], 0); assert_eq!(v["text"], "a!", "cumulative per item"); assert_eq!(v["meta_info"]["finish_reason"]["type"], "length"); + assert!(v["meta_info"]["e2e_latency"].as_f64().unwrap() >= 0.010); tx1.send(done(11, "?")).await.unwrap(); let v = parse(&stream.next().await.unwrap()); assert_eq!(v["index"], 1); assert_eq!(v["text"], "b?"); + assert!(v["meta_info"]["e2e_latency"].as_f64().unwrap() >= 0.010); assert_eq!(stream.next().await.unwrap(), "[DONE]"); assert!(stream.next().await.is_none()); @@ -497,7 +718,7 @@ mod tests { async fn per_item_error_carries_index() { let (tx0, rx0) = mpsc::channel(8); let (tx1, rx1) = mpsc::channel(8); - let receivers = vec![("10".into(), rx0), ("11".into(), rx1)]; + let receivers = vec![timed_receiver(10, rx0), timed_receiver(11, rx1)]; let stream = generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true); futures::pin_mut!(stream); @@ -523,7 +744,7 @@ mod tests { #[tokio::test] async fn incremental_emits_deltas_with_cumulative_count() { let (tx, rx) = mpsc::channel(8); - let receivers = vec![("10".into(), rx)]; + let receivers = vec![timed_receiver(10, rx)]; let stream = generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, true); futures::pin_mut!(stream); @@ -532,6 +753,7 @@ mod tests { let v = parse(&stream.next().await.unwrap()); assert_eq!(v["text"], "Hello"); assert_eq!(v["meta_info"]["completion_tokens"], 1); + assert!(v["meta_info"].get("e2e_latency").is_none()); tx.send(frame(10, " world")).await.unwrap(); let v = parse(&stream.next().await.unwrap()); @@ -540,12 +762,14 @@ mod tests { v["meta_info"]["completion_tokens"], 2, "count stays cumulative" ); + assert!(v["meta_info"].get("e2e_latency").is_none()); tx.send(done(10, "!")).await.unwrap(); let v = parse(&stream.next().await.unwrap()); assert_eq!(v["text"], "!"); assert_eq!(v["meta_info"]["completion_tokens"], 3); assert_eq!(v["meta_info"]["finish_reason"]["type"], "length"); + assert!(v["meta_info"]["e2e_latency"].as_f64().unwrap() >= 0.010); assert_eq!(stream.next().await.unwrap(), "[DONE]"); } @@ -555,7 +779,7 @@ mod tests { #[tokio::test] async fn single_shape_omits_index() { let (tx, rx) = mpsc::channel(8); - let receivers = vec![("10".into(), rx)]; + let receivers = vec![timed_receiver(10, rx)]; let stream = generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false); futures::pin_mut!(stream); @@ -564,6 +788,7 @@ mod tests { let v = parse(&stream.next().await.unwrap()); assert_eq!(v["text"], "hi"); assert!(v.get("index").is_none(), "single response has no index"); + assert!(v["meta_info"]["e2e_latency"].as_f64().unwrap() >= 0.010); assert_eq!(stream.next().await.unwrap(), "[DONE]"); } @@ -575,7 +800,7 @@ mod tests { #[tokio::test] async fn cumulative_backlog_coalesces_to_latest() { let (tx, rx) = mpsc::channel(8); - let receivers = vec![("10".into(), rx)]; + let receivers = vec![timed_receiver(10, rx)]; let stream = generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false); futures::pin_mut!(stream); @@ -602,7 +827,7 @@ mod tests { #[tokio::test] async fn incremental_backlog_emits_every_delta() { let (tx, rx) = mpsc::channel(8); - let receivers = vec![("10".into(), rx)]; + let receivers = vec![timed_receiver(10, rx)]; let stream = generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, false); futures::pin_mut!(stream);