[model-gateway]: optimize metrics for minimal CPU and memory overhead (#16041)

This commit is contained in:
Simo Lin
2025-12-28 22:28:08 -08:00
committed by GitHub
parent f13949e5a9
commit ec8c831d5a

View File

@@ -1,18 +1,62 @@
use std::{
borrow::Cow,
net::{IpAddr, Ipv4Addr, SocketAddr},
sync::Arc,
time::Duration,
};
use dashmap::DashMap;
use metrics::{counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram};
use metrics_exporter_prometheus::{Matcher, PrometheusBuilder};
use once_cell::sync::Lazy;
// =============================================================================
// STRING INTERNING
// =============================================================================
//
// Dynamic strings (model_id, worker URLs, paths) are interned to avoid repeated
// heap allocations. The interner uses Arc<str> which is cheap to clone and
// allows the metrics crate to store references without repeated allocations.
//
// Performance characteristics:
// - First occurrence: One allocation + DashMap insert
// - Subsequent occurrences: DashMap lookup + Arc::clone (very cheap)
// - Memory: Strings are never freed (acceptable for bounded label cardinality)
/// Global string interner for metric labels.
/// Uses DashMap for lock-free concurrent access.
static STRING_INTERNER: Lazy<DashMap<String, Arc<str>>> = Lazy::new(DashMap::new);
/// Intern a string, returning a cheaply-cloneable Arc<str>.
///
/// This function is designed for high-throughput scenarios where the same
/// strings (model IDs, worker URLs) appear repeatedly. The first call allocates,
pub fn intern_string(s: &str) -> Arc<str> {
// Fast path: check if already interned
if let Some(entry) = STRING_INTERNER.get(s) {
return Arc::clone(entry.value());
}
// Slow path: intern the string
// Use entry API to avoid TOCTOU race
STRING_INTERNER
.entry(s.to_string())
.or_insert_with(|| Arc::from(s))
.clone()
}
pub fn interner_size() -> usize {
STRING_INTERNER.len()
}
// =============================================================================
// STATIC STRING CONSTANTS
// =============================================================================
/// Static string constants for boolean labels to avoid allocations.
pub const STREAMING_TRUE: &str = "true";
pub const STREAMING_FALSE: &str = "false";
/// Convert a bool to a static string reference (zero-cost).
#[inline]
pub const fn bool_to_static_str(b: bool) -> &'static str {
if b {
STREAMING_TRUE
@@ -25,6 +69,7 @@ pub const fn bool_to_static_str(b: bool) -> &'static str {
/// Returns a static string for known codes, or None for unknown codes.
#[inline]
pub fn status_code_to_static_str(code: u16) -> Option<&'static str> {
// Using a match with explicit arms is faster than a lookup table for this size
match code {
200 => Some("200"),
201 => Some("201"),
@@ -392,7 +437,11 @@ pub mod metrics_labels {
pub const ERROR_INTERNAL: &str = "internal_error";
}
/// SMG Metrics helper struct for the new layered metrics architecture
/// SMG Metrics helper struct for the new layered metrics architecture.
///
/// Design principles for low overhead:
/// - Dynamic labels use string interning (single allocation per unique value)
/// - Static labels use the metrics crate's internal caching
pub struct Metrics;
/// Parameters for recording streaming metrics.
@@ -420,10 +469,11 @@ impl Metrics {
/// Here we want a metric to directly reflect user's experience ("I am sending a request")
/// when viewing the router as a blackbox, and is bumped immediately when the request arrives.
pub fn record_http_request(method: &'static str, path: &str) {
let path_interned = intern_string(path);
counter!(
"smg_http_requests_total",
"method" => method,
"path" => path.to_string(),
"path" => path_interned,
)
.increment(1);
}
@@ -431,37 +481,33 @@ impl Metrics {
/// Record HTTP request duration.
/// For best performance, pass static strings for method.
pub fn record_http_duration(method: &'static str, path: &str, duration: Duration) {
let path_interned = intern_string(path);
histogram!(
"smg_http_request_duration_seconds",
"method" => method,
"path" => path.to_string()
"path" => path_interned
)
.record(duration.as_secs_f64());
}
/// Set active HTTP connections count
/// Set active HTTP connections count.
pub fn set_http_connections_active(count: usize) {
gauge!("smg_http_connections_active").set(count as f64);
}
/// Record HTTP response.
/// Uses static strings for common status codes to avoid allocations.
pub fn record_http_response(status_code: u16, error_code: &str) {
// Use static string for common codes, allocate only for rare ones
let status_str: Cow<'static, str> = match status_code_to_static_str(status_code) {
Some(s) => Cow::Borrowed(s),
None => Cow::Owned(status_code.to_string()),
};
// metrics crate accepts Into<SharedString> which handles Cow efficiently
let status_str: Cow<'static, str> = status_code_to_cow(status_code);
let error_interned = intern_string(error_code);
counter!(
"smg_http_responses_total",
"status_code" => status_str,
"error_code" => error_code.to_string()
"error_code" => error_interned
)
.increment(1);
}
/// Record rate limit decision
/// Record rate limit decision.
pub fn record_http_rate_limit(result: &'static str) {
counter!(
"smg_http_rate_limit_total",
@@ -476,9 +522,10 @@ impl Metrics {
/// Record a routed request.
///
/// Uses string interning for model_id to avoid repeated allocations.
///
/// # Arguments
/// * `streaming` - Use `bool_to_static_str(request.stream)` or the constants
/// `STREAMING_TRUE`/`STREAMING_FALSE` to avoid allocation.
pub fn record_router_request(
router_type: &'static str,
backend_type: &'static str,
@@ -487,19 +534,21 @@ impl Metrics {
endpoint: &'static str,
streaming: &'static str,
) {
let model = intern_string(model_id);
counter!(
"smg_router_requests_total",
"router_type" => router_type,
"backend_type" => backend_type,
"connection_mode" => connection_mode,
"model" => model_id.to_string(),
"model" => model,
"endpoint" => endpoint,
"streaming" => streaming
)
.increment(1);
}
/// Record router request duration
/// Record router request duration.
/// Uses string interning for model_id.
pub fn record_router_duration(
router_type: &'static str,
backend_type: &'static str,
@@ -508,18 +557,20 @@ impl Metrics {
endpoint: &'static str,
duration: Duration,
) {
let model = intern_string(model_id);
histogram!(
"smg_router_request_duration_seconds",
"router_type" => router_type,
"backend_type" => backend_type,
"connection_mode" => connection_mode,
"model" => model_id.to_string(),
"model" => model,
"endpoint" => endpoint
)
.record(duration.as_secs_f64());
}
/// Record a router error
/// Record a router error.
/// Uses string interning for model_id.
pub fn record_router_error(
router_type: &'static str,
backend_type: &'static str,
@@ -528,19 +579,21 @@ impl Metrics {
endpoint: &'static str,
error_type: &'static str,
) {
let model = intern_string(model_id);
counter!(
"smg_router_request_errors_total",
"router_type" => router_type,
"backend_type" => backend_type,
"connection_mode" => connection_mode,
"model" => model_id.to_string(),
"model" => model,
"endpoint" => endpoint,
"error_type" => error_type
)
.increment(1);
}
/// Record pipeline stage duration (gRPC only)
/// Record pipeline stage duration (gRPC only).
/// All labels are static, so this is very fast.
pub fn record_router_stage_duration(
router_type: &'static str,
stage: &'static str,
@@ -555,18 +608,19 @@ impl Metrics {
}
/// Record upstream backend response.
/// Uses static strings for common status codes to avoid allocations.
/// Uses static strings for common status codes and interning for error_code.
pub fn record_router_upstream_response(
router_type: &'static str,
status_code: u16,
error_code: &str,
) {
let status_str: Cow<'static, str> = status_code_to_cow(status_code);
let error_interned = intern_string(error_code);
counter!(
"smg_router_upstream_responses_total",
"router_type" => router_type,
"status_code" => status_str,
"error_code" => error_code.to_string()
"error_code" => error_interned
)
.increment(1);
}
@@ -575,7 +629,8 @@ impl Metrics {
// Layer 2: Router inference metrics (gRPC only)
// ========================================================================
/// Record time to first token
/// Record time to first token.
/// Uses string interning for model_id.
pub fn record_router_ttft(
router_type: &'static str,
backend_type: &'static str,
@@ -583,11 +638,12 @@ impl Metrics {
endpoint: &'static str,
duration: Duration,
) {
let model = intern_string(model_id);
histogram!(
"smg_router_ttft_seconds",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model_id.to_string(),
"model" => model,
"endpoint" => endpoint
)
.record(duration.as_secs_f64());
@@ -601,11 +657,12 @@ impl Metrics {
endpoint: &'static str,
duration: Duration,
) {
let model = intern_string(model_id);
histogram!(
"smg_router_tpot_seconds",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model_id.to_string(),
"model" => model,
"endpoint" => endpoint
)
.record(duration.as_secs_f64());
@@ -620,18 +677,20 @@ impl Metrics {
token_type: &'static str,
count: u64,
) {
let model = intern_string(model_id);
counter!(
"smg_router_tokens_total",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model_id.to_string(),
"model" => model,
"endpoint" => endpoint,
"token_type" => token_type
)
.increment(count);
}
/// Record total generation duration
/// Record total generation duration.
/// Uses string interning for model_id.
pub fn record_router_generation_duration(
router_type: &'static str,
backend_type: &'static str,
@@ -639,11 +698,12 @@ impl Metrics {
endpoint: &'static str,
duration: Duration,
) {
let model = intern_string(model_id);
histogram!(
"smg_router_generation_duration_seconds",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model_id.to_string(),
"model" => model,
"endpoint" => endpoint
)
.record(duration.as_secs_f64());
@@ -665,8 +725,8 @@ impl Metrics {
output_tokens,
} = params;
// Allocate model string once, clone as needed for each metric
let model = model_id.to_string();
// Intern model string once - Arc::clone is just a ref count increment
let model = intern_string(model_id);
// TTFT and TPOT (only if we have a first token time)
if let Some(ttft_duration) = ttft {
@@ -674,7 +734,7 @@ impl Metrics {
"smg_router_ttft_seconds",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model.clone(),
"model" => Arc::clone(&model),
"endpoint" => endpoint
)
.record(ttft_duration.as_secs_f64());
@@ -687,7 +747,7 @@ impl Metrics {
"smg_router_tpot_seconds",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model.clone(),
"model" => Arc::clone(&model),
"endpoint" => endpoint
)
.record(tpot.as_secs_f64());
@@ -699,7 +759,7 @@ impl Metrics {
"smg_router_generation_duration_seconds",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model.clone(),
"model" => Arc::clone(&model),
"endpoint" => endpoint
)
.record(generation_duration.as_secs_f64());
@@ -710,7 +770,7 @@ impl Metrics {
"smg_router_tokens_total",
"router_type" => router_type,
"backend_type" => backend_type,
"model" => model.clone(),
"model" => Arc::clone(&model),
"endpoint" => endpoint,
"token_type" => metrics_labels::TOKEN_INPUT
)
@@ -740,11 +800,12 @@ impl Metrics {
model_id: &str,
size: usize,
) {
let model = intern_string(model_id);
gauge!(
"smg_worker_pool_size",
"worker_type" => worker_type,
"connection_mode" => connection_mode,
"model" => model_id.to_string()
"model" => model
)
.set(size as f64);
}
@@ -780,11 +841,12 @@ impl Metrics {
model_id: &str,
policy: &'static str,
) {
let model = intern_string(model_id);
counter!(
"smg_worker_selection_total",
"worker_type" => worker_type,
"connection_mode" => connection_mode,
"model" => model_id.to_string(),
"model" => model,
"policy" => policy
)
.increment(1);
@@ -839,18 +901,20 @@ impl Metrics {
/// Set running requests per worker
pub fn set_worker_requests_active(worker: &str, count: usize) {
let worker_interned = intern_string(worker);
gauge!(
"smg_worker_requests_active",
"worker" => worker.to_string()
"worker" => worker_interned
)
.set(count as f64);
}
/// Set worker health status
pub fn set_worker_health(worker_url: &str, healthy: bool) {
let worker_interned = intern_string(worker_url);
gauge!(
"smg_worker_health",
"worker" => worker_url.to_string()
"worker" => worker_interned
)
.set(if healthy { 1.0 } else { 0.0 });
}
@@ -861,18 +925,20 @@ impl Metrics {
/// Set circuit breaker state (0=closed, 1=open, 2=half_open)
pub fn set_worker_cb_state(worker: &str, state_code: u8) {
let worker_interned = intern_string(worker);
gauge!(
"smg_worker_cb_state",
"worker" => worker.to_string()
"worker" => worker_interned
)
.set(state_code as f64);
}
/// Record circuit breaker state transition
pub fn record_worker_cb_transition(worker: &str, from: &'static str, to: &'static str) {
let worker_interned = intern_string(worker);
counter!(
"smg_worker_cb_transitions_total",
"worker" => worker.to_string(),
"worker" => worker_interned,
"from" => from,
"to" => to
)
@@ -881,9 +947,10 @@ impl Metrics {
/// Record circuit breaker outcome
pub fn record_worker_cb_outcome(worker: &str, outcome: &'static str) {
let worker_interned = intern_string(worker);
counter!(
"smg_worker_cb_outcomes_total",
"worker" => worker.to_string(),
"worker" => worker_interned,
"outcome" => outcome
)
.increment(1);
@@ -891,18 +958,20 @@ impl Metrics {
/// Set circuit breaker consecutive failures
pub fn set_worker_cb_consecutive_failures(worker: &str, count: u32) {
let worker_interned = intern_string(worker);
gauge!(
"smg_worker_cb_consecutive_failures",
"worker" => worker.to_string()
"worker" => worker_interned
)
.set(count as f64);
}
/// Set circuit breaker consecutive successes
pub fn set_worker_cb_consecutive_successes(worker: &str, count: u32) {
let worker_interned = intern_string(worker);
gauge!(
"smg_worker_cb_consecutive_successes",
"worker" => worker.to_string()
"worker" => worker_interned
)
.set(count as f64);
}
@@ -932,7 +1001,6 @@ impl Metrics {
}
/// Record retry backoff duration.
/// Uses static strings for common attempt numbers (1-5).
pub fn record_worker_retry_backoff(attempt: u32, duration: Duration) {
let attempt_str: Cow<'static, str> = match attempt {
1 => Cow::Borrowed("1"),
@@ -997,10 +1065,12 @@ impl Metrics {
/// Record MCP tool call
pub fn record_mcp_tool_call(model_id: &str, tool_name: &str, result: &'static str) {
let model = intern_string(model_id);
let tool = intern_string(tool_name);
counter!(
"smg_mcp_tool_calls_total",
"model" => model_id.to_string(),
"tool_name" => tool_name.to_string(),
"model" => model,
"tool_name" => tool,
"result" => result
)
.increment(1);
@@ -1008,10 +1078,12 @@ impl Metrics {
/// Record MCP tool execution duration
pub fn record_mcp_tool_duration(model_id: &str, tool_name: &str, duration: Duration) {
let model = intern_string(model_id);
let tool = intern_string(tool_name);
histogram!(
"smg_mcp_tool_duration_seconds",
"model" => model_id.to_string(),
"tool_name" => tool_name.to_string()
"model" => model,
"tool_name" => tool
)
.record(duration.as_secs_f64());
}
@@ -1023,9 +1095,10 @@ impl Metrics {
/// Record MCP tool loop iteration
pub fn record_mcp_tool_iteration(model_id: &str) {
let model = intern_string(model_id);
counter!(
"smg_mcp_tool_iterations_total",
"model" => model_id.to_string()
"model" => model
)
.increment(1);
}
@@ -1086,14 +1159,17 @@ impl Metrics {
// ========================================================================
pub fn remove_worker_metrics(worker_url: &str) {
gauge!("smg_worker_cb_consecutive_failures", "worker" => worker_url.to_string()).set(0.0);
gauge!("smg_worker_cb_consecutive_successes", "worker" => worker_url.to_string()).set(0.0);
gauge!("smg_worker_requests_active", "worker" => worker_url.to_string()).set(0.0);
// Intern once, clone (cheap) for each metric
let worker = intern_string(worker_url);
gauge!("smg_worker_cb_consecutive_failures", "worker" => Arc::clone(&worker)).set(0.0);
gauge!("smg_worker_cb_consecutive_successes", "worker" => Arc::clone(&worker)).set(0.0);
gauge!("smg_worker_requests_active", "worker" => Arc::clone(&worker)).set(0.0);
// Zero for these metrics have special valid meaning, thus we set to -1 temporarily
// (and will remove them completely after https://github.com/metrics-rs/metrics/issues/653)
gauge!("smg_worker_cb_state", "worker" => worker_url.to_string()).set(-1.0);
gauge!("smg_worker_health","worker" => worker_url.to_string()).set(-1.0);
gauge!("smg_worker_cb_state", "worker" => Arc::clone(&worker)).set(-1.0);
gauge!("smg_worker_health", "worker" => worker).set(-1.0);
}
}
@@ -1334,4 +1410,87 @@ mod tests {
assert_eq!(socket_addr.to_string(), "127.0.0.1:29000");
}
// ========================================================================
// String interning tests
// ========================================================================
#[test]
fn test_intern_string_returns_same_arc() {
let s1 = intern_string("test_model");
let s2 = intern_string("test_model");
// Should return the same Arc (pointer equality)
assert!(Arc::ptr_eq(&s1, &s2));
assert_eq!(&*s1, "test_model");
}
#[test]
fn test_intern_string_different_strings() {
let s1 = intern_string("model_a");
let s2 = intern_string("model_b");
// Different strings should have different Arcs
assert!(!Arc::ptr_eq(&s1, &s2));
assert_eq!(&*s1, "model_a");
assert_eq!(&*s2, "model_b");
}
#[test]
fn test_intern_string_empty() {
let s1 = intern_string("");
let s2 = intern_string("");
assert!(Arc::ptr_eq(&s1, &s2));
assert_eq!(&*s1, "");
}
#[test]
fn test_interner_size_grows() {
let initial_size = interner_size();
// Intern some unique strings
let unique = format!("unique_test_string_{}", initial_size);
intern_string(&unique);
assert!(interner_size() > initial_size);
}
#[test]
fn test_bool_to_static_str() {
assert_eq!(bool_to_static_str(true), "true");
assert_eq!(bool_to_static_str(false), "false");
}
#[test]
fn test_status_code_to_static_str() {
// Common codes should return static strings
assert_eq!(status_code_to_static_str(200), Some("200"));
assert_eq!(status_code_to_static_str(404), Some("404"));
assert_eq!(status_code_to_static_str(500), Some("500"));
// Uncommon codes should return None
assert_eq!(status_code_to_static_str(418), None);
assert_eq!(status_code_to_static_str(999), None);
}
#[test]
fn test_status_code_to_cow() {
// Common codes should be borrowed
let cow_200 = status_code_to_cow(200);
assert!(matches!(cow_200, Cow::Borrowed(_)));
assert_eq!(cow_200, "200");
// Uncommon codes should be owned
let cow_418 = status_code_to_cow(418);
assert!(matches!(cow_418, Cow::Owned(_)));
assert_eq!(cow_418, "418");
}
#[test]
fn test_method_to_static_str() {
assert_eq!(method_to_static_str("GET"), "GET");
assert_eq!(method_to_static_str("POST"), "POST");
assert_eq!(method_to_static_str("UNKNOWN"), "OTHER");
}
}