[model-gateway] feat(metrics): implement Layer 2 router metrics (smg_router_*) (#15124)

This commit is contained in:
Simo Lin
2025-12-14 10:34:12 -08:00
committed by GitHub
parent f9bceea064
commit b11af135eb
6 changed files with 401 additions and 31 deletions
@@ -765,6 +765,9 @@ pub mod smg_labels {
pub const ENDPOINT_CHAT: &str = "chat";
pub const ENDPOINT_GENERATE: &str = "generate";
pub const ENDPOINT_RESPONSES: &str = "responses";
pub const ENDPOINT_COMPLETIONS: &str = "completions";
pub const ENDPOINT_RERANK: &str = "rerank";
pub const ENDPOINT_EMBEDDINGS: &str = "embeddings";
// Worker types
pub const WORKER_REGULAR: &str = "regular";
@@ -812,6 +815,22 @@ pub mod smg_labels {
// Circuit breaker outcomes
pub const CB_SUCCESS: &str = "success";
pub const CB_FAILURE: &str = "failure";
// Router error types
pub const ERROR_NO_WORKERS: &str = "no_workers";
pub const ERROR_TIMEOUT: &str = "timeout";
pub const ERROR_BACKEND: &str = "backend_error";
pub const ERROR_VALIDATION: &str = "validation_error";
pub const ERROR_INTERNAL: &str = "internal_error";
// Pipeline stages (gRPC router)
pub const STAGE_PREPARATION: &str = "preparation";
pub const STAGE_WORKER_SELECTION: &str = "worker_selection";
pub const STAGE_CLIENT_ACQUISITION: &str = "client_acquisition";
pub const STAGE_REQUEST_BUILDING: &str = "request_building";
pub const STAGE_DISPATCH_METADATA: &str = "dispatch_metadata";
pub const STAGE_REQUEST_EXECUTION: &str = "request_execution";
pub const STAGE_RESPONSE_PROCESSING: &str = "response_processing";
}
/// SMG Metrics helper struct for the new layered metrics architecture
+133 -19
View File
@@ -3,7 +3,7 @@
//! This module defines the RequestPipeline orchestrator that coordinates
//! the execution of pipeline stages from request preparation to response delivery.
use std::sync::Arc;
use std::{sync::Arc, time::Instant};
use axum::response::{IntoResponse, Response};
use tracing::error;
@@ -13,9 +13,11 @@ use super::{
context::*,
harmony,
regular::{processor, stages::*, streaming},
utils::error_type_from_status,
};
use crate::{
core::WorkerRegistry,
observability::metrics::{smg_labels, SmgMetrics},
policies::PolicyRegistry,
protocols::{
chat::{ChatCompletionRequest, ChatCompletionResponse},
@@ -34,6 +36,8 @@ use crate::{
#[derive(Clone)]
pub struct RequestPipeline {
stages: Arc<Vec<Box<dyn PipelineStage>>>,
/// Backend type for metrics labeling
backend_type: &'static str,
}
impl RequestPipeline {
@@ -79,6 +83,7 @@ impl RequestPipeline {
Self {
stages: Arc::new(stages),
backend_type: smg_labels::BACKEND_REGULAR,
}
}
@@ -108,6 +113,7 @@ impl RequestPipeline {
Self {
stages: Arc::new(stages),
backend_type: smg_labels::BACKEND_REGULAR,
}
}
@@ -137,6 +143,7 @@ impl RequestPipeline {
Self {
stages: Arc::new(stages),
backend_type: smg_labels::BACKEND_PD,
}
}
@@ -182,6 +189,7 @@ impl RequestPipeline {
Self {
stages: Arc::new(stages),
backend_type: smg_labels::BACKEND_PD,
}
}
@@ -193,22 +201,49 @@ impl RequestPipeline {
model_id: Option<String>,
components: Arc<SharedComponents>,
) -> Response {
let start = Instant::now();
// Clone Arc for metrics (cheap atomic increment) to avoid borrow issues
let request_for_metrics = Arc::clone(&request);
let streaming = request.stream;
// Record request start
SmgMetrics::record_router_request(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
&request_for_metrics.model,
smg_labels::ENDPOINT_CHAT,
streaming,
);
let mut ctx = RequestContext::for_chat(request, headers, model_id, components);
for (idx, stage) in self.stages.iter().enumerate() {
for stage in self.stages.iter() {
match stage.execute(&mut ctx).await {
Ok(Some(response)) => {
// Stage completed successfully with a response (e.g., streaming)
// Stage completed with streaming response - record success and return
SmgMetrics::record_router_duration(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
&request_for_metrics.model,
smg_labels::ENDPOINT_CHAT,
start.elapsed(),
);
return response;
}
Ok(None) => {
continue;
}
Ok(None) => continue,
Err(response) => {
// Error occurred
SmgMetrics::record_router_error(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
&request_for_metrics.model,
smg_labels::ENDPOINT_CHAT,
error_type_from_status(response.status()),
);
error!(
"Stage {} ({}) failed with status {}",
idx + 1,
"Stage {} failed with status {}",
stage.name(),
response.status()
);
@@ -218,12 +253,30 @@ impl RequestPipeline {
}
match ctx.state.response.final_response {
Some(FinalResponse::Chat(response)) => axum::Json(response).into_response(),
Some(FinalResponse::Chat(response)) => {
SmgMetrics::record_router_duration(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
&request_for_metrics.model,
smg_labels::ENDPOINT_CHAT,
start.elapsed(),
);
axum::Json(response).into_response()
}
Some(FinalResponse::Generate(_)) => {
error!(
function = "execute_chat",
"Wrong response type: expected Chat, got Generate"
);
SmgMetrics::record_router_error(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
&request_for_metrics.model,
smg_labels::ENDPOINT_CHAT,
smg_labels::ERROR_INTERNAL,
);
error::internal_error("wrong_response_type", "Internal error: wrong response type")
}
None => {
@@ -231,6 +284,14 @@ impl RequestPipeline {
function = "execute_chat",
"No response produced by pipeline"
);
SmgMetrics::record_router_error(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
&request_for_metrics.model,
smg_labels::ENDPOINT_CHAT,
smg_labels::ERROR_INTERNAL,
);
error::internal_error("no_response_produced", "No response produced")
}
}
@@ -244,22 +305,49 @@ impl RequestPipeline {
model_id: Option<String>,
components: Arc<SharedComponents>,
) -> Response {
let start = Instant::now();
// Clone model_id for metrics before moving into context
// GenerateRequest doesn't have a model field, so we use model_id
let model_for_metrics = model_id.clone();
let streaming = request.stream;
// Record request start
SmgMetrics::record_router_request(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
model_for_metrics.as_deref().unwrap_or("unknown"),
smg_labels::ENDPOINT_GENERATE,
streaming,
);
let mut ctx = RequestContext::for_generate(request, headers, model_id, components);
for (idx, stage) in self.stages.iter().enumerate() {
for stage in self.stages.iter() {
match stage.execute(&mut ctx).await {
Ok(Some(response)) => {
// Stage completed successfully with a response (e.g., streaming)
SmgMetrics::record_router_duration(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
model_for_metrics.as_deref().unwrap_or("unknown"),
smg_labels::ENDPOINT_GENERATE,
start.elapsed(),
);
return response;
}
Ok(None) => {
continue;
}
Ok(None) => continue,
Err(response) => {
// Error occurred
SmgMetrics::record_router_error(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
model_for_metrics.as_deref().unwrap_or("unknown"),
smg_labels::ENDPOINT_GENERATE,
error_type_from_status(response.status()),
);
error!(
"Stage {} ({}) failed with status {}",
idx + 1,
"Stage {} failed with status {}",
stage.name(),
response.status()
);
@@ -269,12 +357,30 @@ impl RequestPipeline {
}
match ctx.state.response.final_response {
Some(FinalResponse::Generate(response)) => axum::Json(response).into_response(),
Some(FinalResponse::Generate(response)) => {
SmgMetrics::record_router_duration(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
model_for_metrics.as_deref().unwrap_or("unknown"),
smg_labels::ENDPOINT_GENERATE,
start.elapsed(),
);
axum::Json(response).into_response()
}
Some(FinalResponse::Chat(_)) => {
error!(
function = "execute_generate",
"Wrong response type: expected Generate, got Chat"
);
SmgMetrics::record_router_error(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
model_for_metrics.as_deref().unwrap_or("unknown"),
smg_labels::ENDPOINT_GENERATE,
smg_labels::ERROR_INTERNAL,
);
error::internal_error("wrong_response_type", "Internal error: wrong response type")
}
None => {
@@ -282,6 +388,14 @@ impl RequestPipeline {
function = "execute_generate",
"No response produced by pipeline"
);
SmgMetrics::record_router_error(
smg_labels::ROUTER_GRPC,
self.backend_type,
smg_labels::CONNECTION_GRPC,
model_for_metrics.as_deref().unwrap_or("unknown"),
smg_labels::ENDPOINT_GENERATE,
smg_labels::ERROR_INTERNAL,
);
error::internal_error("no_response_produced", "No response produced")
}
}
@@ -3,6 +3,7 @@
use std::{collections::HashMap, sync::Arc};
use axum::response::Response;
use http::StatusCode;
use serde_json::{json, Map, Value};
use tracing::{error, warn};
use uuid::Uuid;
@@ -15,6 +16,7 @@ use super::{
use crate::{
core::Worker,
grpc_client::sglang_proto::{InputLogProbs, OutputLogProbs},
observability::metrics::smg_labels,
protocols::{
chat::{ChatCompletionRequest, ChatMessage},
common::{
@@ -957,6 +959,33 @@ pub fn parse_finish_reason(reason_str: &str, completion_tokens: i32) -> Generate
}
}
// ============================================================================
// Metrics helper functions (shared by HTTP routers and gRPC pipeline)
// ============================================================================
/// Map route path to endpoint label for metrics
pub fn route_to_endpoint(route: &str) -> &'static str {
match route {
"/v1/chat/completions" => smg_labels::ENDPOINT_CHAT,
"/generate" => smg_labels::ENDPOINT_GENERATE,
"/v1/completions" => smg_labels::ENDPOINT_COMPLETIONS,
"/v1/rerank" => smg_labels::ENDPOINT_RERANK,
"/v1/responses" => smg_labels::ENDPOINT_RESPONSES,
_ => "other",
}
}
/// Map HTTP status code to error type label for metrics
pub fn error_type_from_status(status: StatusCode) -> &'static str {
match status.as_u16() {
400 => smg_labels::ERROR_VALIDATION,
404 => smg_labels::ERROR_NO_WORKERS,
408 | 504 => smg_labels::ERROR_TIMEOUT,
500..=599 => smg_labels::ERROR_BACKEND,
_ => smg_labels::ERROR_INTERNAL,
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
@@ -22,7 +22,7 @@ use crate::{
},
observability::{
events::{self, Event},
metrics::RouterMetrics,
metrics::{smg_labels, RouterMetrics, SmgMetrics},
otel_trace::inject_trace_context_http,
},
policies::{LoadBalancingPolicy, PolicyRegistry},
@@ -33,7 +33,11 @@ use crate::{
generate::GenerateRequest,
rerank::RerankRequest,
},
routers::{error, header_utils, RouterTrait},
routers::{
error,
grpc::utils::{error_type_from_status, route_to_endpoint},
header_utils, RouterTrait,
},
};
#[derive(Debug)]
@@ -275,10 +279,22 @@ impl PDRouter {
let start_time = Instant::now();
let route = context.route;
let model = context.model_id.unwrap_or("default");
let endpoint = route_to_endpoint(route);
// Record request start (Layer 2)
SmgMetrics::record_router_request(
smg_labels::ROUTER_HTTP,
smg_labels::BACKEND_PD,
smg_labels::CONNECTION_HTTP,
model,
endpoint,
context.is_stream,
);
// Clone request once outside the retry loop, then use Arc to share across attempts
// This avoids O(retries) clones by sharing the same data
let shared_request = Arc::new(original_request.clone());
RetryExecutor::execute_response_with_retry(
let response = RetryExecutor::execute_response_with_retry(
&self.retry_config,
{
move |attempt: u32| {
@@ -345,7 +361,31 @@ impl PDRouter {
},
|| RouterMetrics::record_retries_exhausted(route),
)
.await
.await;
// Record Layer 2 metrics
let duration = start_time.elapsed();
if response.status().is_success() {
SmgMetrics::record_router_duration(
smg_labels::ROUTER_HTTP,
smg_labels::BACKEND_PD,
smg_labels::CONNECTION_HTTP,
model,
endpoint,
duration,
);
} else if !is_retryable_status(response.status()) {
SmgMetrics::record_router_error(
smg_labels::ROUTER_HTTP,
smg_labels::BACKEND_PD,
smg_labels::CONNECTION_HTTP,
model,
endpoint,
error_type_from_status(response.status()),
);
}
response
}
async fn handle_decode_error_response(
+34 -2
View File
@@ -23,7 +23,7 @@ use crate::{
},
observability::{
events::{self, Event},
metrics::RouterMetrics,
metrics::{smg_labels, RouterMetrics, SmgMetrics},
otel_trace::inject_trace_context_http,
},
policies::PolicyRegistry,
@@ -37,7 +37,11 @@ use crate::{
rerank::{RerankRequest, RerankResponse, RerankResult},
responses::{ResponsesGetParams, ResponsesRequest},
},
routers::{error, header_utils, RouterTrait},
routers::{
error,
grpc::utils::{error_type_from_status, route_to_endpoint},
header_utils, RouterTrait,
},
};
/// Regular router that uses injected load balancing policies
@@ -165,6 +169,18 @@ impl Router {
let start = Instant::now();
let is_stream = typed_req.is_stream();
let text = typed_req.extract_text_for_routing();
let model = model_id.unwrap_or("default");
let endpoint = route_to_endpoint(route);
// Record request start (Layer 2)
SmgMetrics::record_router_request(
smg_labels::ROUTER_HTTP,
smg_labels::BACKEND_REGULAR,
smg_labels::CONNECTION_HTTP,
model,
endpoint,
is_stream,
);
let response = RetryExecutor::execute_response_with_retry(
&self.retry_config,
@@ -199,8 +215,24 @@ impl Router {
let duration = start.elapsed();
RouterMetrics::record_request(route);
RouterMetrics::record_generate_duration(duration);
SmgMetrics::record_router_duration(
smg_labels::ROUTER_HTTP,
smg_labels::BACKEND_REGULAR,
smg_labels::CONNECTION_HTTP,
model,
endpoint,
duration,
);
} else if !is_retryable_status(response.status()) {
RouterMetrics::record_request_error(route, "non_retryable_error");
SmgMetrics::record_router_error(
smg_labels::ROUTER_HTTP,
smg_labels::BACKEND_REGULAR,
smg_labels::CONNECTION_HTTP,
model,
endpoint,
error_type_from_status(response.status()),
);
}
response
+142 -6
View File
@@ -2,6 +2,7 @@ use std::{
any::Any,
collections::HashSet,
sync::{atomic::AtomicBool, Arc},
time::Instant,
};
use axum::{
@@ -35,6 +36,7 @@ use crate::{
app_context::AppContext,
core::{model_type::Endpoint, ModelCard, ProviderType, RuntimeType, Worker, WorkerRegistry},
data_connector::{ConversationId, ListParams, ResponseId, SortOrder},
observability::metrics::{smg_labels, SmgMetrics},
protocols::{
chat::ChatCompletionRequest,
responses::{
@@ -576,6 +578,20 @@ impl crate::routers::RouterTrait for OpenAIRouter {
body: &ChatCompletionRequest,
model_id: Option<&str>,
) -> Response {
let start = Instant::now();
let model = model_id.unwrap_or(body.model.as_str());
let streaming = body.stream;
// Record request start
SmgMetrics::record_router_request(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
streaming,
);
let auth_header = extract_auth_header(headers, &None);
let worker = match self
@@ -583,18 +599,44 @@ impl crate::routers::RouterTrait for OpenAIRouter {
.await
{
Ok(w) => w,
Err(response) => return response,
Err(response) => {
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
smg_labels::ERROR_NO_WORKERS,
);
return response;
}
};
let mut payload = match to_value(body) {
Ok(v) => v,
Err(e) => {
return error_responses::bad_request(format!("Failed to serialize request: {}", e))
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
smg_labels::ERROR_VALIDATION,
);
return error_responses::bad_request(format!("Failed to serialize request: {}", e));
}
};
let provider = self.get_provider_arc_for_worker(worker.as_ref(), model_id);
if let Err(e) = provider.transform_request(&mut payload, Endpoint::Chat) {
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
smg_labels::ERROR_VALIDATION,
);
return error_responses::bad_request(format!("Provider transform error: {}", e));
}
@@ -630,6 +672,14 @@ impl crate::routers::RouterTrait for OpenAIRouter {
Ok(r) => r,
Err(e) => {
worker.circuit_breaker().record_failure();
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
smg_labels::ERROR_BACKEND,
);
return (
StatusCode::SERVICE_UNAVAILABLE,
format!("Failed to contact upstream: {}", e),
@@ -646,6 +696,14 @@ impl crate::routers::RouterTrait for OpenAIRouter {
match resp.bytes().await {
Ok(body) => {
worker.circuit_breaker().record_success();
SmgMetrics::record_router_duration(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
start.elapsed(),
);
let mut response = Response::new(Body::from(body));
*response.status_mut() = status;
if let Some(ct) = content_type {
@@ -655,6 +713,14 @@ impl crate::routers::RouterTrait for OpenAIRouter {
}
Err(e) => {
worker.circuit_breaker().record_failure();
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
smg_labels::ERROR_BACKEND,
);
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to read response: {}", e),
@@ -663,6 +729,15 @@ impl crate::routers::RouterTrait for OpenAIRouter {
}
}
} else {
// For streaming, record duration at start since we can't track completion
SmgMetrics::record_router_duration(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_CHAT,
start.elapsed(),
);
let stream = resp.bytes_stream();
let (tx, rx) = mpsc::unbounded_channel();
tokio::spawn(async move {
@@ -696,15 +771,38 @@ impl crate::routers::RouterTrait for OpenAIRouter {
body: &ResponsesRequest,
model_id: Option<&str>,
) -> Response {
let start = Instant::now();
let model = model_id.unwrap_or(body.model.as_str());
let streaming = body.stream.unwrap_or(false);
// Record request start
SmgMetrics::record_router_request(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_RESPONSES,
streaming,
);
let auth_header = extract_auth_header(headers, &None);
let model = model_id.unwrap_or(body.model.as_str());
let worker = match self
.select_worker_for_model(model, auth_header.as_ref())
.await
{
Ok(w) => w,
Err(response) => return response,
Err(response) => {
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_RESPONSES,
smg_labels::ERROR_NO_WORKERS,
);
return response;
}
};
let mut request_body = body.clone();
@@ -755,6 +853,14 @@ impl crate::routers::RouterTrait for OpenAIRouter {
.get_conversation(&conv_id)
.await
{
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_RESPONSES,
smg_labels::ERROR_VALIDATION,
);
return error_responses::not_found("conversation", &conv_id.0);
}
@@ -864,12 +970,28 @@ impl crate::routers::RouterTrait for OpenAIRouter {
let mut payload = match to_value(&request_body) {
Ok(v) => v,
Err(e) => {
return error_responses::bad_request(format!("Failed to serialize request: {}", e))
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_RESPONSES,
smg_labels::ERROR_VALIDATION,
);
return error_responses::bad_request(format!("Failed to serialize request: {}", e));
}
};
let provider = self.get_provider_arc_for_worker(worker.as_ref(), model_id);
if let Err(e) = provider.transform_request(&mut payload, Endpoint::Responses) {
SmgMetrics::record_router_error(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_RESPONSES,
smg_labels::ERROR_VALIDATION,
);
return error_responses::bad_request(format!("Provider transform error: {}", e));
}
@@ -891,11 +1013,25 @@ impl crate::routers::RouterTrait for OpenAIRouter {
previous_response_id: original_previous_response_id,
});
if ctx.is_streaming() {
let response = if ctx.is_streaming() {
handle_streaming_response(ctx).await
} else {
self.handle_non_streaming_response(ctx).await
};
// Record duration only for successful requests (errors tracked inside handlers)
if response.status().is_success() {
SmgMetrics::record_router_duration(
smg_labels::ROUTER_OPENAI,
smg_labels::BACKEND_EXTERNAL,
smg_labels::CONNECTION_HTTP,
model,
smg_labels::ENDPOINT_RESPONSES,
start.elapsed(),
);
}
response
}
async fn get_response(