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

sgl-router/src/core/mod.rs

@@ -22,7 +22,7 @@ pub use error::{WorkerError, WorkerResult};
 pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor};
 pub use worker::{
     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_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats};

sgl-router/src/core/worker_builder.rs

-use super::circuit_breaker::CircuitBreakerConfig;
-use super::worker::{BasicWorker, ConnectionMode, DPAwareWorker, HealthConfig, WorkerType};
+use super::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
+use super::worker::{
+    BasicWorker, ConnectionMode, DPAwareWorker, HealthConfig, WorkerMetadata, WorkerType,
+};
 use crate::grpc::client::SglangSchedulerClient;
 use std::collections::HashMap;
 
@@ -88,23 +90,30 @@ impl BasicWorkerBuilder {
 
     /// Build the BasicWorker instance
     pub fn build(self) -> BasicWorker {
-        // Use the existing constructor methods for now
-        let mut worker =
-            BasicWorker::with_connection_mode(self.url, self.worker_type, self.connection_mode);
-
-        // Apply optional configurations using existing methods
-        if !self.labels.is_empty() {
-            worker = worker.with_labels(self.labels);
-        }
-
-        worker = worker.with_health_config(self.health_config);
-        worker = worker.with_circuit_breaker_config(self.circuit_breaker_config);
+        use std::sync::{
+            atomic::{AtomicBool, AtomicUsize},
+            Arc,
+        };
+        use tokio::sync::Mutex;
+
+        let metadata = WorkerMetadata {
+            url: self.url.clone(),
+            worker_type: self.worker_type,
+            connection_mode: self.connection_mode,
+            labels: self.labels,
+            health_config: self.health_config,
+        };
 
-        if let Some(client) = self.grpc_client {
-            worker = worker.with_grpc_client(client);
+        BasicWorker {
+            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
     }
 }
 

sgl-router/src/routers/http/pd_router.rs

@@ -4,7 +4,7 @@ use super::pd_types::{api_path, PDRouterError};
 use crate::config::types::RetryConfig;
 use crate::core::{
     is_retryable_status, BasicWorkerBuilder, CircuitBreakerConfig, HealthConfig, RetryExecutor,
-    Worker, WorkerFactory, WorkerLoadGuard, WorkerRegistry, WorkerType,
+    Worker, WorkerLoadGuard, WorkerRegistry, WorkerType,
 };
 use crate::metrics::RouterMetrics;
 use crate::policies::{LoadBalancingPolicy, PolicyRegistry};
@@ -220,13 +220,12 @@ impl PDRouter {
 
         // 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
-        let worker = WorkerFactory::create_prefill_with_config(
-            url.clone(),
-            bootstrap_port,
-            self.circuit_breaker_config.clone(),
-        );
+        let worker = BasicWorkerBuilder::new(url.clone())
+            .worker_type(WorkerType::Prefill { bootstrap_port })
+            .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
         self.worker_registry.register(worker_arc.clone());
@@ -261,12 +260,12 @@ impl PDRouter {
 
         // 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
-        let worker = WorkerFactory::create_decode_with_config(
-            url.clone(),
-            self.circuit_breaker_config.clone(),
-        );
+        let worker = BasicWorkerBuilder::new(url.clone())
+            .worker_type(WorkerType::Decode)
+            .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
         self.worker_registry.register(worker_arc.clone());