[Rust Server] Add e2e latency metadata and fix Sarashina import (#35125)

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

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<Duration>,
+    e2e_latency: Option<Duration>,
+}
+
+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<AppState> {
     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<EgressItem>,
     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)
+/// then either (unary) drain them concurrently into a request-ordered JSON array,
-/// multiplex their streams into one SSE response, each frame carrying its `index`.
+/// 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<GenerateRequest>,
     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).
+        // Unary: poll every item concurrently, as Python's gather does. `join_all`
-        let mut results = Vec::with_capacity(receivers.len());
+        // preserves input order for the final JSON array, while each drain observes
-        for (rid_str, mut rx) in receivers {
+        // its own terminal output promptly (important for per-item e2e_latency).
-            let (_status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await;
+        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<EgressItem>)>,
+    receivers: Vec<(Rid, mpsc::Receiver<EgressItem>, 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<Rid> = receivers.iter().map(|(rid, _)| rid.clone()).collect();
+        let rid_strs: Vec<Rid> = receivers
+            .iter()
+            .map(|(rid, _, _)| rid.clone())
+            .collect();
+        let mut timings: Vec<RequestTiming> = receivers
+            .iter()
+            .map(|(_, _, timing)| timing.clone())
+            .collect();
         let mut accs: Vec<OutputAccumulator> =
             (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<usize>,
+    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<EgressItem>,
+    ) -> (Rid, mpsc::Receiver<EgressItem>, 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);