chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[model-gateway] add mcp and discovery metrics (#15156)

Browse Source
This commit is contained in:
Simo Lin
2025-12-14 22:54:53 -08:00
committed by GitHub
parent 21cfebac65
commit 7b9156c773
5 changed files with 163 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
sgl-model-gateway/src/core/steps/mcp_registration.rs
View File

@@ -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)

12
sgl-model-gateway/src/observability/metrics.rs
View File

@@ -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";

64
sgl-model-gateway/src/routers/grpc/harmony/responses.rs
View File

@@ -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(&current_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,
&current_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(&current_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,
&current_request.model,
)
.await
{
Ok(results) => results,
Err(err_response) => {
@@ -1156,6 +1174,7 @@ async fn execute_mcp_tools(
mcp_manager: &Arc<McpManager>,
tool_calls: &[ToolCall],
tracking: &mut McpCallTracking,
model_id: &str,
) -> Result<Vec<ToolResult>, 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(),

44
sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs
View File

@@ -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(&current_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(
&current_request.model,
&tool_name,
tool_duration,
);
SmgMetrics::record_mcp_tool_call(
&current_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(

54
sgl-model-gateway/src/service_discovery.rs
View File

@@ -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!(