[model-gateway] add mcp and discovery metrics (#15156)
This commit is contained in:
@@ -7,6 +7,7 @@ use tracing::{debug, error, info, warn};
|
|||||||
use crate::{
|
use crate::{
|
||||||
app_context::AppContext,
|
app_context::AppContext,
|
||||||
mcp::{config::McpServerConfig, manager::McpManager},
|
mcp::{config::McpServerConfig, manager::McpManager},
|
||||||
|
observability::metrics::SmgMetrics,
|
||||||
workflow::*,
|
workflow::*,
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -151,6 +152,9 @@ impl StepExecutor for RegisterMcpServerStep {
|
|||||||
// Register the client in the manager's client map
|
// Register the client in the manager's client map
|
||||||
mcp_manager.register_static_server(config_request.name.clone(), mcp_client);
|
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);
|
info!("Registered MCP server: {}", config_request.name);
|
||||||
|
|
||||||
Ok(StepResult::Success)
|
Ok(StepResult::Success)
|
||||||
|
|||||||
@@ -804,6 +804,18 @@ pub mod smg_labels {
|
|||||||
pub const DISCOVERY_CONSUL: &str = "consul";
|
pub const DISCOVERY_CONSUL: &str = "consul";
|
||||||
pub const DISCOVERY_MANUAL: &str = "manual";
|
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
|
// Rate limit results
|
||||||
pub const RATE_LIMIT_ALLOWED: &str = "allowed";
|
pub const RATE_LIMIT_ALLOWED: &str = "allowed";
|
||||||
pub const RATE_LIMIT_REJECTED: &str = "rejected";
|
pub const RATE_LIMIT_REJECTED: &str = "rejected";
|
||||||
|
|||||||
@@ -32,7 +32,7 @@
|
|||||||
//! ```
|
//! ```
|
||||||
use std::{
|
use std::{
|
||||||
sync::Arc,
|
sync::Arc,
|
||||||
time::{SystemTime, UNIX_EPOCH},
|
time::{Instant, SystemTime, UNIX_EPOCH},
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::response::Response;
|
use axum::response::Response;
|
||||||
@@ -45,6 +45,7 @@ use uuid::Uuid;
|
|||||||
use crate::{
|
use crate::{
|
||||||
data_connector::{ConversationItemStorage, ConversationStorage, ResponseId, ResponseStorage},
|
data_connector::{ConversationItemStorage, ConversationStorage, ResponseId, ResponseStorage},
|
||||||
mcp::{self, McpManager},
|
mcp::{self, McpManager},
|
||||||
|
observability::metrics::{smg_labels, SmgMetrics},
|
||||||
protocols::{
|
protocols::{
|
||||||
common::{Function, ToolCall, ToolChoice, ToolChoiceValue, Usage},
|
common::{Function, ToolCall, ToolChoice, ToolChoiceValue, Usage},
|
||||||
responses::{
|
responses::{
|
||||||
@@ -324,6 +325,9 @@ async fn execute_with_mcp_loop(
|
|||||||
loop {
|
loop {
|
||||||
iteration_count += 1;
|
iteration_count += 1;
|
||||||
|
|
||||||
|
// Record tool loop iteration metric
|
||||||
|
SmgMetrics::record_mcp_tool_iteration(¤t_request.model);
|
||||||
|
|
||||||
// Safety check: prevent infinite loops
|
// Safety check: prevent infinite loops
|
||||||
if iteration_count > MAX_TOOL_ITERATIONS {
|
if iteration_count > MAX_TOOL_ITERATIONS {
|
||||||
error!(
|
error!(
|
||||||
@@ -434,7 +438,13 @@ async fn execute_with_mcp_loop(
|
|||||||
|
|
||||||
// Execute MCP tools (if any)
|
// Execute MCP tools (if any)
|
||||||
let mcp_results = if !mcp_tool_calls.is_empty() {
|
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 {
|
} else {
|
||||||
Vec::new()
|
Vec::new()
|
||||||
};
|
};
|
||||||
@@ -759,6 +769,9 @@ async fn execute_mcp_tool_loop_streaming(
|
|||||||
loop {
|
loop {
|
||||||
iteration_count += 1;
|
iteration_count += 1;
|
||||||
|
|
||||||
|
// Record tool loop iteration metric
|
||||||
|
SmgMetrics::record_mcp_tool_iteration(¤t_request.model);
|
||||||
|
|
||||||
// Safety check: prevent infinite loops
|
// Safety check: prevent infinite loops
|
||||||
if iteration_count > MAX_TOOL_ITERATIONS {
|
if iteration_count > MAX_TOOL_ITERATIONS {
|
||||||
emitter.emit_error(
|
emitter.emit_error(
|
||||||
@@ -869,8 +882,13 @@ async fn execute_mcp_tool_loop_streaming(
|
|||||||
|
|
||||||
// Execute MCP tools (if any)
|
// Execute MCP tools (if any)
|
||||||
let mcp_results = if !mcp_tool_calls.is_empty() {
|
let mcp_results = if !mcp_tool_calls.is_empty() {
|
||||||
match execute_mcp_tools(&ctx.mcp_manager, &mcp_tool_calls, &mut mcp_tracking)
|
match execute_mcp_tools(
|
||||||
.await
|
&ctx.mcp_manager,
|
||||||
|
&mcp_tool_calls,
|
||||||
|
&mut mcp_tracking,
|
||||||
|
¤t_request.model,
|
||||||
|
)
|
||||||
|
.await
|
||||||
{
|
{
|
||||||
Ok(results) => results,
|
Ok(results) => results,
|
||||||
Err(err_response) => {
|
Err(err_response) => {
|
||||||
@@ -1156,6 +1174,7 @@ async fn execute_mcp_tools(
|
|||||||
mcp_manager: &Arc<McpManager>,
|
mcp_manager: &Arc<McpManager>,
|
||||||
tool_calls: &[ToolCall],
|
tool_calls: &[ToolCall],
|
||||||
tracking: &mut McpCallTracking,
|
tracking: &mut McpCallTracking,
|
||||||
|
model_id: &str,
|
||||||
) -> Result<Vec<ToolResult>, Response> {
|
) -> Result<Vec<ToolResult>, Response> {
|
||||||
let mut results = Vec::new();
|
let mut results = Vec::new();
|
||||||
|
|
||||||
@@ -1192,10 +1211,13 @@ async fn execute_mcp_tools(
|
|||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|
||||||
match mcp_manager
|
let tool_start = Instant::now();
|
||||||
|
let tool_result = mcp_manager
|
||||||
.call_tool(&tool_call.function.name, args_map)
|
.call_tool(&tool_call.function.name, args_map)
|
||||||
.await
|
.await;
|
||||||
{
|
let tool_duration = tool_start.elapsed();
|
||||||
|
|
||||||
|
match tool_result {
|
||||||
Ok(mcp_result) => {
|
Ok(mcp_result) => {
|
||||||
debug!(
|
debug!(
|
||||||
tool_name = %tool_call.function.name,
|
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 {
|
results.push(ToolResult {
|
||||||
call_id: tool_call.id.clone(),
|
call_id: tool_call.id.clone(),
|
||||||
tool_name: tool_call.function.name.clone(),
|
tool_name: tool_call.function.name.clone(),
|
||||||
@@ -1262,6 +1300,18 @@ async fn execute_mcp_tools(
|
|||||||
Some(error_msg),
|
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)
|
// Return error result to model (let it handle gracefully)
|
||||||
results.push(ToolResult {
|
results.push(ToolResult {
|
||||||
call_id: tool_call.id.clone(),
|
call_id: tool_call.id.clone(),
|
||||||
|
|||||||
@@ -3,7 +3,7 @@
|
|||||||
use std::{
|
use std::{
|
||||||
collections::HashMap,
|
collections::HashMap,
|
||||||
sync::Arc,
|
sync::Arc,
|
||||||
time::{SystemTime, UNIX_EPOCH},
|
time::{Instant, SystemTime, UNIX_EPOCH},
|
||||||
};
|
};
|
||||||
|
|
||||||
use axum::{
|
use axum::{
|
||||||
@@ -22,6 +22,7 @@ use uuid::Uuid;
|
|||||||
use super::conversions;
|
use super::conversions;
|
||||||
use crate::{
|
use crate::{
|
||||||
mcp::{self, McpManager},
|
mcp::{self, McpManager},
|
||||||
|
observability::metrics::{smg_labels, SmgMetrics},
|
||||||
protocols::{
|
protocols::{
|
||||||
chat::{
|
chat::{
|
||||||
ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse,
|
ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse,
|
||||||
@@ -283,6 +284,9 @@ pub(super) async fn execute_tool_loop(
|
|||||||
if !tool_calls.is_empty() {
|
if !tool_calls.is_empty() {
|
||||||
state.iteration += 1;
|
state.iteration += 1;
|
||||||
|
|
||||||
|
// Record tool loop iteration metric
|
||||||
|
SmgMetrics::record_mcp_tool_iteration(¤t_request.model);
|
||||||
|
|
||||||
debug!(
|
debug!(
|
||||||
"Tool loop iteration {}: found {} tool call(s)",
|
"Tool loop iteration {}: found {} tool call(s)",
|
||||||
state.iteration,
|
state.iteration,
|
||||||
@@ -378,6 +382,7 @@ pub(super) async fn execute_tool_loop(
|
|||||||
tool_name, call_id, args_json_str
|
tool_name, call_id, args_json_str
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let tool_start = Instant::now();
|
||||||
let (output_str, success, error) = match ctx
|
let (output_str, success, error) = match ctx
|
||||||
.mcp_manager
|
.mcp_manager
|
||||||
.call_tool(tool_name.as_str(), args_json_str.as_str())
|
.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))
|
(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
|
// Record the call in state
|
||||||
state.record_call(
|
state.record_call(
|
||||||
@@ -619,7 +641,7 @@ async fn execute_tool_loop_streaming_internal(
|
|||||||
.duration_since(UNIX_EPOCH)
|
.duration_since(UNIX_EPOCH)
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.as_secs();
|
.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());
|
emitter.set_original_request(original_request.clone());
|
||||||
|
|
||||||
// Emit initial response.created and response.in_progress events
|
// Emit initial response.created and response.in_progress events
|
||||||
@@ -641,6 +663,10 @@ async fn execute_tool_loop_streaming_internal(
|
|||||||
|
|
||||||
loop {
|
loop {
|
||||||
state.iteration += 1;
|
state.iteration += 1;
|
||||||
|
|
||||||
|
// Record tool loop iteration metric
|
||||||
|
SmgMetrics::record_mcp_tool_iteration(&model);
|
||||||
|
|
||||||
if state.iteration > MAX_ITERATIONS {
|
if state.iteration > MAX_ITERATIONS {
|
||||||
return Err(format!(
|
return Err(format!(
|
||||||
"Tool loop exceeded maximum iterations ({})",
|
"Tool loop exceeded maximum iterations ({})",
|
||||||
@@ -815,6 +841,7 @@ async fn execute_tool_loop_streaming_internal(
|
|||||||
"Calling MCP tool '{}' with args: {}",
|
"Calling MCP tool '{}' with args: {}",
|
||||||
tool_name, args_json_str
|
tool_name, args_json_str
|
||||||
);
|
);
|
||||||
|
let tool_start = Instant::now();
|
||||||
let (output_str, success, error) = match ctx
|
let (output_str, success, error) = match ctx
|
||||||
.mcp_manager
|
.mcp_manager
|
||||||
.call_tool(tool_name.as_str(), args_json_str.as_str())
|
.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))
|
(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
|
// Record the call in state
|
||||||
state.record_call(
|
state.record_call(
|
||||||
|
|||||||
@@ -19,7 +19,9 @@ use tokio::{task, time};
|
|||||||
use tracing::{debug, error, info, warn};
|
use tracing::{debug, error, info, warn};
|
||||||
|
|
||||||
use crate::{
|
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,
|
protocols::worker_spec::WorkerConfigRequest,
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -333,7 +335,8 @@ async fn handle_pod_event(
|
|||||||
let worker_url = pod_info.worker_url(port);
|
let worker_url = pod_info.worker_url(port);
|
||||||
|
|
||||||
if pod_info.is_healthy() {
|
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() {
|
let mut tracker = match tracked_pods.lock() {
|
||||||
Ok(tracker) => tracker,
|
Ok(tracker) => tracker,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -343,10 +346,10 @@ async fn handle_pod_event(
|
|||||||
};
|
};
|
||||||
|
|
||||||
if tracker.contains(pod_info) {
|
if tracker.contains(pod_info) {
|
||||||
false
|
(false, tracker.len())
|
||||||
} else {
|
} else {
|
||||||
tracker.insert(pod_info.clone());
|
tracker.insert(pod_info.clone());
|
||||||
true
|
(true, tracker.len())
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -410,12 +413,31 @@ async fn handle_pod_event(
|
|||||||
Ok(_) => {
|
Ok(_) => {
|
||||||
debug!("Worker addition job submitted for: {}", worker_url);
|
debug!("Worker addition job submitted for: {}", worker_url);
|
||||||
RouterMetrics::record_discovery_update(1, 0);
|
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) => {
|
Err(e) => {
|
||||||
error!(
|
error!(
|
||||||
"Failed to submit worker addition job for {}: {}",
|
"Failed to submit worker addition job for {}: {}",
|
||||||
worker_url, e
|
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() {
|
if let Ok(mut tracker) = tracked_pods.lock() {
|
||||||
tracker.remove(pod_info);
|
tracker.remove(pod_info);
|
||||||
}
|
}
|
||||||
@@ -427,6 +449,12 @@ async fn handle_pod_event(
|
|||||||
worker_url
|
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 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() {
|
let mut tracked = match tracked_pods.lock() {
|
||||||
Ok(tracked) => tracked,
|
Ok(tracked) => tracked,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -447,7 +476,8 @@ async fn handle_pod_deletion(
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
tracked.remove(pod_info)
|
let removed = tracked.remove(pod_info);
|
||||||
|
(removed, tracked.len())
|
||||||
};
|
};
|
||||||
|
|
||||||
if was_tracked {
|
if was_tracked {
|
||||||
@@ -469,6 +499,18 @@ async fn handle_pod_deletion(
|
|||||||
} else {
|
} else {
|
||||||
debug!("Submitted worker removal job for {}", worker_url);
|
debug!("Submitted worker removal job for {}", worker_url);
|
||||||
RouterMetrics::record_discovery_update(0, 1);
|
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 {
|
} else {
|
||||||
error!(
|
error!(
|
||||||
|
|||||||
Reference in New Issue
Block a user