[model-gateway] Add Layer 3 worker metrics (smg_worker_*) (#15130)

This commit is contained in:
Simo Lin
2025-12-14 20:25:46 -08:00
committed by GitHub
parent 3912ee4991
commit fb96669ff9
7 changed files with 227 additions and 15 deletions

View File

@@ -8,7 +8,7 @@ use std::{
use tracing::info;
use crate::observability::metrics::RouterMetrics;
use crate::observability::metrics::{RouterMetrics, SmgMetrics};
/// Circuit breaker configuration
#[derive(Debug, Clone)]
@@ -96,6 +96,9 @@ impl CircuitBreaker {
/// Create a new circuit breaker with custom configuration and metric label
pub fn with_config_and_label(config: CircuitBreakerConfig, metric_label: String) -> Self {
let init_state = CircuitState::Closed;
// New metrics
SmgMetrics::set_worker_cb_state(&metric_label, init_state.to_int());
// Legacy metrics
RouterMetrics::set_cb_state(&metric_label, init_state.to_int());
Self {
state: Arc::new(RwLock::new(init_state)),
@@ -153,6 +156,9 @@ impl CircuitBreaker {
}
let outcome_str = if success { "success" } else { "failure" };
// New metrics
SmgMetrics::record_worker_cb_outcome(&self.metric_label, outcome_str);
// Legacy metrics
RouterMetrics::record_cb_outcome(&self.metric_label, outcome_str);
self.publish_gauge_metrics();
}
@@ -232,6 +238,10 @@ impl CircuitBreaker {
let from = old_state.as_str();
let to = new_state.as_str();
info!("Circuit breaker state transition: {} -> {}", from, to);
// New metrics
SmgMetrics::record_worker_cb_transition(&self.metric_label, from, to);
SmgMetrics::set_worker_cb_state(&self.metric_label, new_state.to_int());
// Legacy metrics
RouterMetrics::record_cb_state_transition(&self.metric_label, from, to);
RouterMetrics::set_cb_state(&self.metric_label, new_state.to_int());
self.publish_gauge_metrics();
@@ -315,6 +325,10 @@ impl CircuitBreaker {
// TODO maybe publish whenever the variable is changed
fn publish_gauge_metrics(&self) {
// New metrics
SmgMetrics::set_worker_cb_consecutive_failures(&self.metric_label, self.failure_count());
SmgMetrics::set_worker_cb_consecutive_successes(&self.metric_label, self.success_count());
// Legacy metrics
RouterMetrics::set_cb_consecutive_failures(&self.metric_label, self.failure_count());
RouterMetrics::set_cb_consecutive_successes(&self.metric_label, self.success_count());
}

View File

@@ -1,13 +1,13 @@
//! Step to remove workers from worker registry.
use std::sync::Arc;
use std::{collections::HashSet, sync::Arc};
use async_trait::async_trait;
use tracing::{debug, warn};
use crate::{
app_context::AppContext,
observability::metrics::RouterMetrics,
observability::metrics::{RouterMetrics, SmgMetrics},
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -27,6 +27,20 @@ impl StepExecutor for RemoveFromWorkerRegistryStep {
worker_urls.len()
);
// Collect unique worker configurations before removal for pool size updates
let unique_configs: HashSet<_> = worker_urls
.iter()
.filter_map(|url| app_context.worker_registry.get_by_url(url))
.map(|w| {
let meta = w.metadata();
(
meta.worker_type.clone(),
meta.connection_mode.clone(),
w.model_id().to_string(),
)
})
.collect();
let mut removed_count = 0;
for worker_url in worker_urls.iter() {
if app_context
@@ -49,9 +63,34 @@ impl StepExecutor for RemoveFromWorkerRegistryStep {
debug!("Removed {} worker(s) from registry", removed_count);
}
// Update active workers metric
// Update active workers metric (legacy)
RouterMetrics::set_active_workers(app_context.worker_registry.len());
// Update Layer 3 worker pool size metrics for unique configurations
for (worker_type, connection_mode, model_id) in unique_configs {
// Get labels before moving values into get_workers_filtered
let worker_type_label = worker_type.as_metric_label();
let connection_mode_label = connection_mode.as_metric_label();
let pool_size = app_context
.worker_registry
.get_workers_filtered(
Some(&model_id),
Some(worker_type),
Some(connection_mode),
None,
false,
)
.len();
SmgMetrics::set_worker_pool_size(
worker_type_label,
connection_mode_label,
&model_id,
pool_size,
);
}
Ok(StepResult::Success)
}

View File

@@ -1,6 +1,6 @@
//! Unified worker registration step.
use std::sync::Arc;
use std::{collections::HashSet, sync::Arc};
use async_trait::async_trait;
use tracing::debug;
@@ -8,7 +8,7 @@ use tracing::debug;
use crate::{
app_context::AppContext,
core::Worker,
observability::metrics::RouterMetrics,
observability::metrics::{RouterMetrics, SmgMetrics},
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
};
@@ -37,9 +37,47 @@ impl StepExecutor for RegisterWorkersStep {
worker_ids.push(worker_id);
}
// Update active workers metric
// Update active workers metric (legacy)
RouterMetrics::set_active_workers(app_context.worker_registry.len());
// Collect unique worker configurations to avoid redundant metric updates
let unique_configs: HashSet<_> = workers
.iter()
.map(|w| {
let meta = w.metadata();
(
meta.worker_type.clone(),
meta.connection_mode.clone(),
w.model_id().to_string(),
)
})
.collect();
// Update Layer 3 worker pool size metrics per unique type/connection/model
for (worker_type, connection_mode, model_id) in unique_configs {
// Get labels before moving values into get_workers_filtered
let worker_type_label = worker_type.as_metric_label();
let connection_mode_label = connection_mode.as_metric_label();
let pool_size = app_context
.worker_registry
.get_workers_filtered(
Some(&model_id),
Some(worker_type),
Some(connection_mode),
None,
false,
)
.len();
SmgMetrics::set_worker_pool_size(
worker_type_label,
connection_mode_label,
&model_id,
pool_size,
);
}
context.set("worker_ids", worker_ids);
Ok(StepResult::Success)
}

View File

@@ -17,7 +17,7 @@ use super::{
};
use crate::{
core::{BasicWorkerBuilder, DPAwareWorkerBuilder},
observability::metrics::RouterMetrics,
observability::metrics::{smg_labels, RouterMetrics, SmgMetrics},
protocols::worker_spec::WorkerInfo,
routers::grpc::client::GrpcClient,
};
@@ -310,6 +310,14 @@ impl ConnectionMode {
_ => false,
}
}
/// Get the metric label for this connection mode
pub fn as_metric_label(&self) -> &'static str {
match self {
ConnectionMode::Http => smg_labels::CONNECTION_HTTP,
ConnectionMode::Grpc { .. } => smg_labels::CONNECTION_GRPC,
}
}
}
impl fmt::Display for ConnectionMode {
@@ -392,6 +400,17 @@ impl fmt::Display for WorkerType {
}
}
impl WorkerType {
/// Get the metric label for this worker type
pub fn as_metric_label(&self) -> &'static str {
match self {
WorkerType::Regular => smg_labels::WORKER_REGULAR,
WorkerType::Prefill { .. } => smg_labels::WORKER_PREFILL,
WorkerType::Decode => smg_labels::WORKER_DECODE,
}
}
}
/// Health check configuration
#[derive(Debug, Clone)]
pub struct HealthConfig {
@@ -537,7 +556,9 @@ impl BasicWorker {
}
fn update_running_requests_metrics(&self) {
RouterMetrics::set_running_requests(self.url(), self.load());
let load = self.load();
RouterMetrics::set_running_requests(self.url(), load);
SmgMetrics::set_worker_requests_active(self.url(), load);
}
}
@@ -574,10 +595,16 @@ impl Worker for BasicWorker {
ConnectionMode::Grpc { .. } => self.grpc_health_check().await?,
};
// Get worker type label for metrics
let worker_type_str = self.metadata.worker_type.as_metric_label();
if health_result {
self.consecutive_failures.store(0, Ordering::Release);
let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1;
// Record health check success metric
SmgMetrics::record_worker_health_check(worker_type_str, smg_labels::CB_SUCCESS);
if !self.is_healthy()
&& successes >= self.metadata.health_config.success_threshold as usize
{
@@ -589,6 +616,9 @@ impl Worker for BasicWorker {
self.consecutive_successes.store(0, Ordering::Release);
let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1;
// Record health check failure metric
SmgMetrics::record_worker_health_check(worker_type_str, smg_labels::CB_FAILURE);
if self.is_healthy()
&& failures >= self.metadata.health_config.failure_threshold as usize
{

View File

@@ -9,6 +9,7 @@ use tracing::{error, warn};
use super::PipelineStage;
use crate::{
core::{ConnectionMode, Worker, WorkerRegistry, WorkerType},
observability::metrics::{smg_labels, SmgMetrics},
policies::PolicyRegistry,
routers::{
error,
@@ -146,7 +147,17 @@ impl WorkerSelectionStage {
// Select worker using the policy
let idx = policy.select_worker(&available, text)?;
Some(available[idx].clone())
let selected = available[idx].clone();
// Record worker selection metric
SmgMetrics::record_worker_selection(
smg_labels::WORKER_REGULAR,
smg_labels::CONNECTION_GRPC,
model_id.unwrap_or("default"),
policy.name(),
);
Some(selected)
}
fn select_pd_pair(
@@ -195,6 +206,23 @@ impl WorkerSelectionStage {
let prefill_idx = policy.select_worker(&available_prefill, text)?;
let decode_idx = policy.select_worker(&available_decode, text)?;
let model = model_id.unwrap_or("default");
let policy_name = policy.name();
// Record worker selection metrics for both prefill and decode
SmgMetrics::record_worker_selection(
smg_labels::WORKER_PREFILL,
smg_labels::CONNECTION_GRPC,
model,
policy_name,
);
SmgMetrics::record_worker_selection(
smg_labels::WORKER_DECODE,
smg_labels::CONNECTION_GRPC,
model,
policy_name,
);
Some((
available_prefill[prefill_idx].clone(),
available_decode[decode_idx].clone(),

View File

@@ -345,11 +345,26 @@ impl PDRouter {
)
.await;
let _status = response.status();
let not_error = _status.is_success() || _status.is_client_error();
let status = response.status();
let not_error = status.is_success() || status.is_client_error();
prefill.record_outcome(not_error);
decode.record_outcome(not_error);
// Record worker errors for server errors (5xx)
if status.is_server_error() {
let error_type = error_type_from_status(status);
SmgMetrics::record_worker_error(
smg_labels::WORKER_PREFILL,
smg_labels::CONNECTION_HTTP,
error_type,
);
SmgMetrics::record_worker_error(
smg_labels::WORKER_DECODE,
smg_labels::CONNECTION_HTTP,
error_type,
);
}
response
}
}
@@ -358,8 +373,16 @@ impl PDRouter {
|delay, attempt| {
RouterMetrics::record_retry(route);
RouterMetrics::record_retry_backoff_duration(delay, attempt);
// Layer 3 worker metrics (PD mode uses both prefill and decode workers)
SmgMetrics::record_worker_retry(smg_labels::WORKER_PREFILL, endpoint);
SmgMetrics::record_worker_retry(smg_labels::WORKER_DECODE, endpoint);
SmgMetrics::record_worker_retry_backoff(attempt, delay);
},
|| {
RouterMetrics::record_retries_exhausted(route);
SmgMetrics::record_worker_retries_exhausted(smg_labels::WORKER_PREFILL, endpoint);
SmgMetrics::record_worker_retries_exhausted(smg_labels::WORKER_DECODE, endpoint);
},
|| RouterMetrics::record_retries_exhausted(route),
)
.await;
@@ -729,6 +752,21 @@ impl PDRouter {
"decode",
)?;
// Record worker selection metrics (Layer 3)
let model = model_id.unwrap_or("default");
SmgMetrics::record_worker_selection(
smg_labels::WORKER_PREFILL,
smg_labels::CONNECTION_HTTP,
model,
prefill_policy.name(),
);
SmgMetrics::record_worker_selection(
smg_labels::WORKER_DECODE,
smg_labels::CONNECTION_HTTP,
model,
decode_policy.name(),
);
Ok((prefill, decode))
}

View File

@@ -157,6 +157,15 @@ impl Router {
};
let idx = policy.select_worker(&available, text)?;
// Record worker selection metric (Layer 3)
SmgMetrics::record_worker_selection(
smg_labels::WORKER_REGULAR,
smg_labels::CONNECTION_HTTP,
model_id.unwrap_or("default"),
policy.name(),
);
Some(available[idx].clone())
}
@@ -206,9 +215,15 @@ impl Router {
|delay, attempt| {
RouterMetrics::record_retry(route);
RouterMetrics::record_retry_backoff_duration(delay, attempt);
// Layer 3 worker metrics
SmgMetrics::record_worker_retry(smg_labels::WORKER_REGULAR, endpoint);
SmgMetrics::record_worker_retry_backoff(attempt, delay);
},
// on_exhausted hook
|| RouterMetrics::record_retries_exhausted(route),
|| {
RouterMetrics::record_retries_exhausted(route);
SmgMetrics::record_worker_retries_exhausted(smg_labels::WORKER_REGULAR, endpoint);
},
)
.await;
@@ -290,7 +305,17 @@ impl Router {
events::RequestReceivedEvent {}.emit();
worker.record_outcome(response.status().is_success());
let status = response.status();
worker.record_outcome(status.is_success());
// Record worker errors for server errors (5xx)
if status.is_server_error() {
SmgMetrics::record_worker_error(
smg_labels::WORKER_REGULAR,
smg_labels::CONNECTION_HTTP,
error_type_from_status(status),
);
}
response
}