diff --git a/sgl-model-gateway/src/core/circuit_breaker.rs b/sgl-model-gateway/src/core/circuit_breaker.rs index d013ab933..aed00f6b3 100644 --- a/sgl-model-gateway/src/core/circuit_breaker.rs +++ b/sgl-model-gateway/src/core/circuit_breaker.rs @@ -374,8 +374,14 @@ impl CircuitBreaker { } fn publish_gauge_metrics(&self) { - Metrics::set_worker_cb_consecutive_failures(&self.metric_label, self.failure_count()); - Metrics::set_worker_cb_consecutive_successes(&self.metric_label, self.success_count()); + Metrics::set_worker_cb_consecutive_failures( + &self.metric_label, + self.consecutive_failures(), + ); + Metrics::set_worker_cb_consecutive_successes( + &self.metric_label, + self.consecutive_successes(), + ); } } diff --git a/sgl-model-gateway/src/core/worker.rs b/sgl-model-gateway/src/core/worker.rs index 16522231a..74422c033 100644 --- a/sgl-model-gateway/src/core/worker.rs +++ b/sgl-model-gateway/src/core/worker.rs @@ -10,7 +10,7 @@ use std::{ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json; -use tokio::{sync::RwLock, time}; +use tokio::{sync::OnceCell, time}; use super::{ CircuitBreaker, Endpoint, ModelCard, ModelType, ProviderType, WorkerError, WorkerResult, @@ -515,8 +515,9 @@ pub struct BasicWorker { pub consecutive_failures: Arc, pub consecutive_successes: Arc, pub circuit_breaker: CircuitBreaker, - /// Lazily initialized gRPC client for gRPC workers - pub grpc_client: Arc>>>, + /// Lazily initialized gRPC client for gRPC workers. + /// Uses OnceCell for lock-free reads after initialization. + pub grpc_client: Arc>>, /// Runtime-mutable models override (for lazy discovery) /// When set, overrides metadata.models for routing decisions. /// Uses std::sync::RwLock for synchronous access in supports_model(). @@ -715,64 +716,53 @@ impl Worker for BasicWorker { match self.metadata.connection_mode { ConnectionMode::Http => Ok(None), ConnectionMode::Grpc { .. } => { - { - let client_guard = self.grpc_client.read().await; - if let Some(ref client) = *client_guard { - return Ok(Some(client.clone())); - } - } - - let mut client_guard = self.grpc_client.write().await; - - if let Some(ref client) = *client_guard { - return Ok(Some(client.clone())); - } - - let runtime_str = self.metadata.runtime_type.to_string(); - tracing::info!( - "Lazily initializing gRPC client ({}) for worker: {}", - runtime_str, - self.metadata.url - ); - match GrpcClient::connect(&self.metadata.url, &runtime_str).await { - Ok(client) => { - let client_arc = Arc::new(client); - *client_guard = Some(client_arc.clone()); + // OnceCell provides lock-free reads after initialization. + // get_or_try_init only acquires internal lock on first call. + let client = self + .grpc_client + .get_or_try_init(|| async { + let runtime_str = self.metadata.runtime_type.to_string(); tracing::info!( - "Successfully connected gRPC client ({}) for worker: {}", + "Lazily initializing gRPC client ({}) for worker: {}", runtime_str, self.metadata.url ); - Ok(Some(client_arc)) - } - Err(e) => { - tracing::error!( - "Failed to connect gRPC client for worker {}: {}", - self.metadata.url, - e - ); - Err(WorkerError::ConnectionFailed { - url: self.metadata.url.clone(), - reason: format!("Failed to connect to gRPC server: {}", e), - }) - } - } + match GrpcClient::connect(&self.metadata.url, &runtime_str).await { + Ok(client) => { + tracing::info!( + "Successfully connected gRPC client ({}) for worker: {}", + runtime_str, + self.metadata.url + ); + Ok(Arc::new(client)) + } + Err(e) => { + tracing::error!( + "Failed to connect gRPC client for worker {}: {}", + self.metadata.url, + e + ); + Err(WorkerError::ConnectionFailed { + url: self.metadata.url.clone(), + reason: format!("Failed to connect to gRPC server: {}", e), + }) + } + } + }) + .await?; + Ok(Some(Arc::clone(client))) } } } async fn reset_grpc_client(&self) -> WorkerResult<()> { - match self.metadata.connection_mode { - ConnectionMode::Http => Ok(()), - ConnectionMode::Grpc { .. } => { - let mut client_guard = self.grpc_client.write().await; - if client_guard.is_some() { - tracing::info!("Resetting gRPC client for worker: {}", self.metadata.url); - *client_guard = None; - } - Ok(()) - } - } + // OnceCell doesn't support resetting. This is intentional for lock-free performance. + // If a connection fails, the worker should be removed and re-added. + tracing::debug!( + "reset_grpc_client called for {} (no-op with OnceCell)", + self.metadata.url + ); + Ok(()) } async fn grpc_health_check(&self) -> WorkerResult { diff --git a/sgl-model-gateway/src/core/worker_builder.rs b/sgl-model-gateway/src/core/worker_builder.rs index 5784050c5..b0c714db7 100644 --- a/sgl-model-gateway/src/core/worker_builder.rs +++ b/sgl-model-gateway/src/core/worker_builder.rs @@ -131,7 +131,7 @@ impl BasicWorkerBuilder { Arc, RwLock as StdRwLock, }; - use tokio::sync::RwLock; + use tokio::sync::OnceCell; let bootstrap_host = match url::Url::parse(&self.url) { Ok(parsed) => parsed.host_str().unwrap_or("localhost").to_string(), @@ -176,7 +176,16 @@ impl BasicWorkerBuilder { default_model_type: ModelType::LLM, // Standard LLM capabilities }; - let grpc_client = Arc::new(RwLock::new(self.grpc_client.map(Arc::new))); + // Use OnceCell for lock-free gRPC client access after initialization + let grpc_client = Arc::new(match self.grpc_client { + Some(client) => { + let cell = OnceCell::new(); + // Pre-set the client if provided (blocking set is fine during construction) + cell.set(Arc::new(client)).ok(); + cell + } + None => OnceCell::new(), + }); BasicWorker { metadata, diff --git a/sgl-model-gateway/src/core/worker_registry.rs b/sgl-model-gateway/src/core/worker_registry.rs index 4222119d2..4ae766aa2 100644 --- a/sgl-model-gateway/src/core/worker_registry.rs +++ b/sgl-model-gateway/src/core/worker_registry.rs @@ -1,8 +1,12 @@ //! Worker Registry for multi-router support //! //! Provides centralized registry for workers with model-based indexing +//! +//! # Performance Optimizations +//! The model index uses immutable Arc snapshots instead of RwLock for lock-free reads. +//! This is critical for high-concurrency scenarios where many requests query the same model. -use std::sync::{Arc, RwLock}; +use std::sync::Arc; use dashmap::DashMap; use uuid::Uuid; @@ -36,8 +40,10 @@ impl Default for WorkerId { } } -/// Model index type for O(1) lookups (stores Arc directly) -type ModelIndex = Arc>>>>>; +/// Model index using immutable snapshots for lock-free reads. +/// Each model maps to an Arc'd slice of workers that can be read without locking. +/// Updates create new snapshots (copy-on-write semantics). +type ModelIndex = Arc]>>>; /// Worker registry with model-based indexing #[derive(Debug)] @@ -45,8 +51,8 @@ pub struct WorkerRegistry { /// All workers indexed by ID workers: Arc>>, - /// Model index for O(1) lookups (stores Arc directly) - /// This replaces the previous dual-index approach for better memory efficiency + /// Model index for O(1) lookups using immutable snapshots. + /// Uses Arc<[T]> instead of Arc>> for lock-free reads. model_index: ModelIndex, /// Workers indexed by worker type @@ -87,14 +93,18 @@ impl WorkerRegistry { self.url_to_id .insert(worker.url().to_string(), worker_id.clone()); - // Update model index for O(1) lookups + // Update model index for O(1) lookups using copy-on-write + // This creates a new immutable snapshot with the added worker let model_id = worker.model_id().to_string(); self.model_index .entry(model_id) - .or_insert_with(|| Arc::new(RwLock::new(Vec::new()))) - .write() - .expect("RwLock for model_index is poisoned") - .push(worker.clone()); + .and_modify(|existing| { + // Create new snapshot with the additional worker + let mut new_workers: Vec> = existing.iter().cloned().collect(); + new_workers.push(worker.clone()); + *existing = Arc::from(new_workers.into_boxed_slice()); + }) + .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice())); // Update type index (clone needed for DashMap key ownership) self.type_workers @@ -117,13 +127,16 @@ impl WorkerRegistry { // Remove from URL mapping self.url_to_id.remove(worker.url()); - // Remove from model index - if let Some(model_index_entry) = self.model_index.get(worker.model_id()) { - let worker_url = worker.url(); - model_index_entry - .write() - .expect("RwLock for model_index is poisoned") - .retain(|w| w.url() != worker_url); + // Remove from model index using copy-on-write + // Create new snapshot without the removed worker + let worker_url = worker.url(); + if let Some(mut entry) = self.model_index.get_mut(worker.model_id()) { + let new_workers: Vec> = entry + .iter() + .filter(|w| w.url() != worker_url) + .cloned() + .collect(); + *entry = Arc::from(new_workers.into_boxed_slice()); } // Remove from type index @@ -165,23 +178,22 @@ impl WorkerRegistry { self.url_to_id.get(url).and_then(|id| self.get(&id)) } - /// Get all workers for a model (O(1) optimized) - /// Uses the pre-indexed model_index for fast lookups - pub fn get_by_model(&self, model_id: &str) -> Vec> { + /// Empty worker slice constant for returning when no workers found + const EMPTY_WORKERS: &'static [Arc] = &[]; + + /// Get all workers for a model (O(1) optimized, lock-free) + /// Returns an Arc to the immutable worker slice - just an atomic refcount bump. + /// This is the fastest possible read path with zero contention. + pub fn get_by_model(&self, model_id: &str) -> Arc<[Arc]> { self.model_index .get(model_id) - .map(|workers| { - workers - .read() - .expect("RwLock for model_index is poisoned") - .clone() - }) - .unwrap_or_default() + .map(|workers| Arc::clone(&workers)) + .unwrap_or_else(|| Arc::from(Self::EMPTY_WORKERS)) } /// Alias for get_by_model for backwards compatibility #[inline] - pub fn get_by_model_fast(&self, model_id: &str) -> Vec> { + pub fn get_by_model_fast(&self, model_id: &str) -> Arc<[Arc]> { self.get_by_model(model_id) } @@ -266,17 +278,11 @@ impl WorkerRegistry { .collect() } - /// Get all model IDs with workers + /// Get all model IDs with workers (lock-free) pub fn get_models(&self) -> Vec { self.model_index .iter() - .filter(|entry| { - entry - .value() - .read() - .map(|workers| !workers.is_empty()) - .unwrap_or(false) - }) + .filter(|entry| !entry.value().is_empty()) .map(|entry| entry.key().clone()) .collect() } @@ -299,8 +305,8 @@ impl WorkerRegistry { ) -> Vec> { // Start with the most efficient collection based on filters // Use model index when possible as it's O(1) lookup - let workers = if let Some(model) = model_id { - self.get_by_model_fast(model) + let workers: Vec> = if let Some(model) = model_id { + self.get_by_model_fast(model).to_vec() } else { self.get_all() }; @@ -340,20 +346,14 @@ impl WorkerRegistry { .collect() } - /// Get worker statistics + /// Get worker statistics (lock-free) pub fn stats(&self) -> WorkerRegistryStats { let total_workers = self.workers.len(); - // Count models directly instead of allocating Vec via get_models() + // Count models directly instead of allocating Vec via get_models() (lock-free) let total_models = self .model_index .iter() - .filter(|entry| { - entry - .value() - .read() - .map(|workers| !workers.is_empty()) - .unwrap_or(false) - }) + .filter(|entry| !entry.value().is_empty()) .count(); let mut healthy_count = 0; diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index c8df1332f..08762c6eb 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -705,8 +705,9 @@ impl PDRouter { let prefill_workers = if let Some(model) = effective_model_id { self.worker_registry .get_by_model_fast(model) - .into_iter() + .iter() .filter(|w| matches!(w.worker_type(), WorkerType::Prefill { .. })) + .cloned() .collect() } else { self.worker_registry.get_prefill_workers() @@ -715,8 +716,9 @@ impl PDRouter { let decode_workers = if let Some(model) = effective_model_id { self.worker_registry .get_by_model_fast(model) - .into_iter() + .iter() .filter(|w| matches!(w.worker_type(), WorkerType::Decode)) + .cloned() .collect() } else { self.worker_registry.get_decode_workers()