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

[model-gateway] optimize worker registry and reduce lock contention in grpc client fetch (#15336)

Browse Source
This commit is contained in:
Simo Lin
2025-12-17 08:49:10 -10:00
committed by GitHub
parent d747147a26
commit 53e151945a
5 changed files with 112 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/core/circuit_breaker.rs
+8 -2
View File
@@ -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(),
);
}
}
sgl-model-gateway/src/core/worker.rs
+42 -52
View File
@@ -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<AtomicUsize>,
pub consecutive_successes: Arc<AtomicUsize>,
pub circuit_breaker: CircuitBreaker,
/// Lazily initialized gRPC client for gRPC workers
pub grpc_client: Arc<RwLock<Option<Arc<GrpcClient>>>>,
/// Lazily initialized gRPC client for gRPC workers.
/// Uses OnceCell for lock-free reads after initialization.
pub grpc_client: Arc<OnceCell<Arc<GrpcClient>>>,
/// 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<bool> {
sgl-model-gateway/src/core/worker_builder.rs
+11 -2
View File
@@ -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,
sgl-model-gateway/src/core/worker_registry.rs
+47 -47
View File
@@ -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<dyn Worker> directly)
type ModelIndex = Arc<DashMap<String, Arc<RwLock<Vec<Arc<dyn Worker>>>>>>;
/// 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<DashMap<String, Arc<[Arc<dyn Worker>]>>>;
/// Worker registry with model-based indexing
#[derive(Debug)]
@@ -45,8 +51,8 @@ pub struct WorkerRegistry {
/// All workers indexed by ID
workers: Arc<DashMap<WorkerId, Arc<dyn Worker>>>,
/// Model index for O(1) lookups (stores Arc<dyn Worker> 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<RwLock<Vec<T>>> 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<Arc<dyn Worker>> = 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<Arc<dyn Worker>> = 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<Arc<dyn Worker>> {
/// Empty worker slice constant for returning when no workers found
const EMPTY_WORKERS: &'static [Arc<dyn Worker>] = &[];
/// 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<dyn Worker>]> {
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<Arc<dyn Worker>> {
pub fn get_by_model_fast(&self, model_id: &str) -> Arc<[Arc<dyn Worker>]> {
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<String> {
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<Arc<dyn Worker>> {
// 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<Arc<dyn Worker>> = 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;
sgl-model-gateway/src/routers/http/pd_router.rs
+4 -2
View File
@@ -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()