[model-gateway] Add streaming metrics for harmony gRPC router (#15147)

This commit is contained in:
Simo Lin
2025-12-14 20:54:43 -08:00
committed by GitHub
parent 1e641ee4a7
commit bd9c3a47d6
2 changed files with 53 additions and 6 deletions

View File

@@ -756,6 +756,7 @@ pub mod smg_labels {
pub const BACKEND_REGULAR: &str = "regular";
pub const BACKEND_PD: &str = "pd";
pub const BACKEND_EXTERNAL: &str = "external";
pub const BACKEND_HARMONY: &str = "harmony";
// Connection modes
pub const CONNECTION_HTTP: &str = "http";

View File

@@ -4,6 +4,7 @@ use std::{
collections::{hash_map::Entry::Vacant, HashMap},
io,
sync::Arc,
time::Instant,
};
use axum::{body::Body, http::StatusCode, response::Response};
@@ -19,6 +20,7 @@ use super::{
};
use crate::{
grpc_client::sglang_proto::generate_complete::MatchedStop::{MatchedStopStr, MatchedTokenId},
observability::metrics::{smg_labels, SmgMetrics, StreamingMetricsParams},
protocols::{
chat::{
ChatCompletionRequest, ChatCompletionStreamResponse, ChatMessageDelta, ChatStreamChoice,
@@ -183,6 +185,10 @@ impl HarmonyStreamingProcessor {
original_request: Arc<ChatCompletionRequest>,
tx: &mpsc::UnboundedSender<Result<Bytes, io::Error>>,
) -> Result<(), String> {
// Timing for metrics
let start_time = Instant::now();
let mut first_token_time: Option<Instant> = None;
// Per-index state management (for n>1 support)
let mut parsers: HashMap<u32, HarmonyParserAdapter> = HashMap::new();
let mut is_firsts: HashMap<u32, bool> = HashMap::new();
@@ -202,6 +208,11 @@ impl HarmonyStreamingProcessor {
let chunk = chunk_wrapper.as_sglang();
let index = chunk.index;
// Track first token time for TTFT metric
if first_token_time.is_none() {
first_token_time = Some(Instant::now());
}
// Initialize parser for this index if needed
if let Vacant(e) = parsers.entry(index) {
e.insert(
@@ -285,11 +296,12 @@ impl HarmonyStreamingProcessor {
}
}
// Compute totals once for both usage chunk and metrics
let total_prompt: u32 = prompt_tokens.values().sum();
let total_completion: u32 = completion_tokens.values().sum();
// Emit final usage if requested
if let Some(true) = stream_options.as_ref().and_then(|so| so.include_usage) {
let total_prompt: u32 = prompt_tokens.values().sum();
let total_completion: u32 = completion_tokens.values().sum();
Self::emit_usage_chunk(
total_prompt,
total_completion,
@@ -302,6 +314,18 @@ impl HarmonyStreamingProcessor {
// Mark stream as completed successfully to prevent abort on drop
grpc_stream.mark_completed();
// Record streaming metrics
SmgMetrics::record_streaming_metrics(StreamingMetricsParams {
router_type: smg_labels::ROUTER_GRPC,
backend_type: smg_labels::BACKEND_HARMONY,
model_id: &original_request.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(())
}
@@ -313,6 +337,10 @@ impl HarmonyStreamingProcessor {
original_request: Arc<ChatCompletionRequest>,
tx: &mpsc::UnboundedSender<Result<Bytes, io::Error>>,
) -> Result<(), String> {
// Timing for metrics
let start_time = Instant::now();
let mut first_token_time: Option<Instant> = None;
// Phase 1: Process prefill stream (collect metadata)
let mut prompt_tokens: HashMap<u32, u32> = HashMap::new();
@@ -342,6 +370,11 @@ impl HarmonyStreamingProcessor {
let chunk = chunk_wrapper.as_sglang();
let index = chunk.index;
// Track first token time for TTFT metric
if first_token_time.is_none() {
first_token_time = Some(Instant::now());
}
// Initialize parser for this index if needed
if let Vacant(e) = parsers.entry(index) {
e.insert(
@@ -425,11 +458,12 @@ impl HarmonyStreamingProcessor {
// This ensures that if client disconnects during decode, BOTH streams send abort
prefill_stream.mark_completed();
// Compute totals once for both usage chunk and metrics
let total_prompt: u32 = prompt_tokens.values().sum();
let total_completion: u32 = completion_tokens.values().sum();
// Emit final usage if requested
if let Some(true) = stream_options.as_ref().and_then(|so| so.include_usage) {
let total_prompt: u32 = prompt_tokens.values().sum();
let total_completion: u32 = completion_tokens.values().sum();
Self::emit_usage_chunk(
total_prompt,
total_completion,
@@ -439,6 +473,18 @@ impl HarmonyStreamingProcessor {
)?;
}
// Record streaming metrics
SmgMetrics::record_streaming_metrics(StreamingMetricsParams {
router_type: smg_labels::ROUTER_GRPC,
backend_type: smg_labels::BACKEND_HARMONY,
model_id: &original_request.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(())
}