[model-gateway] add --disable-health-check option to skip worker health probes (#17002)

This commit is contained in:
Ziwen Zhao
2026-01-13 17:28:32 -08:00
committed by GitHub
parent 5938c3b06a
commit cf25852a1d
17 changed files with 159 additions and 14 deletions
+1 -1
View File
@@ -729,7 +729,7 @@ Router flags map to these values:
- **Retries**: Default max retries = 5 with exponential backoff (`--retry-max-retries`, `--retry-initial-backoff-ms`, `--retry-max-backoff-ms`, `--retry-backoff-multiplier`, `--retry-jitter-factor`). Retries trigger on 408/429/500/502/503/504.
- **Circuit Breakers**: Per worker thresholds (`--cb-failure-threshold`, `--cb-success-threshold`, `--cb-timeout-duration-secs`, `--cb-window-duration-secs`). Disable via `--disable-circuit-breaker`.
- **Rate Limiting**: Token bucket driven by `--max-concurrent-requests`. Set `--rate-limit-tokens-per-second` to override refill rate. Configure request queue via `--queue-size` and `--queue-timeout-secs`; queued requests observe FIFO order and respect cancellation.
- **Health Checks**: Runtime probes via `--health-check-interval-secs`, `--health-check-timeout-secs`, failure/success thresholds, and `--health-check-endpoint`.
- **Health Checks**: Runtime probes via `--health-check-interval-secs`, `--health-check-timeout-secs`, failure/success thresholds, and `--health-check-endpoint`. Use `--disable-health-check` to skip health checks entirely.
- **Cache Management**: `/flush_cache` ensures LRU eviction when redeploying PD workers.
## Load Balancing Policies
@@ -391,6 +391,7 @@ struct Router {
health_check_timeout_secs: u64,
health_check_interval_secs: u64,
health_check_endpoint: String,
disable_health_check: bool,
enable_igw: bool,
queue_size: usize,
queue_timeout_secs: u64,
@@ -591,6 +592,7 @@ impl Router {
timeout_secs: self.health_check_timeout_secs,
check_interval_secs: self.health_check_interval_secs,
endpoint: self.health_check_endpoint.clone(),
disable_health_check: self.disable_health_check,
})
.tokenizer_cache(config::TokenizerCacheConfig {
enable_l0: self.tokenizer_cache_enable_l0,
@@ -692,6 +694,7 @@ impl Router {
health_check_timeout_secs = 5,
health_check_interval_secs = 60,
health_check_endpoint = String::from("/health"),
disable_health_check = false,
enable_igw = false,
queue_size = 100,
queue_timeout_secs = 60,
@@ -777,6 +780,7 @@ impl Router {
health_check_timeout_secs: u64,
health_check_interval_secs: u64,
health_check_endpoint: String,
disable_health_check: bool,
enable_igw: bool,
queue_size: usize,
queue_timeout_secs: u64,
@@ -875,6 +879,7 @@ impl Router {
health_check_timeout_secs,
health_check_interval_secs,
health_check_endpoint,
disable_health_check,
enable_igw,
queue_size,
queue_timeout_secs,
@@ -86,6 +86,7 @@ class RouterArgs:
health_check_timeout_secs: int = 5
health_check_interval_secs: int = 60
health_check_endpoint: str = "/health"
disable_health_check: bool = False
# Circuit breaker configuration
cb_failure_threshold: int = 10
cb_success_threshold: int = 3
@@ -599,6 +600,12 @@ class RouterArgs:
default=RouterArgs.health_check_endpoint,
help="Health check endpoint path",
)
health_group.add_argument(
f"--{prefix}disable-health-check",
action="store_true",
default=RouterArgs.disable_health_check,
help="Disable all worker health checks at startup",
)
# Tokenizer configuration
tokenizer_group.add_argument(
f"--{prefix}model-path",
@@ -157,3 +157,64 @@ class TestIGWMode:
logger.info("Worker: id=%s, url=%s", w.id, w.url)
finally:
gateway.shutdown()
@pytest.mark.e2e
class TestDisableHealthCheck:
"""Tests for --disable-health-check CLI option."""
def test_disable_health_check_workers_immediately_healthy(
self, model_pool: ModelPool
):
"""Test that workers are immediately healthy when health checks are disabled."""
http_instance = model_pool.get("llama-8b", ConnectionMode.HTTP)
gateway = Gateway()
gateway.start(
igw_mode=True,
extra_args=["--disable-health-check"],
)
try:
# Add worker - should be immediately healthy since health checks are disabled
success, worker_id = gateway.add_worker(
http_instance.worker_url,
wait_ready=True,
ready_timeout=10, # Short timeout since it should be immediate
)
assert success, f"Failed to add worker: {worker_id}"
logger.info("Added worker with health checks disabled: %s", worker_id)
# Verify worker is healthy
workers = gateway.list_workers()
assert len(workers) >= 1, "Expected at least one worker"
for worker in workers:
logger.info(
"Worker: id=%s, status=%s, disable_health_check=%s",
worker.id,
worker.status,
worker.metadata.get("disable_health_check"),
)
# Worker should be healthy immediately
assert (
worker.status == "healthy"
), "Worker should be healthy when health checks disabled"
finally:
gateway.shutdown()
def test_disable_health_check_gateway_starts_without_health_checker(
self, model_pool: ModelPool
):
"""Test that gateway starts successfully with health checks disabled."""
gateway = Gateway()
gateway.start(
igw_mode=True,
extra_args=["--disable-health-check"],
)
try:
assert gateway.health(), "Gateway should be healthy"
logger.info("Gateway started with health checks disabled")
finally:
gateway.shutdown()
+2
View File
@@ -554,6 +554,7 @@ pub struct HealthCheckConfig {
pub timeout_secs: u64,
pub check_interval_secs: u64,
pub endpoint: String,
pub disable_health_check: bool,
}
impl Default for HealthCheckConfig {
@@ -564,6 +565,7 @@ impl Default for HealthCheckConfig {
timeout_secs: 5,
check_interval_secs: 60,
endpoint: "/health".to_string(),
disable_health_check: false,
}
}
}
+4
View File
@@ -590,6 +590,9 @@ impl JobQueue {
health_failure_threshold: router_config
.health_check
.failure_threshold,
disable_health_check: router_config
.health_check
.disable_health_check,
max_connection_attempts: router_config
.health_check
.success_threshold
@@ -652,6 +655,7 @@ impl JobQueue {
health_check_interval_secs: router_config.health_check.check_interval_secs,
health_success_threshold: router_config.health_check.success_threshold,
health_failure_threshold: router_config.health_check.failure_threshold,
disable_health_check: router_config.health_check.disable_health_check,
max_connection_attempts: router_config.health_check.success_threshold * 10,
dp_aware: router_config.dp_aware,
};
@@ -60,6 +60,7 @@ impl StepExecutor<ExternalWorkerWorkflowData> for CreateExternalWorkersStep {
endpoint: cfg.endpoint.clone(),
failure_threshold: cfg.failure_threshold,
success_threshold: cfg.success_threshold,
disable_health_check: cfg.disable_health_check || config.disable_health_check,
}
};
@@ -98,7 +99,11 @@ impl StepExecutor<ExternalWorkerWorkflowData> for CreateExternalWorkersStep {
}
let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
worker.set_healthy(false);
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)",
@@ -132,7 +137,11 @@ impl StepExecutor<ExternalWorkerWorkflowData> for CreateExternalWorkersStep {
}
let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
worker.set_healthy(false);
if health_config.disable_health_check {
worker.set_healthy(true);
} else {
worker.set_healthy(false);
}
debug!(
"Created external worker for model {} at {}",
@@ -106,7 +106,7 @@ impl StepExecutor<LocalWorkerWorkflowData> for CreateLocalWorkerStep {
let circuit_breaker_config = build_circuit_breaker_config(app_context);
// Build health config
let health_config = build_health_config(app_context);
let health_config = build_health_config(app_context, config);
// Normalize URL
let normalized_url = normalize_url(&config.url, connection_mode);
@@ -270,7 +270,7 @@ fn build_circuit_breaker_config(app_context: &AppContext) -> CircuitBreakerConfi
}
}
fn build_health_config(app_context: &AppContext) -> HealthConfig {
fn build_health_config(app_context: &AppContext, config: &WorkerConfigRequest) -> HealthConfig {
let cfg = &app_context.router_config.health_check;
HealthConfig {
timeout_secs: cfg.timeout_secs,
@@ -278,6 +278,7 @@ fn build_health_config(app_context: &AppContext) -> HealthConfig {
endpoint: cfg.endpoint.clone(),
failure_threshold: cfg.failure_threshold,
success_threshold: cfg.success_threshold,
disable_health_check: cfg.disable_health_check || config.disable_health_check,
}
}
@@ -334,7 +335,11 @@ fn create_dp_aware_workers(
}
let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
worker.set_healthy(false);
if health_config.disable_health_check {
worker.set_healthy(true);
} else {
worker.set_healthy(false);
}
workers.push(worker);
debug!(
@@ -358,6 +363,8 @@ fn create_single_worker(
config: &WorkerConfigRequest,
final_labels: &HashMap<String, String>,
) -> Vec<Arc<dyn Worker>> {
let health_check_disabled = health_config.disable_health_check;
let mut builder = BasicWorkerBuilder::new(normalized_url.to_string())
.model(model_card)
.worker_type(worker_type)
@@ -374,7 +381,11 @@ fn create_single_worker(
}
let worker = Arc::new(builder.build()) as Arc<dyn Worker>;
worker.set_healthy(false);
if health_check_disabled {
worker.set_healthy(true);
} else {
worker.set_healthy(false);
}
debug!(
"Created worker object for {} ({:?}) with {} labels",
@@ -80,6 +80,9 @@ impl StepExecutor<WorkerUpdateWorkflowData> for UpdateWorkerPropertiesStep {
success_threshold: request
.health_success_threshold
.unwrap_or(existing_health.success_threshold),
disable_health_check: request
.disable_health_check
.unwrap_or(existing_health.disable_health_check),
};
// Determine API key: use new one if provided, otherwise keep existing
+15
View File
@@ -531,6 +531,8 @@ pub struct HealthConfig {
pub failure_threshold: u32,
/// Number of consecutive successes before marking healthy
pub success_threshold: u32,
/// Whether to disable health checks for this worker
pub disable_health_check: bool,
}
impl Default for HealthConfig {
@@ -541,6 +543,7 @@ impl Default for HealthConfig {
endpoint: "/health".to_string(),
failure_threshold: 3,
success_threshold: 2,
disable_health_check: false,
}
}
}
@@ -698,6 +701,13 @@ 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);
}
return Ok(());
}
let health_result = match &self.metadata.connection_mode {
ConnectionMode::Http => self.http_health_check().await?,
ConnectionMode::Grpc { .. } => self.grpc_health_check().await?,
@@ -1249,6 +1259,7 @@ pub fn worker_to_info(worker: &Arc<dyn Worker>) -> WorkerInfo {
chat_template: worker.chat_template(model_id).map(String::from),
bootstrap_port,
metadata: worker.metadata().labels.clone(),
disable_health_check: worker.metadata().health_config.disable_health_check,
job_status: None,
}
}
@@ -1322,6 +1333,7 @@ mod tests {
assert_eq!(config.endpoint, "/health");
assert_eq!(config.failure_threshold, 3);
assert_eq!(config.success_threshold, 2);
assert!(!config.disable_health_check);
}
#[test]
@@ -1332,12 +1344,14 @@ mod tests {
endpoint: "/healthz".to_string(),
failure_threshold: 5,
success_threshold: 3,
disable_health_check: true,
};
assert_eq!(config.timeout_secs, 10);
assert_eq!(config.check_interval_secs, 60);
assert_eq!(config.endpoint, "/healthz");
assert_eq!(config.failure_threshold, 5);
assert_eq!(config.success_threshold, 3);
assert!(config.disable_health_check);
}
#[test]
@@ -1376,6 +1390,7 @@ mod tests {
endpoint: "/custom-health".to_string(),
failure_threshold: 4,
success_threshold: 2,
disable_health_check: false,
};
use crate::core::BasicWorkerBuilder;
@@ -397,6 +397,7 @@ mod tests {
check_interval_secs: 60,
failure_threshold: 3,
success_threshold: 2,
disable_health_check: false,
};
let cb_config = CircuitBreakerConfig {
@@ -489,6 +490,7 @@ mod tests {
check_interval_secs: 45,
failure_threshold: 5,
success_threshold: 3,
disable_health_check: false,
};
let worker = DPAwareWorkerBuilder::new("http://localhost:8080", 3, 16)
@@ -671,6 +671,7 @@ impl WorkerRegistry {
// This is especially important when there are many workers
let health_futures: Vec<_> = workers
.iter()
.filter(|worker| !worker.metadata().health_config.disable_health_check)
.map(|worker| {
let worker = worker.clone();
async move {
+5
View File
@@ -378,6 +378,10 @@ struct CliArgs {
#[arg(long, default_value = "/health", help_heading = "Health Checks")]
health_check_endpoint: String,
/// Disable all worker health checks at startup
#[arg(long, default_value_t = false, help_heading = "Health Checks")]
disable_health_check: bool,
// ==================== Tokenizer ====================
/// Model path for loading tokenizer (HuggingFace ID or local path)
#[arg(long, help_heading = "Tokenizer")]
@@ -987,6 +991,7 @@ impl CliArgs {
timeout_secs: self.health_check_timeout_secs,
check_interval_secs: self.health_check_interval_secs,
endpoint: self.health_check_endpoint.clone(),
disable_health_check: self.disable_health_check,
})
.tokenizer_cache(TokenizerCacheConfig {
enable_l0: self.tokenizer_cache_enable_l0,
@@ -80,6 +80,10 @@ pub struct WorkerConfigRequest {
#[serde(default = "default_health_failure_threshold")]
pub health_failure_threshold: u32,
/// Disable periodic health checks for this worker (default: false)
#[serde(default)]
pub disable_health_check: bool,
/// Maximum connection attempts during worker registration (default: 20)
#[serde(default = "default_max_connection_attempts")]
pub max_connection_attempts: u32,
@@ -165,6 +169,9 @@ pub struct WorkerInfo {
#[serde(skip_serializing_if = "HashMap::is_empty")]
pub metadata: HashMap<String, String>,
/// Whether health checks are disabled for this worker
pub disable_health_check: bool,
/// Job status for async operations (if available)
#[serde(skip_serializing_if = "Option::is_none")]
pub job_status: Option<JobStatus>,
@@ -191,6 +198,7 @@ impl WorkerInfo {
chat_template: None,
bootstrap_port: None,
metadata: HashMap::new(),
disable_health_check: false,
job_status,
}
}
@@ -271,6 +279,10 @@ pub struct WorkerUpdateRequest {
/// Update health failure threshold
#[serde(skip_serializing_if = "Option::is_none")]
pub health_failure_threshold: Option<u32>,
/// Disable periodic health checks for this worker
#[serde(skip_serializing_if = "Option::is_none")]
pub disable_health_check: Option<bool>,
}
/// Generic API response
+11 -7
View File
@@ -880,13 +880,17 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
let router_manager = RouterManager::from_config(&config, &app_context).await?;
let router: Arc<dyn RouterTrait> = router_manager.clone();
let _health_checker = app_context
.worker_registry
.start_health_checker(config.router_config.health_check.check_interval_secs);
debug!(
"Started health checker for workers with {}s interval",
config.router_config.health_check.check_interval_secs
);
if !config.router_config.health_check.disable_health_check {
let _health_checker = app_context
.worker_registry
.start_health_checker(config.router_config.health_check.check_interval_secs);
debug!(
"Started health checker for workers with {}s interval",
config.router_config.health_check.check_interval_secs
);
} else {
info!("Global health checks disabled via CLI/config; skipping health checker");
}
if let Some(ref load_monitor) = app_context.load_monitor {
load_monitor.start().await;
@@ -470,6 +470,7 @@ async fn handle_pod_event(
.check_interval_secs,
health_success_threshold: app_context.router_config.health_check.success_threshold,
health_failure_threshold: app_context.router_config.health_check.failure_threshold,
disable_health_check: app_context.router_config.health_check.disable_health_check,
max_connection_attempts: app_context.router_config.health_check.success_threshold
* 20,
dp_aware: app_context.router_config.dp_aware,
@@ -41,6 +41,7 @@ async fn test_policy_registry_with_router_manager() {
health_check_interval_secs: 60,
health_success_threshold: 2,
health_failure_threshold: 3,
disable_health_check: false,
max_connection_attempts: 20,
dp_aware: false,
};
@@ -73,6 +74,7 @@ async fn test_policy_registry_with_router_manager() {
health_check_interval_secs: 60,
health_success_threshold: 2,
health_failure_threshold: 3,
disable_health_check: false,
max_connection_attempts: 20,
dp_aware: false,
};
@@ -101,6 +103,7 @@ async fn test_policy_registry_with_router_manager() {
health_check_interval_secs: 60,
health_success_threshold: 2,
health_failure_threshold: 3,
disable_health_check: false,
max_connection_attempts: 20,
dp_aware: false,
};