feat: gate Kubernetes workers until weight update

src/core/worker.rs

@@ -78,6 +78,12 @@ pub trait Worker: Send + Sync + fmt::Debug {
     /// Set the worker's health status
     fn set_healthy(&self, healthy: bool);
 
+    /// Check whether the worker has the current policy weights and may receive traffic.
+    fn is_usable(&self) -> bool;
+
+    /// Allow or prevent routing traffic to this worker.
+    fn set_usable(&self, usable: bool);
+
     /// Perform an async health check on the worker
     async fn check_health_async(&self) -> WorkerResult<()>;
 
@@ -121,9 +127,9 @@ pub trait Worker: Send + Sync + fmt::Debug {
     /// Get the circuit breaker for this worker
     fn circuit_breaker(&self) -> &CircuitBreaker;
 
-    /// Check if the worker is available (healthy + circuit closed/half-open)
+    /// Check if the worker is available for inference traffic.
     fn is_available(&self) -> bool {
-        self.is_healthy() && self.circuit_breaker().can_execute()
+        self.is_healthy() && self.is_usable() && self.circuit_breaker().can_execute()
     }
 
     /// Record the outcome of a request to this worker
@@ -357,6 +363,7 @@ pub struct BasicWorker {
     load_counter: Arc<AtomicUsize>,
     processed_counter: Arc<AtomicUsize>,
     healthy: Arc<AtomicBool>,
+    usable: Arc<AtomicBool>,
     consecutive_failures: Arc<AtomicUsize>,
     consecutive_successes: Arc<AtomicUsize>,
     circuit_breaker: CircuitBreaker,
@@ -369,6 +376,7 @@ impl fmt::Debug for BasicWorker {
         f.debug_struct("BasicWorker")
             .field("metadata", &self.metadata)
             .field("healthy", &self.healthy.load(Ordering::Relaxed))
+            .field("usable", &self.usable.load(Ordering::Relaxed))
             .field("circuit_breaker", &self.circuit_breaker)
             .field("has_grpc_client", &self.grpc_client.is_some())
             .finish()
@@ -398,6 +406,7 @@ impl BasicWorker {
             load_counter: Arc::new(AtomicUsize::new(0)),
             processed_counter: Arc::new(AtomicUsize::new(0)),
             healthy: Arc::new(AtomicBool::new(true)),
+            usable: Arc::new(AtomicBool::new(true)),
             consecutive_failures: Arc::new(AtomicUsize::new(0)),
             consecutive_successes: Arc::new(AtomicUsize::new(0)),
             circuit_breaker: CircuitBreaker::new(),
@@ -410,6 +419,11 @@ impl BasicWorker {
         self
     }
 
+    pub fn with_usable(self, usable: bool) -> Self {
+        self.usable.store(usable, Ordering::Release);
+        self
+    }
+
     pub fn with_health_config(mut self, config: HealthConfig) -> Self {
         self.metadata.health_config = config;
         self
@@ -471,6 +485,14 @@ impl Worker for BasicWorker {
         RouterMetrics::set_worker_health(self.url(), healthy);
     }
 
+    fn is_usable(&self) -> bool {
+        self.usable.load(Ordering::Acquire)
+    }
+
+    fn set_usable(&self, usable: bool) {
+        self.usable.store(usable, Ordering::Release);
+    }
+
     async fn check_health_async(&self) -> WorkerResult<()> {
         use std::time::Duration;
 
@@ -631,6 +653,12 @@ impl DPAwareWorker {
         self.base_worker = self.base_worker.with_labels(labels);
         self
     }
+
+    /// Configure whether this worker may receive inference traffic.
+    pub fn with_usable(mut self, usable: bool) -> Self {
+        self.base_worker = self.base_worker.with_usable(usable);
+        self
+    }
 }
 
 #[async_trait]
@@ -655,6 +683,14 @@ impl Worker for DPAwareWorker {
         self.base_worker.set_healthy(healthy);
     }
 
+    fn is_usable(&self) -> bool {
+        self.base_worker.is_usable()
+    }
+
+    fn set_usable(&self, usable: bool) {
+        self.base_worker.set_usable(usable);
+    }
+
     async fn check_health_async(&self) -> WorkerResult<()> {
         // Delegate to the base worker's health check logic
         self.base_worker.check_health_async().await

src/policies/cache_aware.rs

@@ -317,7 +317,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
             workers
                 .iter()
                 .position(|w| w.url() == tenant_url)
-                .filter(|&idx| workers[idx].is_healthy())
+                .filter(|&idx| workers[idx].is_available())
         } else {
             // Low cache match: use worker with minimum load
             healthy_indices

src/policies/consistent_hash.rs

@@ -246,7 +246,12 @@ impl ConsistentHashPolicy {
 
     /// Update the hash ring when workers change
     fn update_hash_ring(&self, workers: &[Arc<dyn Worker>]) {
-        let worker_urls: Vec<String> = workers.iter().map(|w| w.url().to_string()).collect();
+        let worker_urls: Vec<String> = workers
+            .iter()
+            .filter(|worker| worker.is_available())
+            .map(|worker| worker.url().to_string())
+            .collect();
+        let worker_count = worker_urls.len();
 
         // Check if workers have changed
         {
@@ -280,8 +285,8 @@ impl ConsistentHashPolicy {
 
         info!(
             "Updated consistent hash ring with {} workers and {} virtual nodes",
-            workers.len(),
-            workers.len() as u32 * VIRTUAL_NODES_PER_WORKER
+            worker_count,
+            worker_count as u32 * VIRTUAL_NODES_PER_WORKER
         );
     }
 
@@ -638,7 +643,7 @@ impl LoadBalancingPolicy for ConsistentHashPolicy {
         match selected_idx {
             Some(idx) => {
                 // Verify the worker is healthy
-                if workers[idx].is_healthy() && workers[idx].circuit_breaker().can_execute() {
+                if workers[idx].is_available() {
                     let worker_url = workers[idx].url();
                     debug!(
                         "CONSISTENT_HASH_DEBUG: Selected worker at index {}: {}",

src/policies/mod.rs

@@ -164,7 +164,7 @@ pub(crate) fn get_healthy_worker_indices(workers: &[Arc<dyn Worker>]) -> Vec<usi
     workers
         .iter()
         .enumerate()
-        .filter(|(_, w)| w.is_healthy() && w.circuit_breaker().can_execute())
+        .filter(|(_, w)| w.is_available())
         .map(|(idx, _)| idx)
         .collect()
 }

src/routers/http/router.rs

@@ -25,9 +25,9 @@ use axum::{
     response::{IntoResponse, Response},
     Json,
 };
-use futures_util::StreamExt;
+use futures_util::{future::join_all, StreamExt};
 use reqwest::Client;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio_stream::wrappers::UnboundedReceiverStream;
@@ -357,17 +357,16 @@ impl Router {
                 }
             }
 
-            if healthy_host_count > 0 {
+            if healthy_host_count == unique_hosts_vec.len() {
                 info!(
-                    "{} out of {} unique hosts are healthy (representing {} workers)",
+                    "All {} unique hosts are healthy (representing {} workers)",
                     healthy_host_count,
-                    unique_hosts_vec.len(),
                     worker_urls.len()
                 );
                 return Ok(());
             } else {
                 debug!(
-                   "Waiting for at least 1 of {} unique hosts to become healthy ({} unhealthy: {:?})",
+                    "Waiting for all {} unique hosts to become healthy ({} unhealthy: {:?})",
                     unique_hosts_vec.len(),
                     unhealthy_hosts.len(),
                     unhealthy_hosts
@@ -1199,6 +1198,70 @@ impl Router {
     }
 
     pub async fn add_worker(&self, worker_url: &str) -> Result<String, String> {
+        self.add_worker_with_usability(worker_url, true).await
+    }
+
+    /// Register a Kubernetes-discovered worker without routing requests to it until its first
+    /// successful weight update.
+    pub async fn add_discovered_worker(&self, worker_url: &str) -> Result<String, String> {
+        if self.has_registered_endpoint(worker_url).await {
+            return Ok(format!("Worker endpoint already registered: {worker_url}"));
+        }
+
+        self.add_worker_with_usability(worker_url, false).await
+    }
+
+    async fn has_registered_endpoint(&self, worker_url: &str) -> bool {
+        let Ok(candidate_url) = reqwest::Url::parse(worker_url) else {
+            return false;
+        };
+        let (Some(candidate_host), Some(candidate_port)) =
+            (candidate_url.host_str(), candidate_url.port_or_known_default())
+        else {
+            return false;
+        };
+        let Ok(candidate_addrs) = tokio::net::lookup_host((candidate_host, candidate_port)).await
+        else {
+            return false;
+        };
+        let candidate_ips = candidate_addrs.map(|address| address.ip()).collect::<HashSet<_>>();
+
+        for worker in self.worker_registry.get_all() {
+            let registered_url = strip_dp_rank(worker.url()).trim_end_matches('/');
+            let registered_url = registered_url.strip_suffix("/v1").unwrap_or(registered_url);
+            let Ok(registered_url) = reqwest::Url::parse(registered_url) else {
+                continue;
+            };
+            let (Some(registered_host), Some(registered_port)) = (
+                registered_url.host_str(),
+                registered_url.port_or_known_default(),
+            ) else {
+                continue;
+            };
+            if registered_port != candidate_port {
+                continue;
+            }
+            let Ok(registered_addrs) =
+                tokio::net::lookup_host((registered_host, registered_port)).await
+            else {
+                continue;
+            };
+            if registered_addrs
+                .map(|address| address.ip())
+                .any(|ip| candidate_ips.contains(&ip))
+            {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    async fn add_worker_with_usability(
+        &self,
+        worker_url: &str,
+        usable: bool,
+    ) -> Result<String, String> {
         let start_time = std::time::Instant::now();
         let client = reqwest::Client::builder()
             .timeout(Duration::from_secs(self.worker_startup_timeout_secs))
@@ -1255,7 +1318,8 @@ impl Router {
                                     WorkerType::Regular,
                                 )
                                 .with_circuit_breaker_config(self.circuit_breaker_config.clone())
-                                .with_labels(labels.clone());
+                                .with_labels(labels.clone())
+                                .with_usable(usable);
 
                                 let worker_arc: Arc<dyn Worker> = Arc::new(new_worker);
                                 self.worker_registry.register(worker_arc.clone());
@@ -1296,7 +1360,8 @@ impl Router {
                                     .with_circuit_breaker_config(
                                         self.circuit_breaker_config.clone(),
                                     )
-                                    .with_labels(labels);
+                                    .with_labels(labels)
+                                    .with_usable(usable);
 
                             let worker_arc = Arc::new(new_worker);
                             self.worker_registry.register(worker_arc.clone());
@@ -1615,21 +1680,30 @@ impl RouterTrait for Router {
 
     async fn health(&self, _req: Request<Body>) -> Response {
         let workers = self.worker_registry.get_all();
-        let unhealthy_servers: Vec<_> = workers
-            .iter()
-            .filter(|w| !w.is_healthy())
-            .map(|w| w.url().to_string())
-            .collect();
+        let health = join_all(workers.iter().map(|worker| worker.check_health_async())).await;
+        for (worker, result) in workers.iter().zip(health) {
+            worker.set_healthy(result.is_ok());
+        }
 
-        if unhealthy_servers.is_empty() {
-            (StatusCode::OK, "All servers healthy").into_response()
-        } else {
-            (
-                StatusCode::SERVICE_UNAVAILABLE,
-                format!("Unhealthy servers: {:?}", unhealthy_servers),
+        let available_workers = workers.iter().filter(|worker| worker.is_available()).count();
+        if available_workers > 0 {
+            return (
+                StatusCode::OK,
+                format!("{available_workers} workers available"),
             )
-                .into_response()
+                .into_response();
         }
+
+        let unavailable_servers: Vec<_> = workers
+            .iter()
+            .filter(|worker| !worker.is_available())
+            .map(|worker| worker.url().to_string())
+            .collect();
+        (
+            StatusCode::SERVICE_UNAVAILABLE,
+            format!("No available workers: {unavailable_servers:?}"),
+        )
+            .into_response()
     }
 
     async fn health_generate(&self, req: Request<Body>) -> Response {
@@ -1839,9 +1913,9 @@ impl RouterTrait for Router {
     }
 
     fn readiness(&self) -> Response {
-        // Regular router is ready if it has at least one healthy worker
+        // Regular router is ready if it has at least one routable worker.
         let workers = self.worker_registry.get_all();
-        let healthy_count = workers.iter().filter(|w| w.is_healthy()).count();
+        let healthy_count = workers.iter().filter(|w| w.is_available()).count();
         let total_workers = workers.len();
 
         if healthy_count > 0 {
@@ -2648,12 +2722,11 @@ mod tests {
     #[tokio::test]
     async fn test_wait_for_healthy_workers_partial_health() {
         // One healthy server + one unreachable URL.
-        // The new behaviour succeeds when at least one host is healthy.
         let (healthy_url, _handle) = start_healthy_mock_server().await;
         let unreachable_url = "http://127.0.0.1:1".to_string(); // port 1 is unreachable
 
-        let result = Router::wait_for_healthy_workers(&[healthy_url, unreachable_url], 5, 1).await;
-        assert!(result.is_ok());
+        let result = Router::wait_for_healthy_workers(&[healthy_url, unreachable_url], 2, 1).await;
+        assert!(result.is_err());
     }
 
     /// Helper: start a mock server that returns 503 on /health for a given
@@ -2729,14 +2802,13 @@ mod tests {
         assert_eq!(resp.status().as_u16(), 503);
 
         // ── Step 1: wait_for_healthy_workers (mirrors PDRouter::new startup) ──
-        // This succeeds because the healthy worker responds immediately,
-        // even though the delayed worker is still returning 503.
+        // Startup waits for the delayed worker as well as the immediately healthy worker.
         let result =
             Router::wait_for_healthy_workers(&[healthy_url.clone(), delayed_url.clone()], 10, 1)
                 .await;
         assert!(
             result.is_ok(),
-            "Startup should succeed with at least one healthy worker"
+            "Startup should succeed once all workers are healthy"
         );
 
         // ── Step 2: register workers in the registry (mirrors PDRouter::new) ──