feat(gateway): add --remove-unhealthy-workers

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,

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",

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 |
 
 ---
 

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` |
 
 ---
 

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,
         }
     }
 }

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<dyn Worker> = 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<dyn Worker> = 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"
+        );
+    }
 }

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,

model_gateway/src/server.rs

@@ -1044,9 +1044,10 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
         info!("Global health checks disabled via CLI/config; skipping health checker");
         None
     } else {
-        let hc = app_context
-            .worker_registry
-            .start_health_checker(config.router_config.health_check.check_interval_secs);
+        let hc = app_context.worker_registry.start_health_checker(
+            config.router_config.health_check.check_interval_secs,
+            config.router_config.health_check.remove_unhealthy_workers,
+        );
         debug!(
             "Started health checker for workers with {}s interval",
             config.router_config.health_check.check_interval_secs