From ad1f4b126567a833b9700446896e61437912174a Mon Sep 17 00:00:00 2001 From: Anirudh Ravichandran Date: Sat, 26 Sep 2026 09:58:51 -0700 Subject: [PATCH] fix: discover late workers' models on the next health check A worker that isn't serving when the router starts is registered under the 'unknown' placeholder model, and requests that name a model only reach workers indexed under it. The late worker got none of that traffic until the model refresh on every 5th health check (4 minutes after startup at the default 60s interval), while sticky sessions pinned to the early workers; in prime-rl job 1498 all 335 opening sessions landed on 2 of 16 servers. The health checker now re-reads /v1/models on every check for workers still under the placeholder; the rest keep the 5-check refresh for LoRA changes. GET /workers also reports each worker's live models, since model_id is fixed at registration and kept showing 'unknown' after routing recovered. Model refreshes run concurrently, at most 32 at a time, and each result is applied as it arrives, so unreachable workers cost a health-check cycle one 5s request timeout in total rather than one each. Registering a worker at a URL that is already registered now removes the old worker from every index first, so the models discovered for it don't carry over to the replacement. Registration, removal and model syncs now take one registry lock, so a model sync that races a replacement or removal can't put the dropped worker back into the model index. Replacing a worker now hands back the worker it replaced, and the router manager releases that worker's model from its policy counts before counting the replacement; removal releases the model of the worker actually removed. A worker indexed only under the placeholder, even more than once, is retried on every check. --- src/core/mod.rs | 2 +- src/core/worker.rs | 6 +- src/core/worker_registry.rs | 183 +++++++++++++++++++++++++++------- src/protocols/worker_spec.rs | 4 + src/routers/http/router.rs | 64 ++++++++++-- src/routers/router_manager.rs | 82 ++++++++++++--- src/server.rs | 1 + 7 files changed, 277 insertions(+), 65 deletions(-) diff --git a/src/core/mod.rs b/src/core/mod.rs index 5938779be..34cc6f62b 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -22,6 +22,6 @@ pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecuto pub use worker::{ fetch_models_from_worker, start_health_checker, BasicWorker, ConnectionMode, DPAwareWorker, HealthChecker, HealthConfig, Worker, WorkerCollection, WorkerFactory, WorkerLoadGuard, - WorkerType, + WorkerType, UNKNOWN_MODEL_ID, }; pub use worker_registry::{strip_dp_rank, WorkerId, WorkerRegistry, WorkerRegistryStats}; diff --git a/src/core/worker.rs b/src/core/worker.rs index 28a9303e8..4cb930be0 100644 --- a/src/core/worker.rs +++ b/src/core/worker.rs @@ -15,6 +15,10 @@ static WORKER_CLIENT: LazyLock = LazyLock::new(|| { .expect("Failed to create worker HTTP client") }); +/// Placeholder `model_id` for a worker registered without one, e.g. because its +/// `/v1/models` could not be read yet. +pub const UNKNOWN_MODEL_ID: &str = "unknown"; + /// Query a worker's `/v1/models` endpoint and return the list of model IDs. /// /// Returns an empty vec if the endpoint is unreachable or returns an unexpected format. @@ -217,7 +221,7 @@ pub trait Worker: Send + Sync + fmt::Debug { .labels .get("model_id") .map(|s| s.as_str()) - .unwrap_or("unknown") + .unwrap_or(UNKNOWN_MODEL_ID) } /// Get the priority of this worker (higher value = higher priority) diff --git a/src/core/worker_registry.rs b/src/core/worker_registry.rs index 69a2290ee..d1666175c 100644 --- a/src/core/worker_registry.rs +++ b/src/core/worker_registry.rs @@ -2,8 +2,9 @@ //! //! Provides centralized registry for workers with model-based indexing -use crate::core::{ConnectionMode, Worker, WorkerType}; +use crate::core::{ConnectionMode, Worker, WorkerType, UNKNOWN_MODEL_ID}; use dashmap::DashMap; +use parking_lot::Mutex; use std::sync::{Arc, RwLock}; use tracing::{info, warn}; use uuid::Uuid; @@ -78,6 +79,11 @@ pub struct WorkerRegistry { /// the Worker trait's model_id() reads from immutable metadata. /// A worker can serve multiple models (e.g. base model + LoRA adapters). known_models: Arc>>, + + /// Serializes the writers (`register`, `remove`, `sync_worker_models`). Each + /// updates several maps, so a model sync interleaved with a replacement or + /// removal would re-index the worker that was just dropped. Reads don't lock. + write_lock: Arc>, } impl WorkerRegistry { @@ -91,16 +97,32 @@ impl WorkerRegistry { connection_workers: Arc::new(DashMap::new()), url_to_id: Arc::new(DashMap::new()), known_models: Arc::new(DashMap::new()), + write_lock: Arc::new(Mutex::new(())), } } - /// Register a new worker + /// Register a worker, replacing any worker already registered at its URL pub fn register(&self, worker: Arc) -> WorkerId { - let worker_id = if let Some(existing_id) = self.url_to_id.get(worker.url()) { - // Worker with this URL already exists, update it - existing_id.clone() - } else { - WorkerId::new() + self.register_replacing(worker).0 + } + + /// `register`, also returning the worker it replaced, so callers that track + /// per-model state (e.g. `PolicyRegistry` counts) can release the old one. + pub fn register_replacing( + &self, + worker: Arc, + ) -> (WorkerId, Option>) { + let _guard = self.write_lock.lock(); + // A replacement keeps the existing ID, but the old worker leaves every + // index first, including the models discovered for it, so none of them + // outlive it. + let existing_id = self.url_to_id.get(worker.url()).map(|id| id.clone()); + let (worker_id, replaced) = match existing_id { + Some(id) => { + let replaced = self.remove_unlocked(&id); + (id, replaced) + } + None => (WorkerId::new(), None), }; // Store worker @@ -137,11 +159,17 @@ impl WorkerRegistry { .or_default() .push(worker_id.clone()); - worker_id + (worker_id, replaced) } /// Remove a worker by ID pub fn remove(&self, worker_id: &WorkerId) -> Option> { + let _guard = self.write_lock.lock(); + self.remove_unlocked(worker_id) + } + + /// `remove` for callers that already hold `write_lock`. + fn remove_unlocked(&self, worker_id: &WorkerId) -> Option> { if let Some((_, worker)) = self.workers.remove(worker_id) { // Remove from URL mapping self.url_to_id.remove(worker.url()); @@ -192,11 +220,9 @@ impl WorkerRegistry { /// Remove a worker by URL pub fn remove_by_url(&self, url: &str) -> Option> { - if let Some((_, worker_id)) = self.url_to_id.remove(url) { - self.remove(&worker_id) - } else { - None - } + let _guard = self.write_lock.lock(); + let (_, worker_id) = self.url_to_id.remove(url)?; + self.remove_unlocked(&worker_id) } /// Get a worker by ID @@ -305,6 +331,7 @@ impl WorkerRegistry { /// it to indexes for newly discovered models. This supports workers that /// serve multiple models (e.g. a base model + LoRA adapters). pub fn sync_worker_models(&self, url: &str, new_models: &[String]) { + let _guard = self.write_lock.lock(); let Some(worker) = self.get_by_url(url) else { return; }; @@ -364,6 +391,25 @@ impl WorkerRegistry { .insert(url.to_string(), new_models.to_vec()); } + /// Models a worker is indexed under: the set last discovered from its + /// `/v1/models`, else the `model_id` label it was registered with. + pub fn worker_models(&self, url: &str) -> Vec { + if let Some(models) = self.known_models.get(url) { + return models.clone(); + } + self.get_by_url(url) + .map(|worker| vec![worker.model_id().to_string()]) + .unwrap_or_default() + } + + /// Whether a worker is in no real model's pool: every model it is indexed + /// under is the placeholder. + fn has_unknown_model(&self, url: &str) -> bool { + self.worker_models(url) + .iter() + .all(|model| model == UNKNOWN_MODEL_ID) + } + /// Get all model IDs with workers pub fn get_models(&self) -> Vec { self.model_workers @@ -462,9 +508,12 @@ impl WorkerRegistry { /// /// Periodically checks `/health` on every worker and refreshes the model index /// by querying `/v1/models`. If a worker's loaded model changes (e.g. `LoRA` - /// load/evict), the model index is updated automatically. + /// load/evict), the model index is updated automatically. Workers whose models + /// are still unknown are queried on every check. #[must_use] pub fn start_health_checker(&self, check_interval_secs: u64) -> crate::core::HealthChecker { + use futures::stream::{self, StreamExt}; + use std::collections::HashMap; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -475,6 +524,7 @@ impl WorkerRegistry { let handle = tokio::spawn(async move { const LOAD_RESET_INTERVAL: u64 = 10; const MODEL_REFRESH_INTERVAL: u64 = 5; + const MAX_CONCURRENT_MODEL_FETCHES: usize = 32; let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs)); @@ -528,33 +578,42 @@ impl WorkerRegistry { check_count += 1; - // Periodically refresh model discovery + // Refresh model discovery: every MODEL_REFRESH_INTERVAL checks + // to pick up LoRA adapter loads and evictions, and on every + // check for workers still under the placeholder model (not + // serving yet when registered), so they join their model's + // pool as soon as they are up instead of minutes later. // TODO: notify PolicyRegistry on model changes so per-model // policies stay in sync (requires a callback or moving this // logic to a layer that has access to both registries). - if check_count.is_multiple_of(MODEL_REFRESH_INTERVAL) { - // Deduplicate by base URL so DP workers (@0, @1, …) - // sharing the same endpoint only trigger one fetch. - let mut fetched: std::collections::HashMap> = - std::collections::HashMap::new(); - - for worker in &workers { - if !worker.is_healthy() { - continue; - } - let base_url = strip_dp_rank(worker.url()).to_string(); - let new_models = match fetched.entry(base_url.clone()) { - std::collections::hash_map::Entry::Occupied(e) => e.get().clone(), - std::collections::hash_map::Entry::Vacant(e) => { - let models = - crate::core::worker::fetch_models_from_worker(&base_url).await; - e.insert(models.clone()); - models - } - }; - if !new_models.is_empty() { - registry.sync_worker_models(worker.url(), &new_models); - } + let full_refresh = check_count.is_multiple_of(MODEL_REFRESH_INTERVAL); + // Group the URLs of workers due for a refresh by base URL, so DP + // workers (@0, @1, …) sharing an endpoint share one fetch. + let mut due: HashMap> = HashMap::new(); + for worker in &workers { + if worker.is_healthy() + && (full_refresh || registry.has_unknown_model(worker.url())) + { + due.entry(strip_dp_rank(worker.url()).to_string()) + .or_default() + .push(worker.url().to_string()); + } + } + // Fetch concurrently (bounded) and apply each result as it + // arrives: unreachable workers cost the cycle about one request + // timeout in total and don't hold up the rest. + let mut fetches = stream::iter(due) + .map(|(base_url, urls)| async move { + let models = crate::core::worker::fetch_models_from_worker(&base_url).await; + (urls, models) + }) + .buffer_unordered(MAX_CONCURRENT_MODEL_FETCHES); + while let Some((urls, models)) = fetches.next().await { + if models.is_empty() { + continue; + } + for url in urls { + registry.sync_worker_models(&url, &models); } } @@ -849,4 +908,52 @@ mod tests { "LoRA index should be empty after worker removal" ); } + + /// Re-registering a URL (e.g. POST /workers for a worker that now serves + /// another model) must not carry over what was discovered for the old one. + #[test] + fn test_reregistering_url_replaces_discovered_models() { + let registry = WorkerRegistry::new(); + let url = "http://worker1:8000"; + let worker_serving = |model: &str| -> Arc { + let labels = HashMap::from([("model_id".to_string(), model.to_string())]); + Arc::from(WorkerFactory::create_regular_with_labels( + url.to_string(), + labels, + CircuitBreakerConfig::default(), + )) + }; + + let old_id = registry.register(worker_serving("old-model")); + registry.sync_worker_models(url, &["old-model".to_string(), "old-lora".to_string()]); + + let (new_id, replaced) = registry.register_replacing(worker_serving("new-model")); + assert_eq!(new_id, old_id); + assert_eq!(replaced.unwrap().model_id(), "old-model"); + assert_eq!(registry.worker_models(url), ["new-model"]); + assert!(registry.get_by_model_fast("old-model").is_empty()); + assert!(registry.get_by_model_fast("old-lora").is_empty()); + assert_eq!(registry.get_by_model_fast("new-model").len(), 1); + assert_eq!(registry.get_by_type(&WorkerType::Regular).len(), 1); + } + + /// A worker whose models list only the placeholder, even more than once, is + /// still in no real model's pool and must be retried on every check. + #[test] + fn test_has_unknown_model_ignores_duplicate_placeholders() { + let registry = WorkerRegistry::new(); + let url = "http://worker1:8000"; + registry.register(Arc::new(crate::core::BasicWorker::new( + url.to_string(), + WorkerType::Regular, + ))); + assert!(registry.has_unknown_model(url)); + + let placeholder = UNKNOWN_MODEL_ID.to_string(); + registry.sync_worker_models(url, &[placeholder.clone(), placeholder]); + assert!(registry.has_unknown_model(url)); + + registry.sync_worker_models(url, &["base-model".to_string()]); + assert!(!registry.has_unknown_model(url)); + } } diff --git a/src/protocols/worker_spec.rs b/src/protocols/worker_spec.rs index 58287909a..dc9663459 100644 --- a/src/protocols/worker_spec.rs +++ b/src/protocols/worker_spec.rs @@ -48,6 +48,10 @@ pub struct WorkerInfo { /// Model ID this worker serves pub model_id: String, + /// Models the worker is indexed under: last discovered from its + /// `/v1/models`, else `model_id`, which is fixed at registration + pub models: Vec, + /// Worker priority pub priority: u32, diff --git a/src/routers/http/router.rs b/src/routers/http/router.rs index 71635b463..c4f4a6c61 100644 --- a/src/routers/http/router.rs +++ b/src/routers/http/router.rs @@ -2494,27 +2494,43 @@ mod tests { (format!("http://{}", addr), handle) } + /// Helper: start a mock worker that lists `models` on `/v1/models`. Until + /// `ready_after` elapses, `/health` and `/v1/models` return 503, like an + /// engine that is still loading. async fn start_model_listing_worker( models: Vec<&'static str>, + ready_after: Duration, ) -> (String, tokio::task::JoinHandle<()>) { use axum::{routing::get, Json, Router as AxumRouter}; use tokio::net::TcpListener; - let models: Vec = models.into_iter().map(String::from).collect(); + let start = Instant::now(); + let data: Vec = models + .iter() + .map(|id| serde_json::json!({"id": id, "object": "model"})) + .collect(); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let app = AxumRouter::new() - .route("/health", get(|| async { StatusCode::OK })) + .route( + "/health", + get(move || async move { + if start.elapsed() < ready_after { + StatusCode::SERVICE_UNAVAILABLE + } else { + StatusCode::OK + } + }), + ) .route( "/v1/models", get(move || { - let models = models.clone(); + let data = data.clone(); async move { - let data: Vec = models - .iter() - .map(|id| serde_json::json!({"id": id, "object": "model"})) - .collect(); - Json(serde_json::json!({"object": "list", "data": data})) + if start.elapsed() < ready_after { + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + Json(serde_json::json!({"object": "list", "data": data})).into_response() } }), ); @@ -2592,7 +2608,8 @@ mod tests { #[tokio::test] async fn test_router_new_indexes_all_discovered_models_immediately() { - let (url, _handle) = start_model_listing_worker(vec!["base-model", "rft-run-1"]).await; + let (url, _handle) = + start_model_listing_worker(vec!["base-model", "rft-run-1"], Duration::ZERO).await; let config = crate::config::types::RouterConfig { mode: crate::config::types::RoutingMode::Regular { worker_urls: vec![url.clone()], @@ -2626,6 +2643,35 @@ mod tests { assert_eq!(worker.url(), url); } + #[tokio::test] + async fn test_late_worker_joins_model_index_on_next_health_check() { + // Still loading when registered, so the worker is filed under the + // placeholder model, as when it misses Router::new's model discovery. + let (url, _handle) = + start_model_listing_worker(vec!["base-model"], Duration::from_millis(500)).await; + let registry = WorkerRegistry::new(); + registry.register(Arc::new(BasicWorker::new(url.clone(), WorkerType::Regular))); + assert!(registry.get_by_model_fast("base-model").is_empty()); + + // Checks run every second and the periodic full refresh is only due on + // the 5th (~4s), so the model has to come from the per-check retry. + let health_checker = registry.start_health_checker(1); + let deadline = Instant::now() + Duration::from_secs(3); + while registry.get_by_model_fast("base-model").is_empty() { + assert!( + Instant::now() < deadline, + "late worker was not indexed under its model before the periodic refresh" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + assert!(registry + .get_by_model_fast(crate::core::UNKNOWN_MODEL_ID) + .is_empty()); + assert_eq!(registry.worker_models(&url), ["base-model"]); + + health_checker.shutdown().await; + } + #[test] fn test_inline_header_conversion_matches_headers_to_request_headers() { // Verify that the inline header conversion pattern used in vllm_pd_router diff --git a/src/routers/router_manager.rs b/src/routers/router_manager.rs index 939e3350f..2eacd7137 100644 --- a/src/routers/router_manager.rs +++ b/src/routers/router_manager.rs @@ -211,7 +211,13 @@ impl RouterManager { // Register worker let worker_arc: Arc = Arc::from(worker); - let worker_id = self.worker_registry.register(worker_arc.clone()); + let (worker_id, replaced) = self.worker_registry.register_replacing(worker_arc.clone()); + + // A worker replaced at the same URL leaves its model's policy count, as + // on removal, before the replacement joins its own model's + if let Some(replaced) = replaced { + self.policy_registry.on_worker_removed(replaced.model_id()); + } // Notify PolicyRegistry about the new worker // Extract policy hint from labels if provided @@ -252,21 +258,12 @@ impl RouterManager { &self, url: &str, ) -> Result { - // Get worker to extract model_id before removing - let model_id = self - .worker_registry - .get_by_url(url) - .map(|worker| worker.model_id().to_string()); - - if let Some(_worker) = self.worker_registry.remove_by_url(url) { - // Notify PolicyRegistry about worker removal - if let Some(ref model_id) = model_id { - self.policy_registry.on_worker_removed(model_id); - - info!("Removed worker with URL {} for model {}", url, model_id); - } else { - info!("Removed worker with URL {}", url); - } + // Take the model from the worker actually removed, so a replacement that + // lands in between can't make us release the wrong model's count + if let Some(worker) = self.worker_registry.remove_by_url(url) { + let model_id = worker.model_id(); + self.policy_registry.on_worker_removed(model_id); + info!("Removed worker with URL {} for model {}", url, model_id); Ok(WorkerApiResponse { success: true, @@ -348,6 +345,7 @@ impl RouterManager { id: id.to_string(), url: worker.url().to_string(), model_id: worker.model_id().to_string(), + models: self.worker_registry.worker_models(worker.url()), priority: worker.priority(), cost: worker.cost(), worker_type: match worker.worker_type() { @@ -846,3 +844,55 @@ impl std::fmt::Debug for RouterManager { .finish() } } + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + fn worker_request(url: &str, model: &str) -> WorkerConfigRequest { + WorkerConfigRequest { + url: url.to_string(), + model_id: Some(model.to_string()), + worker_type: None, + priority: None, + cost: None, + labels: HashMap::new(), + bootstrap_port: None, + } + } + + /// Replacing a worker at the same URL releases the old model's policy count, + /// and removing the replacement releases its own, so nothing leaks. + #[tokio::test] + async fn test_replacing_worker_rebalances_policy_counts() { + let policy_registry = Arc::new(crate::policies::PolicyRegistry::new( + crate::config::PolicyConfig::RoundRobin, + )); + let manager = RouterManager::new( + RouterConfig::default(), + reqwest::Client::new(), + Arc::new(WorkerRegistry::new()), + policy_registry.clone(), + ); + let url = "http://worker1:8000"; + + manager + .add_worker(worker_request(url, "old-model")) + .await + .unwrap(); + manager + .add_worker(worker_request(url, "new-model")) + .await + .unwrap(); + assert_eq!( + policy_registry.get_worker_counts(), + HashMap::from([("new-model".to_string(), 1)]) + ); + assert!(policy_registry.get_policy("old-model").is_none()); + + manager.remove_worker_from_registry(url).unwrap(); + assert!(policy_registry.get_worker_counts().is_empty()); + assert!(policy_registry.get_policy("new-model").is_none()); + } +} diff --git a/src/server.rs b/src/server.rs index e90b73e75..1dac205b5 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1094,6 +1094,7 @@ async fn list_workers_rest( let mut worker_info = serde_json::json!({ "url": worker.url(), "model_id": worker.model_id(), + "models": state.context.worker_registry.worker_models(worker.url()), "worker_type": match worker.worker_type() { WorkerType::Regular => "regular", WorkerType::Prefill { .. } => "prefill",