diff --git a/sgl-model-gateway/src/core/circuit_breaker.rs b/sgl-model-gateway/src/core/circuit_breaker.rs index 3138c057c..7b3a44b5a 100644 --- a/sgl-model-gateway/src/core/circuit_breaker.rs +++ b/sgl-model-gateway/src/core/circuit_breaker.rs @@ -95,8 +95,10 @@ 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; + RouterMetrics::set_cb_state(&metric_label, init_state.to_int()); Self { - state: Arc::new(RwLock::new(CircuitState::Closed)), + 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)), diff --git a/sgl-model-gateway/src/core/mod.rs b/sgl-model-gateway/src/core/mod.rs index 8e0f7db28..388c75404 100644 --- a/sgl-model-gateway/src/core/mod.rs +++ b/sgl-model-gateway/src/core/mod.rs @@ -33,7 +33,7 @@ pub use model_type::{Endpoint, ModelType}; pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor}; pub use worker::{ worker_to_info, BasicWorker, ConnectionMode, DPAwareWorker, HealthChecker, HealthConfig, - RuntimeType, Worker, WorkerFactory, WorkerLoadGuard, WorkerType, + RuntimeType, Worker, WorkerFactory, WorkerLoadGuard, WorkerLoadGuardV2, WorkerType, }; pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder}; pub use worker_manager::{LoadMonitor, WorkerManager}; diff --git a/sgl-model-gateway/src/core/worker.rs b/sgl-model-gateway/src/core/worker.rs index 2e1eb0e1c..14635dc49 100644 --- a/sgl-model-gateway/src/core/worker.rs +++ b/sgl-model-gateway/src/core/worker.rs @@ -130,26 +130,6 @@ pub trait Worker: Send + Sync + fmt::Debug { /// Record the outcome of a request to this worker fn record_outcome(&self, success: bool) { self.circuit_breaker().record_outcome(success); - let after = self.circuit_breaker().state(); - - if before != after { - let from = before.as_str(); - let to = after.as_str(); - RouterMetrics::record_cb_state_transition(self.url(), from, to); - } - - let state_code = self.circuit_breaker().state().to_int(); - RouterMetrics::set_cb_state(self.url(), state_code); - - // Update consecutive failures/successes gauges - RouterMetrics::set_cb_consecutive_failures( - self.url(), - self.circuit_breaker().failure_count(), - ); - RouterMetrics::set_cb_consecutive_successes( - self.url(), - self.circuit_breaker().success_count(), - ); } /// Check if this worker is DP-aware @@ -1054,6 +1034,24 @@ pub fn workers_to_urls(workers: &[Box]) -> Vec { workers.iter().map(|w| w.url().to_string()).collect() } +// TODO migrate code to V2 (and then remove this name suffix) +pub struct WorkerLoadGuardV2 { + worker: Arc, +} + +impl WorkerLoadGuardV2 { + pub fn new(worker: Arc) -> Self { + worker.increment_load(); + Self { worker } + } +} + +impl Drop for WorkerLoadGuardV2 { + fn drop(&mut self) { + self.worker.decrement_load(); + } +} + /// RAII guard for worker load management pub struct WorkerLoadGuard<'a> { workers: Vec<&'a dyn Worker>, diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 09b7026de..ad2af1e85 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -19,7 +19,8 @@ use tracing::{debug, error}; use crate::{ config::types::RetryConfig, core::{ - is_retryable_status, ConnectionMode, RetryExecutor, Worker, WorkerRegistry, WorkerType, + is_retryable_status, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuardV2, + WorkerRegistry, WorkerType, }, observability::{ events::{self, Event}, @@ -265,19 +266,8 @@ impl Router { None => self.policy_registry.get_default_policy(), }; - let load_incremented = if policy.name() == "cache_aware" { - worker.increment_load(); - true - } else { - false - }; - - // Keep a clone for potential cleanup on retry - let worker_for_cleanup = if load_incremented { - Some(worker.clone()) - } else { - None - }; + let load_guard = + (policy.name() == "cache_aware").then(|| WorkerLoadGuardV2::new(worker.clone())); events::RequestSentEvent { url: worker.url().to_string(), @@ -294,7 +284,7 @@ impl Router { route, worker.url(), is_stream, - load_incremented, + load_guard, ) .await; @@ -302,14 +292,6 @@ impl Router { worker.record_outcome(response.status().is_success()); - // For retryable failures, we need to decrement load since send_typed_request - // won't have done it (it only decrements on success or non-retryable failures) - if is_retryable_status(response.status()) && load_incremented { - if let Some(cleanup_worker) = worker_for_cleanup { - cleanup_worker.decrement_load(); - } - } - response } @@ -453,7 +435,7 @@ impl Router { route: &'static str, worker_url: &str, is_stream: bool, - load_incremented: bool, // Whether load was incremented for this request + mut load_guard: Option, ) -> Response { // Get the worker once and reuse for API key and load tracking let worker = self.worker_registry.get_by_url(worker_url); @@ -536,13 +518,6 @@ impl Router { worker_url, route, e ); - // Decrement load on error if it was incremented - if load_incremented { - if let Some(ref w) = worker { - w.decrement_load(); - } - } - return convert_reqwest_error(e); } }; @@ -567,19 +542,9 @@ impl Router { } }; - // Decrement load counter for non-streaming requests if it was incremented - if load_incremented { - if let Some(ref w) = worker { - w.decrement_load(); - } - } - + drop(load_guard); response - } else if load_incremented { - // For streaming with load tracking, we need to manually decrement when done - // Clone the worker Arc for the async block instead of looking it up again - let stream_worker = worker.clone(); - + } else { // Preserve headers for streaming response let mut response_headers = header_utils::preserve_response_headers(res.headers()); // Ensure we set the correct content-type for SSE @@ -591,16 +556,14 @@ impl Router { // Spawn task to forward stream and detect completion tokio::spawn(async move { let mut stream = stream; - let mut decremented = false; while let Some(chunk) = stream.next().await { match chunk { Ok(bytes) => { // Check for stream end marker using memmem for efficiency - if memmem::find(&bytes, b"data: [DONE]").is_some() { - if let Some(ref w) = stream_worker { - w.decrement_load(); - decremented = true; - } + if load_guard.is_some() + && memmem::find(&bytes, b"data: [DONE]").is_some() + { + load_guard = None; } if tx.send(Ok(bytes)).is_err() { break; @@ -612,46 +575,7 @@ impl Router { } } } - if !decremented { - if let Some(ref w) = stream_worker { - w.decrement_load(); - } - } - }); - - let stream = UnboundedReceiverStream::new(rx); - let body = Body::from_stream(stream); - - let mut response = Response::new(body); - *response.status_mut() = status; - *response.headers_mut() = response_headers; - response - } else { - // For requests without load tracking, just stream - // Preserve headers for streaming response - let mut response_headers = header_utils::preserve_response_headers(res.headers()); - // Ensure we set the correct content-type for SSE - response_headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); - - let stream = res.bytes_stream(); - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - - // Spawn task to forward stream - tokio::spawn(async move { - let mut stream = stream; - while let Some(chunk) = stream.next().await { - match chunk { - Ok(bytes) => { - if tx.send(Ok(bytes)).is_err() { - break; - } - } - Err(e) => { - let _ = tx.send(Err(format!("Stream error: {}", e))); - break; - } - } - } + drop(load_guard); }); let stream = UnboundedReceiverStream::new(rx);