diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index 5b72f22f7..4fb170994 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -836,6 +836,26 @@ pub mod smg_labels { /// SMG Metrics helper struct for the new layered metrics architecture pub struct SmgMetrics; +/// Parameters for recording streaming metrics. +pub struct StreamingMetricsParams<'a> { + /// Router type label (e.g., "grpc", "http") + pub router_type: &'static str, + /// Backend type label (e.g., "regular", "pd") + pub backend_type: &'static str, + /// Model identifier (will be converted to owned String for metrics) + pub model_id: &'a str, + /// Endpoint label (e.g., "chat", "generate") + pub endpoint: &'static str, + /// Time to first token (None if no tokens were generated) + pub ttft: Option, + /// Total generation time + pub generation_duration: Duration, + /// Input token count (None for endpoints that don't track this) + pub input_tokens: Option, + /// Output token count + pub output_tokens: u64, +} + impl SmgMetrics { /// Record an HTTP request pub fn record_http_request(method: &str, path: &str, status_class: &str) { @@ -1055,6 +1075,86 @@ impl SmgMetrics { .record(duration.as_secs_f64()); } + /// Record all streaming metrics in a single batch call. + /// + /// This consolidates TTFT, TPOT, generation duration, and token metrics + /// into one function, handling TPOT calculation internally. + pub fn record_streaming_metrics(params: StreamingMetricsParams<'_>) { + let StreamingMetricsParams { + router_type, + backend_type, + model_id, + endpoint, + ttft, + generation_duration, + input_tokens, + output_tokens, + } = params; + // metrics-rs requires owned strings for dynamic labels (uses Cow<'static, str>). + // We allocate once and clone for each metric - unavoidable with this API. + let model = model_id.to_string(); + + // TTFT and TPOT (only if we have a first token time) + if let Some(ttft_duration) = ttft { + histogram!( + "smg_router_ttft_seconds", + "router_type" => router_type, + "backend_type" => backend_type, + "model" => model.clone(), + "endpoint" => endpoint + ) + .record(ttft_duration.as_secs_f64()); + + // TPOT - only meaningful with >1 output token + if output_tokens > 1 { + let time_after_first = generation_duration.saturating_sub(ttft_duration); + let tpot = time_after_first / (output_tokens as u32 - 1); + histogram!( + "smg_router_tpot_seconds", + "router_type" => router_type, + "backend_type" => backend_type, + "model" => model.clone(), + "endpoint" => endpoint + ) + .record(tpot.as_secs_f64()); + } + } + + // Generation duration + histogram!( + "smg_router_generation_duration_seconds", + "router_type" => router_type, + "backend_type" => backend_type, + "model" => model.clone(), + "endpoint" => endpoint + ) + .record(generation_duration.as_secs_f64()); + + // Input tokens (if available) + if let Some(input) = input_tokens { + counter!( + "smg_router_tokens_total", + "router_type" => router_type, + "backend_type" => backend_type, + "model" => model.clone(), + "endpoint" => endpoint, + "token_type" => smg_labels::TOKEN_INPUT + ) + .increment(input); + } + + // Output tokens + counter!( + "smg_router_tokens_total", + "router_type" => router_type, + "backend_type" => backend_type, + "model" => model, + "endpoint" => endpoint, + "token_type" => smg_labels::TOKEN_OUTPUT + ) + .increment(output_tokens); + } + // ======================================================================== // Layer 3: Worker metrics // ======================================================================== diff --git a/sgl-model-gateway/src/routers/grpc/pipeline.rs b/sgl-model-gateway/src/routers/grpc/pipeline.rs index d0a858451..6a218ffeb 100644 --- a/sgl-model-gateway/src/routers/grpc/pipeline.rs +++ b/sgl-model-gateway/src/routers/grpc/pipeline.rs @@ -65,6 +65,7 @@ impl RequestPipeline { reasoning_parser_factory, configured_tool_parser, configured_reasoning_parser, + smg_labels::BACKEND_REGULAR, )); let stages: Vec> = vec![ @@ -171,6 +172,7 @@ impl RequestPipeline { reasoning_parser_factory, configured_tool_parser, configured_reasoning_parser, + smg_labels::BACKEND_PD, )); let stages: Vec> = vec![ diff --git a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs index afc42211b..74f291604 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs @@ -14,6 +14,7 @@ use tracing::{debug, error, warn}; use crate::{ grpc_client::sglang_proto::generate_complete::MatchedStop::{MatchedStopStr, MatchedTokenId}, + observability::metrics::{smg_labels, SmgMetrics, StreamingMetricsParams}, protocols::{ chat::{ChatCompletionRequest, ChatCompletionStreamResponse}, common::{ @@ -43,6 +44,16 @@ pub struct StreamingProcessor { reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, configured_reasoning_parser: Option, + backend_type: &'static str, +} + +/// Context for generate endpoint streaming - groups config params to reduce function arguments +struct GenerateStreamContext { + request_id: String, + weight_version: String, + return_logprob: bool, + backend_type: &'static str, + model: String, } impl StreamingProcessor { @@ -52,6 +63,7 @@ impl StreamingProcessor { reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, configured_reasoning_parser: Option, + backend_type: &'static str, ) -> Self { Self { tokenizer, @@ -59,6 +71,7 @@ impl StreamingProcessor { reasoning_parser_factory, configured_tool_parser, configured_reasoning_parser, + backend_type, } } @@ -164,6 +177,10 @@ impl StreamingProcessor { original_request: Arc, tx: &UnboundedSender>, ) -> Result<(), String> { + // Metrics timing + let start_time = Instant::now(); + let mut first_token_time: Option = None; + // Extract request parameters let separate_reasoning = original_request.separate_reasoning; let tool_choice = &original_request.tool_choice; @@ -246,6 +263,11 @@ impl StreamingProcessor { match gen_response.into_response() { ProtoResponseVariant::Chunk(chunk) => { + // Track TTFT immediately on first chunk received from backend + if first_token_time.is_none() { + first_token_time = Some(Instant::now()); + } + let index = chunk.index(); // For vLLM, accumulate completion tokens (vLLM sends deltas) @@ -548,6 +570,20 @@ impl StreamingProcessor { // Mark stream as completed successfully to prevent abort on drop grpc_stream.mark_completed(); + // Record streaming metrics + let total_prompt: u32 = prompt_tokens.values().sum(); + let total_completion: u32 = completion_tokens.values().sum(); + SmgMetrics::record_streaming_metrics(StreamingMetricsParams { + router_type: smg_labels::ROUTER_GRPC, + backend_type: self.backend_type, + model_id: model, + endpoint: smg_labels::ENDPOINT_CHAT, + ttft: first_token_time.map(|t| t.duration_since(start_time)), + generation_duration: start_time.elapsed(), + input_tokens: Some(total_prompt as u64), + output_tokens: total_completion as u64, + }); + Ok(()) } @@ -603,30 +639,28 @@ impl StreamingProcessor { generate_request: Arc, dispatch: context::DispatchMetadata, ) -> Response { - let return_logprob = generate_request.return_logprob.unwrap_or(false); - // Create SSE channel let (tx, rx) = mpsc::unbounded_channel::>(); + // Build context once, clone for spawned task + let ctx = GenerateStreamContext { + request_id: dispatch.request_id.clone(), + weight_version: dispatch + .weight_version + .clone() + .unwrap_or_else(|| "default".to_string()), + return_logprob: generate_request.return_logprob.unwrap_or(false), + backend_type: self.backend_type, + model: dispatch.model.clone(), + }; + // Spawn background task based on execution mode match execution_result { context::ExecutionResult::Single { stream } => { let tokenizer = self.tokenizer.clone(); - let request_id = dispatch.request_id.clone(); - let weight_version = dispatch - .weight_version - .clone() - .unwrap_or_else(|| "default".to_string()); tokio::spawn(async move { - let result = Self::process_generate_streaming( - tokenizer, - stream, - request_id, - weight_version, - return_logprob, - &tx, - ) - .await; + let result = + Self::process_generate_streaming(tokenizer, stream, ctx, &tx).await; if let Err(e) = result { let error_chunk = format!("data: {{\"error\": \"{}\"}}\n\n", e); @@ -637,22 +671,10 @@ impl StreamingProcessor { }); } context::ExecutionResult::Dual { prefill, decode } => { - // For PD mode, need to handle prefill stream for input_logprobs let tokenizer = self.tokenizer.clone(); - let request_id = dispatch.request_id.clone(); - let weight_version = dispatch - .weight_version - .clone() - .unwrap_or_else(|| "default".to_string()); tokio::spawn(async move { let result = Self::process_generate_streaming_dual( - tokenizer, - prefill, - *decode, - request_id, - weight_version, - return_logprob, - &tx, + tokenizer, prefill, *decode, ctx, &tx, ) .await; @@ -675,12 +697,11 @@ impl StreamingProcessor { async fn process_generate_streaming( tokenizer: Arc, mut stream: ProtoStream, - request_id: String, - weight_version: String, - _include_logprobs: bool, + ctx: GenerateStreamContext, tx: &UnboundedSender>, ) -> Result<(), String> { let start_time = Instant::now(); + let mut first_token_time: Option = None; // Track state per index for n>1 case let mut accumulated_texts: HashMap = HashMap::new(); @@ -691,6 +712,11 @@ impl StreamingProcessor { match gen_response.into_response() { ProtoResponseVariant::Chunk(chunk) => { + // Track TTFT immediately on first chunk received from backend + if first_token_time.is_none() { + first_token_time = Some(Instant::now()); + } + let index = chunk.index(); // Both backends send delta token_ids, so accumulate for both @@ -708,7 +734,7 @@ impl StreamingProcessor { accumulated_text.push_str(&chunk_text); // Generate unique ID per index - let index_id = format!("{}-{}", request_id, index); + let index_id = format!("{}-{}", ctx.request_id, index); // Build streaming response chunk (SGLang format) let chunk_response = serde_json::json!({ @@ -718,7 +744,7 @@ impl StreamingProcessor { "id": index_id, "finish_reason": null, "prompt_tokens": chunk.prompt_tokens(), - "weight_version": &weight_version, + "weight_version": &ctx.weight_version, "completion_tokens": current_completion_tokens, "cached_tokens": chunk.cached_tokens() }, @@ -737,7 +763,7 @@ impl StreamingProcessor { let accumulated_text = accumulated_texts.get(&index).cloned().unwrap_or_default(); let completion_tokens = *completion_tokens_map.get(&index).unwrap_or(&0); - let index_id = format!("{}-{}", request_id, index); + let index_id = format!("{}-{}", ctx.request_id, index); let e2e_latency = start_time.elapsed().as_secs_f64(); // Send final chunk with finish_reason @@ -748,7 +774,7 @@ impl StreamingProcessor { "id": index_id, "finish_reason": complete.finish_reason(), "prompt_tokens": complete.prompt_tokens(), - "weight_version": &weight_version, + "weight_version": &ctx.weight_version, "completion_tokens": completion_tokens, "cached_tokens": complete.cached_tokens(), "e2e_latency": e2e_latency @@ -775,6 +801,10 @@ impl StreamingProcessor { // Mark stream as completed successfully to prevent abort on drop stream.mark_completed(); + // Record streaming metrics + let total_completion: u32 = completion_tokens_map.values().sum(); + Self::record_generate_metrics(start_time, first_token_time, total_completion, &ctx); + Ok(()) } @@ -783,13 +813,11 @@ impl StreamingProcessor { tokenizer: Arc, mut prefill_stream: ProtoStream, decode_stream: ProtoStream, - request_id: String, - weight_version: String, - return_logprob: bool, + ctx: GenerateStreamContext, tx: &UnboundedSender>, ) -> Result<(), String> { // Collect input_logprobs from prefill stream if requested - let input_token_logprobs = if return_logprob { + let input_token_logprobs = if ctx.return_logprob { let mut input_logprobs = None; while let Some(response) = prefill_stream.next().await { let gen_response = response.map_err(|e| format!("Prefill stream error: {}", e))?; @@ -817,9 +845,7 @@ impl StreamingProcessor { let result = Self::process_generate_streaming_with_input_logprobs( tokenizer, decode_stream, - request_id, - weight_version, - return_logprob, + ctx, input_token_logprobs, tx, ) @@ -838,13 +864,12 @@ impl StreamingProcessor { async fn process_generate_streaming_with_input_logprobs( tokenizer: Arc, mut stream: ProtoStream, - request_id: String, - weight_version: String, - _include_logprobs: bool, + ctx: GenerateStreamContext, input_token_logprobs: Option>>>, tx: &UnboundedSender>, ) -> Result<(), String> { let start_time = Instant::now(); + let mut first_token_time: Option = None; // Track state per index for n>1 case let mut accumulated_texts: HashMap = HashMap::new(); @@ -857,6 +882,11 @@ impl StreamingProcessor { match gen_response.into_response() { ProtoResponseVariant::Chunk(chunk) => { + // Track TTFT immediately on first chunk received from backend + if first_token_time.is_none() { + first_token_time = Some(Instant::now()); + } + let index = chunk.index(); // Both backends send delta token_ids, so accumulate for both @@ -880,7 +910,7 @@ impl StreamingProcessor { } // Generate unique ID per index - let index_id = format!("{}-{}", request_id, index); + let index_id = format!("{}-{}", ctx.request_id, index); // Build streaming response chunk with cumulative logprobs let current_output_logprobs = accumulated_output_logprobs @@ -894,7 +924,7 @@ impl StreamingProcessor { "id": index_id, "finish_reason": null, "prompt_tokens": chunk.prompt_tokens(), - "weight_version": &weight_version, + "weight_version": &ctx.weight_version, "input_token_logprobs": input_token_logprobs.as_ref(), "output_token_logprobs": current_output_logprobs, "completion_tokens": current_completion_tokens, @@ -921,7 +951,7 @@ impl StreamingProcessor { let final_output_logprobs = accumulated_output_logprobs .get(&index) .and_then(|o| o.as_ref()); - let index_id = format!("{}-{}", request_id, index); + let index_id = format!("{}-{}", ctx.request_id, index); let e2e_latency = start_time.elapsed().as_secs_f64(); // Parse finish_reason @@ -938,7 +968,7 @@ impl StreamingProcessor { "id": index_id, "finish_reason": finish_reason, "prompt_tokens": complete.prompt_tokens(), - "weight_version": &weight_version, + "weight_version": &ctx.weight_version, "input_token_logprobs": input_token_logprobs.as_ref(), "output_token_logprobs": final_output_logprobs, "completion_tokens": completion_tokens, @@ -967,6 +997,10 @@ impl StreamingProcessor { // Mark stream as completed successfully to prevent abort on drop stream.mark_completed(); + // Record streaming metrics + let total_completion: u32 = completion_tokens_map.values().sum(); + Self::record_generate_metrics(start_time, first_token_time, total_completion, &ctx); + Ok(()) } @@ -974,6 +1008,25 @@ impl StreamingProcessor { // Helper Methods // ======================================================================== + /// Record streaming metrics for generate endpoint + fn record_generate_metrics( + start_time: Instant, + first_token_time: Option, + total_completion: u32, + ctx: &GenerateStreamContext, + ) { + SmgMetrics::record_streaming_metrics(StreamingMetricsParams { + router_type: smg_labels::ROUTER_GRPC, + backend_type: ctx.backend_type, + model_id: &ctx.model, + endpoint: smg_labels::ENDPOINT_GENERATE, + ttft: first_token_time.map(|t| t.duration_since(start_time)), + generation_duration: start_time.elapsed(), + input_tokens: None, // generate endpoint doesn't expose prompt tokens in streaming + output_tokens: total_completion as u64, + }); + } + /// Process a chunk of tokens through the stop decoder fn process_chunk_tokens( stop_decoder: &mut StopSequenceDecoder,