diff --git a/Cargo.toml b/Cargo.toml index 9489359c..5a6c041a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -89,6 +89,7 @@ portpicker = "0.1" tempfile = "3.8" lazy_static = "1.4" opentelemetry_sdk = { version = "0.27", features = ["testing"] } +tokio = { version = "1.42.0", features = ["full", "test-util"] } [[bench]] name = "request_processing" diff --git a/src/core/mod.rs b/src/core/mod.rs index 2fe028df..66b3af10 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -20,7 +20,7 @@ pub use circuit_breaker::{ pub use error::{WorkerError, WorkerResult}; pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor}; pub use worker::{ - start_health_checker, BasicWorker, ConnectionMode, DPAwareWorker, HealthChecker, HealthConfig, - Worker, WorkerCollection, WorkerFactory, WorkerLoadGuard, WorkerType, + BasicWorker, ConnectionMode, DPAwareWorker, HealthChecker, HealthConfig, Worker, + WorkerCollection, WorkerFactory, WorkerLoadGuard, WorkerType, }; pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats}; diff --git a/src/core/worker.rs b/src/core/worker.rs index f48a8a2d..98ce7e18 100644 --- a/src/core/worker.rs +++ b/src/core/worker.rs @@ -1,7 +1,6 @@ use super::{CircuitBreaker, CircuitBreakerConfig, WorkerError, WorkerResult}; use crate::metrics::RouterMetrics; use async_trait::async_trait; -use futures; use serde_json; use std::fmt; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; @@ -58,7 +57,10 @@ pub trait Worker: Send + Sync + fmt::Debug { /// Decrement the load counter fn decrement_load(&self); - /// Reset the load counter to 0 (for sync/recovery) + /// Reset the load counter to 0. Administrative operation: only safe when + /// the caller has independently established that no requests are in + /// flight for this worker. The runtime never calls this automatically; + /// health checks and routing must not discard active request accounting. fn reset_load(&self) { // Default implementation - does nothing // Workers that track load should override this @@ -448,18 +450,31 @@ impl Worker for BasicWorker { fn increment_load(&self) { self.load_counter.fetch_add(1, Ordering::Relaxed); + RouterMetrics::set_worker_load(self.url(), self.load()); } fn decrement_load(&self) { - self.load_counter + if self + .load_counter .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { current.checked_sub(1) }) - .ok(); + .is_err() + { + tracing::warn!( + worker_url = %self.metadata.url, + "Attempted to decrement load counter that is already at 0" + ); + } + RouterMetrics::set_worker_load(self.url(), self.load()); } + /// Administrative operation: resets the load counter to 0. Only safe when + /// the caller has independently established that no requests are in + /// flight for this worker; the runtime never calls this automatically. fn reset_load(&self) { self.load_counter.store(0, Ordering::Relaxed); + RouterMetrics::set_worker_load(self.url(), 0); } fn processed_requests(&self) -> usize { @@ -812,36 +827,73 @@ pub fn workers_to_urls(workers: &[Box]) -> Vec { workers.iter().map(|w| w.url().to_string()).collect() } -/// RAII guard for worker load management -pub struct WorkerLoadGuard<'a> { - workers: Vec<&'a dyn Worker>, +/// RAII guard for worker load management. +/// +/// Increments the load counter of the tracked worker(s) on creation and +/// decrements exactly once when released (explicitly or on drop). Handles can +/// be shared via [`WorkerLoadGuard::share`] so that a streaming forward task +/// can own the guard while the routing task releases it early on retryable +/// failures; only the first release decrements. +pub struct WorkerLoadGuard { + inner: Arc, } -impl<'a> WorkerLoadGuard<'a> { +struct WorkerLoadGuardInner { + workers: Vec>, + released: AtomicBool, +} + +impl WorkerLoadGuard { /// Create a new load guard for a single worker - pub fn new(worker: &'a dyn Worker) -> Self { + pub fn new(worker: Arc) -> Self { worker.increment_load(); + RouterMetrics::set_running_requests(worker.url(), worker.load()); Self { - workers: vec![worker], + inner: Arc::new(WorkerLoadGuardInner { + workers: vec![worker], + released: AtomicBool::new(false), + }), } } /// Create a new load guard for multiple workers - pub fn new_multi(workers: Vec<&'a dyn Worker>) -> Self { + pub fn new_multi(workers: Vec>) -> Self { // Increment load counters for all workers for worker in &workers { worker.increment_load(); + RouterMetrics::set_running_requests(worker.url(), worker.load()); + } + Self { + inner: Arc::new(WorkerLoadGuardInner { + workers, + released: AtomicBool::new(false), + }), + } + } + + /// Share the guard without touching load counters. All shared handles + /// refer to the same release state: only the first release decrements. + pub fn share(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + } + } + + /// Decrement the tracked workers' load counters exactly once. Safe to + /// call multiple times; subsequent calls are no-ops. + pub fn release(&self) { + if !self.inner.released.swap(true, Ordering::AcqRel) { + for worker in &self.inner.workers { + worker.decrement_load(); + RouterMetrics::set_running_requests(worker.url(), worker.load()); + } } - Self { workers } } } -impl<'a> Drop for WorkerLoadGuard<'a> { +impl Drop for WorkerLoadGuard { fn drop(&mut self) { - // Decrement load counters for all workers - for worker in &self.workers { - worker.decrement_load(); - } + self.release(); } } @@ -872,95 +924,90 @@ impl HealthChecker { } } -/// Start an async background health checker for a collection of workers -pub fn start_health_checker( - workers: std::sync::Arc>>>, - check_interval_secs: u64, -) -> HealthChecker { - let shutdown = Arc::new(AtomicBool::new(false)); - let shutdown_clone = shutdown.clone(); - - let handle = tokio::spawn(async move { - let mut interval = - tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs)); - - // Counter for periodic load reset (every 10 health check cycles) - let mut check_count = 0u64; - const LOAD_RESET_INTERVAL: u64 = 10; - - loop { - interval.tick().await; - - // Check for shutdown signal - if shutdown_clone.load(Ordering::Acquire) { - tracing::debug!("Health checker shutting down"); - break; - } +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + use std::time::Duration; - check_count += 1; + fn test_worker() -> Arc { + Arc::new(BasicWorker::new( + "http://test:8080".to_string(), + WorkerType::Regular, + )) + } - // Check health of all workers - let workers_to_check = match workers.read() { - Ok(guard) => guard.clone(), - Err(poisoned) => { - tracing::error!("Worker lock poisoned: {}", poisoned); - continue; - } - }; + #[test] + fn test_worker_load_guard_drop_decrements() { + let worker = test_worker(); + { + let _guard = WorkerLoadGuard::new(worker.clone()); + assert_eq!(worker.load(), 1); + } + assert_eq!(worker.load(), 0); + } - // Periodically reset load counters to prevent drift - // Only do this when we believe all workers should be idle - if check_count.is_multiple_of(LOAD_RESET_INTERVAL) { - let max_load = workers_to_check.iter().map(|w| w.load()).max().unwrap_or(0); - // Only reset if load appears to be very low (likely drift) - if max_load <= 2 { - tracing::debug!( - "Resetting load counters to prevent drift (max_load: {})", - max_load - ); - for worker in &workers_to_check { - worker.reset_load(); - } - } - } + #[test] + fn test_worker_load_guard_release_is_idempotent() { + let worker = test_worker(); + let guard = WorkerLoadGuard::new(worker.clone()); + assert_eq!(worker.load(), 1); + guard.release(); + assert_eq!(worker.load(), 0); + guard.release(); + assert_eq!(worker.load(), 0); + drop(guard); + assert_eq!(worker.load(), 0); + } - // Perform health checks concurrently - let health_checks = workers_to_check.iter().map(|worker| { - let worker_url = worker.url().to_string(); - let was_healthy = worker.is_healthy(); - - async move { - match worker.check_health_async().await { - Ok(_) => { - if !was_healthy { - tracing::info!("Worker {} is now healthy", worker_url); - } - } - Err(e) => { - if was_healthy { - tracing::warn!("Worker {} health check failed: {}", worker_url, e); - } else { - // Worker was already unhealthy, log at debug level - tracing::debug!("Worker {} remains unhealthy: {}", worker_url, e); - } - } - } - } - }); + #[test] + fn test_worker_load_guard_shared_handle_decrements_once() { + let worker = test_worker(); + let guard = WorkerLoadGuard::new(worker.clone()); + let shared = guard.share(); + assert_eq!(worker.load(), 1); + // First release wins regardless of which handle triggers it. + drop(shared); + assert_eq!(worker.load(), 0); + drop(guard); + assert_eq!(worker.load(), 0, "second release must be a no-op"); + } - // Execute all health checks concurrently - futures::future::join_all(health_checks).await; - } - }); + #[test] + fn test_worker_load_guard_release_then_shared_drop_is_noop() { + // Streaming retryable path: release immediately, then hand a shared + // handle to the forward task; its drop must not decrement again. + let worker = test_worker(); + let guard = WorkerLoadGuard::new(worker.clone()); + guard.release(); + assert_eq!(worker.load(), 0); + let shared = guard.share(); + drop(shared); + drop(guard); + assert_eq!(worker.load(), 0); + } - HealthChecker { handle, shutdown } -} + #[test] + fn test_worker_load_guard_multi() { + let worker1 = test_worker(); + let worker2 = Arc::new(BasicWorker::new( + "http://test:8081".to_string(), + WorkerType::Regular, + )) as Arc; + let guard = WorkerLoadGuard::new_multi(vec![worker1.clone(), worker2.clone()]); + assert_eq!(worker1.load(), 1); + assert_eq!(worker2.load(), 1); + guard.release(); + assert_eq!(worker1.load(), 0); + assert_eq!(worker2.load(), 0); + } -#[cfg(test)] -mod tests { - use super::*; - use std::thread; - use std::time::Duration; + #[test] + fn test_decrement_at_zero_clamps() { + let worker = test_worker(); + worker.decrement_load(); + assert_eq!(worker.load(), 0); + } // Test WorkerType #[test] @@ -1398,11 +1445,14 @@ mod tests { // Test WorkerLoadGuard #[test] fn test_load_guard_single_worker() { - let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + let worker: Arc = Arc::new(BasicWorker::new( + "http://test:8080".to_string(), + WorkerType::Regular, + )); assert_eq!(worker.load(), 0); { - let _guard = WorkerLoadGuard::new(&worker); + let _guard = WorkerLoadGuard::new(worker.clone()); assert_eq!(worker.load(), 1); } @@ -1412,16 +1462,23 @@ mod tests { #[test] fn test_load_guard_multiple_workers() { - let workers: Vec> = vec![ - WorkerFactory::create_regular("http://w1:8080".to_string()), - WorkerFactory::create_regular("http://w2:8080".to_string()), - WorkerFactory::create_regular("http://w3:8080".to_string()), + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8080".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8080".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w3:8080".to_string(), + WorkerType::Regular, + )), ]; - let worker_refs: Vec<&dyn Worker> = workers.iter().map(|w| w.as_ref()).collect(); - { - let _guard = WorkerLoadGuard::new_multi(worker_refs); + let _guard = WorkerLoadGuard::new_multi(workers.clone()); // All loads incremented assert_eq!(workers[0].load(), 1); assert_eq!(workers[1].load(), 1); @@ -1436,23 +1493,20 @@ mod tests { #[test] fn test_load_guard_panic_safety() { - let worker = Arc::new(BasicWorker::new( + let worker: Arc = Arc::new(BasicWorker::new( "http://test:8080".to_string(), WorkerType::Regular, )); assert_eq!(worker.load(), 0); - // Clone for use inside catch_unwind - let worker_clone = Arc::clone(&worker); - // Use AssertUnwindSafe wrapper for the test // This is safe because we're only testing the load counter behavior. use std::panic::AssertUnwindSafe; // This will panic, but the guard should still clean up let result = std::panic::catch_unwind(AssertUnwindSafe(|| { - let _guard = WorkerLoadGuard::new(worker_clone.as_ref()); - assert_eq!(worker_clone.load(), 1); + let _guard = WorkerLoadGuard::new(worker.clone()); + assert_eq!(worker.load(), 1); panic!("Test panic"); })); diff --git a/src/core/worker_registry.rs b/src/core/worker_registry.rs index ec9e3409..b1d23356 100644 --- a/src/core/worker_registry.rs +++ b/src/core/worker_registry.rs @@ -364,10 +364,6 @@ impl WorkerRegistry { let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs)); - // Counter for periodic load reset (every 10 health check cycles) - let mut check_count = 0u64; - const LOAD_RESET_INTERVAL: u64 = 10; - loop { interval.tick().await; @@ -383,19 +379,25 @@ impl WorkerRegistry { .map(|entry| entry.value().clone()) .collect(); - // Perform health checks - for worker in &workers { - let _ = worker.check_health_async().await; // Use async version directly - } - - // Reset loads periodically - check_count += 1; - if check_count.is_multiple_of(LOAD_RESET_INTERVAL) { - tracing::debug!("Resetting worker loads (cycle {})", check_count); - for worker in &workers { - worker.reset_load(); - } - } + // Perform health checks in parallel. Health checking must not + // mutate load accounting: a failed /health probe does not + // cancel in-flight requests, so a worker marked unhealthy can + // still hold valid load when it recovers. Reset on recovery + // would erase that load (or erase a newly routed request's + // increment racing the reset), recreating the undercount that + // this checker previously caused. WorkerLoadGuard keeps + // increments and decrements paired on every request path, so + // no reset is needed. + let health_futures: Vec<_> = workers + .iter() + .map(|worker| { + let worker = worker.clone(); + async move { + let _ = worker.check_health_async().await; + } + }) + .collect(); + futures::future::join_all(health_futures).await; } }); @@ -424,9 +426,76 @@ pub struct WorkerRegistryStats { #[cfg(test)] mod tests { use super::*; - use crate::core::{CircuitBreakerConfig, WorkerFactory}; + use crate::core::error::WorkerResult; + use crate::core::worker::WorkerMetadata; + use crate::core::{BasicWorker, CircuitBreakerConfig, HealthConfig, WorkerFactory}; + use async_trait::async_trait; use std::collections::HashMap; + /// Test worker that skips real HTTP health checks so the health-checker + /// loop can be driven deterministically under paused Tokio time (no I/O). + #[derive(Debug)] + struct NoopHealthWorker(Arc); + + #[async_trait] + impl Worker for NoopHealthWorker { + fn url(&self) -> &str { + self.0.url() + } + + fn worker_type(&self) -> WorkerType { + self.0.worker_type() + } + + fn connection_mode(&self) -> ConnectionMode { + self.0.connection_mode() + } + + fn is_healthy(&self) -> bool { + self.0.is_healthy() + } + + fn set_healthy(&self, healthy: bool) { + self.0.set_healthy(healthy); + } + + async fn check_health_async(&self) -> WorkerResult<()> { + Ok(()) + } + + fn load(&self) -> usize { + self.0.load() + } + + fn increment_load(&self) { + self.0.increment_load(); + } + + fn decrement_load(&self) { + self.0.decrement_load(); + } + + fn reset_load(&self) { + self.0.reset_load(); + } + + fn processed_requests(&self) -> usize { + self.0.processed_requests() + } + + fn increment_processed(&self) { + self.0.increment_processed(); + } + + fn metadata(&self) -> &WorkerMetadata { + self.0.metadata() + } + + fn circuit_breaker(&self) -> &crate::core::CircuitBreaker { + self.0.circuit_breaker() + } + } + #[test] fn test_worker_registry() { let registry = WorkerRegistry::new(); @@ -523,4 +592,103 @@ mod tests { assert_eq!(llama_workers_after.len(), 1); assert_eq!(llama_workers_after[0].url(), "http://worker2:8080"); } + + /// Start a mock HTTP server whose /health endpoint always returns 200. + async fn start_healthy_mock_server() -> (String, tokio::task::JoinHandle<()>) { + use axum::{http::StatusCode, routing::get, Router as AxumRouter}; + use tokio::net::TcpListener; + + 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 })); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (format!("http://{}", addr), handle) + } + + fn registry_worker(url: &str) -> Arc { + Arc::new( + BasicWorker::new(url.to_string(), WorkerType::Regular).with_health_config( + HealthConfig { + timeout_secs: 2, + check_interval_secs: 1, + endpoint: "/health".to_string(), + failure_threshold: 3, + success_threshold: 1, + }, + ), + ) + } + + /// Regression test for issue #197: health checks must not reset the load + /// counters of workers with requests in flight. + /// + /// Uses paused Tokio time and a no-I/O health check to advance through + /// well over 10 health-check cycles, the cadence at which the previous + /// implementation reset load counters. This test fails against that + /// implementation and passes with the fix. + #[tokio::test(start_paused = true)] + async fn test_health_checker_preserves_inflight_load() { + let registry = WorkerRegistry::new(); + let worker: Arc = Arc::new(NoopHealthWorker(Arc::new(BasicWorker::new( + "http://worker:8080".to_string(), + WorkerType::Regular, + )))); + registry.register(worker.clone()); + + // Simulate 5 requests in flight. + for _ in 0..5 { + worker.increment_load(); + } + assert_eq!(worker.load(), 5); + + // Advance through 12 health-check cycles (the first tick fires + // immediately, so this covers check counts well past the old + // 10-cycle reset interval). + let checker = registry.start_health_checker(1); + for _ in 0..12 { + tokio::time::advance(std::time::Duration::from_secs(1)).await; + tokio::task::yield_now().await; + } + + assert_eq!( + worker.load(), + 5, + "health checker must not reset the load counter of a worker with in-flight requests" + ); + + // Under paused time the checker task cannot observe a shutdown flag + // until the next tick fires, so drop the handle instead of awaiting + // shutdown; the test runtime aborts the task. + drop(checker); + } + + /// Load must survive an unhealthy->healthy recovery transition: a failed + /// /health probe does not cancel in-flight requests, so the load a worker + /// reports at recovery time may be real work, not drift. + #[tokio::test] + async fn test_health_checker_preserves_load_across_recovery() { + let (url, _server) = start_healthy_mock_server().await; + let registry = WorkerRegistry::new(); + let worker = registry_worker(&url); + worker.set_healthy(false); + // Load accumulated while the worker was down. Some of it may still + // be in flight when the worker recovers. + for _ in 0..3 { + worker.increment_load(); + } + registry.register(worker.clone()); + + let checker = registry.start_health_checker(1); + tokio::time::sleep(std::time::Duration::from_millis(2500)).await; + checker.shutdown().await; + + assert!(worker.is_healthy(), "worker should have recovered"); + assert_eq!( + worker.load(), + 3, + "load must be preserved across the unhealthy->healthy transition" + ); + } } diff --git a/src/policies/cache_aware.rs b/src/policies/cache_aware.rs index dc209a19..6ac63acd 100644 --- a/src/policies/cache_aware.rs +++ b/src/policies/cache_aware.rs @@ -171,57 +171,16 @@ impl CacheAwarePolicy { } } - fn select_worker_min_load( - &self, - workers: &[Arc], - request_text: Option<&str>, - healthy_indices: &[usize], - model_id: &str, - max_load: usize, - min_load: usize, - ) -> Option { - // Log load balancing trigger (only compute worker loads if debug enabled) - if tracing::enabled!(tracing::Level::DEBUG) { - let worker_loads: Vec<(&str, usize)> = - workers.iter().map(|w| (w.url(), w.load())).collect(); - debug!( - "Load balancing triggered | max: {} | min: {} | workers: {:?}", - max_load, min_load, worker_loads - ); - } - - RouterMetrics::record_load_balancing_event(); - RouterMetrics::set_load_range(max_load, min_load); - - // Use shortest queue when imbalanced - let min_load_idx = healthy_indices - .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied()?; - - // Even in imbalanced mode, update the tree to maintain cache state - if let Some(text) = request_text { - // Get the tree reference without locking the entire HashMap - // DashMap only locks the specific shard containing this key - let tree = self.trees.get(model_id).map(|entry| entry.value().clone()); - - if let Some(tree) = tree { - let worker_url = workers[min_load_idx].url(); - tree.insert(text, worker_url); - } else { - debug!( - "Warning: No tree found for model '{}', skipping cache update", - model_id - ); - } - } - - // Increment processed counter - workers[min_load_idx].increment_processed(); - RouterMetrics::record_processed_request(workers[min_load_idx].url()); - RouterMetrics::record_policy_decision(self.name(), workers[min_load_idx].url()); - - Some(min_load_idx) + /// Is this specific worker hot enough that a request should be steered away + /// from it, even at the cost of a cache miss? + /// + /// Same thresholds as before, but the comparison is `this worker` against the + /// least-loaded worker, not `the busiest worker` against the least-loaded one. + /// A worker that is merely busier than its peers keeps serving its cached + /// prefixes; only one that is genuinely running away gives them up. + fn is_worker_overloaded(&self, worker_load: usize, min_load: usize) -> bool { + worker_load.saturating_sub(min_load) > self.config.balance_abs_threshold + && (worker_load as f32) > (min_load as f32 * self.config.balance_rel_threshold) } } @@ -249,27 +208,28 @@ impl LoadBalancingPolicy for CacheAwarePolicy { }); let min_load = if min_load == usize::MAX { 0 } else { min_load }; - // Check if load is imbalanced - let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold - && (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold); - debug!( - "Load status for model: max_load={}, min_load={}, is_imbalanced={}", - max_load, min_load, is_imbalanced + "Load status for model: max_load={}, min_load={}", + max_load, min_load ); - if is_imbalanced { - return self.select_worker_min_load( - workers, - request_text, - &healthy_indices, - model_id, - max_load, - min_load, - ); - } - - // Use cache-aware routing when balanced + // NOTE: the load check is applied per request, against the worker this + // request actually wants, rather than fleet-wide. The previous behaviour + // asked "is any pair of workers imbalanced?" and, if so, discarded cache + // affinity for *every* request -- including requests whose preferred + // worker was idle. Under prefill/decode disaggregation that is + // catastrophic: prefill worker load includes queued requests, so at high + // concurrency the fleet spread clears the threshold almost permanently, + // routing degenerates to shortest-queue, and the prefix cache stops being + // used at all. Measured on SWE-bench Pro (4 prefill / 6 decode, 1024 in + // parallel), that took the prefill prefix cache hit rate from 88.9% to + // 68.5% -- a 2.8x increase in prompt tokens actually computed -- and + // halved end-to-end throughput. + // + // Shedding load only matters for the worker that is actually hot, so that + // is the only case where affinity is given up. See is_worker_overloaded. + + // Use cache-aware routing let text = request_text.unwrap_or(""); // Get the tree reference without locking the entire HashMap @@ -314,10 +274,31 @@ impl LoadBalancingPolicy for CacheAwarePolicy { let selected_idx = if match_rate > self.config.cache_threshold { // Cache hit path: find worker by URL (compare &str directly, no allocation) let tenant_url: &str = &result.tenant; - workers + let cached_idx = workers .iter() .position(|w| w.url() == tenant_url) - .filter(|&idx| workers[idx].is_healthy()) + .filter(|&idx| workers[idx].is_healthy()); + + match cached_idx { + // The worker holding this prefix is running away from the rest of + // the fleet, so pay the cache miss and shed the load. This is the + // hot-spot case that load balancing exists for. + Some(idx) if self.is_worker_overloaded(workers[idx].load(), min_load) => { + RouterMetrics::record_load_balancing_event(); + RouterMetrics::set_load_range(max_load, min_load); + debug!( + "Cached worker {} overloaded (load={}, min={}), steering away", + workers[idx].url(), + workers[idx].load(), + min_load + ); + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied() + } + other => other, + } } else { // Low cache match: use worker with minimum load healthy_indices @@ -542,6 +523,140 @@ mod tests { } } + /// Helper: pin a worker's load counter to an exact value. + fn set_load(worker: &Arc, n: usize) { + worker.reset_load(); + for _ in 0..n { + worker.increment_load(); + } + } + + fn three_workers() -> Vec> { + vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w3:8000".to_string(), + WorkerType::Regular, + )), + ] + } + + fn balance_policy() -> CacheAwarePolicy { + CacheAwarePolicy::with_config(CacheAwareConfig { + cache_threshold: 0.5, + balance_abs_threshold: 5, + balance_rel_threshold: 2.0, + eviction_interval_secs: 0, + max_tree_size: 10000, + }) + } + + /// Regression test: a hot worker elsewhere in the fleet must not cost cache + /// affinity for requests that want an idle worker. + /// + /// The previous implementation gated on `max_load - min_load` across the whole + /// fleet, so one runaway worker disabled prefix routing for every request. Under + /// P/D disaggregation, where prefill load counters include queued requests, that + /// gate is essentially always open at high concurrency and the prefix cache stops + /// being used at all. + #[test] + fn test_affinity_survives_unrelated_hot_worker() { + let policy = balance_policy(); + let workers = three_workers(); + policy.init_workers(&workers); + + // Make w3 strictly the least loaded so the first request tenants to it. + set_load(&workers[0], 2); + set_load(&workers[1], 1); + set_load(&workers[2], 0); + let primed = policy + .select_worker(&workers, Some("conversation-A")) + .unwrap(); + assert_eq!(primed, 2, "setup: first request should tenant to w3"); + + // Now w1 runs away. Fleet is wildly imbalanced (100 vs 0), which the old + // fleet-wide gate would have treated as "ignore the cache entirely". + // w3 -- the worker actually holding this prefix -- stays cold. + set_load(&workers[0], 100); + set_load(&workers[1], 0); + set_load(&workers[2], 1); + + for _ in 0..5 { + let idx = policy + .select_worker(&workers, Some("conversation-A")) + .unwrap(); + assert_eq!( + idx, 2, + "affinity must be preserved: the cached worker w3 is not the hot one" + ); + } + } + + /// The other half of the contract: when the worker holding the prefix is itself + /// the runaway, the cache miss is worth paying and the request is steered away. + /// This is the hot-spot case load balancing exists for. + #[test] + fn test_affinity_yields_when_cached_worker_is_the_hot_one() { + let policy = balance_policy(); + let workers = three_workers(); + policy.init_workers(&workers); + + // Tenant "conversation-B" to w1. + set_load(&workers[0], 0); + set_load(&workers[1], 1); + set_load(&workers[2], 2); + let primed = policy + .select_worker(&workers, Some("conversation-B")) + .unwrap(); + assert_eq!(primed, 0, "setup: first request should tenant to w1"); + + // w1 is now the runaway and holds the prefix. + set_load(&workers[0], 100); + set_load(&workers[1], 0); + set_load(&workers[2], 0); + + for _ in 0..5 { + let idx = policy + .select_worker(&workers, Some("conversation-B")) + .unwrap(); + assert_ne!(idx, 0, "must steer away from the overloaded cached worker"); + } + } + + /// A worker that is merely busier than its peers keeps serving its cached + /// prefixes -- the override is for runaways, not for ordinary variation. + #[test] + fn test_affinity_survives_mild_load_difference() { + let policy = balance_policy(); + let workers = three_workers(); + policy.init_workers(&workers); + + set_load(&workers[0], 2); + set_load(&workers[1], 1); + set_load(&workers[2], 0); + let primed = policy + .select_worker(&workers, Some("conversation-C")) + .unwrap(); + assert_eq!(primed, 2); + + // w3 is busier than the others but well inside balance_abs_threshold (5). + set_load(&workers[0], 0); + set_load(&workers[1], 0); + set_load(&workers[2], 4); + + let idx = policy + .select_worker(&workers, Some("conversation-C")) + .unwrap(); + assert_eq!(idx, 2, "a 4-request lead is not a hot spot"); + } + #[test] fn test_cache_aware_worker_removal() { let config = CacheAwareConfig { diff --git a/src/routers/http/router.rs b/src/routers/http/router.rs index f68e4b08..51618efc 100644 --- a/src/routers/http/router.rs +++ b/src/routers/http/router.rs @@ -1,7 +1,7 @@ use crate::config::types::RetryConfig; use crate::core::{ is_retryable_status, BasicWorker, CircuitBreakerConfig, DPAwareWorker, HealthConfig, - RetryExecutor, Worker, WorkerRegistry, WorkerType, + RetryExecutor, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType, }; use crate::metrics::RouterMetrics; use crate::otel_http::{self, ClientRequestOptions}; @@ -571,17 +571,11 @@ impl Router { None => self.policy_registry.get_default_policy(), }; - let load_incremented = if policy.name() == "cache_aware" { - worker.increment_load(); - RouterMetrics::set_running_requests(worker.url(), worker.load()); - true - } else { - false - }; - - // Keep a clone for potential cleanup on retry - let worker_for_cleanup = if load_incremented { - Some(worker.clone()) + // The load guard owns the load accounting for this attempt: + // it increments on creation and decrements exactly once on + // release (drop, retry transfer, or stream completion). + let load_guard = if policy.name() == "cache_aware" { + Some(WorkerLoadGuard::new(worker.clone())) } else { None }; @@ -593,7 +587,7 @@ impl Router { route, worker.url(), is_stream, - load_incremented, + load_guard, ) .await; @@ -602,18 +596,6 @@ impl Router { let status = response.status(); worker.record_outcome(status.is_success() || status.is_client_error()); - // For retryable failures, we need to decrement load since send_typed_request - // won't have done it (it only decrements on success or non-retryable failures) - if is_retryable_status(response.status()) && load_incremented { - if let Some(cleanup_worker) = worker_for_cleanup { - cleanup_worker.decrement_load(); - RouterMetrics::set_running_requests( - cleanup_worker.url(), - cleanup_worker.load(), - ); - } - } - response }, // should_retry predicate @@ -771,7 +753,7 @@ impl Router { route: &str, worker_url: &str, is_stream: bool, - load_incremented: bool, // Whether load was incremented for this request + load_guard: Option, // Load guard for cache-aware policy ) -> Response { let (mut request_builder, extracted_dp_rank, request_url) = if self.intra_node_data_parallel_size > 1 { @@ -856,14 +838,7 @@ impl Router { worker_url, route, e ); - // Decrement load on error if it was incremented - if load_incremented { - if let Some(worker) = self.worker_registry.get_by_url(worker_url) { - worker.decrement_load(); - RouterMetrics::set_running_requests(worker_url, worker.load()); - } - } - + // The load guard is dropped on return, decrementing load. return ( StatusCode::INTERNAL_SERVER_ERROR, format!("Request failed: {}", e), @@ -887,32 +862,29 @@ impl Router { response } Err(e) => { - // IMPORTANT: Decrement load on error before returning - if load_incremented { - if let Some(worker) = self.worker_registry.get_by_url(worker_url) { - worker.decrement_load(); - RouterMetrics::set_running_requests(worker_url, worker.load()); - } - } - + // The load guard is dropped on return, decrementing load. let error_msg = format!("Failed to get response body: {}", e); (StatusCode::INTERNAL_SERVER_ERROR, error_msg).into_response() } }; - // Decrement load counter for non-streaming requests if it was incremented - if load_incremented { - if let Some(worker) = self.worker_registry.get_by_url(worker_url) { - worker.decrement_load(); - RouterMetrics::set_running_requests(worker_url, worker.load()); - } - } + // The load guard is dropped when this function returns, after the + // response body has been fully buffered, decrementing load once. response - } else if load_incremented { - // For streaming with load tracking, we need to manually decrement when done - let registry = Arc::clone(&self.worker_registry); - let worker_url = worker_url.to_string(); + } else if let Some(guard) = load_guard { + // Streaming with load tracking: the guard moves into the forward + // task and is released when the stream ends (completion, error, + // or client disconnect). For retryable statuses the load is + // released immediately, before the retry executor selects the + // next worker. + let status_is_retryable = is_retryable_status(status); + let task_guard = if status_is_retryable { + guard.release(); + guard.share() + } else { + guard + }; // Preserve headers for streaming response let mut response_headers = header_utils::preserve_response_headers(res.headers()); @@ -922,25 +894,14 @@ impl Router { let stream = res.bytes_stream(); let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - // Spawn task to forward stream and detect completion + // Spawn task to forward stream; the guard is released on drop + // when the task finishes. tokio::spawn(async move { + let _guard = task_guard; let mut stream = stream; - let mut decremented = false; while let Some(chunk) = stream.next().await { match chunk { Ok(bytes) => { - // Check for stream end marker - if bytes - .as_ref() - .windows(12) - .any(|window| window == b"data: [DONE]") - { - if let Some(worker) = registry.get_by_url(&worker_url) { - worker.decrement_load(); - RouterMetrics::set_running_requests(&worker_url, worker.load()); - decremented = true; - } - } if tx.send(Ok(bytes)).is_err() { break; } @@ -951,12 +912,6 @@ impl Router { } } } - if !decremented { - if let Some(worker) = registry.get_by_url(&worker_url) { - worker.decrement_load(); - RouterMetrics::set_running_requests(&worker_url, worker.load()); - } - } }); let stream = UnboundedReceiverStream::new(rx); diff --git a/src/routers/http/vllm_pd_router.rs b/src/routers/http/vllm_pd_router.rs index 52012c60..9d06c124 100644 --- a/src/routers/http/vllm_pd_router.rs +++ b/src/routers/http/vllm_pd_router.rs @@ -6,7 +6,7 @@ use super::pd_router::PdRouterBase; use super::pd_types::{error_chain, PDRouterError}; use super::vllm_service_discovery::{MoriIOTransferMode, ServiceRegistry, ServiceType}; use crate::config::KvConnector; -use crate::core::{BasicWorker, Worker, WorkerType}; +use crate::core::{BasicWorker, Worker, WorkerLoadGuard, WorkerType}; use crate::metrics::RouterMetrics; use crate::otel_http::{self, ClientRequestOptions}; use crate::policies::PolicyRegistry; @@ -1174,8 +1174,11 @@ impl VllmPDRouter { path ); - // Increment prefill load at the start of the prefill phase - prefill_worker.increment_load(); + // Increment prefill load at the start of the prefill phase. The guard + // releases the load on every exit path, including task cancellation + // (e.g. client disconnect dropping the handler future), which the + // previous manual decrements leaked. + let prefill_guard = WorkerLoadGuard::new(prefill_worker.clone()); let prefill_zmq_addr = self.get_zmq_address(prefill_worker.base_url(), ServiceType::Prefill); @@ -1270,7 +1273,6 @@ impl VllmPDRouter { { Ok(resp) => resp, Err(e) => { - prefill_worker.decrement_load(); let full_error = error_chain(&e); let duration = start_time.elapsed(); RouterMetrics::record_pd_prefill_error(&prefill_base_url); @@ -1292,7 +1294,6 @@ impl VllmPDRouter { let prefill_bytes = match prefill_response.bytes().await { Ok(bytes) => bytes, Err(e) => { - prefill_worker.decrement_load(); let full_error = error_chain(&e); let duration = start_time.elapsed(); RouterMetrics::record_pd_prefill_error(&prefill_base_url); @@ -1322,7 +1323,6 @@ impl VllmPDRouter { let prefill_response_json: Value = match serde_json::from_slice(&prefill_bytes) { Ok(json) => json, Err(e) => { - prefill_worker.decrement_load(); let duration = start_time.elapsed(); RouterMetrics::record_pd_prefill_error(&prefill_base_url); RouterMetrics::record_pd_request(path); @@ -1348,9 +1348,10 @@ impl VllmPDRouter { // Stop profiling on prefill server after its work is done self.stop_profiling(&prefill_base_url).await; - // Prefill phase complete: decrement prefill load, increment decode load - prefill_worker.decrement_load(); - decode_worker.increment_load(); + // Prefill phase complete: release prefill load, track decode load. + // The decode guard releases on every exit path, including cancellation. + prefill_guard.release(); + let decode_guard = WorkerLoadGuard::new(decode_worker.clone()); debug!("✅ vLLM Stage 1 completed, starting Stage 2 - Decode"); @@ -1451,7 +1452,6 @@ impl VllmPDRouter { { Ok(resp) => resp, Err(e) => { - decode_worker.decrement_load(); let full_error = error_chain(&e); let duration = start_time.elapsed(); RouterMetrics::record_pd_decode_error(&decode_base_url); @@ -1467,8 +1467,8 @@ impl VllmPDRouter { // Stop profiling on decode server after response received self.stop_profiling(&decode_base_url).await; - // Decode phase complete: decrement decode load - decode_worker.decrement_load(); + // Decode phase complete: release decode load + decode_guard.release(); let status = decode_response.status(); let headers = decode_response.headers().clone(); @@ -2418,6 +2418,7 @@ impl WorkerManagement for VllmPDRouter { #[cfg(test)] mod tests { use super::*; + use axum::extract::Json; use serde_json::json; #[test] @@ -2689,4 +2690,221 @@ mod tests { let params = moriio_write_decode_params(Some("tx-abc"), 4, Some(2)); assert_eq!(params["remote_dp_rank"], 2); } + + // --- Load-accounting cancellation tests (regression for #197 review P1) --- + + /// Build a VllmPDRouter without AppContext for focused tests. + fn test_vllm_pd_router() -> VllmPDRouter { + let worker_registry = Arc::new(crate::core::WorkerRegistry::new()); + let policy_registry = + Arc::new(PolicyRegistry::new(crate::config::PolicyConfig::RoundRobin)); + let pd_router = PdRouterBase { + worker_registry, + policy_registry: policy_registry.clone(), + worker_startup_timeout_secs: 5, + worker_startup_check_interval_secs: 1, + worker_loads: Arc::new(tokio::sync::watch::channel(HashMap::new()).1), + load_monitor_handle: None, + client: reqwest::Client::new(), + circuit_breaker_config: crate::core::CircuitBreakerConfig::default(), + dp_size: 1, + }; + VllmPDRouter { + pd_router, + service_registry: Arc::new(ServiceRegistry::new()), + http_client: reqwest::Client::new(), + policy_registry, + use_discovery: false, + enable_profiling: false, + profile_timeout_secs: 0, + profiling_tasks: Arc::new(Mutex::new(HashMap::new())), + intra_node_data_parallel_size: 1, + prefill_dp_round_robin: Arc::new(AtomicUsize::new(0)), + kv_connector: KvConnector::Nixl, + mooncake_prefill_info: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// Mock server whose POST /v1/chat/completions handler signals `entered` + /// and then waits for `release` before responding. + async fn start_pending_mock_server() -> ( + String, + Arc, + Arc, + tokio::task::JoinHandle<()>, + ) { + use axum::routing::post; + use tokio::net::TcpListener; + + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let entered_clone = entered.clone(); + let release_clone = release.clone(); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let app = axum::Router::new().route( + "/v1/chat/completions", + post(move |Json(_body): Json| { + let entered = entered_clone.clone(); + let release = release_clone.clone(); + async move { + entered.notify_one(); + release.notified().await; + (axum::http::StatusCode::OK, "{}") + } + }), + ); + + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (format!("http://{}", addr), entered, release, handle) + } + + /// Mock server whose POST /v1/chat/completions responds immediately with a + /// JSON body containing kv_transfer_params (prefill stage). + async fn start_immediate_prefill_mock_server() -> (String, tokio::task::JoinHandle<()>) { + use axum::routing::post; + use tokio::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let app = axum::Router::new().route( + "/v1/chat/completions", + post(|Json(_body): Json| async move { + ( + axum::http::StatusCode::OK, + "{\"kv_transfer_params\":{\"dummy\":true}}", + ) + }), + ); + + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (format!("http://{}", addr), handle) + } + + async fn wait_for_load(worker: &Arc, expected: usize) { + for _ in 0..200 { + if worker.load() == expected { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + panic!( + "timed out waiting for worker load to reach {}, got {}", + expected, + worker.load() + ); + } + + /// Cancelling a two-stage request while the prefill response is pending + /// must release the prefill worker's load. + #[tokio::test] + async fn test_pd_request_cancellation_releases_prefill_load() { + let router = test_vllm_pd_router(); + let (prefill_url, entered, _release, _server) = start_pending_mock_server().await; + + let prefill_worker: Arc = Arc::new(BasicWorker::new( + prefill_url.clone(), + WorkerType::Prefill { + bootstrap_port: None, + }, + )); + let decode_worker: Arc = Arc::new(BasicWorker::new( + "http://127.0.0.1:1".to_string(), + WorkerType::Decode, + )); + + let prefill_task_worker = prefill_worker.clone(); + let decode_task_worker = decode_worker.clone(); + let handle = tokio::spawn(async move { + let request = json!({ + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + }); + let _ = router + .process_vllm_two_stage_request( + request, + prefill_task_worker, + decode_task_worker, + "/v1/chat/completions", + None, + ) + .await; + }); + + // Wait until the request reached the prefill worker and its load is + // counted, then cancel the whole request task. + entered.notified().await; + wait_for_load(&prefill_worker, 1).await; + handle.abort(); + let _ = handle.await; + + assert_eq!( + prefill_worker.load(), + 0, + "cancelling a P/D request must release its prefill load" + ); + assert_eq!(decode_worker.load(), 0); + } + + /// Cancelling a two-stage request while the decode response is pending + /// must release both the decode and prefill workers' load. + #[tokio::test] + async fn test_pd_request_cancellation_releases_decode_load() { + let router = test_vllm_pd_router(); + let (prefill_url, _prefill_server) = start_immediate_prefill_mock_server().await; + let (decode_url, decode_entered, _release, _decode_server) = + start_pending_mock_server().await; + + let prefill_worker: Arc = Arc::new(BasicWorker::new( + prefill_url.clone(), + WorkerType::Prefill { + bootstrap_port: None, + }, + )); + let decode_worker: Arc = + Arc::new(BasicWorker::new(decode_url.clone(), WorkerType::Decode)); + + let prefill_task_worker = prefill_worker.clone(); + let decode_task_worker = decode_worker.clone(); + let handle = tokio::spawn(async move { + let request = json!({ + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + }); + let _ = router + .process_vllm_two_stage_request( + request, + prefill_task_worker, + decode_task_worker, + "/v1/chat/completions", + None, + ) + .await; + }); + + // Wait until the prefill stage completed and the request reached the + // decode worker, then cancel the whole request task. + decode_entered.notified().await; + wait_for_load(&decode_worker, 1).await; + handle.abort(); + let _ = handle.await; + + assert_eq!( + decode_worker.load(), + 0, + "cancelling a P/D request must release its decode load" + ); + assert_eq!( + prefill_worker.load(), + 0, + "the prefill load must have been released at the transition" + ); + } }