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

[model-gateway] reduce cpu overhead (#15316)

Browse Source
This commit is contained in:
Simo Lin
2025-12-17 07:42:01 -10:00
committed by GitHub
parent 0c00220795
commit d747147a26
7 changed files with 251 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

169
sgl-model-gateway/src/core/circuit_breaker.rs
View File

@@ -1,8 +1,5 @@
use std::{
sync::{
atomic::{AtomicU32, AtomicU64, Ordering},
Arc, RwLock,
},
sync::atomic::{AtomicU32, AtomicU64, AtomicU8, Ordering},
time::{Duration, Instant},
};
@@ -34,6 +31,11 @@ impl Default for CircuitBreakerConfig {
}
}
/// Circuit breaker state constants for atomic storage
const STATE_CLOSED: u8 = 0;
const STATE_OPEN: u8 = 1;
const STATE_HALF_OPEN: u8 = 2;
/// Circuit breaker state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
@@ -66,23 +68,49 @@ impl CircuitState {
pub fn to_int(&self) -> u8 {
match self {
CircuitState::Closed => 0u8,
CircuitState::Open => 1u8,
CircuitState::HalfOpen => 2u8,
CircuitState::Closed => STATE_CLOSED,
CircuitState::Open => STATE_OPEN,
CircuitState::HalfOpen => STATE_HALF_OPEN,
}
}
fn from_int(v: u8) -> Self {
match v {
STATE_CLOSED => CircuitState::Closed,
STATE_OPEN => CircuitState::Open,
STATE_HALF_OPEN => CircuitState::HalfOpen,
_ => CircuitState::Closed, // Default to closed for safety
}
}
}
/// Circuit breaker implementation
/// Get current time as milliseconds since an arbitrary epoch.
/// Uses Instant for monotonic time, converting to ms for atomic storage.
#[inline]
fn now_ms() -> u64 {
// Use a static reference point for consistent timing
static START: std::sync::OnceLock<Instant> = std::sync::OnceLock::new();
let start = START.get_or_init(Instant::now);
start.elapsed().as_millis() as u64
}
/// Circuit breaker implementation using lock-free atomics for hot paths.
///
/// This implementation avoids RwLock contention by using atomic operations
/// for state checks (the most common operation). Only state transitions
/// use compare-and-swap which is still lock-free.
#[derive(Debug)]
pub struct CircuitBreaker {
state: Arc<RwLock<CircuitState>>,
consecutive_failures: Arc<AtomicU32>,
consecutive_successes: Arc<AtomicU32>,
total_failures: Arc<AtomicU64>,
total_successes: Arc<AtomicU64>,
last_failure_time: Arc<RwLock<Option<Instant>>>,
last_state_change: Arc<RwLock<Instant>>,
/// Circuit state stored as atomic u8 (0=Closed, 1=Open, 2=HalfOpen)
state: AtomicU8,
consecutive_failures: AtomicU32,
consecutive_successes: AtomicU32,
total_failures: AtomicU64,
total_successes: AtomicU64,
/// Last failure time in milliseconds (from now_ms())
last_failure_time_ms: AtomicU64,
/// Last state change time in milliseconds (from now_ms())
last_state_change_ms: AtomicU64,
config: CircuitBreakerConfig,
metric_label: String,
}
@@ -98,13 +126,13 @@ impl CircuitBreaker {
let init_state = CircuitState::Closed;
Metrics::set_worker_cb_state(&metric_label, init_state.to_int());
Self {
state: Arc::new(RwLock::new(init_state)),
consecutive_failures: Arc::new(AtomicU32::new(0)),
consecutive_successes: Arc::new(AtomicU32::new(0)),
total_failures: Arc::new(AtomicU64::new(0)),
total_successes: Arc::new(AtomicU64::new(0)),
last_failure_time: Arc::new(RwLock::new(None)),
last_state_change: Arc::new(RwLock::new(Instant::now())),
state: AtomicU8::new(STATE_CLOSED),
consecutive_failures: AtomicU32::new(0),
consecutive_successes: AtomicU32::new(0),
total_failures: AtomicU64::new(0),
total_successes: AtomicU64::new(0),
last_failure_time_ms: AtomicU64::new(0),
last_state_change_ms: AtomicU64::new(now_ms()),
config,
metric_label,
}
@@ -115,7 +143,8 @@ impl CircuitBreaker {
&self.metric_label
}
/// Check if a request can be executed
/// Check if a request can be executed (lock-free hot path)
#[inline]
pub fn can_execute(&self) -> bool {
let state = self.state();
match state {
@@ -125,20 +154,47 @@ impl CircuitBreaker {
}
}
/// Get the current state
/// Get the current state (lock-free)
#[inline]
pub fn state(&self) -> CircuitState {
self.check_and_update_state_returning()
}
/// Check and update state, returning the current state to avoid double lock
/// Check and update state, returning the current state (lock-free)
#[inline]
fn check_and_update_state_returning(&self) -> CircuitState {
let current_state = *self.state.read().unwrap();
let current_state_int = self.state.load(Ordering::Acquire);
let current_state = CircuitState::from_int(current_state_int);
if current_state == CircuitState::Open {
let last_change = *self.last_state_change.read().unwrap();
if last_change.elapsed() >= self.config.timeout_duration {
self.transition_to(CircuitState::HalfOpen);
return CircuitState::HalfOpen;
let last_change_ms = self.last_state_change_ms.load(Ordering::Acquire);
let elapsed_ms = now_ms().saturating_sub(last_change_ms);
let timeout_ms = self.config.timeout_duration.as_millis() as u64;
if elapsed_ms >= timeout_ms {
// Try to transition to HalfOpen using CAS
if self
.state
.compare_exchange(
STATE_OPEN,
STATE_HALF_OPEN,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
self.last_state_change_ms.store(now_ms(), Ordering::Release);
self.consecutive_failures.store(0, Ordering::Release);
self.consecutive_successes.store(0, Ordering::Release);
info!("Circuit breaker state transition: open -> half_open");
Metrics::record_worker_cb_transition(&self.metric_label, "open", "half_open");
Metrics::set_worker_cb_state(&self.metric_label, STATE_HALF_OPEN);
self.publish_gauge_metrics();
return CircuitState::HalfOpen;
}
// Another thread already transitioned, re-read the state
return CircuitState::from_int(self.state.load(Ordering::Acquire));
}
}
current_state
@@ -163,7 +219,7 @@ impl CircuitBreaker {
self.consecutive_failures.store(0, Ordering::Release);
let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1;
let current_state = *self.state.read().unwrap();
let current_state = CircuitState::from_int(self.state.load(Ordering::Acquire));
match current_state {
CircuitState::HalfOpen => {
@@ -184,12 +240,10 @@ impl CircuitBreaker {
self.consecutive_successes.store(0, Ordering::Release);
let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1;
{
let mut last_failure = self.last_failure_time.write().unwrap();
*last_failure = Some(Instant::now());
}
// Update last failure time atomically
self.last_failure_time_ms.store(now_ms(), Ordering::Release);
let current_state = *self.state.read().unwrap();
let current_state = CircuitState::from_int(self.state.load(Ordering::Acquire));
match current_state {
CircuitState::Closed => {
@@ -204,16 +258,14 @@ impl CircuitBreaker {
}
}
/// Transition to a new state
/// Transition to a new state (uses CAS for lock-free operation)
fn transition_to(&self, new_state: CircuitState) {
let mut state = self.state.write().unwrap();
let old_state = *state;
let new_state_int = new_state.to_int();
let old_state_int = self.state.swap(new_state_int, Ordering::AcqRel);
let old_state = CircuitState::from_int(old_state_int);
if old_state != new_state {
*state = new_state;
let mut last_change = self.last_state_change.write().unwrap();
*last_change = Instant::now();
self.last_state_change_ms.store(now_ms(), Ordering::Release);
match new_state {
CircuitState::Closed => {
@@ -260,12 +312,20 @@ impl CircuitBreaker {
/// Get time since last failure
pub fn time_since_last_failure(&self) -> Option<Duration> {
self.last_failure_time.read().unwrap().map(|t| t.elapsed())
let last_ms = self.last_failure_time_ms.load(Ordering::Acquire);
if last_ms == 0 {
None
} else {
let elapsed_ms = now_ms().saturating_sub(last_ms);
Some(Duration::from_millis(elapsed_ms))
}
}
/// Get time since last state change
pub fn time_since_last_state_change(&self) -> Duration {
self.last_state_change.read().unwrap().elapsed()
let last_ms = self.last_state_change_ms.load(Ordering::Acquire);
let elapsed_ms = now_ms().saturating_sub(last_ms);
Duration::from_millis(elapsed_ms)
}
/// Check if the circuit is in a half-open state
@@ -322,13 +382,15 @@ impl CircuitBreaker {
impl Clone for CircuitBreaker {
fn clone(&self) -> Self {
Self {
state: Arc::clone(&self.state),
consecutive_failures: Arc::clone(&self.consecutive_failures),
consecutive_successes: Arc::clone(&self.consecutive_successes),
total_failures: Arc::clone(&self.total_failures),
total_successes: Arc::clone(&self.total_successes),
last_failure_time: Arc::clone(&self.last_failure_time),
last_state_change: Arc::clone(&self.last_state_change),
state: AtomicU8::new(self.state.load(Ordering::Acquire)),
consecutive_failures: AtomicU32::new(self.consecutive_failures.load(Ordering::Acquire)),
consecutive_successes: AtomicU32::new(
self.consecutive_successes.load(Ordering::Acquire),
),
total_failures: AtomicU64::new(self.total_failures.load(Ordering::Relaxed)),
total_successes: AtomicU64::new(self.total_successes.load(Ordering::Relaxed)),
last_failure_time_ms: AtomicU64::new(self.last_failure_time_ms.load(Ordering::Acquire)),
last_state_change_ms: AtomicU64::new(self.last_state_change_ms.load(Ordering::Acquire)),
config: self.config.clone(),
metric_label: self.metric_label.clone(),
}
@@ -528,7 +590,8 @@ mod tests {
assert_eq!(cb2.consecutive_failures(), 1);
cb1.record_failure();
assert_eq!(cb2.consecutive_failures(), 2);
assert_eq!(cb1.consecutive_failures(), 2);
assert_eq!(cb2.consecutive_failures(), 1); // cb2 is unchanged
}
#[test]

9
sgl-model-gateway/src/middleware.rs
View File

@@ -26,7 +26,7 @@ use tracing::{debug, error, field::Empty, info, info_span, warn, Span};
pub use crate::core::token_bucket::TokenBucket;
use crate::{
observability::metrics::{metrics_labels, Metrics},
observability::metrics::{method_to_static_str, metrics_labels, Metrics},
routers::error::extract_error_code_from_response,
server::AppState,
wasm::{
@@ -641,7 +641,8 @@ where
}
fn call(&mut self, req: Request) -> Self::Future {
let method = req.method().as_str().to_owned();
// Convert method to static string to avoid allocation
let method = method_to_static_str(req.method().as_str());
let path = normalize_path_for_metrics(req.uri().path());
let start = Instant::now();
@@ -664,8 +665,8 @@ where
let duration = start.elapsed();
let status_class = status_to_class(response.status().as_u16());
Metrics::record_http_request(&method, &path, status_class);
Metrics::record_http_duration(&method, &path, duration);
Metrics::record_http_request(method, &path, status_class);
Metrics::record_http_duration(method, &path, duration);
Ok(response)
})

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

@@ -1,4 +1,5 @@
use std::{
borrow::Cow,
net::{IpAddr, Ipv4Addr, SocketAddr},
time::Duration,
};
@@ -6,6 +7,80 @@ use std::{
use metrics::{counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram};
use metrics_exporter_prometheus::{Matcher, PrometheusBuilder};
/// 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
} else {
STREAMING_FALSE
}
}
/// Static lookup table for common HTTP status codes to avoid allocations.
/// 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> {
match code {
200 => Some("200"),
201 => Some("201"),
204 => Some("204"),
400 => Some("400"),
401 => Some("401"),
403 => Some("403"),
404 => Some("404"),
408 => Some("408"),
422 => Some("422"),
429 => Some("429"),
500 => Some("500"),
502 => Some("502"),
503 => Some("503"),
504 => Some("504"),
_ => None,
}
}
/// Static HTTP method strings to avoid allocations on every request.
pub mod http_methods {
pub const GET: &str = "GET";
pub const POST: &str = "POST";
pub const PUT: &str = "PUT";
pub const DELETE: &str = "DELETE";
pub const PATCH: &str = "PATCH";
pub const HEAD: &str = "HEAD";
pub const OPTIONS: &str = "OPTIONS";
}
/// Convert HTTP method to static string. Returns the method as-is for unknown methods.
#[inline]
pub fn method_to_static_str(method: &str) -> &'static str {
match method {
"GET" => http_methods::GET,
"POST" => http_methods::POST,
"PUT" => http_methods::PUT,
"DELETE" => http_methods::DELETE,
"PATCH" => http_methods::PATCH,
"HEAD" => http_methods::HEAD,
"OPTIONS" => http_methods::OPTIONS,
// For unknown methods, we return a static "OTHER" to avoid allocation
// This is acceptable since unknown methods are rare in practice
_ => "OTHER",
}
}
/// Get status code as Cow - static for common codes, allocated for rare ones.
#[inline]
pub fn status_code_to_cow(code: u16) -> Cow<'static, str> {
match status_code_to_static_str(code) {
Some(s) => Cow::Borrowed(s),
None => Cow::Owned(code.to_string()),
}
}
#[derive(Debug, Clone)]
pub struct PrometheusConfig {
pub port: u16,
@@ -333,22 +408,25 @@ pub struct StreamingMetricsParams<'a> {
}
impl Metrics {
/// Record an HTTP request
pub fn record_http_request(method: &str, path: &str, status_class: &str) {
/// Record an HTTP request.
/// For best performance, pass static strings (use `method_to_static_str` for method,
/// `status_to_class` for status_class returns static, and cache normalized paths).
pub fn record_http_request(method: &'static str, path: &str, status_class: &'static str) {
counter!(
"smg_http_requests_total",
"method" => method.to_string(),
"method" => method,
"path" => path.to_string(),
"status" => status_class.to_string()
"status" => status_class
)
.increment(1);
}
/// Record HTTP request duration
pub fn record_http_duration(method: &str, path: &str, duration: Duration) {
/// Record HTTP request duration.
/// For best performance, pass static strings for method.
pub fn record_http_duration(method: &'static str, path: &str, duration: Duration) {
histogram!(
"smg_http_request_duration_seconds",
"method" => method.to_string(),
"method" => method,
"path" => path.to_string()
)
.record(duration.as_secs_f64());
@@ -359,11 +437,18 @@ impl Metrics {
gauge!("smg_http_connections_active").set(count as f64);
}
/// Record HTTP response
/// 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
counter!(
"smg_http_responses_total",
"status_code" => status_code.to_string(),
"status_code" => status_str,
"error_code" => error_code.to_string()
)
.increment(1);
@@ -382,14 +467,18 @@ impl Metrics {
// Layer 2: Router metrics
// ========================================================================
/// Record a routed request
/// Record a routed request.
///
/// # 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,
connection_mode: &'static str,
model_id: &str,
endpoint: &'static str,
streaming: bool,
streaming: &'static str,
) {
counter!(
"smg_router_requests_total",
@@ -398,7 +487,7 @@ impl Metrics {
"connection_mode" => connection_mode,
"model" => model_id.to_string(),
"endpoint" => endpoint,
"streaming" => streaming.to_string()
"streaming" => streaming
)
.increment(1);
}
@@ -458,16 +547,18 @@ impl Metrics {
.record(duration.as_secs_f64());
}
/// Record upstream backend response
/// Record upstream backend response.
/// Uses static strings for common status codes to avoid allocations.
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);
counter!(
"smg_router_upstream_responses_total",
"router_type" => router_type,
"status_code" => status_code.to_string(),
"status_code" => status_str,
"error_code" => error_code.to_string()
)
.increment(1);
@@ -566,8 +657,8 @@ impl Metrics {
input_tokens,
output_tokens,
} = params;
// metrics-rs requires owned strings for dynamic labels (uses Cow<'static, str>).
// We allocate once and clone for each metric - unavoidable with this API.
// Allocate model string once, clone as needed for each metric
let model = model_id.to_string();
// TTFT and TPOT (only if we have a first token time)
@@ -596,7 +687,7 @@ impl Metrics {
}
}
// Generation duration
// Generation duration (always recorded)
histogram!(
"smg_router_generation_duration_seconds",
"router_type" => router_type,
@@ -619,7 +710,7 @@ impl Metrics {
.increment(input);
}
// Output tokens
// Output tokens (always recorded - move model on final use)
counter!(
"smg_router_tokens_total",
"router_type" => router_type,
@@ -792,11 +883,20 @@ impl Metrics {
.increment(1);
}
/// Record retry backoff duration
/// 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"),
2 => Cow::Borrowed("2"),
3 => Cow::Borrowed("3"),
4 => Cow::Borrowed("4"),
5 => Cow::Borrowed("5"),
_ => Cow::Owned(attempt.to_string()),
};
histogram!(
"smg_worker_retry_backoff_seconds",
"attempt" => attempt.to_string()
"attempt" => attempt_str
)
.record(duration.as_secs_f64());
}

6
sgl-model-gateway/src/routers/grpc/pipeline.rs
View File

@@ -17,7 +17,7 @@ use super::{
};
use crate::{
core::WorkerRegistry,
observability::metrics::{metrics_labels, Metrics},
observability::metrics::{bool_to_static_str, metrics_labels, Metrics},
policies::PolicyRegistry,
protocols::{
chat::{ChatCompletionRequest, ChatCompletionResponse},
@@ -215,7 +215,7 @@ impl RequestPipeline {
metrics_labels::CONNECTION_GRPC,
&request_for_metrics.model,
metrics_labels::ENDPOINT_CHAT,
streaming,
bool_to_static_str(streaming),
);
let mut ctx = RequestContext::for_chat(request, headers, model_id, components);
@@ -320,7 +320,7 @@ impl RequestPipeline {
metrics_labels::CONNECTION_GRPC,
model_for_metrics.as_deref().unwrap_or("unknown"),
metrics_labels::ENDPOINT_GENERATE,
streaming,
bool_to_static_str(streaming),
);
let mut ctx = RequestContext::for_generate(request, headers, model_id, components);

4
sgl-model-gateway/src/routers/http/pd_router.rs
View File

@@ -22,7 +22,7 @@ use crate::{
},
observability::{
events::{self, Event},
metrics::{metrics_labels, Metrics},
metrics::{bool_to_static_str, metrics_labels, Metrics},
otel_trace::inject_trace_context_http,
},
policies::{LoadBalancingPolicy, PolicyRegistry},
@@ -288,7 +288,7 @@ impl PDRouter {
metrics_labels::CONNECTION_HTTP,
model,
endpoint,
context.is_stream,
bool_to_static_str(context.is_stream),
);
// Clone request once outside the retry loop, then use Arc to share across attempts
// This avoids O(retries) clones by sharing the same data

4
sgl-model-gateway/src/routers/http/router.rs
View File

@@ -24,7 +24,7 @@ use crate::{
},
observability::{
events::{self, Event},
metrics::{metrics_labels, Metrics},
metrics::{bool_to_static_str, metrics_labels, Metrics},
otel_trace::inject_trace_context_http,
},
policies::PolicyRegistry,
@@ -189,7 +189,7 @@ impl Router {
metrics_labels::CONNECTION_HTTP,
model,
endpoint,
is_stream,
bool_to_static_str(is_stream),
);
let response = RetryExecutor::execute_response_with_retry(

6
sgl-model-gateway/src/routers/openai/router.rs
View File

@@ -36,7 +36,7 @@ use crate::{
app_context::AppContext,
core::{model_type::Endpoint, ModelCard, ProviderType, RuntimeType, Worker, WorkerRegistry},
data_connector::{ConversationId, ListParams, ResponseId, SortOrder},
observability::metrics::{metrics_labels, Metrics},
observability::metrics::{bool_to_static_str, metrics_labels, Metrics},
protocols::{
chat::ChatCompletionRequest,
responses::{
@@ -589,7 +589,7 @@ impl crate::routers::RouterTrait for OpenAIRouter {
metrics_labels::CONNECTION_HTTP,
model,
metrics_labels::ENDPOINT_CHAT,
streaming,
bool_to_static_str(streaming),
);
let auth_header = extract_auth_header(headers, &None);
@@ -782,7 +782,7 @@ impl crate::routers::RouterTrait for OpenAIRouter {
metrics_labels::CONNECTION_HTTP,
model,
metrics_labels::ENDPOINT_RESPONSES,
streaming,
bool_to_static_str(streaming),
);
let auth_header = extract_auth_header(headers, &None);