diff --git a/model_gateway/src/policies/bucket.rs b/model_gateway/src/policies/bucket.rs index f16cabc9f0..b62a10c7ab 100644 --- a/model_gateway/src/policies/bucket.rs +++ b/model_gateway/src/policies/bucket.rs @@ -562,9 +562,18 @@ impl Bucket { #[cfg(test)] mod tests { + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{BasicWorkerBuilder, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + #[tokio::test] async fn test_load_balancing_conditions() { // Test 1: Basic load balancing trigger @@ -579,18 +588,21 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), ]; @@ -793,18 +805,21 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), ]; @@ -996,18 +1011,21 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), ]; @@ -1104,18 +1122,21 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), ]; diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index 711cb73e6e..cdb280b66e 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -918,10 +918,18 @@ impl Default for CacheAwarePolicy { #[cfg(test)] mod tests { use kv_index::{compute_content_hash, SequenceHash, StoredBlock, WorkerBlockMap}; + use openai_protocol::worker::HealthCheckConfig; use super::*; use crate::worker::{BasicWorkerBuilder, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + #[test] fn test_cache_aware_with_balanced_load() { // Create policy without eviction thread for testing @@ -935,12 +943,14 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), ]; @@ -997,9 +1007,11 @@ mod tests { let worker1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let worker2 = BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); // Create significant load imbalance @@ -1033,11 +1045,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -1097,6 +1111,7 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), )]; @@ -1268,6 +1283,7 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), )]; @@ -1331,11 +1347,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -1370,6 +1388,7 @@ mod tests { let workers: Vec> = vec![Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), )]; policy.init_workers(&workers); @@ -1394,9 +1413,11 @@ mod tests { let w1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let w2 = BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); // Give w1 higher load @@ -1440,11 +1461,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -1489,6 +1512,7 @@ mod tests { let workers: Vec> = vec![Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), )]; @@ -1506,11 +1530,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -1574,11 +1600,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -1610,9 +1638,11 @@ mod tests { let w1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let w2 = BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); // Give w1 higher load so min-load picks w2 for _ in 0..3 { @@ -1647,9 +1677,11 @@ mod tests { let w1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let w2 = BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); for _ in 0..3 { w1.increment_load(); @@ -1683,11 +1715,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -1743,11 +1777,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -1799,9 +1835,11 @@ mod tests { let w1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let w2 = BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); // Create heavy imbalance: w1 has 20 load, w2 has 0 @@ -1839,11 +1877,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; diff --git a/model_gateway/src/policies/consistent_hashing.rs b/model_gateway/src/policies/consistent_hashing.rs index d6cf0100da..d50e5a385a 100644 --- a/model_gateway/src/policies/consistent_hashing.rs +++ b/model_gateway/src/policies/consistent_hashing.rs @@ -198,9 +198,18 @@ impl LoadBalancingPolicy for ConsistentHashingPolicy { mod tests { use std::collections::HashMap; + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{BasicWorkerBuilder, HashRing, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + fn headers_with_routing_key(key: &str) -> http::HeaderMap { let mut headers = http::HeaderMap::new(); headers.insert("x-smg-routing-key", key.parse().unwrap()); @@ -219,6 +228,7 @@ mod tests { Arc::new( BasicWorkerBuilder::new(*url) .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ) as Arc }) diff --git a/model_gateway/src/policies/manual.rs b/model_gateway/src/policies/manual.rs index c5a50d0473..a8186c43f3 100644 --- a/model_gateway/src/policies/manual.rs +++ b/model_gateway/src/policies/manual.rs @@ -301,15 +301,25 @@ fn min_group_select(workers: &[Arc], healthy_indices: &[usize]) -> u mod tests { use std::collections::HashMap; + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{BasicWorkerBuilder, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + fn create_workers(urls: &[&str]) -> Vec> { urls.iter() .map(|url| { Arc::new( BasicWorkerBuilder::new(*url) .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ) as Arc }) diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index f25020d572..872c526a0b 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -182,9 +182,18 @@ pub struct SelectWorkerInfo<'a> { #[cfg(test)] mod tests { + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{BasicWorkerBuilder, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + #[test] fn test_get_healthy_worker_indices() { let workers: Vec> = vec![ @@ -192,18 +201,21 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key2") + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ), ]; diff --git a/model_gateway/src/policies/power_of_two.rs b/model_gateway/src/policies/power_of_two.rs index a6c76c2873..0e545ba15d 100644 --- a/model_gateway/src/policies/power_of_two.rs +++ b/model_gateway/src/policies/power_of_two.rs @@ -132,11 +132,18 @@ impl Default for PowerOfTwoPolicy { #[cfg(test)] mod tests { - use openai_protocol::worker::SchedulerLoadSnapshot; + use openai_protocol::worker::{HealthCheckConfig, SchedulerLoadSnapshot}; use super::*; use crate::worker::{BasicWorkerBuilder, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + /// Create a `WorkerLoadResponse` with a single DP rank at the given token_usage ratio. fn make_load(token_usage: f64) -> WorkerLoadResponse { WorkerLoadResponse { @@ -163,12 +170,15 @@ mod tests { let policy = PowerOfTwoPolicy::new(); let worker1 = BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let worker2 = BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let worker3 = BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); // Set different loads @@ -204,11 +214,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -240,6 +252,7 @@ mod tests { let workers: Vec> = vec![Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), )]; @@ -263,11 +276,13 @@ mod tests { // 2. Create Worker A: Idle (0 reqs), but has high token usage in cache let worker_a = BasicWorkerBuilder::new("http://worker_a:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); // 3. Create Worker B: Busy (5 reqs), but missing from cache let worker_b = BasicWorkerBuilder::new("http://worker_b:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); // Manually increment load on Worker B to simulate active requests @@ -321,6 +336,7 @@ mod tests { let create_worker = |url: &str, reqs: usize| { let w = BasicWorkerBuilder::new(url) .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); for _ in 0..reqs { w.increment_load(); diff --git a/model_gateway/src/policies/prefix_hash.rs b/model_gateway/src/policies/prefix_hash.rs index bf25f57d27..a5f6d4c7ab 100644 --- a/model_gateway/src/policies/prefix_hash.rs +++ b/model_gateway/src/policies/prefix_hash.rs @@ -239,6 +239,8 @@ impl LoadBalancingPolicy for PrefixHashPolicy { #[cfg(test)] mod tests { + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{BasicWorkerBuilder, HashRing, WorkerType}; @@ -248,6 +250,10 @@ mod tests { Arc::new( BasicWorkerBuilder::new(*url) .worker_type(WorkerType::Regular) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) .build(), ) as Arc }) diff --git a/model_gateway/src/policies/random.rs b/model_gateway/src/policies/random.rs index 79cd990889..2ac49dfe9d 100644 --- a/model_gateway/src/policies/random.rs +++ b/model_gateway/src/policies/random.rs @@ -50,9 +50,18 @@ impl LoadBalancingPolicy for RandomPolicy { mod tests { use std::collections::HashMap; + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{BasicWorkerBuilder, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + #[test] fn test_random_selection() { let policy = RandomPolicy::new(); @@ -60,16 +69,19 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -93,11 +105,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -120,6 +134,7 @@ mod tests { let workers: Vec> = vec![Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), )]; diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 08913a5436..dfcb3650ae 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -689,6 +689,7 @@ impl std::fmt::Debug for PolicyRegistry { mod tests { use std::sync::Arc; + use openai_protocol::worker::HealthCheckConfig; use smg_mesh::{MeshSyncManager, StateStores}; use super::*; @@ -697,6 +698,13 @@ mod tests { worker::{BasicWorkerBuilder, Worker, WorkerType, UNKNOWN_MODEL_ID}, }; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + #[test] fn test_policy_registry_basic() { let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); @@ -782,6 +790,7 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), )]; @@ -826,12 +835,14 @@ mod tests { BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ); let worker2: Arc = Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(), ); let workers = vec![worker1, worker2]; diff --git a/model_gateway/src/policies/round_robin.rs b/model_gateway/src/policies/round_robin.rs index 1aed1c0680..2ab0c21946 100644 --- a/model_gateway/src/policies/round_robin.rs +++ b/model_gateway/src/policies/round_robin.rs @@ -58,9 +58,18 @@ impl LoadBalancingPolicy for RoundRobinPolicy { #[cfg(test)] mod tests { + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{BasicWorkerBuilder, WorkerType}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + #[test] fn test_round_robin_selection() { let policy = RoundRobinPolicy::new(); @@ -68,16 +77,19 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -98,16 +110,19 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w3:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; @@ -130,11 +145,13 @@ mod tests { Arc::new( BasicWorkerBuilder::new("http://w1:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), Arc::new( BasicWorkerBuilder::new("http://w2:8000") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(), ), ]; diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 614f6bb3a8..cc4547201e 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -783,9 +783,18 @@ impl RouterTrait for Router { #[cfg(test)] mod tests { + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::{config::types::PolicyConfig, worker::BasicWorkerBuilder}; + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + fn create_test_regular_router() -> Router { // Create registries let worker_registry = Arc::new(WorkerRegistry::new()); @@ -794,9 +803,11 @@ mod tests { // Register test workers let worker1 = BasicWorkerBuilder::new("http://worker1:8080") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); let worker2 = BasicWorkerBuilder::new("http://worker2:8080") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); worker_registry.register_or_replace(Arc::new(worker1)); worker_registry.register_or_replace(Arc::new(worker2)); diff --git a/model_gateway/src/worker/builder.rs b/model_gateway/src/worker/builder.rs index 94c28a9477..ec1241dee7 100644 --- a/model_gateway/src/worker/builder.rs +++ b/model_gateway/src/worker/builder.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use arc_swap::ArcSwap; use openai_protocol::{ model_card::ModelCard, - worker::{HealthCheckConfig, WorkerModels, WorkerSpec}, + worker::{HealthCheckConfig, WorkerModels, WorkerSpec, WorkerStatus}, }; use super::{ @@ -32,6 +32,11 @@ pub struct BasicWorkerBuilder { http_client: Option, /// Resolved resilience config (if not set, defaults are used). resilience: Option, + /// Initial lifecycle status. If unset, defaults to `Pending` for + /// health-checked workers and `Ready` for `disable_health_check == true`. + /// Callers replacing an existing worker (e.g. metadata updates) should + /// pass the old worker's status to avoid kicking it back to Pending. + initial_status: Option, } impl BasicWorkerBuilder { @@ -45,6 +50,7 @@ impl BasicWorkerBuilder { grpc_client: None, http_client: None, resilience: None, + initial_status: None, } } @@ -58,6 +64,7 @@ impl BasicWorkerBuilder { grpc_client: None, http_client: None, resilience: None, + initial_status: None, } } @@ -73,6 +80,7 @@ impl BasicWorkerBuilder { grpc_client: None, http_client: None, resilience: None, + initial_status: None, } } @@ -127,6 +135,19 @@ impl BasicWorkerBuilder { self } + /// Override the initial lifecycle status. + /// + /// By default, `build()` chooses `Pending` for health-checked workers and + /// `Ready` for workers with `disable_health_check == true`. Callers that + /// replace an existing worker (e.g. metadata-only updates via + /// `register_or_replace`) should pass the old worker's status here to + /// avoid kicking a healthy worker back to Pending and causing avoidable + /// 503s while it re-proves itself. + pub fn status(mut self, status: WorkerStatus) -> Self { + self.initial_status = Some(status); + self + } + /// Set health check endpoint path (internal-only, from router config). pub fn health_endpoint(mut self, endpoint: impl Into) -> Self { self.health_endpoint = endpoint.into(); @@ -239,10 +260,19 @@ impl BasicWorkerBuilder { None => OnceCell::new(), }); - // Workers start Ready (routable). PR 6b will change this to Pending - // for health-checked workers. - let initial_status = openai_protocol::worker::WorkerStatus::Ready; - Metrics::set_worker_health(&metadata.spec.url, true); + // Caller can override the initial status (e.g. when replacing an + // existing worker, to preserve its prior status). Otherwise: + // - workers with health checks disabled start Ready (routable) + // - workers with health checks enabled start Pending (not routable + // until the health checker promotes them after success_threshold) + let initial_status = + self.initial_status + .unwrap_or(if metadata.health_config.disable_health_check { + WorkerStatus::Ready + } else { + WorkerStatus::Pending + }); + Metrics::set_worker_health(&metadata.spec.url, initial_status == WorkerStatus::Ready); let http_client = self.http_client.unwrap_or_else(|| { reqwest::Client::builder() @@ -263,6 +293,7 @@ impl BasicWorkerBuilder { status: Arc::new(AtomicU8::new(initial_status as u8)), consecutive_failures: Arc::new(AtomicUsize::new(0)), consecutive_successes: Arc::new(AtomicUsize::new(0)), + total_pending_probes: Arc::new(AtomicUsize::new(0)), circuit_breaker: CircuitBreaker::with_config_and_label( self.circuit_breaker_config, metadata.spec.url.clone(), @@ -329,7 +360,9 @@ mod tests { assert_eq!(worker.url(), "http://localhost:8080"); assert_eq!(worker.worker_type(), &WorkerType::Regular); assert_eq!(worker.connection_mode(), &ConnectionMode::Http); - assert!(worker.is_healthy()); + // Health-checked workers start Pending (not routable until health checker promotes) + assert!(!worker.is_healthy()); + assert_eq!(worker.status(), WorkerStatus::Pending); } #[test] @@ -341,7 +374,7 @@ mod tests { assert_eq!(worker.url(), "http://localhost:8080"); assert_eq!(worker.worker_type(), &WorkerType::Decode); assert_eq!(worker.connection_mode(), &ConnectionMode::Http); - assert!(worker.is_healthy()); + assert!(!worker.is_healthy()); } #[test] diff --git a/model_gateway/src/worker/registry.rs b/model_gateway/src/worker/registry.rs index 06d983b4e2..5233a61da2 100644 --- a/model_gateway/src/worker/registry.rs +++ b/model_gateway/src/worker/registry.rs @@ -17,6 +17,7 @@ use std::{ }; use dashmap::{mapref::entry::Entry, DashMap}; +use openai_protocol::worker::WorkerStatus; use parking_lot::RwLock; use smg_mesh::OptionalMeshSyncManager; use tokio::sync::broadcast; @@ -991,13 +992,13 @@ impl WorkerRegistry { .collect(); let checked_workers = futures::future::join_all(futs).await; - // Only remove workers whose health check actually failed - // this tick. Workers that are unhealthy but passing checks - // (e.g. mesh-synced, pre-activation) are recovering — leave - // them alone until they either become healthy or truly fail. + // Only remove Failed workers — not Pending (still starting) + // or NotReady (may recover). Failed is terminal, so we + // don't gate on *failed (a worker can reach Failed via + // max_pending_probes on a probe that returned Ok(false)). if let Some(ref job_queue) = job_queue { - for (worker, failed) in &checked_workers { - if !worker.is_healthy() && *failed { + for (worker, _failed) in &checked_workers { + if worker.status() == WorkerStatus::Failed { let url = worker.url().to_string(); tracing::warn!( worker_url = %url, @@ -1229,17 +1230,19 @@ mod tests { } #[tokio::test] - async fn test_health_checker_removes_unhealthy_workers() { + async fn test_pending_worker_stays_pending_under_failed_probes() { use openai_protocol::worker::HealthCheckConfig; + // Pending workers do NOT transition to NotReady on failed probes — + // they stay Pending until either `success_threshold` consecutive + // successes (→Ready) or `max_pending_probes` total attempts (→Failed). + // This test verifies the Pending state is sticky during early failures. 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( + let worker: Arc = Arc::new( BasicWorkerBuilder::new("http://127.0.0.1:1") .worker_type(WorkerType::Regular) .labels(labels) @@ -1252,29 +1255,73 @@ mod tests { }) .build(), ); + assert_eq!(worker.status(), WorkerStatus::Pending); - registry.register(Arc::from(worker)).unwrap(); - assert_eq!(registry.stats().total_workers, 1); - - // Start health checker with remove_unhealthy=true but no job queue. - // Without a job queue the removal job can't be submitted, but the worker - // should still be marked unhealthy by the health check. + registry.register(worker.clone()).unwrap(); let _hc = registry.start_health_checker(1, true, None); - // Wait for the health check to run and mark the worker unhealthy + // Wait for a few probe attempts. With check_interval=1 and a + // non-existent URL, probes will fail. With failure_threshold=1, + // max_pending_probes is 10 — so after ~3s we expect Pending still. tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; + let stats = registry.stats(); + assert_eq!(stats.total_workers, 1); + assert_eq!(stats.healthy_workers, 0, "Pending workers are not healthy"); + // Worker should still be Pending — not yet Failed (max_pending_probes + // is 10, we've only run ~3 probes) and not NotReady (Pending workers + // don't transition to NotReady on failure). + let after = registry.get_by_url("http://127.0.0.1:1").unwrap(); + assert_eq!( + after.status(), + WorkerStatus::Pending, + "Pending should be sticky under early failures" + ); + } + + #[test] + fn test_failed_worker_can_be_removed_from_registry() { + // Drives the removal-on-Failed logic without needing the async + // health checker loop or a JobQueue. Verifies that the registry + // accepts removal of a Failed worker via remove() and that the + // worker is gone afterward. + use openai_protocol::worker::HealthCheckConfig; + + let registry = WorkerRegistry::new(); + + let worker: Arc = Arc::new( + BasicWorkerBuilder::new("http://failed:8080") + .worker_type(WorkerType::Regular) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..HealthCheckConfig::default() + }) + .build(), + ); + let worker_id = registry.register(worker.clone()).unwrap(); + assert_eq!(worker.status(), WorkerStatus::Ready); + + // Simulate the state machine reaching Failed. + worker.set_status(WorkerStatus::Failed); assert_eq!( - registry.stats().healthy_workers, - 0, - "Worker should be marked unhealthy after failed health checks" + registry.get(&worker_id).unwrap().status(), + WorkerStatus::Failed ); + + // The health checker's removal loop only acts on workers in Failed state. + // Verify we can remove a Failed worker through the standard remove() API + // (which is what Job::RemoveWorker eventually invokes). + assert!(registry.remove(&worker_id).is_some()); + assert!(registry.get(&worker_id).is_none()); + assert_eq!(registry.stats().total_workers, 0); } #[tokio::test] - async fn test_health_checker_keeps_unhealthy_workers_when_disabled() { + async fn test_health_checker_keeps_workers_when_remove_unhealthy_disabled() { use openai_protocol::worker::HealthCheckConfig; + // When --remove-unhealthy-workers is false, workers in any state + // (Pending, NotReady, Failed) should remain in the registry. let registry = WorkerRegistry::new(); let mut labels = HashMap::new(); @@ -1297,7 +1344,7 @@ mod tests { registry.register(Arc::from(worker)).unwrap(); assert_eq!(registry.stats().total_workers, 1); - // Start health checker with remove_unhealthy=false (default behavior) + // Start health checker with remove_unhealthy=false let _hc = registry.start_health_checker(1, false, None); tokio::time::sleep(tokio::time::Duration::from_secs(3)).await; @@ -1421,6 +1468,65 @@ mod tests { assert_eq!(registry.get_by_model("o4-mini").len(), 1); } + #[test] + fn test_builder_status_override_on_replace() { + // Regression test: metadata-only updates must not kick a healthy + // worker back to Pending. The builder exposes a `.status()` setter + // so callers (e.g. UpdateWorkerPropertiesStep) can pass the old + // worker's status when constructing the replacement. + let registry = WorkerRegistry::new(); + + // First worker starts Pending (health checks enabled by default), + // then gets promoted to Ready (simulating what the health checker does). + let first: Arc = Arc::new( + BasicWorkerBuilder::new("http://worker:8080") + .worker_type(WorkerType::Regular) + .model(ModelCard::new("llama-3")) + .build(), + ); + assert_eq!(first.status(), WorkerStatus::Pending); + first.set_status(WorkerStatus::Ready); + + let first_id = registry.register(first.clone()).unwrap(); + assert_eq!( + registry.get(&first_id).unwrap().status(), + WorkerStatus::Ready + ); + + // Caller (e.g. UpdateWorkerPropertiesStep) reads the old status and + // passes it to the builder. The builder honors the override instead + // of defaulting to Pending. + let preserved_status = first.status(); + let second: Arc = Arc::new( + BasicWorkerBuilder::new("http://worker:8080") + .worker_type(WorkerType::Regular) + .model(ModelCard::new("llama-3")) + .priority(99) + .status(preserved_status) + .build(), + ); + assert_eq!( + second.status(), + WorkerStatus::Ready, + "builder must honor explicit status override" + ); + + assert!(registry.replace(&first_id, second)); + + let after = registry.get(&first_id).unwrap(); + assert_eq!(after.status(), WorkerStatus::Ready); + assert_eq!(after.priority(), 99, "new priority should be applied"); + } + + #[test] + fn test_builder_default_status_is_pending() { + // Without an explicit override, health-checked workers start Pending. + let worker = BasicWorkerBuilder::new("http://worker:8080") + .worker_type(WorkerType::Regular) + .build(); + assert_eq!(worker.status(), WorkerStatus::Pending); + } + #[test] fn test_register_rejects_duplicate_url() { let registry = WorkerRegistry::new(); diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index df6778b902..d2409cfd79 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -527,6 +527,10 @@ pub struct BasicWorker { pub status: Arc, pub consecutive_failures: Arc, pub consecutive_successes: Arc, + /// Total health check probes attempted while in Pending state. + /// Unlike consecutive counters, this counts ALL attempts (not just consecutive). + /// Used to detect misconfigured workers stuck in Pending. + pub total_pending_probes: Arc, pub circuit_breaker: CircuitBreaker, /// Lazily initialized gRPC client for gRPC workers. /// Uses OnceCell for lock-free reads after initialization. @@ -591,46 +595,118 @@ impl Worker for BasicWorker { async fn check_health_async(&self) -> WorkerResult<()> { if self.metadata.health_config.disable_health_check { - if !self.is_healthy() { - self.set_healthy(true); + if self.status() != WorkerStatus::Ready { + self.set_status(WorkerStatus::Ready); } return Ok(()); } - let health_result = match &self.metadata.spec.connection_mode { - ConnectionMode::Http => self.http_health_check().await?, - ConnectionMode::Grpc => self.grpc_health_check().await?, - }; - - // Get worker type label for metrics + let current_status = self.status(); + let health_config = &self.metadata.health_config; let worker_type_str = self.metadata.spec.worker_type.as_metric_label(); + // Failed is terminal — skip the probe entirely. The health checker + // loop will eventually remove the worker (if --remove-unhealthy-workers + // is set) or it stays Failed until explicit re-registration. + if current_status == WorkerStatus::Failed { + return Err(WorkerError::HealthCheckFailed { + url: self.metadata.spec.url.clone(), + reason: "worker is in Failed state".to_string(), + }); + } + + // Track total probes in Pending for startup timeout detection + if current_status == WorkerStatus::Pending { + self.total_pending_probes.fetch_add(1, Ordering::Relaxed); + } + + // Transport-level errors (e.g. gRPC connect failure, TLS handshake + // failure, DNS resolution failure) are treated as failed probes + // rather than short-circuiting, so the Pending timeout and + // NotReady→Failed paths can observe them. The error is logged with + // the worker URL so operators can debug why a worker won't come Ready. + let health_result = match &self.metadata.spec.connection_mode { + ConnectionMode::Http => self.http_health_check().await, + ConnectionMode::Grpc => self.grpc_health_check().await, + } + .unwrap_or_else(|err| { + tracing::warn!( + worker_url = %self.metadata.spec.url, + error = %err, + "Health check probe transport error (treating as failed probe)" + ); + false + }); + if health_result { self.consecutive_failures.store(0, Ordering::Release); let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1; - - // Record health check success metric Metrics::record_worker_health_check(worker_type_str, metrics_labels::CB_SUCCESS); - if !self.is_healthy() - && successes >= self.metadata.health_config.success_threshold as usize + // Pending → Ready or NotReady → Ready on success_threshold. + // Failed is terminal — workers that reached Failed must be removed + // and re-registered. (replace() preserves the old status, so a + // metadata update on a Failed worker would not recover it.) + if matches!( + current_status, + WorkerStatus::Pending | WorkerStatus::NotReady + ) && successes >= health_config.success_threshold as usize { - self.set_healthy(true); + self.set_status(WorkerStatus::Ready); self.consecutive_successes.store(0, Ordering::Release); + // Reset pending probe counter on successful promotion + self.total_pending_probes.store(0, Ordering::Relaxed); + } else if current_status == WorkerStatus::Pending { + // Even on a successful probe, enforce the Pending cap. A + // worker that flaps F,S,F,S,... never reaches success_threshold + // and would otherwise grow `total_pending_probes` without bound. + let max_pending = health_config.failure_threshold as usize * 10; + if self.total_pending_probes.load(Ordering::Relaxed) >= max_pending { + self.set_status(WorkerStatus::Failed); + self.consecutive_successes.store(0, Ordering::Release); + self.consecutive_failures.store(0, Ordering::Release); + } } Ok(()) } else { self.consecutive_successes.store(0, Ordering::Release); let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1; - - // Record health check failure metric Metrics::record_worker_health_check(worker_type_str, metrics_labels::CB_FAILURE); - if self.is_healthy() - && failures >= self.metadata.health_config.failure_threshold as usize - { - self.set_healthy(false); - self.consecutive_failures.store(0, Ordering::Release); + match current_status { + WorkerStatus::Ready => { + // Ready → NotReady on failure_threshold + if failures >= health_config.failure_threshold as usize { + self.set_status(WorkerStatus::NotReady); + self.consecutive_failures.store(0, Ordering::Release); + } + } + WorkerStatus::NotReady => { + // NotReady → Failed on liveness_failure_threshold (3× failure_threshold) + let liveness_threshold = health_config.failure_threshold as usize * 3; + if failures >= liveness_threshold { + self.set_status(WorkerStatus::Failed); + self.consecutive_failures.store(0, Ordering::Release); + } + } + WorkerStatus::Pending => { + // Pending → Failed on max_pending_probes (10× failure_threshold). + // Note: this uses `total_pending_probes` (total attempts), not + // `consecutive_failures`, because a worker flapping between + // success/failure in Pending without reaching success_threshold + // is also suspect. Reset consecutive_failures for consistency + // with other failure-path transitions (Failed is terminal, but + // this keeps counter state clean). + let max_pending = health_config.failure_threshold as usize * 10; + let total = self.total_pending_probes.load(Ordering::Relaxed); + if total >= max_pending { + self.set_status(WorkerStatus::Failed); + self.consecutive_failures.store(0, Ordering::Release); + } + } + WorkerStatus::Failed => { + // Already failed — no further transitions + } } Err(WorkerError::HealthCheckFailed { @@ -698,6 +774,14 @@ impl Worker for BasicWorker { &self.http_client } + fn supports_model(&self, model_id: &str) -> bool { + let overridden = self.models_override.load(); + if !overridden.is_wildcard() { + return overridden.supports(model_id); + } + self.metadata.supports_model(model_id) + } + fn models(&self) -> Vec { let overridden = self.models_override.load(); let source = if overridden.is_wildcard() { @@ -708,14 +792,6 @@ impl Worker for BasicWorker { source.to_vec() } - fn supports_model(&self, model_id: &str) -> bool { - let overridden = self.models_override.load(); - if !overridden.is_wildcard() { - return overridden.supports(model_id); - } - self.metadata.supports_model(model_id) - } - fn set_models(&self, models: Vec) { tracing::debug!( "Setting {} models for worker {} via lazy discovery", @@ -988,22 +1064,24 @@ impl Drop for HealthChecker { } } -/// Helper to convert Worker trait object to WorkerInfo struct +/// Helper to convert Worker trait object to WorkerInfo struct. +/// +/// Both `is_healthy` and `status` are derived from the same atomic snapshot +/// to avoid TOCTOU between the two fields. The `status` field exposes the +/// real lifecycle state (Pending/Ready/NotReady/Failed) so API consumers +/// can distinguish "starting up" from "broken" — `is_healthy` collapses +/// everything to a routability bool for backwards compatibility. pub fn worker_to_info(worker: &Arc) -> WorkerInfo { let metadata = worker.metadata(); let spec = metadata.spec.clone(); - let is_healthy = worker.is_healthy(); + let status = worker.status(); WorkerInfo { id: worker.url().to_string(), model_id: spec.models.primary().map(|m| m.id.clone()), spec, - is_healthy, - status: Some(if is_healthy { - WorkerStatus::Ready - } else { - WorkerStatus::NotReady - }), + is_healthy: status == WorkerStatus::Ready, + status: Some(status), load: worker.load(), job_status: None, } @@ -1013,12 +1091,23 @@ pub fn worker_to_info(worker: &Arc) -> WorkerInfo { mod tests { use std::{thread, time::Duration}; + use openai_protocol::worker::HealthCheckConfig; + use super::*; use crate::worker::{ circuit_breaker::{CircuitBreakerConfig, CircuitState}, BasicWorkerBuilder, }; + /// Health config that skips health checks — workers start Ready immediately. + /// Use in tests that don't test the health check lifecycle. + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..HealthCheckConfig::default() + } + } + #[test] fn test_worker_type_display() { assert_eq!(WorkerType::Regular.to_string(), "regular"); @@ -1070,9 +1159,9 @@ mod tests { #[test] fn test_basic_worker_creation() { - use crate::worker::BasicWorkerBuilder; let worker = BasicWorkerBuilder::new("http://test:8080") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); assert_eq!(worker.url(), "http://test:8080"); assert_eq!(worker.worker_type(), &WorkerType::Regular); @@ -1151,17 +1240,41 @@ mod tests { #[test] fn test_health_status() { - use crate::worker::BasicWorkerBuilder; let worker = BasicWorkerBuilder::new("http://test:8080") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); assert!(worker.is_healthy()); + assert_eq!(worker.status(), WorkerStatus::Ready); worker.set_healthy(false); assert!(!worker.is_healthy()); + assert_eq!(worker.status(), WorkerStatus::NotReady); + + worker.set_healthy(true); + assert!(worker.is_healthy()); + assert_eq!(worker.status(), WorkerStatus::Ready); + } + + #[test] + fn test_pending_worker_not_routable() { + // Default health config: health checks enabled → starts Pending + let worker = BasicWorkerBuilder::new("http://test:8080") + .worker_type(WorkerType::Regular) + .build(); + + assert_eq!(worker.status(), WorkerStatus::Pending); + assert!(!worker.is_healthy()); // Pending is not routable + assert!(!worker.is_available()); // Pending is not available + + // set_healthy(false) on Pending is a no-op (hasn't proven itself) + worker.set_healthy(false); + assert_eq!(worker.status(), WorkerStatus::Pending); + // set_healthy(true) promotes Pending → Ready worker.set_healthy(true); + assert_eq!(worker.status(), WorkerStatus::Ready); assert!(worker.is_healthy()); } @@ -1483,6 +1596,7 @@ mod tests { let dp_worker = BasicWorkerBuilder::new("http://worker1:8080") .dp_config(0, 2) .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); assert!(dp_worker.is_healthy()); @@ -1502,9 +1616,9 @@ mod tests { #[test] fn test_worker_circuit_breaker() { - use crate::worker::BasicWorkerBuilder; let worker = BasicWorkerBuilder::new("http://test:8080") .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); assert!(worker.is_available()); @@ -1533,10 +1647,10 @@ mod tests { window_duration: Duration::from_secs(60), }; - use crate::worker::BasicWorkerBuilder; let worker = BasicWorkerBuilder::new("http://test:8080") .worker_type(WorkerType::Regular) .circuit_breaker_config(config) + .health_config(no_health_check()) .build(); worker.record_outcome(503); @@ -1558,6 +1672,7 @@ mod tests { let dp_worker = BasicWorkerBuilder::new("http://worker:8080") .dp_config(0, 2) .worker_type(WorkerType::Regular) + .health_config(no_health_check()) .build(); assert!(dp_worker.is_available()); @@ -1572,27 +1687,31 @@ mod tests { #[tokio::test] async fn test_mixed_worker_types() { - use crate::worker::BasicWorkerBuilder; + let hc = no_health_check(); let regular: Box = Box::new( BasicWorkerBuilder::new("http://regular:8080") .worker_type(WorkerType::Regular) + .health_config(hc.clone()) .build(), ); let prefill: Box = Box::new( BasicWorkerBuilder::new("http://prefill:8080") .worker_type(WorkerType::Prefill) .bootstrap_port(Some(9090)) + .health_config(hc.clone()) .build(), ); let decode: Box = Box::new( BasicWorkerBuilder::new("http://decode:8080") .worker_type(WorkerType::Decode) + .health_config(hc.clone()) .build(), ); let dp_aware_regular: Box = Box::new( BasicWorkerBuilder::new("http://dp:8080") .dp_config(0, 2) .worker_type(WorkerType::Regular) + .health_config(hc.clone()) .api_key("test_api_key") .build(), ); @@ -1600,6 +1719,7 @@ mod tests { BasicWorkerBuilder::new("http://dp-prefill:8080") .dp_config(1, 2) .worker_type(WorkerType::Prefill) + .health_config(hc.clone()) .api_key("test_api_key") .build(), ); @@ -1607,6 +1727,7 @@ mod tests { BasicWorkerBuilder::new("http://dp-decode:8080") .dp_config(0, 4) .worker_type(WorkerType::Decode) + .health_config(hc.clone()) .api_key("test_api_key") .build(), ); diff --git a/model_gateway/src/workflow/steps/external/create_workers.rs b/model_gateway/src/workflow/steps/external/create_workers.rs index ae7ae84b46..18dec59d80 100644 --- a/model_gateway/src/workflow/steps/external/create_workers.rs +++ b/model_gateway/src/workflow/steps/external/create_workers.rs @@ -120,12 +120,8 @@ impl StepExecutor for CreateExternalWorkersStep { builder = builder.labels(labels.clone()); } + // Builder sets initial status: Pending if health-checked, Ready if not. let worker = Arc::new(builder.build()) as Arc; - if health_config.disable_health_check { - worker.set_healthy(true); - } else { - worker.set_healthy(false); - } info!( "Created wildcard worker at {} (accepts any model, user auth forwarded)", @@ -161,12 +157,8 @@ impl StepExecutor for CreateExternalWorkersStep { builder = builder.labels(labels.clone()); } + // Builder sets initial status: Pending if health-checked, Ready if not. let worker = Arc::new(builder.build()) as Arc; - if health_config.disable_health_check { - worker.set_healthy(true); - } else { - worker.set_healthy(false); - } info!( "Created external worker at {} with {} discovered models", diff --git a/model_gateway/src/workflow/steps/local/create_worker.rs b/model_gateway/src/workflow/steps/local/create_worker.rs index 5860503adc..bb3fd3ad0e 100644 --- a/model_gateway/src/workflow/steps/local/create_worker.rs +++ b/model_gateway/src/workflow/steps/local/create_worker.rs @@ -169,9 +169,8 @@ impl StepExecutor for CreateLocalWorkerStep { builder = builder.kv_role(r); } - let worker = Arc::new(builder.build()) as Arc; - worker.set_healthy(health_config.disable_health_check); - worker + // Builder sets initial status: Pending if health-checked, Ready if not. + Arc::new(builder.build()) as Arc }) .collect(); diff --git a/model_gateway/src/workflow/steps/local/update_worker_properties.rs b/model_gateway/src/workflow/steps/local/update_worker_properties.rs index 7a40237767..c89a9ac733 100644 --- a/model_gateway/src/workflow/steps/local/update_worker_properties.rs +++ b/model_gateway/src/workflow/steps/local/update_worker_properties.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use async_trait::async_trait; +use openai_protocol::worker::WorkerStatus; use tracing::{debug, info, warn}; use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; @@ -71,8 +72,21 @@ impl StepExecutor for UpdateWorkerPropertiesStep { .clone() .or_else(|| worker.metadata().spec.api_key.clone()); - // Create a new worker with updated properties - // Use base_url() so DP workers start from the un-suffixed URL + // Create a new worker with updated properties. + // Use base_url() so DP workers start from the un-suffixed URL. + // + // Status: a metadata-only update is not a re-registration — the + // worker is the same endpoint. Preserve the old status so a + // healthy worker stays routable through the update. The one + // exception is when the update flips `disable_health_check` to + // `true`: the health checker skips disabled workers, so a stale + // Pending/NotReady status would never recover. Force Ready in + // that case. + let next_status = if updated_health_config.disable_health_check { + WorkerStatus::Ready + } else { + worker.status() + }; let mut builder = BasicWorkerBuilder::new(worker.base_url()) .worker_type(*worker.worker_type()) .connection_mode(*worker.connection_mode()) @@ -84,7 +98,8 @@ impl StepExecutor for UpdateWorkerPropertiesStep { .http_client(worker.http_client().clone()) .resilience(worker.resilience().clone()) .priority(updated_priority) - .cost(updated_cost); + .cost(updated_cost) + .status(next_status); if let Some(ref api_key) = updated_api_key { builder = builder.api_key(api_key.clone()); diff --git a/model_gateway/src/workflow/steps/shared/activate.rs b/model_gateway/src/workflow/steps/shared/activate.rs index 8a6687af7a..00725ff557 100644 --- a/model_gateway/src/workflow/steps/shared/activate.rs +++ b/model_gateway/src/workflow/steps/shared/activate.rs @@ -8,10 +8,7 @@ use wfaas::{ use crate::workflow::data::WorkerRegistrationData; -/// Unified step to activate workers by marking them as healthy. -/// -/// This is the final step in any worker registration workflow. -/// Works with any workflow data type that implements `WorkerRegistrationData`. +/// Final step in any worker registration workflow: flip Pending → Ready. pub struct ActivateWorkersStep; #[async_trait] @@ -23,10 +20,12 @@ impl StepExecutor for ActivateWorke .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; for worker in workers { - worker.set_healthy(true); + if !worker.is_healthy() { + worker.set_healthy(true); + } } - info!("Activated {} worker(s) (marked as healthy)", workers.len()); + info!("Activated {} worker(s)", workers.len()); Ok(StepResult::Success) } diff --git a/model_gateway/tests/api/api_endpoints_test.rs b/model_gateway/tests/api/api_endpoints_test.rs index 47dfc82cd0..d5357c27d5 100644 --- a/model_gateway/tests/api/api_endpoints_test.rs +++ b/model_gateway/tests/api/api_endpoints_test.rs @@ -887,7 +887,7 @@ mod responses_endpoint_tests { async fn test_v1_responses_input_items() { // This test uses OpenAI mode because the input_items endpoint // is only implemented in OpenAIRouter and reads from storage (no workers needed) - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .openai_mode(vec!["http://dummy.local".to_string()]) // Dummy URL (won't be called) .random_policy() .host("127.0.0.1") @@ -900,6 +900,7 @@ mod responses_endpoint_tests { .queue_size(0) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -1010,7 +1011,7 @@ mod responses_endpoint_tests { async fn create_openai_ctx(port: u16) -> AppTestContext { use smg::config::RouterConfig; - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .openai_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -1023,6 +1024,7 @@ mod responses_endpoint_tests { .queue_size(0) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; AppTestContext::new_with_config( config, @@ -1276,7 +1278,7 @@ mod error_tests { #[tokio::test] async fn test_payload_too_large() { // Create context with small payload limit - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -1288,6 +1290,7 @@ mod error_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -1576,7 +1579,7 @@ mod pd_mode_tests { .and_then(|p| p.trim_end_matches('/').parse::().ok()) .unwrap_or(9000); - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .prefill_decode_mode(vec![(prefill_url, Some(prefill_port))], vec![decode_url]) .random_policy() .host("127.0.0.1") @@ -1588,6 +1591,7 @@ mod pd_mode_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; // Create app context let app_context = crate::common::create_test_context(config).await; @@ -1713,7 +1717,7 @@ mod request_id_tests { #[tokio::test] async fn test_request_id_with_custom_headers() { // Create config with custom request ID headers - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -1726,6 +1730,7 @@ mod request_id_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/api/parser_endpoints_test.rs b/model_gateway/tests/api/parser_endpoints_test.rs index d8fa2528d4..e119554055 100644 --- a/model_gateway/tests/api/parser_endpoints_test.rs +++ b/model_gateway/tests/api/parser_endpoints_test.rs @@ -31,7 +31,7 @@ struct ParserTestContext { impl ParserTestContext { async fn new(worker_configs: Vec) -> Self { // Create router config with parser support enabled - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -43,6 +43,7 @@ impl ParserTestContext { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; Self::new_with_config(config, worker_configs).await } diff --git a/model_gateway/tests/common/mod.rs b/model_gateway/tests/common/mod.rs index 140e6e127d..f93d5d9c27 100644 --- a/model_gateway/tests/common/mod.rs +++ b/model_gateway/tests/common/mod.rs @@ -177,7 +177,7 @@ impl AppTestContext { worker_configs: Vec, ) -> Pin + Send>> { Box::pin(async move { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -189,6 +189,10 @@ impl AppTestContext { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + // Test mock workers don't need health checks — start Ready immediately. + // Tests that need to exercise Pending/Failed semantics should use + // `new_with_config()` with an explicit config. + config.health_check.disable_health_check = true; Self::new_with_config(config, worker_configs).await }) @@ -410,6 +414,10 @@ pub fn create_test_context( .worker_type(WorkerType::Regular) .runtime_type(RuntimeType::External) .models(models) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) .build(), ); app_context.worker_registry.register(worker); @@ -545,6 +553,10 @@ pub fn create_test_context_with_parsers( .worker_type(WorkerType::Regular) .runtime_type(RuntimeType::External) .models(models) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) .build(), ); app_context.worker_registry.register(worker); @@ -679,6 +691,10 @@ pub fn create_test_context_with_mcp_config( .worker_type(WorkerType::Regular) .runtime_type(RuntimeType::External) .models(models) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) .build(), ); app_context.worker_registry.register(worker); diff --git a/model_gateway/tests/common/test_app.rs b/model_gateway/tests/common/test_app.rs index 41e49a066c..5846bb3352 100644 --- a/model_gateway/tests/common/test_app.rs +++ b/model_gateway/tests/common/test_app.rs @@ -239,6 +239,10 @@ pub fn register_external_worker(ctx: &Arc, url: &str, models: Option .worker_type(WorkerType::Regular) .runtime_type(RuntimeType::External) .models(model_list) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) .build(), ); @@ -257,6 +261,10 @@ pub fn register_external_worker_with_card(ctx: &Arc, url: &str, mode .worker_type(WorkerType::Regular) .runtime_type(RuntimeType::External) .model(model_card) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) .build(), ); diff --git a/model_gateway/tests/common/test_config.rs b/model_gateway/tests/common/test_config.rs index bc79d2cd1d..50b9216f36 100644 --- a/model_gateway/tests/common/test_config.rs +++ b/model_gateway/tests/common/test_config.rs @@ -20,78 +20,98 @@ pub mod defaults { pub const QUEUE_TIMEOUT_SECS: u64 = 60; } +/// Apply default test overrides to a freshly-built RouterConfig. +/// +/// Test mock workers don't expose real health endpoints, so we disable +/// health checks by default. Tests that need to exercise Pending/NotReady/ +/// Failed semantics should override this by setting +/// `config.health_check.disable_health_check = false` after calling +/// `TestRouterConfig::*`. +fn apply_test_defaults(mut config: RouterConfig) -> RouterConfig { + config.health_check.disable_health_check = true; + config +} + /// Builder for common test RouterConfig patterns pub struct TestRouterConfig; impl TestRouterConfig { /// Create a basic round-robin config for routing tests pub fn round_robin(port: u16) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .round_robin_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .round_robin_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .build_unchecked(), + ) } /// Create a random load balancing config pub fn random(port: u16) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .random_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .random_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .build_unchecked(), + ) } /// Create a cache-aware config for routing tests pub fn cache_aware(port: u16) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .cache_aware_policy( - 0.5, // cache_threshold - 32, // balance_abs_threshold - 1.5, // balance_rel_threshold - 60, // eviction_interval_secs - 1000, // max_tree_size - ) - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .cache_aware_policy( + 0.5, // cache_threshold + 32, // balance_abs_threshold + 1.5, // balance_rel_threshold + 60, // eviction_interval_secs + 1000, // max_tree_size + ) + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .build_unchecked(), + ) } /// Create a power-of-two config pub fn power_of_two(port: u16) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .power_of_two_policy(5) // load_check_interval_secs - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .power_of_two_policy(5) // load_check_interval_secs + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .build_unchecked(), + ) } /// Create a manual routing config (for sticky routing tests) @@ -106,87 +126,97 @@ impl TestRouterConfig { /// Create a manual routing config with specified assignment mode pub fn manual_with_mode(port: u16, assignment_mode: ManualAssignmentMode) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .policy(PolicyConfig::Manual { - eviction_interval_secs: 60, - max_idle_secs: 3600, - assignment_mode, - }) - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .policy(PolicyConfig::Manual { + eviction_interval_secs: 60, + max_idle_secs: 3600, + assignment_mode, + }) + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .build_unchecked(), + ) } /// Create a config with custom concurrent request limit (for rate limiting tests) pub fn with_concurrency(port: u16, max_concurrent: i32) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .round_robin_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(max_concurrent) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .round_robin_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(max_concurrent) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .build_unchecked(), + ) } /// Create a config with custom payload size limit pub fn with_payload_limit(port: u16, max_payload_size: usize) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .round_robin_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(max_payload_size) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .round_robin_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(max_payload_size) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .build_unchecked(), + ) } /// Create a config with short timeouts (for timeout/retry tests) pub fn with_short_timeouts(port: u16) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .round_robin_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(5) - .worker_startup_timeout_secs(2) - .worker_startup_check_interval_secs(1) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(5) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .round_robin_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(5) + .worker_startup_timeout_secs(2) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(5) + .build_unchecked(), + ) } /// Create a round-robin config with retry settings pub fn round_robin_with_retry(port: u16, retry_config: RetryConfig) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .round_robin_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .retry_config(retry_config) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .round_robin_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .retry_config(retry_config) + .build_unchecked(), + ) } /// Create a round-robin config with circuit breaker @@ -194,19 +224,21 @@ impl TestRouterConfig { port: u16, circuit_breaker: CircuitBreakerConfig, ) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .round_robin_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .circuit_breaker_config(circuit_breaker) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .round_robin_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .circuit_breaker_config(circuit_breaker) + .build_unchecked(), + ) } /// Create a round-robin config with both retry and circuit breaker @@ -215,20 +247,22 @@ impl TestRouterConfig { retry_config: RetryConfig, circuit_breaker: CircuitBreakerConfig, ) -> RouterConfig { - RouterConfig::builder() - .regular_mode(vec![]) - .round_robin_policy() - .host(defaults::HOST) - .port(port) - .max_payload_size(defaults::MAX_PAYLOAD_SIZE) - .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) - .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) - .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) - .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) - .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) - .retry_config(retry_config) - .circuit_breaker_config(circuit_breaker) - .build_unchecked() + apply_test_defaults( + RouterConfig::builder() + .regular_mode(vec![]) + .round_robin_policy() + .host(defaults::HOST) + .port(port) + .max_payload_size(defaults::MAX_PAYLOAD_SIZE) + .request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS) + .worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS) + .worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS) + .max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS) + .queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS) + .retry_config(retry_config) + .circuit_breaker_config(circuit_breaker) + .build_unchecked(), + ) } } diff --git a/model_gateway/tests/load_guard_raii_test.rs b/model_gateway/tests/load_guard_raii_test.rs index 163275939a..8b48f4a49d 100644 --- a/model_gateway/tests/load_guard_raii_test.rs +++ b/model_gateway/tests/load_guard_raii_test.rs @@ -24,7 +24,14 @@ fn create_sse_response(rx: mpsc::UnboundedReceiver) -> Response { /// Helper to create a test worker fn create_test_worker() -> Arc { - Arc::new(BasicWorkerBuilder::new("http://localhost:8000").build()) + Arc::new( + BasicWorkerBuilder::new("http://localhost:8000") + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(), + ) } #[tokio::test] diff --git a/model_gateway/tests/otel_tracing_test.rs b/model_gateway/tests/otel_tracing_test.rs index 084c75eecf..fefa6cea65 100644 --- a/model_gateway/tests/otel_tracing_test.rs +++ b/model_gateway/tests/otel_tracing_test.rs @@ -120,7 +120,7 @@ async fn test_router_with_tracing() { println!("Mock worker started on: {worker_url}"); // 3. create router config and enable tracing - let router_config = RouterConfig::builder() + let mut router_config = RouterConfig::builder() .regular_mode(vec![worker_url.clone()]) .random_policy() .host("0.0.0.0") @@ -133,6 +133,7 @@ async fn test_router_with_tracing() { .queue_timeout_secs(60) .enable_trace(&collector_endpoint) .build_unchecked(); + router_config.health_check.disable_health_check = true; // 4. Initialize the OTLP client (check if already initialized by another test) let otel_initialized_by_this_test = if otel_trace::is_otel_enabled() { diff --git a/model_gateway/tests/reliability/circuit_breaker_test.rs b/model_gateway/tests/reliability/circuit_breaker_test.rs index 2e48676192..54c9f57299 100644 --- a/model_gateway/tests/reliability/circuit_breaker_test.rs +++ b/model_gateway/tests/reliability/circuit_breaker_test.rs @@ -75,7 +75,7 @@ mod circuit_breaker_tests { /// Test circuit breaker with disabled flag #[tokio::test] async fn test_circuit_breaker_disabled() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -89,6 +89,7 @@ mod circuit_breaker_tests { .disable_circuit_breaker() .disable_retries() .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -182,7 +183,7 @@ mod circuit_breaker_tests { /// Test circuit breaker with retries enabled #[tokio::test] async fn test_circuit_breaker_with_retries() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -206,6 +207,7 @@ mod circuit_breaker_tests { window_duration_secs: 10, }) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/reliability/rate_limiting_test.rs b/model_gateway/tests/reliability/rate_limiting_test.rs index cfda89b8d0..544dee09ba 100644 --- a/model_gateway/tests/reliability/rate_limiting_test.rs +++ b/model_gateway/tests/reliability/rate_limiting_test.rs @@ -81,7 +81,7 @@ mod rate_limiting_tests { /// Test rate limit tokens per second #[tokio::test] async fn test_rate_limit_tokens() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -94,6 +94,7 @@ mod rate_limiting_tests { .rate_limit_tokens_per_second(50) // 50 tokens/sec .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19301)]).await; @@ -189,7 +190,7 @@ mod rate_limiting_tests { #[tokio::test] #[expect(clippy::disallowed_methods)] async fn test_queue_behavior() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -202,6 +203,7 @@ mod rate_limiting_tests { .queue_size(5) // Small queue .queue_timeout_secs(1) // Short timeout .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/reliability/retries_test.rs b/model_gateway/tests/reliability/retries_test.rs index df609ee19c..5e3c5c42ca 100644 --- a/model_gateway/tests/reliability/retries_test.rs +++ b/model_gateway/tests/reliability/retries_test.rs @@ -65,7 +65,7 @@ mod retry_tests { /// Test that retries are disabled when configured #[tokio::test] async fn test_retries_disabled() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -78,6 +78,7 @@ mod retry_tests { .queue_timeout_secs(60) .disable_retries() .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/routing/cache_aware_backward_compat_test.rs b/model_gateway/tests/routing/cache_aware_backward_compat_test.rs index e04ff50ba6..a891eb4c18 100644 --- a/model_gateway/tests/routing/cache_aware_backward_compat_test.rs +++ b/model_gateway/tests/routing/cache_aware_backward_compat_test.rs @@ -1,10 +1,20 @@ use std::{collections::HashMap, sync::Arc}; +use openai_protocol::worker::HealthCheckConfig; use smg::{ policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, SelectWorkerInfo}, worker::{BasicWorkerBuilder, Worker, WorkerType}, }; +/// Test helper: health check config that skips health checks. +/// Workers built with this config start Ready (routable) immediately. +fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } +} + #[test] fn test_backward_compatibility_with_empty_model_id() { let config = CacheAwareConfig { @@ -22,6 +32,7 @@ fn test_backward_compatibility_with_empty_model_id() { let worker1 = BasicWorkerBuilder::new("http://worker1:8080") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(); // No model_id label - should default to "unknown" @@ -31,6 +42,7 @@ fn test_backward_compatibility_with_empty_model_id() { .worker_type(WorkerType::Regular) .api_key("test_api_key") .labels(labels2) + .health_config(no_health_check()) .build(); // Add workers - should both go to "default" tree @@ -72,6 +84,7 @@ fn test_mixed_model_ids() { let worker1 = BasicWorkerBuilder::new("http://worker1:8080") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(); // No model_id label - defaults to "unknown" which goes to "default" tree @@ -81,6 +94,7 @@ fn test_mixed_model_ids() { .worker_type(WorkerType::Regular) .labels(labels2) .api_key("test_api_key") + .health_config(no_health_check()) .build(); let mut labels3 = HashMap::new(); @@ -88,6 +102,7 @@ fn test_mixed_model_ids() { let worker3 = BasicWorkerBuilder::new("http://worker3:8080") .worker_type(WorkerType::Regular) .labels(labels3) + .health_config(no_health_check()) .build(); let mut labels4 = HashMap::new(); @@ -95,6 +110,7 @@ fn test_mixed_model_ids() { let worker4 = BasicWorkerBuilder::new("http://worker4:8080") .worker_type(WorkerType::Regular) .labels(labels4) + .health_config(no_health_check()) .build(); // Add all workers @@ -139,11 +155,13 @@ fn test_remove_worker_by_url_backward_compat() { .worker_type(WorkerType::Regular) .labels(labels1) .api_key("test_api_key") + .health_config(no_health_check()) .build(); let worker2 = BasicWorkerBuilder::new("http://worker2:8080") .worker_type(WorkerType::Regular) .api_key("test_api_key") + .health_config(no_health_check()) .build(); // No model_id label - defaults to "unknown" diff --git a/model_gateway/tests/routing/header_forwarding_test.rs b/model_gateway/tests/routing/header_forwarding_test.rs index 93d108c8a9..9d52ebe6a0 100644 --- a/model_gateway/tests/routing/header_forwarding_test.rs +++ b/model_gateway/tests/routing/header_forwarding_test.rs @@ -66,7 +66,7 @@ mod header_forwarding_tests { /// Test custom request ID headers #[tokio::test] async fn test_custom_request_id_headers() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .random_policy() .host("127.0.0.1") @@ -82,6 +82,7 @@ mod header_forwarding_tests { "x-correlation-id".to_string(), ]) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/routing/payload_size_test.rs b/model_gateway/tests/routing/payload_size_test.rs index a7301a4217..b9941f393d 100644 --- a/model_gateway/tests/routing/payload_size_test.rs +++ b/model_gateway/tests/routing/payload_size_test.rs @@ -23,7 +23,7 @@ mod payload_size_tests { /// Test that small payloads are handled correctly #[tokio::test] async fn test_small_payload() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -35,6 +35,7 @@ mod payload_size_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -75,7 +76,7 @@ mod payload_size_tests { /// Test that payloads within limit are accepted #[tokio::test] async fn test_payload_within_limit() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -87,6 +88,7 @@ mod payload_size_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -129,7 +131,7 @@ mod payload_size_tests { /// Test that payloads exceeding limit are rejected #[tokio::test] async fn test_payload_exceeds_limit() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -141,6 +143,7 @@ mod payload_size_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -188,7 +191,7 @@ mod payload_size_tests { // Use a more reasonable limit for this test let limit_bytes = 10 * 1024; // 10KB limit - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -200,6 +203,7 @@ mod payload_size_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -244,7 +248,7 @@ mod payload_size_tests { /// Test default payload size limit (256MB) #[tokio::test] async fn test_default_payload_limit() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -256,6 +260,7 @@ mod payload_size_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/routing/pd_routing_test.rs b/model_gateway/tests/routing/pd_routing_test.rs index 1247d1ac7c..4fbf9a7270 100644 --- a/model_gateway/tests/routing/pd_routing_test.rs +++ b/model_gateway/tests/routing/pd_routing_test.rs @@ -23,7 +23,7 @@ mod pd_routing_tests { /// Test basic PD mode routing with prefill and decode workers #[tokio::test] async fn test_pd_mode_basic_routing() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .prefill_decode_mode( vec![ ("http://127.0.0.1:19800".to_string(), None), @@ -44,6 +44,7 @@ mod pd_routing_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; // Note: For PD mode tests, we need to start prefill and decode workers separately // The test context will need to handle this specially @@ -90,7 +91,7 @@ mod pd_routing_tests { /// Test PD mode with round robin policy #[tokio::test] async fn test_pd_mode_round_robin() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .prefill_decode_mode( vec![("http://127.0.0.1:19810".to_string(), None)], vec![ @@ -108,6 +109,7 @@ mod pd_routing_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -154,7 +156,7 @@ mod pd_routing_tests { async fn test_pd_mode_with_failing_decode_worker() { use smg::config::RetryConfig; - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .prefill_decode_mode( vec![("http://127.0.0.1:19820".to_string(), None)], vec![ @@ -178,6 +180,7 @@ mod pd_routing_tests { ..Default::default() }) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/routing/power_of_two_test.rs b/model_gateway/tests/routing/power_of_two_test.rs index 05fd63c2c7..197397bfe4 100644 --- a/model_gateway/tests/routing/power_of_two_test.rs +++ b/model_gateway/tests/routing/power_of_two_test.rs @@ -141,7 +141,7 @@ mod power_of_two_tests { window_duration_secs: 10, }; - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .power_of_two_policy(1) .host("127.0.0.1") @@ -155,6 +155,7 @@ mod power_of_two_tests { .retry_config(retry_config) .circuit_breaker_config(circuit_breaker) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/routing/service_discovery_test.rs b/model_gateway/tests/routing/service_discovery_test.rs index d997a6403c..be9ce900a3 100644 --- a/model_gateway/tests/routing/service_discovery_test.rs +++ b/model_gateway/tests/routing/service_discovery_test.rs @@ -23,7 +23,7 @@ mod service_discovery_tests { /// Test service discovery endpoint responds correctly #[tokio::test] async fn test_service_discovery_endpoint() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -35,6 +35,7 @@ mod service_discovery_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -71,7 +72,7 @@ mod service_discovery_tests { /// Test worker registration via discovery shim #[tokio::test] async fn test_worker_registration() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -83,6 +84,7 @@ mod service_discovery_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -129,7 +131,7 @@ mod service_discovery_tests { /// Test worker deregistration via discovery shim #[tokio::test] async fn test_worker_deregistration() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -141,6 +143,7 @@ mod service_discovery_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, @@ -214,7 +217,7 @@ mod service_discovery_tests { /// Test health status reporting for discovery #[tokio::test] async fn test_health_status_endpoint() { - let config = RouterConfig::builder() + let mut config = RouterConfig::builder() .regular_mode(vec![]) .round_robin_policy() .host("127.0.0.1") @@ -226,6 +229,7 @@ mod service_discovery_tests { .max_concurrent_requests(64) .queue_timeout_secs(60) .build_unchecked(); + config.health_check.disable_health_check = true; let ctx = AppTestContext::new_with_config( config, diff --git a/model_gateway/tests/routing/test_openai_routing.rs b/model_gateway/tests/routing/test_openai_routing.rs index d15afa8694..c3ed94cf92 100644 --- a/model_gateway/tests/routing/test_openai_routing.rs +++ b/model_gateway/tests/routing/test_openai_routing.rs @@ -839,6 +839,10 @@ async fn test_openai_router_models_from_registry() { .models(vec![openai_protocol::model_card::ModelCard::new( "gpt-3.5-turbo", )]) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) .build(), ); ctx.worker_registry.register(worker);