[model-gateway] add streaming metrics (TTFT, TPOT, tokens, duration) for gRPC router (#15125)
This commit is contained in:
@@ -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<Duration>,
|
||||
/// Total generation time
|
||||
pub generation_duration: Duration,
|
||||
/// Input token count (None for endpoints that don't track this)
|
||||
pub input_tokens: Option<u64>,
|
||||
/// 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
|
||||
// ========================================================================
|
||||
|
||||
@@ -65,6 +65,7 @@ impl RequestPipeline {
|
||||
reasoning_parser_factory,
|
||||
configured_tool_parser,
|
||||
configured_reasoning_parser,
|
||||
smg_labels::BACKEND_REGULAR,
|
||||
));
|
||||
|
||||
let stages: Vec<Box<dyn PipelineStage>> = vec![
|
||||
@@ -171,6 +172,7 @@ impl RequestPipeline {
|
||||
reasoning_parser_factory,
|
||||
configured_tool_parser,
|
||||
configured_reasoning_parser,
|
||||
smg_labels::BACKEND_PD,
|
||||
));
|
||||
|
||||
let stages: Vec<Box<dyn PipelineStage>> = vec![
|
||||
|
||||
@@ -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<String>,
|
||||
configured_reasoning_parser: Option<String>,
|
||||
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<String>,
|
||||
configured_reasoning_parser: Option<String>,
|
||||
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<ChatCompletionRequest>,
|
||||
tx: &UnboundedSender<Result<Bytes, io::Error>>,
|
||||
) -> Result<(), String> {
|
||||
// Metrics timing
|
||||
let start_time = Instant::now();
|
||||
let mut first_token_time: Option<Instant> = 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<GenerateRequest>,
|
||||
dispatch: context::DispatchMetadata,
|
||||
) -> Response {
|
||||
let return_logprob = generate_request.return_logprob.unwrap_or(false);
|
||||
|
||||
// Create SSE channel
|
||||
let (tx, rx) = mpsc::unbounded_channel::<Result<Bytes, io::Error>>();
|
||||
|
||||
// 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<dyn Tokenizer>,
|
||||
mut stream: ProtoStream,
|
||||
request_id: String,
|
||||
weight_version: String,
|
||||
_include_logprobs: bool,
|
||||
ctx: GenerateStreamContext,
|
||||
tx: &UnboundedSender<Result<Bytes, io::Error>>,
|
||||
) -> Result<(), String> {
|
||||
let start_time = Instant::now();
|
||||
let mut first_token_time: Option<Instant> = None;
|
||||
|
||||
// Track state per index for n>1 case
|
||||
let mut accumulated_texts: HashMap<u32, String> = 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<dyn Tokenizer>,
|
||||
mut prefill_stream: ProtoStream,
|
||||
decode_stream: ProtoStream,
|
||||
request_id: String,
|
||||
weight_version: String,
|
||||
return_logprob: bool,
|
||||
ctx: GenerateStreamContext,
|
||||
tx: &UnboundedSender<Result<Bytes, io::Error>>,
|
||||
) -> 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<dyn Tokenizer>,
|
||||
mut stream: ProtoStream,
|
||||
request_id: String,
|
||||
weight_version: String,
|
||||
_include_logprobs: bool,
|
||||
ctx: GenerateStreamContext,
|
||||
input_token_logprobs: Option<Vec<Vec<Option<f64>>>>,
|
||||
tx: &UnboundedSender<Result<Bytes, io::Error>>,
|
||||
) -> Result<(), String> {
|
||||
let start_time = Instant::now();
|
||||
let mut first_token_time: Option<Instant> = None;
|
||||
|
||||
// Track state per index for n>1 case
|
||||
let mut accumulated_texts: HashMap<u32, String> = 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<Instant>,
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user