Unverified Commit 4f2055ad authored by Simo Lin's avatar Simo Lin Committed by GitHub
Browse files

[router] refactor worker to builder pattern 4/n (#10650)

parent 616a3e20
...@@ -22,7 +22,7 @@ pub use error::{WorkerError, WorkerResult}; ...@@ -22,7 +22,7 @@ pub use error::{WorkerError, WorkerResult};
pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor}; pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor};
pub use worker::{ pub use worker::{
start_health_checker, BasicWorker, ConnectionMode, DPAwareWorker, HealthChecker, HealthConfig, start_health_checker, BasicWorker, ConnectionMode, DPAwareWorker, HealthChecker, HealthConfig,
Worker, WorkerCollection, WorkerFactory, WorkerLoadGuard, WorkerType, Worker, WorkerFactory, WorkerLoadGuard, WorkerType,
}; };
pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder}; pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder};
pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats}; pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats};
This diff is collapsed.
use super::circuit_breaker::CircuitBreakerConfig; use super::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
use super::worker::{BasicWorker, ConnectionMode, DPAwareWorker, HealthConfig, WorkerType}; use super::worker::{
BasicWorker, ConnectionMode, DPAwareWorker, HealthConfig, WorkerMetadata, WorkerType,
};
use crate::grpc::client::SglangSchedulerClient; use crate::grpc::client::SglangSchedulerClient;
use std::collections::HashMap; use std::collections::HashMap;
...@@ -88,23 +90,30 @@ impl BasicWorkerBuilder { ...@@ -88,23 +90,30 @@ impl BasicWorkerBuilder {
/// Build the BasicWorker instance /// Build the BasicWorker instance
pub fn build(self) -> BasicWorker { pub fn build(self) -> BasicWorker {
// Use the existing constructor methods for now use std::sync::{
let mut worker = atomic::{AtomicBool, AtomicUsize},
BasicWorker::with_connection_mode(self.url, self.worker_type, self.connection_mode); Arc,
};
// Apply optional configurations using existing methods use tokio::sync::Mutex;
if !self.labels.is_empty() {
worker = worker.with_labels(self.labels); let metadata = WorkerMetadata {
} url: self.url.clone(),
worker_type: self.worker_type,
worker = worker.with_health_config(self.health_config); connection_mode: self.connection_mode,
worker = worker.with_circuit_breaker_config(self.circuit_breaker_config); labels: self.labels,
health_config: self.health_config,
};
if let Some(client) = self.grpc_client { BasicWorker {
worker = worker.with_grpc_client(client); metadata,
load_counter: Arc::new(AtomicUsize::new(0)),
processed_counter: Arc::new(AtomicUsize::new(0)),
healthy: Arc::new(AtomicBool::new(true)),
consecutive_failures: Arc::new(AtomicUsize::new(0)),
consecutive_successes: Arc::new(AtomicUsize::new(0)),
circuit_breaker: CircuitBreaker::with_config(self.circuit_breaker_config),
grpc_client: self.grpc_client.map(|client| Arc::new(Mutex::new(client))),
} }
worker
} }
} }
......
...@@ -4,7 +4,7 @@ use super::pd_types::{api_path, PDRouterError}; ...@@ -4,7 +4,7 @@ use super::pd_types::{api_path, PDRouterError};
use crate::config::types::RetryConfig; use crate::config::types::RetryConfig;
use crate::core::{ use crate::core::{
is_retryable_status, BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, RetryExecutor, is_retryable_status, BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, RetryExecutor,
Worker, WorkerFactory, WorkerLoadGuard, WorkerRegistry, WorkerType, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType,
}; };
use crate::metrics::RouterMetrics; use crate::metrics::RouterMetrics;
use crate::policies::{LoadBalancingPolicy, PolicyRegistry}; use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
...@@ -220,13 +220,12 @@ impl PDRouter { ...@@ -220,13 +220,12 @@ impl PDRouter {
// Create Worker for the new prefill server with circuit breaker configuration // Create Worker for the new prefill server with circuit breaker configuration
// TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint
let worker = WorkerFactory::create_prefill_with_config( let worker = BasicWorkerBuilder::new(url.clone())
url.clone(), .worker_type(WorkerType::Prefill { bootstrap_port })
bootstrap_port, .circuit_breaker_config(self.circuit_breaker_config.clone())
self.circuit_breaker_config.clone(), .build();
);
let worker_arc: Arc<dyn Worker> = Arc::from(worker); let worker_arc: Arc<dyn Worker> = Arc::new(worker);
// Register the worker in the registry // Register the worker in the registry
self.worker_registry.register(worker_arc.clone()); self.worker_registry.register(worker_arc.clone());
...@@ -261,12 +260,12 @@ impl PDRouter { ...@@ -261,12 +260,12 @@ impl PDRouter {
// Create Worker for the new decode server with circuit breaker configuration // Create Worker for the new decode server with circuit breaker configuration
// TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint
let worker = WorkerFactory::create_decode_with_config( let worker = BasicWorkerBuilder::new(url.clone())
url.clone(), .worker_type(WorkerType::Decode)
self.circuit_breaker_config.clone(), .circuit_breaker_config(self.circuit_breaker_config.clone())
); .build();
let worker_arc: Arc<dyn Worker> = Arc::from(worker); let worker_arc: Arc<dyn Worker> = Arc::new(worker);
// Register the worker in the registry // Register the worker in the registry
self.worker_registry.register(worker_arc.clone()); self.worker_registry.register(worker_arc.clone());
......
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