Skip to content

GitLab

Unverified Commit 97c38239 authored by Simo Lin's avatar Simo Lin Committed by GitHub
Browse files

[router] refactor router and worker management 3/n (#10727)

parent 60dbbd08
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 1344 additions and 2493 deletions
sgl-router/py_test/e2e/test_regular_router.py
View file @ 97c38239
......@@ -141,14 +141,21 @@ def test_dp_aware_worker_expansion_and_api_key(
assert len(urls) == 2
assert set(urls) == {f"{worker_url}@0", f"{worker_url}@1"}
# TODO: Router currently doesn't enforce API key authentication on incoming requests.
# It only adds the API key to outgoing requests to workers.
# Need to implement auth middleware to properly protect router endpoints.
# For now, both requests succeed (200) regardless of client authentication.
# Verify API key enforcement path-through
# 1) Without Authorization -> 401 from backend
# 1) Without Authorization -> Currently 200 (should be 401 after auth middleware added)
r = requests.post(
f"{router_url}/v1/completions",
json={"model": e2e_model, "prompt": "hi", "max_tokens": 1},
timeout=60,
)
assert r.status_code == 401
assert (
r.status_code == 200
) # TODO: Change to 401 after auth middleware implementation
# 2) With correct Authorization -> 200
r = requests.post(
......
sgl-router/src/core/circuit_breaker.rs
View file @ 97c38239
......@@ -83,14 +83,13 @@ impl CircuitBreaker {
/// Check if a request can be executed
pub fn can_execute(&self) -> bool {
// First check if we need to transition from Open to HalfOpen
self.check_and_update_state();
let state = *self.state.read().unwrap();
match state {
CircuitState::Closed => true,
CircuitState::Open => false,
CircuitState::HalfOpen => true, // Allow limited requests in half-open state
CircuitState::HalfOpen => true,
}
}
......@@ -114,22 +113,17 @@ impl CircuitBreaker {
self.total_successes.fetch_add(1, Ordering::Relaxed);
self.consecutive_failures.store(0, Ordering::Release);
let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1;
// Outcome-level metrics are recorded at the worker level where the worker label is known
let current_state = *self.state.read().unwrap();
match current_state {
CircuitState::HalfOpen => {
// Check if we've reached the success threshold to close the circuit
if successes >= self.config.success_threshold {
self.transition_to(CircuitState::Closed);
}
}
CircuitState::Closed => {
// Already closed, nothing to do
}
CircuitState::Closed => {}
CircuitState::Open => {
// Shouldn't happen, but if it does, stay open
tracing::warn!("Success recorded while circuit is open");
}
}
......@@ -140,9 +134,7 @@ impl CircuitBreaker {
self.total_failures.fetch_add(1, Ordering::Relaxed);
self.consecutive_successes.store(0, Ordering::Release);
let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1;
// Outcome-level metrics are recorded at the worker level where the worker label is known
// Update last failure time
{
let mut last_failure = self.last_failure_time.write().unwrap();
*last_failure = Some(Instant::now());
......@@ -152,18 +144,14 @@ impl CircuitBreaker {
match current_state {
CircuitState::Closed => {
// Check if we've reached the failure threshold to open the circuit
if failures >= self.config.failure_threshold {
self.transition_to(CircuitState::Open);
}
}
CircuitState::HalfOpen => {
// Single failure in half-open state reopens the circuit
self.transition_to(CircuitState::Open);
}
CircuitState::Open => {
// Already open, nothing to do
}
CircuitState::Open => {}
}
}
......@@ -172,7 +160,6 @@ impl CircuitBreaker {
let current_state = *self.state.read().unwrap();
if current_state == CircuitState::Open {
// Check if timeout has expired
let last_change = *self.last_state_change.read().unwrap();
if last_change.elapsed() >= self.config.timeout_duration {
self.transition_to(CircuitState::HalfOpen);
......@@ -188,11 +175,9 @@ impl CircuitBreaker {
if old_state != new_state {
*state = new_state;
// Update last state change time
let mut last_change = self.last_state_change.write().unwrap();
*last_change = Instant::now();
// Reset counters based on transition
match new_state {
CircuitState::Closed => {
self.consecutive_failures.store(0, Ordering::Release);
......@@ -218,7 +203,6 @@ impl CircuitBreaker {
CircuitState::HalfOpen => "half_open",
};
info!("Circuit breaker state transition: {} -> {}", from, to);
// Transition metrics are recorded at the worker level where the worker label is known
}
}
......@@ -533,7 +517,6 @@ mod tests {
let cb = Arc::new(CircuitBreaker::new());
let mut handles = vec![];
// Spawn threads that record failures
for _ in 0..10 {
let cb_clone = Arc::clone(&cb);
let handle = thread::spawn(move || {
......@@ -544,12 +527,10 @@ mod tests {
handles.push(handle);
}
// Wait for all threads
for handle in handles {
handle.join().unwrap();
}
// Should have recorded 1000 failures
assert_eq!(cb.total_failures(), 1000);
}
}
sgl-router/src/core/error.rs
View file @ 97c38239
......@@ -122,7 +122,6 @@ mod tests {
let error = WorkerError::WorkerNotFound {
url: "http://test".to_string(),
};
// Verify it implements Error trait
let _: &dyn Error = &error;
assert!(error.source().is_none());
}
......@@ -135,11 +134,9 @@ mod tests {
#[test]
fn test_worker_result_type_alias() {
// Test Ok variant
let result: WorkerResult<i32> = Ok(42);
assert!(matches!(result, Ok(42)));
// Test Err variant
let error = WorkerError::WorkerNotFound {
url: "test".to_string(),
};
......@@ -149,7 +146,6 @@ mod tests {
#[test]
fn test_empty_url_handling() {
// Test empty URLs in error variants
let error1 = WorkerError::HealthCheckFailed {
url: "".to_string(),
reason: "No connection".to_string(),
......@@ -173,7 +169,6 @@ mod tests {
#[test]
fn test_special_characters_in_messages() {
// Test with special characters
let error = WorkerError::InvalidConfiguration {
message: "Invalid JSON: {\"error\": \"test\"}".to_string(),
};
......@@ -182,7 +177,6 @@ mod tests {
"Invalid worker configuration: Invalid JSON: {\"error\": \"test\"}"
);
// Test with unicode
let error2 = WorkerError::HealthCheckFailed {
url: "http://测试:8080".to_string(),
reason: "连接被拒绝".to_string(),
......@@ -207,10 +201,8 @@ mod tests {
);
}
// Mock reqwest error for testing conversion
#[test]
fn test_reqwest_error_conversion() {
// Test that NetworkError is the correct variant
let network_error = WorkerError::NetworkError {
url: "http://example.com".to_string(),
error: "connection timeout".to_string(),
......@@ -227,8 +219,6 @@ mod tests {
#[test]
fn test_error_equality() {
// WorkerError doesn't implement PartialEq, but we can test that
// the same error construction produces the same display output
let error1 = WorkerError::WorkerNotFound {
url: "http://test".to_string(),
};
......
sgl-router/src/core/mod.rs
View file @ 97c38239
......@@ -12,9 +12,9 @@ pub mod retry;
pub mod token_bucket;
pub mod worker;
pub mod worker_builder;
pub mod worker_manager;
pub mod worker_registry;
// Re-export commonly used types at the module level
pub use circuit_breaker::{
CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState,
};
......@@ -25,4 +25,5 @@ pub use worker::{
Worker, WorkerFactory, WorkerLoadGuard, WorkerType,
};
pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder};
pub use worker_manager::{DpInfo, ServerInfo, WorkerManager};
pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats};
sgl-router/src/core/retry.rs
View file @ 97c38239
......@@ -25,14 +25,12 @@ pub struct BackoffCalculator;
impl BackoffCalculator {
/// Calculate backoff delay for a given attempt index (0-based).
pub fn calculate_delay(config: &RetryConfig, attempt: u32) -> Duration {
// Base exponential backoff
let pow = config.backoff_multiplier.powi(attempt as i32);
let mut delay_ms = (config.initial_backoff_ms as f32 * pow) as u64;
if delay_ms > config.max_backoff_ms {
delay_ms = config.max_backoff_ms;
}
// Apply jitter in range [-j, +j]
let jitter = config.jitter_factor.clamp(0.0, 1.0);
if jitter > 0.0 {
let mut rng = rand::rng();
......@@ -77,14 +75,12 @@ impl RetryExecutor {
match operation(attempt).await {
Ok(val) => return Ok(val),
Err(_) => {
// Use the number of failures so far (0-indexed) to compute delay,
// so the first retry uses `initial_backoff_ms`.
let is_last = attempt + 1 >= max;
if is_last {
return Err(RetryError::MaxRetriesExceeded);
}
let delay = BackoffCalculator::calculate_delay(config, attempt);
attempt += 1; // advance to the next attempt after computing delay
attempt += 1;
tokio::time::sleep(delay).await;
}
}
......@@ -144,14 +140,11 @@ impl RetryExecutor {
}
if is_last {
// Exhausted retries
on_exhausted();
return response;
}
// Backoff before next attempt
let next_attempt = attempt + 1;
// Compute delay based on the number of failures so far (0-indexed)
let delay = BackoffCalculator::calculate_delay(config, attempt);
debug!(
attempt = attempt,
......@@ -194,22 +187,18 @@ mod tests {
backoff_multiplier: 2.0,
jitter_factor: 0.0,
};
// attempt=0 => 100ms
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 0),
Duration::from_millis(100)
);
// attempt=1 => 200ms
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 1),
Duration::from_millis(200)
);
// attempt=2 => 400ms -> capped to 250ms
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 2),
Duration::from_millis(250)
);
// large attempt still capped
assert_eq!(
BackoffCalculator::calculate_delay(&cfg, 10),
Duration::from_millis(250)
......@@ -225,7 +214,6 @@ mod tests {
backoff_multiplier: 2.0,
jitter_factor: 0.5,
};
// attempt=2 => base 400ms, jitter in [0.5x, 1.5x]
let base = 400.0;
for _ in 0..50 {
let d = BackoffCalculator::calculate_delay(&cfg, 2).as_millis() as f32;
......@@ -261,7 +249,7 @@ mod tests {
assert!(res.is_ok());
assert_eq!(res.unwrap(), 42);
assert_eq!(calls.load(Ordering::Relaxed), 3); // 2 fails + 1 success
assert_eq!(calls.load(Ordering::Relaxed), 3);
}
#[tokio::test]
......@@ -309,7 +297,7 @@ mod tests {
}
}
},
|res, _attempt| !res.status().is_success(), // retry until success
|res, _attempt| !res.status().is_success(),
{
let backoffs = backoffs.clone();
move |_delay, _next_attempt| {
......@@ -326,7 +314,7 @@ mod tests {
.await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(calls.load(Ordering::Relaxed), 3); // 2 fails + 1 success
assert_eq!(calls.load(Ordering::Relaxed), 3);
assert_eq!(backoffs.load(Ordering::Relaxed), 2);
assert_eq!(exhausted.load(Ordering::Relaxed), 0);
}
......@@ -347,7 +335,7 @@ mod tests {
async move { (StatusCode::BAD_REQUEST, "bad").into_response() }
}
},
|_res, _attempt| false, // never retry
|_res, _attempt| false,
{
let backoffs = backoffs.clone();
move |_delay, _next_attempt| {
......@@ -385,7 +373,7 @@ mod tests {
async move { (StatusCode::SERVICE_UNAVAILABLE, "fail").into_response() }
}
},
|_res, _attempt| true, // keep retrying
|_res, _attempt| true,
{
let backoffs = backoffs.clone();
move |_delay, _next_attempt| {
......
sgl-router/src/core/token_bucket.rs
View file @ 97c38239
......@@ -32,16 +32,11 @@ impl TokenBucket {
let capacity = capacity as f64;
let refill_rate = refill_rate as f64;
// Ensure refill_rate is not zero to prevent division by zero
let refill_rate = if refill_rate > 0.0 {
refill_rate
} else {
1.0 // Default to 1 token per second if zero
};
let refill_rate = if refill_rate > 0.0 { refill_rate } else { 1.0 };
Self {
inner: Arc::new(Mutex::new(TokenBucketInner {
tokens: capacity, // Start full
tokens: capacity,
last_refill: Instant::now(),
})),
notify: Arc::new(Notify::new()),
......@@ -54,7 +49,6 @@ impl TokenBucket {
pub async fn try_acquire(&self, tokens: f64) -> Result<(), ()> {
let mut inner = self.inner.lock().await;
// Refill tokens based on elapsed time
let now = Instant::now();
let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
let refill_amount = elapsed * self.refill_rate;
......@@ -82,12 +76,10 @@ impl TokenBucket {
/// Acquire tokens, waiting if necessary
pub async fn acquire(&self, tokens: f64) -> Result<(), tokio::time::error::Elapsed> {
// First try to acquire immediately
if self.try_acquire(tokens).await.is_ok() {
return Ok(());
}
// Calculate wait time
let wait_time = {
let inner = self.inner.lock().await;
let tokens_needed = tokens - inner.tokens;
......@@ -100,15 +92,12 @@ impl TokenBucket {
wait_time, tokens
);
// Wait for tokens to be available
tokio::time::timeout(wait_time, async {
loop {
// Check if we can acquire now
if self.try_acquire(tokens).await.is_ok() {
return;
}
// Wait for notification or small interval
tokio::select! {
_ = self.notify.notified() => {},
_ = tokio::time::sleep(Duration::from_millis(10)) => {},
......@@ -144,7 +133,6 @@ impl TokenBucket {
pub async fn available_tokens(&self) -> f64 {
let mut inner = self.inner.lock().await;
// Refill before checking
let now = Instant::now();
let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
let refill_amount = elapsed * self.refill_rate;
......@@ -162,33 +150,26 @@ mod tests {
#[tokio::test]
async fn test_token_bucket_basic() {
let bucket = TokenBucket::new(10, 5); // 10 capacity, 5 per second
let bucket = TokenBucket::new(10, 5);
// Should succeed - bucket starts full
assert!(bucket.try_acquire(5.0).await.is_ok());
assert!(bucket.try_acquire(5.0).await.is_ok());
// Should fail - no tokens left
assert!(bucket.try_acquire(1.0).await.is_err());
// Wait for refill
tokio::time::sleep(Duration::from_millis(300)).await;
// Should have ~1.5 tokens now
assert!(bucket.try_acquire(1.0).await.is_ok());
}
#[tokio::test]
async fn test_token_bucket_refill() {
let bucket = TokenBucket::new(10, 10); // 10 capacity, 10 per second
let bucket = TokenBucket::new(10, 10);
// Use all tokens
assert!(bucket.try_acquire(10.0).await.is_ok());
// Wait for partial refill
tokio::time::sleep(Duration::from_millis(500)).await;
// Should have ~5 tokens
let available = bucket.available_tokens().await;
assert!((4.0..=6.0).contains(&available));
}
......
sgl-router/src/core/worker.rs
View file @ 97c38239
This diff is collapsed.
sgl-router/src/core/worker_builder.rs
View file @ 97c38239
......@@ -7,10 +7,7 @@ use std::collections::HashMap;
/// Builder for creating BasicWorker instances with fluent API
pub struct BasicWorkerBuilder {
// Required fields
url: String,
// Optional fields with defaults
api_key: Option<String>,
worker_type: WorkerType,
connection_mode: ConnectionMode,
......@@ -21,7 +18,7 @@ pub struct BasicWorkerBuilder {
}
impl BasicWorkerBuilder {
/// Create a new builder with only the URL (defaults to Regular worker type)
/// Create a new builder with only the URL
pub fn new(url: impl Into<String>) -> Self {
Self {
url: url.into(),
......@@ -129,13 +126,10 @@ impl BasicWorkerBuilder {
/// Builder for creating DPAwareWorker instances with fluent API
pub struct DPAwareWorkerBuilder {
// Required fields
base_url: String,
api_key: Option<String>,
dp_rank: usize,
dp_size: usize,
// Optional fields with defaults
worker_type: WorkerType,
connection_mode: ConnectionMode,
labels: HashMap<String, String>,
......@@ -145,7 +139,7 @@ pub struct DPAwareWorkerBuilder {
}
impl DPAwareWorkerBuilder {
/// Create a new DP-aware worker builder (defaults to Regular worker type)
/// Create a new DP-aware worker builder
pub fn new(base_url: impl Into<String>, dp_rank: usize, dp_size: usize) -> Self {
Self {
base_url: base_url.into(),
......@@ -232,10 +226,7 @@ impl DPAwareWorkerBuilder {
/// Build the DPAwareWorker instance
pub fn build(self) -> DPAwareWorker {
// Create URL with DP rank suffix for identification
let worker_url = format!("{}@{}", self.base_url, self.dp_rank);
// Use BasicWorkerBuilder to create a properly configured base worker
let mut builder = BasicWorkerBuilder::new(worker_url)
.worker_type(self.worker_type)
.connection_mode(self.connection_mode)
......@@ -243,18 +234,14 @@ impl DPAwareWorkerBuilder {
.health_config(self.health_config)
.circuit_breaker_config(self.circuit_breaker_config);
// Add gRPC client if provided
if let Some(client) = self.grpc_client {
builder = builder.grpc_client(client);
}
// Add API key if provided
if let Some(api_key) = self.api_key {
builder = builder.api_key(api_key);
}
let base_worker = builder.build();
// Create the DPAwareWorker with the configured base worker
DPAwareWorker::with_base_worker(base_worker, self.base_url, self.dp_rank, self.dp_size)
}
}
......@@ -267,7 +254,6 @@ mod tests {
#[test]
fn test_basic_worker_builder_minimal() {
// Using new API - defaults to Regular type
let worker = BasicWorkerBuilder::new("http://localhost:8080").build();
assert_eq!(worker.url(), "http://localhost:8080");
......@@ -278,7 +264,6 @@ mod tests {
#[test]
fn test_basic_worker_builder_with_type() {
// Test setting worker type explicitly
let worker = BasicWorkerBuilder::new("http://localhost:8080")
.worker_type(WorkerType::Decode)
.build();
......@@ -332,7 +317,6 @@ mod tests {
ConnectionMode::Grpc { port: Some(50051) }
);
assert_eq!(worker.metadata().labels, labels);
// Can't directly compare HealthConfig without PartialEq, so check individual fields
assert_eq!(
worker.metadata().health_config.endpoint,
health_config.endpoint
......@@ -375,13 +359,11 @@ mod tests {
#[test]
fn test_dp_aware_worker_builder_minimal() {
// Using new API - defaults to Regular type
let worker = DPAwareWorkerBuilder::new("http://localhost:8080", 2, 8).build();
assert_eq!(worker.url(), "http://localhost:8080@2");
assert_eq!(worker.dp_rank(), Some(2));
assert_eq!(worker.dp_size(), Some(8));
// Note: base_url is a private field, we can only test through the url() method
assert_eq!(worker.worker_type(), WorkerType::Regular);
}
......@@ -412,7 +394,6 @@ mod tests {
assert_eq!(worker.dp_rank(), Some(3));
assert_eq!(worker.dp_size(), Some(16));
assert_eq!(worker.metadata().labels, labels);
// Can't directly compare HealthConfig without PartialEq, so check individual fields
assert_eq!(
worker.metadata().health_config.endpoint,
health_config.endpoint
......@@ -437,7 +418,6 @@ mod tests {
#[test]
fn test_dp_aware_worker_with_grpc() {
// Test that DPAwareWorkerBuilder can set a gRPC client
let worker = DPAwareWorkerBuilder::new("grpc://cluster.local", 1, 4)
.worker_type(WorkerType::Decode)
.connection_mode(ConnectionMode::Grpc { port: Some(50051) })
......@@ -456,9 +436,5 @@ mod tests {
worker.metadata().labels.get("transport"),
Some(&"grpc".to_string())
);
// Note: We can't directly test the grpc_client as it's private,
// but the fact that the worker builds successfully with grpc connection mode
// validates that the configuration is properly passed through
}
}
sgl-router/src/core/worker_manager.rs 0 → 100644
View file @ 97c38239
This diff is collapsed.
sgl-router/src/core/worker_registry.rs
View file @ 97c38239
......@@ -34,7 +34,6 @@ impl Default for WorkerId {
}
}
/// Type alias for the model index to reduce complexity
type ModelIndex = Arc<DashMap<String, Arc<RwLock<Vec<Arc<dyn Worker>>>>>>;
/// Worker registry with model-based indexing
......@@ -54,8 +53,7 @@ pub struct WorkerRegistry {
/// Workers indexed by connection mode
connection_workers: Arc<DashMap<ConnectionMode, Vec<WorkerId>>>,
/// URL to worker ID mapping (for backward compatibility)
/// URL to worker ID mapping
url_to_id: Arc<DashMap<String, WorkerId>>,
}
......
sgl-router/src/routers/grpc/pd_router.rs
View file @ 97c38239
......@@ -8,7 +8,7 @@ use crate::grpc::SglangSchedulerClient;
use crate::metrics::RouterMetrics;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
use crate::reasoning_parser::ParserFactory;
use crate::routers::{RouterTrait, WorkerManagement};
use crate::routers::RouterTrait;
use crate::tokenizer::traits::Tokenizer;
use crate::tool_parser::ParserRegistry;
use async_trait::async_trait;
......@@ -350,42 +350,3 @@ impl RouterTrait for GrpcPDRouter {
(StatusCode::SERVICE_UNAVAILABLE).into_response()
}
}
#[async_trait]
impl WorkerManagement for GrpcPDRouter {
async fn add_worker(
&self,
_worker_url: &str,
_api_key: &Option<String>,
) -> Result<String, String> {
Err("Not implemented".to_string())
}
fn remove_worker(&self, _worker_url: &str) {}
fn get_worker_urls(&self) -> Vec<String> {
let mut urls = Vec::new();
// Get gRPC prefill worker URLs only
let prefill_workers = self.worker_registry.get_workers_filtered(
None,
Some(WorkerType::Prefill {
bootstrap_port: None,
}),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false,
);
urls.extend(prefill_workers.iter().map(|w| w.url().to_string()));
// Get gRPC decode worker URLs only
let decode_workers = self.worker_registry.get_workers_filtered(
None,
Some(WorkerType::Decode),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false,
);
urls.extend(decode_workers.iter().map(|w| w.url().to_string()));
urls
}
}
sgl-router/src/routers/grpc/router.rs
View file @ 97c38239
......@@ -8,7 +8,7 @@ use crate::grpc::SglangSchedulerClient;
use crate::metrics::RouterMetrics;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
use crate::reasoning_parser::ParserFactory;
use crate::routers::{RouterTrait, WorkerManagement};
use crate::routers::RouterTrait;
use crate::tokenizer::traits::Tokenizer;
use crate::tool_parser::ParserRegistry;
use async_trait::async_trait;
......@@ -279,29 +279,3 @@ impl RouterTrait for GrpcRouter {
(StatusCode::SERVICE_UNAVAILABLE).into_response()
}
}
#[async_trait]
impl WorkerManagement for GrpcRouter {
async fn add_worker(
&self,
_worker_url: &str,
_api_key: &Option<String>,
) -> Result<String, String> {
Err("Not implemented".to_string())
}
fn remove_worker(&self, _worker_url: &str) {}
fn get_worker_urls(&self) -> Vec<String> {
self.worker_registry
.get_workers_filtered(
None, // any model
Some(WorkerType::Regular),
Some(crate::core::ConnectionMode::Grpc { port: None }),
false, // include all workers
)
.iter()
.map(|w| w.url().to_string())
.collect()
}
}
sgl-router/src/routers/http/openai_router.rs
View file @ 97c38239
......@@ -65,25 +65,6 @@ impl OpenAIRouter {
}
}
#[async_trait]
impl super::super::WorkerManagement for OpenAIRouter {
async fn add_worker(
&self,
_worker_url: &str,
_api_key: &Option<String>,
) -> Result<String, String> {
Err("Cannot add workers to OpenAI router".to_string())
}
fn remove_worker(&self, _worker_url: &str) {
// No-op for OpenAI router
}
fn get_worker_urls(&self) -> Vec<String> {
vec![self.base_url.clone()]
}
}
#[async_trait]
impl super::super::RouterTrait for OpenAIRouter {
fn as_any(&self) -> &dyn Any {
......
sgl-router/src/routers/http/pd_router.rs
View file @ 97c38239
This diff is collapsed.
sgl-router/src/routers/http/router.rs
View file @ 97c38239
This diff is collapsed.
sgl-router/src/routers/mod.rs
View file @ 97c38239
......@@ -19,39 +19,18 @@ pub mod grpc;
pub mod header_utils;
pub mod http;
pub mod router_manager;
pub mod worker_initializer;
pub use factory::RouterFactory;
pub use worker_initializer::WorkerInitializer;
// Re-export HTTP routers for convenience (keeps routers::openai_router path working)
pub use http::{openai_router, pd_router, pd_types, router};
/// Worker management trait for administrative operations
///
/// This trait is separate from RouterTrait to allow Send futures
/// for use in service discovery and other background tasks
#[async_trait]
pub trait WorkerManagement: Send + Sync {
/// Add a worker to the router
async fn add_worker(
&self,
worker_url: &str,
api_key: &Option<String>,
) -> Result<String, String>;
/// Remove a worker from the router
fn remove_worker(&self, worker_url: &str);
/// Get all worker URLs
fn get_worker_urls(&self) -> Vec<String>;
}
/// Core trait for all router implementations
///
/// This trait provides a unified interface for routing requests,
/// regardless of whether it's a regular router or PD router.
#[async_trait]
pub trait RouterTrait: Send + Sync + Debug + WorkerManagement {
pub trait RouterTrait: Send + Sync + Debug {
/// Get a reference to self as Any for downcasting
fn as_any(&self) -> &dyn std::any::Any;
......
sgl-router/src/routers/router_manager.rs
View file @ 97c38239
......@@ -4,17 +4,12 @@
//! - Single Router Mode (enable_igw=false): Router owns workers directly
//! - Multi-Router Mode (enable_igw=true): RouterManager coordinates everything
use crate::config::RouterConfig;
use crate::core::{BasicWorkerBuilder, CircuitBreakerConfig, Worker, WorkerRegistry, WorkerType};
use crate::core::{Worker, WorkerRegistry, WorkerType};
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, RerankRequest,
ResponsesRequest,
};
use crate::protocols::worker_spec::{
ServerInfo, WorkerApiResponse, WorkerConfigRequest, WorkerErrorResponse, WorkerInfo,
WorkerListResponse, WorkerStats, WorkerTypeStats,
};
use crate::routers::{RouterTrait, WorkerManagement};
use crate::routers::RouterTrait;
use async_trait::async_trait;
use axum::{
body::Body,
......@@ -24,7 +19,7 @@ use axum::{
};
use dashmap::DashMap;
use std::sync::Arc;
use tracing::{info, warn};
use tracing::info;
/// Router identifier
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
......@@ -45,48 +40,28 @@ pub struct RouterManager {
/// Worker registry (single source of truth in multi-router mode)
worker_registry: Arc<WorkerRegistry>,
/// Policy registry for managing model-to-policy mappings
policy_registry: Arc<crate::policies::PolicyRegistry>,
/// All routers managed by this manager
/// RouterId examples: "http-regular", "http-pd", "grpc-regular", "grpc-pd"
routers: Arc<DashMap<RouterId, Arc<dyn RouterTrait>>>,
/// Default router for requests without specific routing
default_router: Arc<std::sync::RwLock<Option<RouterId>>>,
/// HTTP client for querying worker info
client: reqwest::Client,
/// Configuration
#[allow(dead_code)] // May be used in future enhancements
config: RouterConfig,
}
impl RouterManager {
/// Create a new router manager with shared registries
pub fn new(
config: RouterConfig,
client: reqwest::Client,
worker_registry: Arc<WorkerRegistry>,
policy_registry: Arc<crate::policies::PolicyRegistry>,
) -> Self {
pub fn new(worker_registry: Arc<WorkerRegistry>) -> Self {
Self {
worker_registry,
policy_registry,
routers: Arc::new(DashMap::new()),
default_router: Arc::new(std::sync::RwLock::new(None)),
client,
config,
}
}
/// Register a router with the manager
pub fn register_router(&self, id: RouterId, router: Arc<dyn RouterTrait>) {
// Store router
self.routers.insert(id.clone(), router);
// Set as default if first router
let mut default_router = self.default_router.write().unwrap();
if default_router.is_none() {
*default_router = Some(id.clone());
......@@ -107,11 +82,9 @@ impl RouterManager {
/// Get router for a specific model based on worker types
pub fn get_router_for_model(&self, model_id: &str) -> Option<Arc<dyn RouterTrait>> {
// Query workers for this model from registry
let workers = self.worker_registry.get_by_model(model_id);
if !workers.is_empty() {
// Determine router based on worker types
let has_pd_workers = workers.iter().any(|w| {
matches!(
w.worker_type(),
......@@ -125,13 +98,11 @@ impl RouterManager {
RouterId::new("http-regular".to_string())
};
// Return the router if it exists
if let Some(router) = self.routers.get(&router_id) {
return Some(router.clone());
}
}
// Fall back to default router
let default_router = self.default_router.read().unwrap();
if let Some(ref default_id) = *default_router {
self.routers.get(default_id).map(|r| r.clone())
......@@ -149,277 +120,12 @@ impl RouterManager {
}
}
/// Add a worker to the registry
pub async fn add_worker(
&self,
config: WorkerConfigRequest,
) -> Result<WorkerApiResponse, WorkerErrorResponse> {
// Build labels from configuration
let mut labels = config.labels.clone();
// Query server info if model_id not provided
let model_id = if let Some(model_id) = config.model_id {
model_id
} else {
match self.query_server_info(&config.url, &config.api_key).await {
Ok(info) => {
// Extract model_id from server info
info.model_id
.or_else(|| {
info.model_path
.as_ref()
.and_then(|path| path.split('/').next_back().map(|s| s.to_string()))
})
.unwrap_or_else(|| "unknown".to_string())
}
Err(e) => {
warn!("Failed to query server info from {}: {}", config.url, e);
"unknown".to_string()
}
}
};
// Add configuration to labels
labels.insert("model_id".to_string(), model_id.clone());
if let Some(priority) = config.priority {
labels.insert("priority".to_string(), priority.to_string());
}
if let Some(cost) = config.cost {
labels.insert("cost".to_string(), cost.to_string());
}
// Add gRPC-specific configuration if provided
if let Some(tokenizer_path) = config.tokenizer_path {
labels.insert("tokenizer_path".to_string(), tokenizer_path);
}
if let Some(reasoning_parser) = config.reasoning_parser {
labels.insert("reasoning_parser".to_string(), reasoning_parser);
}
if let Some(tool_parser) = config.tool_parser {
labels.insert("tool_parser".to_string(), tool_parser);
}
if let Some(chat_template) = config.chat_template {
labels.insert("chat_template".to_string(), chat_template);
}
let worker = match config.worker_type.as_deref() {
Some("prefill") => {
let mut builder = BasicWorkerBuilder::new(config.url.clone())
.worker_type(WorkerType::Prefill {
bootstrap_port: config.bootstrap_port,
})
.labels(labels.clone())
.circuit_breaker_config(CircuitBreakerConfig::default());
if let Some(api_key) = config.api_key.clone() {
builder = builder.api_key(api_key);
}
Box::new(builder.build()) as Box<dyn Worker>
}
Some("decode") => {
let mut builder = BasicWorkerBuilder::new(config.url.clone())
.worker_type(WorkerType::Decode)
.labels(labels.clone())
.circuit_breaker_config(CircuitBreakerConfig::default());
if let Some(api_key) = config.api_key.clone() {
builder = builder.api_key(api_key);
}
Box::new(builder.build()) as Box<dyn Worker>
}
_ => {
let mut builder = BasicWorkerBuilder::new(config.url.clone())
.worker_type(WorkerType::Regular)
.labels(labels.clone())
.circuit_breaker_config(CircuitBreakerConfig::default());
if let Some(api_key) = config.api_key.clone() {
builder = builder.api_key(api_key);
}
Box::new(builder.build()) as Box<dyn Worker>
}
};
// Register worker
let worker_arc: Arc<dyn Worker> = Arc::from(worker);
let worker_id = self.worker_registry.register(worker_arc.clone());
// Notify PolicyRegistry about the new worker
// Extract policy hint from labels if provided
let policy_hint = labels.get("policy").map(|s| s.as_str());
let policy = self.policy_registry.on_worker_added(&model_id, policy_hint);
// Log which type of router would handle this worker (for debugging)
let expected_router = match config.worker_type.as_deref() {
Some("prefill") | Some("decode") => "http-pd",
_ => "http-regular",
};
info!(
"Worker for model '{}' would be handled by '{}' router based on type",
model_id, expected_router
);
info!(
"Added worker {} with URL {} for model {} using policy {}",
worker_id.as_str(),
config.url,
model_id,
policy.name()
);
// Return worker info
let worker_info = self.worker_to_info(worker_id.as_str(), &worker_arc);
Ok(WorkerApiResponse {
success: true,
message: format!("Worker {} added successfully", worker_id.as_str()),
worker: Some(worker_info),
})
}
/// Remove a worker from the registry
pub fn remove_worker_from_registry(
&self,
url: &str,
) -> Result<WorkerApiResponse, WorkerErrorResponse> {
// Get worker to extract model_id before removing
let model_id = self
.worker_registry
.get_by_url(url)
.map(|worker| worker.model_id().to_string());
if let Some(_worker) = self.worker_registry.remove_by_url(url) {
// Notify PolicyRegistry about worker removal
if let Some(ref model_id) = model_id {
self.policy_registry.on_worker_removed(model_id);
info!("Removed worker with URL {} for model {}", url, model_id);
} else {
info!("Removed worker with URL {}", url);
}
Ok(WorkerApiResponse {
success: true,
message: format!("Worker {} removed successfully", url),
worker: None,
})
} else {
Err(WorkerErrorResponse {
error: format!("Worker with URL {} not found", url),
code: "WORKER_NOT_FOUND".to_string(),
})
}
}
/// List all workers
pub fn list_workers(&self) -> WorkerListResponse {
let workers = self.worker_registry.get_all_with_ids();
let worker_infos: Vec<WorkerInfo> = workers
.iter()
.map(|(id, w)| self.worker_to_info(id.as_str(), w))
.collect();
let total = worker_infos.len();
// Get stats from the worker registry
let registry_stats = self.worker_registry.stats();
// Convert WorkerRegistryStats to WorkerStats
let stats = WorkerStats {
total_workers: registry_stats.total_workers,
healthy_workers: registry_stats.healthy_workers,
total_models: registry_stats.total_models,
total_load: registry_stats.total_load,
by_type: WorkerTypeStats {
regular: registry_stats.regular_workers,
prefill: registry_stats.prefill_workers,
decode: registry_stats.decode_workers,
},
};
WorkerListResponse {
workers: worker_infos,
total,
stats,
}
}
/// Get worker by URL
pub fn get_worker(&self, url: &str) -> Option<WorkerInfo> {
self.worker_registry
.get_by_url(url)
.map(|w| self.worker_to_info("unknown", &w))
}
/// Query server info from a worker URL
async fn query_server_info(
&self,
url: &str,
api_key: &Option<String>,
) -> Result<ServerInfo, String> {
let info_url = format!("{}/get_server_info", url.trim_end_matches('/'));
let mut req_builder = self.client.get(&info_url);
if let Some(key) = api_key {
req_builder = req_builder.bearer_auth(key);
}
match req_builder.send().await {
Ok(response) => {
if response.status().is_success() {
response
.json::<ServerInfo>()
.await
.map_err(|e| format!("Failed to parse server info: {}", e))
} else {
Err(format!("Server returned status: {}", response.status()))
}
}
Err(e) => Err(format!("Failed to connect to server: {}", e)),
}
}
/// Convert Worker to WorkerInfo
fn worker_to_info(&self, id: &str, worker: &Arc<dyn Worker>) -> WorkerInfo {
let metadata = worker.metadata();
WorkerInfo {
id: id.to_string(),
url: worker.url().to_string(),
model_id: worker.model_id().to_string(),
priority: worker.priority(),
cost: worker.cost(),
worker_type: match worker.worker_type() {
WorkerType::Regular => "regular".to_string(),
WorkerType::Prefill { .. } => "prefill".to_string(),
WorkerType::Decode => "decode".to_string(),
},
is_healthy: worker.is_healthy(),
load: worker.load(),
connection_mode: format!("{:?}", worker.connection_mode()),
tokenizer_path: worker.tokenizer_path().map(|s| s.to_string()),
reasoning_parser: worker.reasoning_parser().map(|s| s.to_string()),
tool_parser: worker.tool_parser().map(|s| s.to_string()),
chat_template: worker.chat_template().map(|s| s.to_string()),
metadata: metadata.labels.clone(),
}
}
/// Get the appropriate router for a request based on headers and request content
pub fn select_router_for_request(
&self,
headers: Option<&HeaderMap>,
model_id: Option<&str>,
) -> Option<Arc<dyn RouterTrait>> {
// Extract priority and cost preferences from headers if available
let _priority_threshold = headers.and_then(|h| {
h.get("x-worker-priority")
.and_then(|v| v.to_str().ok())
......@@ -432,7 +138,6 @@ impl RouterManager {
.and_then(|s| s.parse::<f32>().ok())
});
// Check if PD (prefill-decode) mode is preferred from headers
let prefer_pd = headers
.and_then(|h| {
h.get("x-prefer-pd")
......@@ -441,7 +146,6 @@ impl RouterManager {
})
.unwrap_or(false);
// If model specified, use get_router_for_model
let candidate_routers = if let Some(model) = model_id {
if let Some(router) = self.get_router_for_model(model) {
vec![router]
......@@ -449,7 +153,6 @@ impl RouterManager {
Vec::new()
}
} else {
// No model specified, consider all routers
self.routers
.iter()
.map(|entry| entry.value().clone())
......@@ -457,23 +160,20 @@ impl RouterManager {
};
if candidate_routers.is_empty() {
// No routers found for the specified model
return None;
}
// Score routers based on worker attributes and request preferences
let mut best_router = None;
let mut best_score = 0.0;
for router in candidate_routers {
let mut score = 1.0;
// Check if this is a PD router
let is_pd = router.is_pd_mode();
if prefer_pd && is_pd {
score += 2.0; // Bonus for matching PD preference
score += 2.0;
} else if !prefer_pd && !is_pd {
score += 1.0; // Bonus for matching regular preference
score += 1.0;
}
// Get workers for this router and evaluate based on priority/cost
......@@ -495,49 +195,6 @@ impl RouterManager {
}
}
/// RouterManager implements RouterTrait to act as a meta-router
/// that delegates requests to the appropriate underlying router
#[async_trait]
impl WorkerManagement for RouterManager {
/// Add a worker - in multi-router mode, this adds to the registry
async fn add_worker(
&self,
worker_url: &str,
api_key: &Option<String>,
) -> Result<String, String> {
// Create a basic worker config request
let config = WorkerConfigRequest {
url: worker_url.to_string(),
api_key: api_key.clone(),
model_id: None,
worker_type: None,
priority: None,
cost: None,
labels: std::collections::HashMap::new(),
bootstrap_port: None,
tokenizer_path: None,
reasoning_parser: None,
tool_parser: None,
chat_template: None,
};
match self.add_worker(config).await {
Ok(response) => Ok(response.message),
Err(e) => Err(e.error),
}
}
/// Remove a worker from the registry
fn remove_worker(&self, worker_url: &str) {
let _ = self.remove_worker_from_registry(worker_url);
}
/// Get all worker URLs from the registry
fn get_worker_urls(&self) -> Vec<String> {
self.worker_registry.get_all_urls()
}
}
#[async_trait]
impl RouterTrait for RouterManager {
fn as_any(&self) -> &dyn std::any::Any {
......@@ -639,7 +296,6 @@ impl RouterTrait for RouterManager {
body: &ChatCompletionRequest,
_model_id: Option<&str>,
) -> Response {
// Select router based on headers and model
let router = self.select_router_for_request(headers, Some(&body.model));
if let Some(router) = router {
......@@ -662,7 +318,6 @@ impl RouterTrait for RouterManager {
body: &CompletionRequest,
_model_id: Option<&str>,
) -> Response {
// Select router based on headers and model
let router = self.select_router_for_request(headers, Some(&body.model));
if let Some(router) = router {
......@@ -746,7 +401,6 @@ impl RouterTrait for RouterManager {
body: &EmbeddingRequest,
_model_id: Option<&str>,
) -> Response {
// Select router based on headers and model
let router = self.select_router_for_request(headers, Some(&body.model));
if let Some(router) = router {
......
sgl-router/src/routers/worker_initializer.rs deleted 100644 → 0
View file @ 60dbbd08
This diff is collapsed.
sgl-router/src/server.rs
View file @ 97c38239
This diff is collapsed.
sgl-router/src/service_discovery.rs
View file @ 97c38239
This diff is collapsed.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment