From 7b9156c773f0467df80cc849637d94a4522a8b5e Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 14 Dec 2025 22:54:53 -0800 Subject: [PATCH] [model-gateway] add mcp and discovery metrics (#15156) --- .../src/core/steps/mcp_registration.rs | 4 ++ .../src/observability/metrics.rs | 12 ++++ .../src/routers/grpc/harmony/responses.rs | 64 +++++++++++++++++-- .../grpc/regular/responses/tool_loop.rs | 44 ++++++++++++- sgl-model-gateway/src/service_discovery.rs | 54 ++++++++++++++-- 5 files changed, 163 insertions(+), 15 deletions(-) diff --git a/sgl-model-gateway/src/core/steps/mcp_registration.rs b/sgl-model-gateway/src/core/steps/mcp_registration.rs index e496fdab1..55fac7a7e 100644 --- a/sgl-model-gateway/src/core/steps/mcp_registration.rs +++ b/sgl-model-gateway/src/core/steps/mcp_registration.rs @@ -7,6 +7,7 @@ use tracing::{debug, error, info, warn}; use crate::{ app_context::AppContext, mcp::{config::McpServerConfig, manager::McpManager}, + observability::metrics::SmgMetrics, workflow::*, }; @@ -151,6 +152,9 @@ impl StepExecutor for RegisterMcpServerStep { // Register the client in the manager's client map mcp_manager.register_static_server(config_request.name.clone(), mcp_client); + // Update active MCP servers metric + SmgMetrics::set_mcp_servers_active(mcp_manager.list_servers().len()); + info!("Registered MCP server: {}", config_request.name); Ok(StepResult::Success) diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index b33b3dff1..98bcaf4a5 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -804,6 +804,18 @@ pub mod smg_labels { pub const DISCOVERY_CONSUL: &str = "consul"; pub const DISCOVERY_MANUAL: &str = "manual"; + // Discovery registration results + pub const REGISTRATION_SUCCESS: &str = "success"; + pub const REGISTRATION_FAILED: &str = "failed"; + pub const REGISTRATION_DUPLICATE: &str = "duplicate"; + + // Deregistration reasons + pub const DEREGISTRATION_HEALTH_CHECK_FAILED: &str = "health_check_failed"; + pub const DEREGISTRATION_TIMEOUT: &str = "timeout"; + pub const DEREGISTRATION_MANUAL: &str = "manual"; + pub const DEREGISTRATION_SHUTDOWN: &str = "shutdown"; + pub const DEREGISTRATION_POD_DELETED: &str = "pod_deleted"; + // Rate limit results pub const RATE_LIMIT_ALLOWED: &str = "allowed"; pub const RATE_LIMIT_REJECTED: &str = "rejected"; diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses.rs index 0193d6cbe..0453aef75 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses.rs @@ -32,7 +32,7 @@ //! ``` use std::{ sync::Arc, - time::{SystemTime, UNIX_EPOCH}, + time::{Instant, SystemTime, UNIX_EPOCH}, }; use axum::response::Response; @@ -45,6 +45,7 @@ use uuid::Uuid; use crate::{ data_connector::{ConversationItemStorage, ConversationStorage, ResponseId, ResponseStorage}, mcp::{self, McpManager}, + observability::metrics::{smg_labels, SmgMetrics}, protocols::{ common::{Function, ToolCall, ToolChoice, ToolChoiceValue, Usage}, responses::{ @@ -324,6 +325,9 @@ async fn execute_with_mcp_loop( loop { iteration_count += 1; + // Record tool loop iteration metric + SmgMetrics::record_mcp_tool_iteration(¤t_request.model); + // Safety check: prevent infinite loops if iteration_count > MAX_TOOL_ITERATIONS { error!( @@ -434,7 +438,13 @@ async fn execute_with_mcp_loop( // Execute MCP tools (if any) let mcp_results = if !mcp_tool_calls.is_empty() { - execute_mcp_tools(&ctx.mcp_manager, &mcp_tool_calls, &mut mcp_tracking).await? + execute_mcp_tools( + &ctx.mcp_manager, + &mcp_tool_calls, + &mut mcp_tracking, + ¤t_request.model, + ) + .await? } else { Vec::new() }; @@ -759,6 +769,9 @@ async fn execute_mcp_tool_loop_streaming( loop { iteration_count += 1; + // Record tool loop iteration metric + SmgMetrics::record_mcp_tool_iteration(¤t_request.model); + // Safety check: prevent infinite loops if iteration_count > MAX_TOOL_ITERATIONS { emitter.emit_error( @@ -869,8 +882,13 @@ async fn execute_mcp_tool_loop_streaming( // Execute MCP tools (if any) let mcp_results = if !mcp_tool_calls.is_empty() { - match execute_mcp_tools(&ctx.mcp_manager, &mcp_tool_calls, &mut mcp_tracking) - .await + match execute_mcp_tools( + &ctx.mcp_manager, + &mcp_tool_calls, + &mut mcp_tracking, + ¤t_request.model, + ) + .await { Ok(results) => results, Err(err_response) => { @@ -1156,6 +1174,7 @@ async fn execute_mcp_tools( mcp_manager: &Arc, tool_calls: &[ToolCall], tracking: &mut McpCallTracking, + model_id: &str, ) -> Result, Response> { let mut results = Vec::new(); @@ -1192,10 +1211,13 @@ async fn execute_mcp_tools( None }; - match mcp_manager + let tool_start = Instant::now(); + let tool_result = mcp_manager .call_tool(&tool_call.function.name, args_map) - .await - { + .await; + let tool_duration = tool_start.elapsed(); + + match tool_result { Ok(mcp_result) => { debug!( tool_name = %tool_call.function.name, @@ -1230,6 +1252,22 @@ async fn execute_mcp_tools( }, ); + // Record MCP tool metrics + SmgMetrics::record_mcp_tool_duration( + model_id, + &tool_call.function.name, + tool_duration, + ); + SmgMetrics::record_mcp_tool_call( + model_id, + &tool_call.function.name, + if is_error { + smg_labels::RESULT_ERROR + } else { + smg_labels::RESULT_SUCCESS + }, + ); + results.push(ToolResult { call_id: tool_call.id.clone(), tool_name: tool_call.function.name.clone(), @@ -1262,6 +1300,18 @@ async fn execute_mcp_tools( Some(error_msg), ); + // Record MCP tool metrics + SmgMetrics::record_mcp_tool_duration( + model_id, + &tool_call.function.name, + tool_duration, + ); + SmgMetrics::record_mcp_tool_call( + model_id, + &tool_call.function.name, + smg_labels::RESULT_ERROR, + ); + // Return error result to model (let it handle gracefully) results.push(ToolResult { call_id: tool_call.id.clone(), diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs index b47423e52..71e97f5a6 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs @@ -3,7 +3,7 @@ use std::{ collections::HashMap, sync::Arc, - time::{SystemTime, UNIX_EPOCH}, + time::{Instant, SystemTime, UNIX_EPOCH}, }; use axum::{ @@ -22,6 +22,7 @@ use uuid::Uuid; use super::conversions; use crate::{ mcp::{self, McpManager}, + observability::metrics::{smg_labels, SmgMetrics}, protocols::{ chat::{ ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse, @@ -283,6 +284,9 @@ pub(super) async fn execute_tool_loop( if !tool_calls.is_empty() { state.iteration += 1; + // Record tool loop iteration metric + SmgMetrics::record_mcp_tool_iteration(¤t_request.model); + debug!( "Tool loop iteration {}: found {} tool call(s)", state.iteration, @@ -378,6 +382,7 @@ pub(super) async fn execute_tool_loop( tool_name, call_id, args_json_str ); + let tool_start = Instant::now(); let (output_str, success, error) = match ctx .mcp_manager .call_tool(tool_name.as_str(), args_json_str.as_str()) @@ -400,6 +405,23 @@ pub(super) async fn execute_tool_loop( (error_json, false, Some(err_str)) } }; + let tool_duration = tool_start.elapsed(); + + // Record MCP tool metrics + SmgMetrics::record_mcp_tool_duration( + ¤t_request.model, + &tool_name, + tool_duration, + ); + SmgMetrics::record_mcp_tool_call( + ¤t_request.model, + &tool_name, + if success { + smg_labels::RESULT_SUCCESS + } else { + smg_labels::RESULT_ERROR + }, + ); // Record the call in state state.record_call( @@ -619,7 +641,7 @@ async fn execute_tool_loop_streaming_internal( .duration_since(UNIX_EPOCH) .unwrap() .as_secs(); - let mut emitter = ResponseStreamEventEmitter::new(response_id, model, created_at); + let mut emitter = ResponseStreamEventEmitter::new(response_id, model.clone(), created_at); emitter.set_original_request(original_request.clone()); // Emit initial response.created and response.in_progress events @@ -641,6 +663,10 @@ async fn execute_tool_loop_streaming_internal( loop { state.iteration += 1; + + // Record tool loop iteration metric + SmgMetrics::record_mcp_tool_iteration(&model); + if state.iteration > MAX_ITERATIONS { return Err(format!( "Tool loop exceeded maximum iterations ({})", @@ -815,6 +841,7 @@ async fn execute_tool_loop_streaming_internal( "Calling MCP tool '{}' with args: {}", tool_name, args_json_str ); + let tool_start = Instant::now(); let (output_str, success, error) = match ctx .mcp_manager .call_tool(tool_name.as_str(), args_json_str.as_str()) @@ -898,6 +925,19 @@ async fn execute_tool_loop_streaming_internal( (error_json, false, Some(err_str)) } }; + let tool_duration = tool_start.elapsed(); + + // Record MCP tool metrics + SmgMetrics::record_mcp_tool_duration(&model, &tool_name, tool_duration); + SmgMetrics::record_mcp_tool_call( + &model, + &tool_name, + if success { + smg_labels::RESULT_SUCCESS + } else { + smg_labels::RESULT_ERROR + }, + ); // Record the call in state state.record_call( diff --git a/sgl-model-gateway/src/service_discovery.rs b/sgl-model-gateway/src/service_discovery.rs index 4e659840b..6024ddb88 100644 --- a/sgl-model-gateway/src/service_discovery.rs +++ b/sgl-model-gateway/src/service_discovery.rs @@ -19,7 +19,9 @@ use tokio::{task, time}; use tracing::{debug, error, info, warn}; use crate::{ - app_context::AppContext, core::Job, observability::metrics::RouterMetrics, + app_context::AppContext, + core::Job, + observability::metrics::{smg_labels, RouterMetrics, SmgMetrics}, protocols::worker_spec::WorkerConfigRequest, }; @@ -333,7 +335,8 @@ async fn handle_pod_event( let worker_url = pod_info.worker_url(port); if pod_info.is_healthy() { - let should_add = { + // Track whether to add and get count in single lock acquisition + let (should_add, tracked_count) = { let mut tracker = match tracked_pods.lock() { Ok(tracker) => tracker, Err(e) => { @@ -343,10 +346,10 @@ async fn handle_pod_event( }; if tracker.contains(pod_info) { - false + (false, tracker.len()) } else { tracker.insert(pod_info.clone()); - true + (true, tracker.len()) } }; @@ -410,12 +413,31 @@ async fn handle_pod_event( Ok(_) => { debug!("Worker addition job submitted for: {}", worker_url); RouterMetrics::record_discovery_update(1, 0); + + // Layer 4: Record successful registration from K8s discovery + SmgMetrics::record_discovery_registration( + smg_labels::DISCOVERY_KUBERNETES, + smg_labels::REGISTRATION_SUCCESS, + ); + + // Update workers discovered gauge (using count from initial lock) + SmgMetrics::set_discovery_workers_discovered( + smg_labels::DISCOVERY_KUBERNETES, + tracked_count, + ); } Err(e) => { error!( "Failed to submit worker addition job for {}: {}", worker_url, e ); + + // Layer 4: Record failed registration + SmgMetrics::record_discovery_registration( + smg_labels::DISCOVERY_KUBERNETES, + smg_labels::REGISTRATION_FAILED, + ); + if let Ok(mut tracker) = tracked_pods.lock() { tracker.remove(pod_info); } @@ -427,6 +449,12 @@ async fn handle_pod_event( worker_url ); } + } else { + // Pod already tracked - this is a duplicate event + SmgMetrics::record_discovery_registration( + smg_labels::DISCOVERY_KUBERNETES, + smg_labels::REGISTRATION_DUPLICATE, + ); } } } @@ -439,7 +467,8 @@ async fn handle_pod_deletion( ) { let worker_url = pod_info.worker_url(port); - let was_tracked = { + // Remove pod and get remaining count in single lock acquisition + let (was_tracked, remaining_count) = { let mut tracked = match tracked_pods.lock() { Ok(tracked) => tracked, Err(e) => { @@ -447,7 +476,8 @@ async fn handle_pod_deletion( return; } }; - tracked.remove(pod_info) + let removed = tracked.remove(pod_info); + (removed, tracked.len()) }; if was_tracked { @@ -469,6 +499,18 @@ async fn handle_pod_deletion( } else { debug!("Submitted worker removal job for {}", worker_url); RouterMetrics::record_discovery_update(0, 1); + + // Layer 4: Record deregistration from K8s pod deletion + SmgMetrics::record_discovery_deregistration( + smg_labels::DISCOVERY_KUBERNETES, + smg_labels::DEREGISTRATION_POD_DELETED, + ); + + // Update workers discovered gauge (using count from initial lock) + SmgMetrics::set_discovery_workers_discovered( + smg_labels::DISCOVERY_KUBERNETES, + remaining_count, + ); } } else { error!(