diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index bb1e8d61c8..07454f3ece 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -424,6 +424,7 @@ struct Router { health_check_interval_secs: u64, health_check_endpoint: String, disable_health_check: bool, + remove_unhealthy_workers: bool, enable_igw: bool, queue_size: usize, queue_timeout_secs: u64, @@ -671,6 +672,7 @@ impl Router { check_interval_secs: self.health_check_interval_secs, endpoint: self.health_check_endpoint.clone(), disable_health_check: self.disable_health_check, + remove_unhealthy_workers: self.remove_unhealthy_workers, }) .tokenizer_cache(config::TokenizerCacheConfig { enable_l0: self.tokenizer_cache_enable_l0, @@ -779,6 +781,7 @@ impl Router { health_check_interval_secs = 60, health_check_endpoint = String::from("/health"), disable_health_check = false, + remove_unhealthy_workers = false, enable_igw = false, queue_size = 100, queue_timeout_secs = 60, @@ -876,6 +879,7 @@ impl Router { health_check_interval_secs: u64, health_check_endpoint: String, disable_health_check: bool, + remove_unhealthy_workers: bool, enable_igw: bool, queue_size: usize, queue_timeout_secs: u64, @@ -982,6 +986,7 @@ impl Router { health_check_interval_secs, health_check_endpoint, disable_health_check, + remove_unhealthy_workers, enable_igw, queue_size, queue_timeout_secs, diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index cb05e971df..b53e9b5ae9 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -92,6 +92,7 @@ class RouterArgs: health_check_interval_secs: int = 60 health_check_endpoint: str = "/health" disable_health_check: bool = False + remove_unhealthy_workers: bool = False # Circuit breaker configuration cb_failure_threshold: int = 10 cb_success_threshold: int = 3 @@ -703,6 +704,12 @@ def add_cli_args( default=RouterArgs.disable_health_check, help="Disable all worker health checks at startup", ) + health_group.add_argument( + f"--{prefix}remove-unhealthy-workers", + action="store_true", + default=RouterArgs.remove_unhealthy_workers, + help="Remove workers from the registry when they are marked unhealthy", + ) # Tokenizer configuration tokenizer_group.add_argument( f"--{prefix}model-path", diff --git a/docs/concepts/reliability/health-checks.md b/docs/concepts/reliability/health-checks.md index c71d5581d9..4b453f767a 100644 --- a/docs/concepts/reliability/health-checks.md +++ b/docs/concepts/reliability/health-checks.md @@ -109,6 +109,7 @@ smg \ | `--health-check-timeout-secs` | `5` | Timeout for each health check request | | `--health-check-endpoint` | `/health` | Endpoint path for health checks | | `--disable-health-check` | `false` | Disable background health checks | +| `--remove-unhealthy-workers` | `false` | Remove workers after being marked unhealthy | --- diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index fc3b88bab8..3b42a78061 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -529,6 +529,7 @@ delay = min(initial_backoff * multiplier^attempt, max_backoff) * (1 + random(0, | `--health-check-interval-secs` | Interval between health checks | `60` | | `--health-check-endpoint` | Health check endpoint path | `/health` | | `--disable-health-check` | Disable all health checks | `false` | +| `--remove-unhealthy-workers` | Remove workers after being marked unhealthy | `false` | --- diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index ebc27985c9..fa4339525a 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -434,6 +434,8 @@ pub struct HealthCheckConfig { pub check_interval_secs: u64, pub endpoint: String, pub disable_health_check: bool, + #[serde(default)] + pub remove_unhealthy_workers: bool, } impl Default for HealthCheckConfig { @@ -445,6 +447,7 @@ impl Default for HealthCheckConfig { check_interval_secs: 60, endpoint: "/health".to_string(), disable_health_check: false, + remove_unhealthy_workers: false, } } } diff --git a/model_gateway/src/core/worker_registry.rs b/model_gateway/src/core/worker_registry.rs index 298bf43621..7bd7c0c835 100644 --- a/model_gateway/src/core/worker_registry.rs +++ b/model_gateway/src/core/worker_registry.rs @@ -221,6 +221,20 @@ impl WorkerRegistry { } } + /// Create a cheap handle that shares all internal Arc state. + /// Unlike `Clone`, this is private so the singleton semantics are preserved. + fn shallow_clone(&self) -> Self { + Self { + workers: self.workers.clone(), + model_index: self.model_index.clone(), + hash_rings: self.hash_rings.clone(), + type_workers: self.type_workers.clone(), + connection_workers: self.connection_workers.clone(), + url_to_id: self.url_to_id.clone(), + mesh_sync: self.mesh_sync.clone(), + } + } + /// Rebuild the hash ring for a model based on current workers in the model index fn rebuild_hash_ring(&self, model_id: &str) { if let Some(workers) = self.model_index.get(model_id) { @@ -660,10 +674,19 @@ impl WorkerRegistry { /// Each worker is checked according to its own `health_config.check_interval_secs`. /// The task sleeps until the next worker is due, so it only wakes when there is /// actual work to do — zero CPU when idle, no polling. - pub(crate) fn start_health_checker(&self, default_interval_secs: u64) -> HealthChecker { + pub(crate) fn start_health_checker( + &self, + default_interval_secs: u64, + remove_unhealthy: bool, + ) -> HealthChecker { let shutdown_notify = Arc::new(tokio::sync::Notify::new()); let shutdown_clone = shutdown_notify.clone(); let workers_ref = self.workers.clone(); + let registry = if remove_unhealthy { + Some(self.shallow_clone()) + } else { + None + }; #[expect( clippy::disallowed_methods, @@ -724,9 +747,25 @@ impl WorkerRegistry { .into_iter() .map(|w| async move { let _ = w.check_health_async().await; + w }) .collect(); - futures::future::join_all(futs).await; + let checked_workers = futures::future::join_all(futs).await; + + // Remove workers that transitioned to unhealthy + if let Some(ref registry) = registry { + for worker in &checked_workers { + if !worker.is_healthy() { + let url = worker.url().to_string(); + tracing::warn!( + worker_url = %url, + "Removing unhealthy worker from registry" + ); + next_check.remove(&url); + registry.remove_by_url(&url); + } + } + } } // Sleep until the earliest deadline or until shutdown is signalled. @@ -890,4 +929,83 @@ mod tests { assert_eq!(llama_workers_after.len(), 1); assert_eq!(llama_workers_after[0].url(), "http://worker2:8080"); } + + #[tokio::test] + async fn test_health_checker_removes_unhealthy_workers() { + use openai_protocol::worker::HealthCheckConfig; + + let registry = WorkerRegistry::new(); + + // Worker pointing at a non-existent URL so health checks fail immediately. + // failure_threshold=1 so it becomes unhealthy after the first failed check. + let mut labels = HashMap::new(); + labels.insert("model_id".to_string(), "test-model".to_string()); + + let worker: Box = Box::new( + BasicWorkerBuilder::new("http://127.0.0.1:1") + .worker_type(WorkerType::Regular) + .labels(labels) + .health_config(HealthCheckConfig { + failure_threshold: 1, + success_threshold: 1, + timeout_secs: 1, + check_interval_secs: 1, + disable_health_check: false, + }) + .build(), + ); + + registry.register(Arc::from(worker)); + assert_eq!(registry.stats().total_workers, 1); + + // Start health checker with remove_unhealthy=true and 1s interval + let _hc = registry.start_health_checker(1, true); + + // Wait for the health check to run and remove the worker + tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; + + assert_eq!( + registry.stats().total_workers, + 0, + "Unhealthy worker should have been removed from the registry" + ); + } + + #[tokio::test] + async fn test_health_checker_keeps_unhealthy_workers_when_disabled() { + use openai_protocol::worker::HealthCheckConfig; + + let registry = WorkerRegistry::new(); + + let mut labels = HashMap::new(); + labels.insert("model_id".to_string(), "test-model".to_string()); + + let worker: Box = Box::new( + BasicWorkerBuilder::new("http://127.0.0.1:1") + .worker_type(WorkerType::Regular) + .labels(labels) + .health_config(HealthCheckConfig { + failure_threshold: 1, + success_threshold: 1, + timeout_secs: 1, + check_interval_secs: 1, + disable_health_check: false, + }) + .build(), + ); + + registry.register(Arc::from(worker)); + assert_eq!(registry.stats().total_workers, 1); + + // Start health checker with remove_unhealthy=false (default behavior) + let _hc = registry.start_health_checker(1, false); + + tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; + + assert_eq!( + registry.stats().total_workers, + 1, + "Worker should remain in the registry when remove_unhealthy is false" + ); + } } diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 6811304b8d..9079fe5bf9 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -406,6 +406,10 @@ struct CliArgs { #[arg(long, default_value_t = false, help_heading = "Health Checks")] disable_health_check: bool, + /// Remove workers from the registry when they are marked unhealthy + #[arg(long, default_value_t = false, help_heading = "Health Checks")] + remove_unhealthy_workers: bool, + // ==================== Tokenizer ==================== /// Model path for loading tokenizer (HuggingFace ID or local path) #[arg(long, alias = "model", help_heading = "Tokenizer")] @@ -1102,6 +1106,7 @@ impl CliArgs { check_interval_secs: self.health_check_interval_secs, endpoint: self.health_check_endpoint.clone(), disable_health_check: self.disable_health_check, + remove_unhealthy_workers: self.remove_unhealthy_workers, }) .tokenizer_cache(TokenizerCacheConfig { enable_l0: self.tokenizer_cache_enable_l0, diff --git a/model_gateway/src/server.rs b/model_gateway/src/server.rs index e17d5f4c54..929db9b607 100644 --- a/model_gateway/src/server.rs +++ b/model_gateway/src/server.rs @@ -1044,9 +1044,10 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box