[router] add move grpc worker management from router to worker manager (#10960)

This commit is contained in:
Simo Lin
2025-09-26 06:57:57 -04:00
committed by GitHub
parent 1b011e68dc
commit 5d4fe1ceee
4 changed files with 232 additions and 288 deletions

View File

@@ -1,7 +1,6 @@
// gRPC Router Implementation
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use axum::{
@@ -13,12 +12,10 @@ use axum::{
use tracing::{debug, error, info, warn};
use crate::config::types::RetryConfig;
use crate::core::{
BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, WorkerRegistry, WorkerType,
};
use crate::core::{WorkerRegistry, WorkerType};
use crate::grpc_client::{proto, SglangSchedulerClient};
use crate::metrics::RouterMetrics;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
use crate::policies::PolicyRegistry;
use crate::protocols::spec::{ChatCompletionRequest, StringOrArray};
use crate::reasoning_parser::ParserFactory;
use crate::routers::RouterTrait;
@@ -38,39 +35,21 @@ pub struct ProcessedMessages {
}
/// gRPC router implementation for SGLang
#[allow(dead_code)] // Fields will be used once implementation is complete
#[allow(dead_code)]
pub struct GrpcRouter {
/// Centralized worker registry
worker_registry: Arc<WorkerRegistry>,
/// Centralized policy registry
policy_registry: Arc<PolicyRegistry>,
/// Load balancing policy
policy: Arc<dyn LoadBalancingPolicy>,
/// Tokenizer for handling text encoding/decoding
tokenizer: Arc<dyn Tokenizer>,
/// Reasoning parser factory for structured reasoning outputs
reasoning_parser_factory: ParserFactory,
/// Tool parser registry for function/tool calls
tool_parser_registry: &'static ParserRegistry,
/// Configuration
timeout_secs: u64,
interval_secs: u64,
dp_aware: bool,
api_key: Option<String>,
retry_config: RetryConfig,
circuit_breaker_config: CircuitBreakerConfig,
}
impl GrpcRouter {
/// Create a new gRPC router
pub async fn new(
worker_urls: Vec<String>,
policy: Arc<dyn LoadBalancingPolicy>,
ctx: &Arc<crate::server::AppContext>,
) -> Result<Self, String> {
// Update metrics
RouterMetrics::set_active_workers(worker_urls.len());
pub async fn new(ctx: &Arc<crate::server::AppContext>) -> Result<Self, String> {
// Extract necessary components from context
let tokenizer = ctx
.tokenizer
@@ -86,77 +65,28 @@ impl GrpcRouter {
.tool_parser_registry
.ok_or_else(|| "gRPC router requires tool parser registry".to_string())?;
// Convert config CircuitBreakerConfig to core CircuitBreakerConfig
let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config();
let core_cb_config = CircuitBreakerConfig {
failure_threshold: circuit_breaker_config.failure_threshold,
success_threshold: circuit_breaker_config.success_threshold,
timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs),
window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs),
};
// Get registries from context
let worker_registry = ctx.worker_registry.clone();
let policy_registry = ctx.policy_registry.clone();
// Create Worker trait objects with gRPC connection mode and register them
// Workers will lazily initialize their gRPC clients on first use
for url in &worker_urls {
let worker = BasicWorkerBuilder::new(url.clone())
.worker_type(WorkerType::Regular)
.connection_mode(crate::core::ConnectionMode::Grpc { port: None })
.circuit_breaker_config(core_cb_config.clone())
.health_config(HealthConfig {
timeout_secs: ctx.router_config.health_check.timeout_secs,
check_interval_secs: ctx.router_config.health_check.check_interval_secs,
endpoint: ctx.router_config.health_check.endpoint.clone(),
failure_threshold: ctx.router_config.health_check.failure_threshold,
success_threshold: ctx.router_config.health_check.success_threshold,
})
.build();
worker_registry.register(Arc::new(worker));
info!(
"Registered gRPC worker at {} (will connect on first use)",
url
);
}
if worker_urls.is_empty() {
return Err("No gRPC workers configured".to_string());
}
// Get only gRPC workers from registry for policy initialization
let workers = worker_registry.get_workers_filtered(
None, // any model
None,
Some(WorkerType::Regular),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false, // include unhealthy workers during initialization
false,
);
// Initialize policy with workers if needed
if let Some(cache_aware) = policy
.as_any()
.downcast_ref::<crate::policies::CacheAwarePolicy>()
{
cache_aware.init_workers(&workers);
}
// No need for local health checkers - WorkerRegistry handles health checking
RouterMetrics::set_active_workers(workers.len());
info!("gRPC router found {} workers in registry", workers.len());
Ok(GrpcRouter {
worker_registry,
policy_registry,
policy,
tokenizer,
reasoning_parser_factory,
tool_parser_registry,
timeout_secs: ctx.router_config.worker_startup_timeout_secs,
interval_secs: ctx.router_config.worker_startup_check_interval_secs,
dp_aware: ctx.router_config.dp_aware,
api_key: ctx.router_config.api_key.clone(),
retry_config: ctx.router_config.effective_retry_config(),
circuit_breaker_config: core_cb_config,
})
}
@@ -576,8 +506,6 @@ impl std::fmt::Debug for GrpcRouter {
let stats = self.worker_registry.stats();
f.debug_struct("GrpcRouter")
.field("workers_count", &stats.total_workers)
.field("timeout_secs", &self.timeout_secs)
.field("interval_secs", &self.interval_secs)
.field("dp_aware", &self.dp_aware)
.finish()
}