From fb96669ff97312d510073ebde79fae123e566a4d Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sun, 14 Dec 2025 20:25:46 -0800 Subject: [PATCH] [model-gateway] Add Layer 3 worker metrics (smg_worker_*) (#15130) --- sgl-model-gateway/src/core/circuit_breaker.rs | 16 ++++++- .../local/remove_from_worker_registry.rs | 45 +++++++++++++++++-- .../src/core/steps/worker/shared/register.rs | 44 ++++++++++++++++-- sgl-model-gateway/src/core/worker.rs | 34 +++++++++++++- .../grpc/common/stages/worker_selection.rs | 30 ++++++++++++- .../src/routers/http/pd_router.rs | 44 ++++++++++++++++-- sgl-model-gateway/src/routers/http/router.rs | 29 +++++++++++- 7 files changed, 227 insertions(+), 15 deletions(-) diff --git a/sgl-model-gateway/src/core/circuit_breaker.rs b/sgl-model-gateway/src/core/circuit_breaker.rs index 7b3a44b5a..b9d0cbe52 100644 --- a/sgl-model-gateway/src/core/circuit_breaker.rs +++ b/sgl-model-gateway/src/core/circuit_breaker.rs @@ -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()); } diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs index 4693687a9..7da3fefce 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs @@ -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) } diff --git a/sgl-model-gateway/src/core/steps/worker/shared/register.rs b/sgl-model-gateway/src/core/steps/worker/shared/register.rs index 4aa01eeb8..ae2e15b77 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/register.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/register.rs @@ -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) } diff --git a/sgl-model-gateway/src/core/worker.rs b/sgl-model-gateway/src/core/worker.rs index 14635dc49..8541c8cfd 100644 --- a/sgl-model-gateway/src/core/worker.rs +++ b/sgl-model-gateway/src/core/worker.rs @@ -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 { diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs index f599fcbcf..7cbf6976b 100644 --- a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -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(), diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 20ab842ac..2c32097f1 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -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)) } diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index ad2af1e85..2a4467f90 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -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 }