From 171912a9e36f510318227577b34446f30ac02d74 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sat, 27 Dec 2025 01:53:03 -0500 Subject: [PATCH] [model-gateway] Add consistent hashing for ManualPolicy routing (#15907) --- .../bindings/python/sglang_router/router.py | 1 + sgl-model-gateway/bindings/python/src/lib.rs | 2 + sgl-model-gateway/src/config/types.rs | 15 +- sgl-model-gateway/src/config/validation.rs | 5 +- sgl-model-gateway/src/core/mod.rs | 2 +- sgl-model-gateway/src/core/worker_registry.rs | 149 ++++- .../src/observability/metrics.rs | 9 + .../src/policies/consistent_hashing.rs | 527 +++++++++++++++ sgl-model-gateway/src/policies/factory.rs | 18 +- sgl-model-gateway/src/policies/manual.rs | 604 ++++++++++++------ sgl-model-gateway/src/policies/mod.rs | 13 +- sgl-model-gateway/src/policies/registry.rs | 7 +- .../grpc/common/stages/worker_selection.rs | 14 +- .../src/routers/http/pd_router.rs | 12 +- sgl-model-gateway/src/routers/http/router.rs | 8 +- 15 files changed, 1180 insertions(+), 206 deletions(-) create mode 100644 sgl-model-gateway/src/policies/consistent_hashing.rs diff --git a/sgl-model-gateway/bindings/python/sglang_router/router.py b/sgl-model-gateway/bindings/python/sglang_router/router.py index e572134ae..95484cb2b 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router.py @@ -26,6 +26,7 @@ def policy_from_str(policy_str: Optional[str]) -> PolicyType: "power_of_two": PolicyType.PowerOfTwo, "bucket": PolicyType.Bucket, "manual": PolicyType.Manual, + "consistent_hashing": PolicyType.ConsistentHashing, } return policy_map[policy_str] diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index b27a5a29e..2ea27573a 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -13,6 +13,7 @@ pub enum PolicyType { PowerOfTwo, Bucket, Manual, + ConsistentHashing, } #[pyclass(eq)] @@ -416,6 +417,7 @@ impl Router { bucket_adjust_interval_secs: self.bucket_adjust_interval_secs, }, PolicyType::Manual => ConfigPolicyConfig::Manual, + PolicyType::ConsistentHashing => ConfigPolicyConfig::ConsistentHashing, } }; diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index c318114f1..fbc7bbc53 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -337,11 +337,19 @@ pub enum PolicyConfig { bucket_adjust_interval_secs: usize, }, - /// Manual routing policy supporting header-based routing: - /// - X-SMG-Target-Worker: Direct routing to a specific worker by URL - /// - X-SMG-Routing-Key: Consistent hash routing for session affinity + /// Manual routing policy with sticky sessions using DashMap. + /// - X-SMG-Routing-Key: Routes to a cached worker or assigns a new one + /// - Provides true sticky sessions with zero key redistribution on worker add + /// - Falls back to random selection if no routing key is provided #[serde(rename = "manual")] Manual, + + /// Consistent hashing policy using hash ring for session affinity: + /// - X-SMG-Target-Worker: Direct routing to a specific worker by URL + /// - X-SMG-Routing-Key: Consistent hash routing for session affinity + /// - Provides O(log n) lookup with minimal redistribution (~1/N keys) on topology change + #[serde(rename = "consistent_hashing")] + ConsistentHashing, } impl PolicyConfig { @@ -353,6 +361,7 @@ impl PolicyConfig { PolicyConfig::PowerOfTwo { .. } => "power_of_two", PolicyConfig::Bucket { .. } => "bucket", PolicyConfig::Manual => "manual", + PolicyConfig::ConsistentHashing => "consistent_hashing", } } } diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index 363bf6723..2209b1c7b 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/sgl-model-gateway/src/config/validation.rs @@ -147,7 +147,10 @@ impl ConfigValidator { fn validate_policy(policy: &PolicyConfig) -> ConfigResult<()> { match policy { - PolicyConfig::Random | PolicyConfig::RoundRobin | PolicyConfig::Manual => {} + PolicyConfig::Random + | PolicyConfig::RoundRobin + | PolicyConfig::Manual + | PolicyConfig::ConsistentHashing => {} PolicyConfig::CacheAware { cache_threshold, balance_abs_threshold: _, diff --git a/sgl-model-gateway/src/core/mod.rs b/sgl-model-gateway/src/core/mod.rs index 8755fd050..cbf6cda9f 100644 --- a/sgl-model-gateway/src/core/mod.rs +++ b/sgl-model-gateway/src/core/mod.rs @@ -41,5 +41,5 @@ pub use worker::{ }; pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder}; pub use worker_manager::{LoadMonitor, WorkerManager}; -pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats}; +pub use worker_registry::{HashRing, WorkerId, WorkerRegistry, WorkerRegistryStats}; pub use worker_service::{WorkerService, WorkerServiceError}; diff --git a/sgl-model-gateway/src/core/worker_registry.rs b/sgl-model-gateway/src/core/worker_registry.rs index 0db95b8b5..cf17fea14 100644 --- a/sgl-model-gateway/src/core/worker_registry.rs +++ b/sgl-model-gateway/src/core/worker_registry.rs @@ -5,6 +5,11 @@ //! # Performance Optimizations //! The model index uses immutable Arc snapshots instead of RwLock for lock-free reads. //! This is critical for high-concurrency scenarios where many requests query the same model. +//! +//! # Consistent Hash Ring +//! The registry maintains a pre-computed hash ring per model for O(log n) consistent hashing. +//! The ring is rebuilt only when workers are added/removed, not per-request. +//! Uses virtual nodes (150 per worker) for even distribution and blake3 for stable hashing. use std::sync::Arc; @@ -16,6 +21,118 @@ use crate::{ observability::metrics::Metrics, }; +/// Number of virtual nodes per physical worker for even distribution. +/// 150 is a common choice that provides good balance between memory and distribution. +const VIRTUAL_NODES_PER_WORKER: usize = 150; + +/// Consistent hash ring for O(log n) worker selection. +/// +/// Each worker is placed at multiple positions (virtual nodes) on the ring +/// based on hash(worker_url + vnode_index). This provides: +/// - Even key distribution across workers +/// - Minimal key redistribution when workers are added/removed (~1/N keys move) +/// - O(log n) lookup via binary search +/// +/// Uses blake3 for stable, fast hashing that's consistent across Rust versions. +#[derive(Debug, Clone)] +pub struct HashRing { + /// Sorted list of (ring_position, worker_url) + /// Multiple entries per worker (virtual nodes) for even distribution. + /// Uses Arc to share URL across all virtual nodes (150 refs vs 150 copies). + entries: Arc<[(u64, Arc)]>, +} + +impl HashRing { + /// Build a hash ring from a list of workers. + /// Creates VIRTUAL_NODES_PER_WORKER entries per worker for even distribution. + pub fn new(workers: &[Arc]) -> Self { + let mut entries: Vec<(u64, Arc)> = + Vec::with_capacity(workers.len() * VIRTUAL_NODES_PER_WORKER); + + for worker in workers { + // Create Arc once per worker, share across all virtual nodes + let url: Arc = Arc::from(worker.url()); + + // Create multiple virtual nodes per worker + for vnode in 0..VIRTUAL_NODES_PER_WORKER { + let vnode_key = format!("{}#{}", url, vnode); + let pos = Self::hash_position(&vnode_key); + entries.push((pos, Arc::clone(&url))); + } + } + + // Sort by ring position for binary search + entries.sort_unstable_by_key(|(pos, _)| *pos); + + Self { + entries: Arc::from(entries.into_boxed_slice()), + } + } + + /// Hash a string to a ring position using blake3 (stable across versions). + #[inline] + fn hash_position(s: &str) -> u64 { + let hash = blake3::hash(s.as_bytes()); + // Take first 8 bytes as u64 + u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap()) + } + + /// Find worker URL for a key using consistent hashing. + /// Returns the first healthy worker URL at or after the key's position (clockwise). + /// + /// - `key`: The routing key to hash + /// - `is_healthy`: Function to check if a worker URL is healthy + pub fn find_healthy_url(&self, key: &str, is_healthy: F) -> Option<&str> + where + F: Fn(&str) -> bool, + { + if self.entries.is_empty() { + return None; + } + + let key_pos = Self::hash_position(key); + + // Binary search to find first entry at or after key_pos + let start = self.entries.partition_point(|(pos, _)| *pos < key_pos); + + // Walk clockwise from start, wrapping around + // Track visited URLs to avoid checking same worker multiple times (virtual nodes) + let mut checked_urls = + std::collections::HashSet::with_capacity(self.worker_count().min(16)); + + for i in 0..self.entries.len() { + let (_, url) = &self.entries[(start + i) % self.entries.len()]; + let url_str: &str = url; + + // Skip if we already checked this worker (from another virtual node) + if !checked_urls.insert(url_str) { + continue; + } + + if is_healthy(url_str) { + return Some(url_str); + } + } + + None + } + + /// Check if the ring is empty + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Get the number of entries in the ring (including virtual nodes) + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Get the number of unique workers in the ring + pub fn worker_count(&self) -> usize { + self.entries.len() / VIRTUAL_NODES_PER_WORKER.max(1) + } +} + /// Unique identifier for a worker #[derive(Debug, Clone, Hash, Eq, PartialEq)] pub struct WorkerId(String); @@ -58,6 +175,10 @@ pub struct WorkerRegistry { /// Uses Arc<[T]> instead of Arc>> for lock-free reads. model_index: ModelIndex, + /// Consistent hash rings per model for O(log n) routing. + /// Rebuilt on worker add/remove (copy-on-write). + hash_rings: Arc>>, + /// Workers indexed by worker type type_workers: Arc>>, @@ -74,12 +195,29 @@ impl WorkerRegistry { Self { workers: Arc::new(DashMap::new()), model_index: Arc::new(DashMap::new()), + hash_rings: Arc::new(DashMap::new()), type_workers: Arc::new(DashMap::new()), connection_workers: Arc::new(DashMap::new()), url_to_id: Arc::new(DashMap::new()), } } + /// Rebuild the hash ring for a model based on current workers in the model index + fn rebuild_hash_ring(&self, model_id: &str) { + if let Some(workers) = self.model_index.get(model_id) { + let ring = HashRing::new(&workers); + self.hash_rings.insert(model_id.to_string(), Arc::new(ring)); + } else { + // No workers for this model, remove the ring + self.hash_rings.remove(model_id); + } + } + + /// Get the hash ring for a model (O(1) lookup) + pub fn get_hash_ring(&self, model_id: &str) -> Option> { + self.hash_rings.get(model_id).map(|r| Arc::clone(&r)) + } + /// Register a new worker pub fn register(&self, worker: Arc) -> WorkerId { let worker_id = if let Some(existing_id) = self.url_to_id.get(worker.url()) { @@ -100,7 +238,7 @@ impl WorkerRegistry { // This creates a new immutable snapshot with the added worker let model_id = worker.model_id().to_string(); self.model_index - .entry(model_id) + .entry(model_id.clone()) .and_modify(|existing| { // Create new snapshot with the additional worker let mut new_workers: Vec> = existing.iter().cloned().collect(); @@ -109,6 +247,9 @@ impl WorkerRegistry { }) .or_insert_with(|| Arc::from(vec![worker.clone()].into_boxed_slice())); + // Rebuild hash ring for this model + self.rebuild_hash_ring(&model_id); + // Update type index (clone needed for DashMap key ownership) self.type_workers .entry(worker.worker_type().clone()) @@ -149,7 +290,8 @@ impl WorkerRegistry { // Remove from model index using copy-on-write // Create new snapshot without the removed worker let worker_url = worker.url(); - if let Some(mut entry) = self.model_index.get_mut(worker.model_id()) { + let model_id = worker.model_id().to_string(); + if let Some(mut entry) = self.model_index.get_mut(&model_id) { let new_workers: Vec> = entry .iter() .filter(|w| w.url() != worker_url) @@ -158,6 +300,9 @@ impl WorkerRegistry { *entry = Arc::from(new_workers.into_boxed_slice()); } + // Rebuild hash ring for this model + self.rebuild_hash_ring(&model_id); + // Remove from type index if let Some(mut type_workers) = self.type_workers.get_mut(worker.worker_type()) { type_workers.retain(|id| id != worker_id); diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index 265b69f9d..8f65390a6 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -810,6 +810,15 @@ impl Metrics { .increment(1); } + /// Record consistent hashing policy execution branch for routing decisions + pub fn record_worker_consistent_hashing_policy_branch(branch: &'static str) { + counter!( + "smg_consistent_hashing_policy_branch_total", + "branch" => branch + ) + .increment(1); + } + /// Set running requests per worker pub fn set_worker_requests_active(worker: &str, count: usize) { gauge!( diff --git a/sgl-model-gateway/src/policies/consistent_hashing.rs b/sgl-model-gateway/src/policies/consistent_hashing.rs new file mode 100644 index 000000000..6da62e328 --- /dev/null +++ b/sgl-model-gateway/src/policies/consistent_hashing.rs @@ -0,0 +1,527 @@ +//! Consistent hashing routing policy with header-based routing support +//! +//! Supports two routing mechanisms via HTTP headers: +//! - `X-SMG-Target-Worker`: Direct routing by worker index (0-based), returns None if unavailable +//! - `X-SMG-Routing-Key`: Consistent hash routing for session affinity +//! +//! ## Consistent Hashing +//! +//! Uses a pre-computed hash ring from WorkerRegistry where: +//! 1. Each worker is placed at a fixed position based on hash(worker_url) +//! 2. Keys are hashed to the ring, then walk clockwise to find first healthy worker +//! 3. When workers scale up/down, only keys in the affected range redistribute (~1/N keys move) +//! +//! The ring is built once when workers are added/removed, not per-request. +//! This ensures O(log n) lookup performance. +//! +//! Complexity: O(log n) binary search + O(k) walk where k = consecutive unhealthy workers. + +use std::sync::Arc; + +use http::header::HeaderName; +use rand::Rng as _; + +use super::{LoadBalancingPolicy, SelectWorkerInfo}; +use crate::{core::Worker, observability::metrics::Metrics}; + +/// Header for direct worker targeting by index (0-based) +static HEADER_TARGET_WORKER: HeaderName = HeaderName::from_static("x-smg-target-worker"); +/// Header for consistent hash routing +static HEADER_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-key"); + +/// Execution branch for metrics +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Branch { + NoHealthyWorkers, + TargetWorkerHit, + TargetWorkerMiss, + RoutingKeyHit, + RandomFallback, +} + +impl Branch { + #[inline] + const fn as_str(&self) -> &'static str { + match self { + Self::NoHealthyWorkers => "no_healthy_workers", + Self::TargetWorkerHit => "target_worker_hit", + Self::TargetWorkerMiss => "target_worker_miss", + Self::RoutingKeyHit => "routing_key_hit", + Self::RandomFallback => "random_fallback", + } + } +} + +#[derive(Debug, Default)] +pub struct ConsistentHashingPolicy; + +impl ConsistentHashingPolicy { + pub fn new() -> Self { + Self + } + + /// Use consistent hashing to find a worker for the given key. + /// Uses pre-computed ring from SelectWorkerInfo if available. + /// + /// The ring returns a worker URL, which we then map to an index in the workers array. + /// This correctly handles filtered worker arrays since we match by URL, not by index. + /// + /// Complexity: O(n) to build healthy URL map + O(log n) ring lookup + O(k) walk + fn find_by_consistent_hash( + workers: &[Arc], + info: &SelectWorkerInfo, + key: &str, + ) -> Option { + // Build URL→index map for healthy workers: O(n) once, O(1) lookups + let healthy_url_to_idx: std::collections::HashMap<&str, usize> = workers + .iter() + .enumerate() + .filter(|(_, w)| w.is_healthy()) + .map(|(i, w)| (w.url(), i)) + .collect(); + + if healthy_url_to_idx.is_empty() { + return None; + } + + // Use pre-computed ring if available + if let Some(ref ring) = info.hash_ring { + // O(1) lookup per URL checked instead of O(n) + let url = ring.find_healthy_url(key, |url| healthy_url_to_idx.contains_key(url))?; + return healthy_url_to_idx.get(url).copied(); + } + + // Fallback: no ring provided, use simple modulo (less optimal but functional) + // This shouldn't happen in normal operation as WorkerSelectionStage provides the ring + let mut healthy_indices: Vec = healthy_url_to_idx.values().copied().collect(); + healthy_indices.sort_unstable(); // Ensure deterministic order + + // Use blake3 for consistent hashing in fallback too + let hash = blake3::hash(key.as_bytes()); + let hash_val = u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap()); + let idx = (hash_val as usize) % healthy_indices.len(); + Some(healthy_indices[idx]) + } + + fn select_worker_impl( + &self, + workers: &[Arc], + info: &SelectWorkerInfo, + ) -> (Option, Branch) { + if workers.is_empty() { + return (None, Branch::NoHealthyWorkers); + } + + // Extract routing headers - to_str() is O(1), just validates ASCII, no allocation + let target_worker = info + .headers + .and_then(|h| h.get(&HEADER_TARGET_WORKER)) + .and_then(|v| v.to_str().ok()) + .filter(|s| !s.is_empty()); + + let routing_key = info + .headers + .and_then(|h| h.get(&HEADER_ROUTING_KEY)) + .and_then(|v| v.to_str().ok()) + .filter(|s| !s.is_empty()); + + // Priority 1: X-SMG-Target-Worker - direct routing by worker index + // O(1) parse + O(1) bounds check + O(1) health check + if let Some(idx_str) = target_worker { + if let Ok(idx) = idx_str.parse::() { + if idx < workers.len() && workers[idx].is_healthy() { + return (Some(idx), Branch::TargetWorkerHit); + } + } + return (None, Branch::TargetWorkerMiss); + } + + // Priority 2: X-SMG-Routing-Key - consistent hash routing (O(log n)) + if let Some(key) = routing_key { + return match Self::find_by_consistent_hash(workers, info, key) { + Some(idx) => (Some(idx), Branch::RoutingKeyHit), + None => (None, Branch::NoHealthyWorkers), + }; + } + + // Priority 3: Implicit routing key from stable headers (session affinity) + let implicit_key = info.headers.and_then(|h| { + h.get("authorization") + .or_else(|| h.get("x-forwarded-for")) + .or_else(|| h.get("cookie")) + .and_then(|v| v.to_str().ok()) + .filter(|s| !s.is_empty()) + }); + + if let Some(key) = implicit_key { + return match Self::find_by_consistent_hash(workers, info, key) { + Some(idx) => (Some(idx), Branch::RoutingKeyHit), + None => (None, Branch::NoHealthyWorkers), + }; + } + + // Fallback: random selection (truly anonymous client) + let healthy_count = workers.iter().filter(|w| w.is_healthy()).count(); + if healthy_count == 0 { + return (None, Branch::NoHealthyWorkers); + } + + let random_healthy_idx = rand::rng().random_range(0..healthy_count); + let idx = workers + .iter() + .enumerate() + .filter(|(_, w)| w.is_healthy()) + .nth(random_healthy_idx) + .map(|(i, _)| i) + .unwrap(); + + (Some(idx), Branch::RandomFallback) + } +} + +impl LoadBalancingPolicy for ConsistentHashingPolicy { + fn select_worker(&self, workers: &[Arc], info: &SelectWorkerInfo) -> Option { + let (result, branch) = self.select_worker_impl(workers, info); + Metrics::record_worker_consistent_hashing_policy_branch(branch.as_str()); + result + } + + fn name(&self) -> &'static str { + "consistent_hashing" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::core::{BasicWorkerBuilder, HashRing, WorkerType}; + + fn headers_with_routing_key(key: &str) -> http::HeaderMap { + let mut headers = http::HeaderMap::new(); + headers.insert("x-smg-routing-key", key.parse().unwrap()); + headers + } + + fn headers_with_target_worker(idx: usize) -> http::HeaderMap { + let mut headers = http::HeaderMap::new(); + headers.insert("x-smg-target-worker", idx.to_string().parse().unwrap()); + headers + } + + fn create_workers(urls: &[&str]) -> Vec> { + urls.iter() + .map(|url| { + Arc::new( + BasicWorkerBuilder::new(*url) + .worker_type(WorkerType::Regular) + .build(), + ) as Arc + }) + .collect() + } + + #[test] + fn test_consistent_routing() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + + let headers = headers_with_routing_key("user-123"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (first_result, _) = policy.select_worker_impl(&workers, &info); + let first_idx = first_result.unwrap(); + + // Same key should always route to same worker + for _ in 0..10 { + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(first_idx)); + assert_eq!(branch, Branch::RoutingKeyHit); + } + } + + #[test] + fn test_different_keys_distribute() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + + let mut distribution = HashMap::new(); + for i in 0..100 { + let headers = headers_with_routing_key(&format!("user-{}", i)); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + let (result, _) = policy.select_worker_impl(&workers, &info); + *distribution.entry(result.unwrap()).or_insert(0) += 1; + } + + assert!(distribution.len() > 1, "Should distribute across workers"); + } + + #[test] + fn test_target_worker_hit() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + let headers = headers_with_target_worker(1); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(1)); + assert_eq!(branch, Branch::TargetWorkerHit); + } + + #[test] + fn test_target_worker_miss_out_of_bounds() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + let headers = headers_with_target_worker(5); // Out of bounds + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, Branch::TargetWorkerMiss); + } + + #[test] + fn test_target_worker_miss_unhealthy() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + workers[1].set_healthy(false); + + let headers = headers_with_target_worker(1); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, Branch::TargetWorkerMiss); + } + + #[test] + fn test_target_worker_priority_over_routing_key() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + let mut headers = http::HeaderMap::new(); + headers.insert("x-smg-target-worker", "1".parse().unwrap()); + headers.insert("x-smg-routing-key", "some-key".parse().unwrap()); + + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(1)); + assert_eq!(branch, Branch::TargetWorkerHit); + } + + #[test] + fn test_fallback_random_distribution() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + + // Without routing headers, should distribute randomly across workers + let mut distribution = HashMap::new(); + for _ in 0..100 { + let info = SelectWorkerInfo::default(); + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert!(result.is_some()); + assert_eq!(branch, Branch::RandomFallback); + *distribution.entry(result.unwrap()).or_insert(0) += 1; + } + + // Should distribute across multiple workers (not always same one) + assert!( + distribution.len() > 1, + "Random fallback should distribute across workers" + ); + } + + #[test] + fn test_no_healthy_workers() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000"]); + workers[0].set_healthy(false); + + let headers = headers_with_routing_key("test"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, Branch::NoHealthyWorkers); + } + + #[test] + fn test_empty_workers() { + let policy = ConsistentHashingPolicy::new(); + let workers: Vec> = vec![]; + + let info = SelectWorkerInfo::default(); + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, Branch::NoHealthyWorkers); + } + + #[test] + fn test_consistent_hash_minimal_redistribution() { + // Test that consistent hashing moves fewer keys than random redistribution + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&[ + "http://w0:8000", + "http://w1:8000", + "http://w2:8000", + "http://w3:8000", + ]); + let ring = Arc::new(HashRing::new(&workers)); + + // Record which worker each key routes to with all workers healthy + let mut key_to_worker_before: HashMap = HashMap::new(); + for i in 0..100 { + let key = format!("user-{}", i); + let headers = headers_with_routing_key(&key); + let info = SelectWorkerInfo { + headers: Some(&headers), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + let (result, _) = policy.select_worker_impl(&workers, &info); + key_to_worker_before.insert(key, result.unwrap()); + } + + // Mark worker 1 as unhealthy + workers[1].set_healthy(false); + + // Record new routing and count how many keys moved + let mut moved_count = 0; + for i in 0..100 { + let key = format!("user-{}", i); + let headers = headers_with_routing_key(&key); + let info = SelectWorkerInfo { + headers: Some(&headers), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + let (result, _) = policy.select_worker_impl(&workers, &info); + let new_worker = result.unwrap(); + let old_worker = key_to_worker_before[&key]; + + if new_worker != old_worker { + moved_count += 1; + } + } + + // With consistent hashing, approximately 1/N keys should move (N = worker count) + // Random redistribution would move approximately (N-1)/N = 75% of keys + // Verify we're significantly better than random (< 50% moved) + let keys_on_failed_worker = key_to_worker_before.values().filter(|&&w| w == 1).count(); + assert!( + moved_count <= keys_on_failed_worker + 5, + "Consistent hashing should only move keys from failed worker (+small variance). \ + Expected ~{}, got {}", + keys_on_failed_worker, + moved_count + ); + assert!( + moved_count < 50, + "Consistent hashing should move fewer than 50% of keys (random would move ~75%), got {}%", + moved_count + ); + } + + #[test] + fn test_routing_key_failover_and_recovery() { + // Test that when a worker fails, keys move to another worker, + // and when it recovers, keys return to the original worker + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w0:8000", "http://w1:8000", "http://w2:8000"]); + let ring = Arc::new(HashRing::new(&workers)); + + // Find which worker a key routes to when all are healthy + let test_key = "session-abc-123"; + let headers = headers_with_routing_key(test_key); + let info = SelectWorkerInfo { + headers: Some(&headers), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + let (result, _) = policy.select_worker_impl(&workers, &info); + let original_idx = result.unwrap(); + + // Mark that worker unhealthy + workers[original_idx].set_healthy(false); + + // Key should now route to a different healthy worker + let (failover_result, _) = policy.select_worker_impl(&workers, &info); + let failover_idx = failover_result.unwrap(); + assert_ne!( + failover_idx, original_idx, + "Should failover to different worker" + ); + assert!( + workers[failover_idx].is_healthy(), + "Failover target should be healthy" + ); + + // Failover should be consistent + for _ in 0..5 { + let (result, _) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(failover_idx), "Failover should be consistent"); + } + + // Recover the original worker + workers[original_idx].set_healthy(true); + + // Key should route back to original worker + let (recovered_result, _) = policy.select_worker_impl(&workers, &info); + assert_eq!( + recovered_result, + Some(original_idx), + "Should return to original worker after recovery" + ); + } + + #[test] + fn test_empty_routing_key_uses_fallback() { + let policy = ConsistentHashingPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + let headers = headers_with_routing_key(""); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert!(result.is_some()); + assert_eq!(branch, Branch::RandomFallback); + } + + #[test] + fn test_policy_name() { + let policy = ConsistentHashingPolicy::new(); + assert_eq!(policy.name(), "consistent_hashing"); + } +} diff --git a/sgl-model-gateway/src/policies/factory.rs b/sgl-model-gateway/src/policies/factory.rs index 3b536709b..3d7c9b4f0 100644 --- a/sgl-model-gateway/src/policies/factory.rs +++ b/sgl-model-gateway/src/policies/factory.rs @@ -3,8 +3,8 @@ use std::sync::Arc; use super::{ - BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, - ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, + BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy, + LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, }; use crate::config::PolicyConfig; @@ -47,6 +47,7 @@ impl PolicyFactory { Arc::new(BucketPolicy::with_config(config)) } PolicyConfig::Manual => Arc::new(ManualPolicy::new()), + PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()), } } @@ -59,6 +60,9 @@ impl PolicyFactory { "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())), "bucket" => Some(Arc::new(BucketPolicy::new())), "manual" => Some(Arc::new(ManualPolicy::new())), + "consistent_hashing" | "consistenthashing" => { + Some(Arc::new(ConsistentHashingPolicy::new())) + } _ => None, } } @@ -96,6 +100,12 @@ mod tests { bucket_adjust_interval_secs: 5, }); assert_eq!(policy.name(), "bucket"); + + let policy = PolicyFactory::create_from_config(&PolicyConfig::Manual); + assert_eq!(policy.name(), "manual"); + + let policy = PolicyFactory::create_from_config(&PolicyConfig::ConsistentHashing); + assert_eq!(policy.name(), "consistent_hashing"); } #[tokio::test] @@ -110,6 +120,10 @@ mod tests { assert!(PolicyFactory::create_by_name("CacheAware").is_some()); assert!(PolicyFactory::create_by_name("bucket").is_some()); assert!(PolicyFactory::create_by_name("Bucket").is_some()); + assert!(PolicyFactory::create_by_name("manual").is_some()); + assert!(PolicyFactory::create_by_name("Manual").is_some()); + assert!(PolicyFactory::create_by_name("consistent_hashing").is_some()); + assert!(PolicyFactory::create_by_name("ConsistentHashing").is_some()); assert!(PolicyFactory::create_by_name("unknown").is_none()); } } diff --git a/sgl-model-gateway/src/policies/manual.rs b/sgl-model-gateway/src/policies/manual.rs index d55cf3eb5..f4d55828d 100644 --- a/sgl-model-gateway/src/policies/manual.rs +++ b/sgl-model-gateway/src/policies/manual.rs @@ -1,105 +1,159 @@ -//! Manual routing policy with header-based routing support +//! Manual routing policy based on routing key header //! -//! Supports two routing mechanisms via HTTP headers: -//! - `X-SMG-Target-Worker`: Direct routing by worker index (0-based), returns None if unavailable -//! - `X-SMG-Routing-Key`: Consistent hash routing for session affinity +//! This policy provides sticky session routing where each unique routing key +//! is consistently mapped to the same worker. Unlike consistent hashing, +//! this policy: +//! - Does NOT redistribute any sessions when workers are added +//! - Only remaps sessions when their assigned worker becomes unhealthy +//! - Maintains up to 2 candidate workers per routing key for fast failover //! -//! Complexity: O(n) for get_healthy_worker_indices (unavoidable), O(1) for routing decisions. +//! Use this when you need stronger stickiness guarantees than consistent hashing, +//! for example with stateful chat sessions where context is stored on the worker. +//! +//! ## Header +//! - `X-SMG-Routing-Key`: The routing key for sticky session routing -use std::{ - hash::{Hash, Hasher}, - sync::Arc, -}; +use std::sync::Arc; +use dashmap::{mapref::entry::Entry, DashMap}; use http::header::HeaderName; -use rand::Rng as _; +use rand::Rng; use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo}; use crate::{core::Worker, observability::metrics::Metrics}; -/// Header for direct worker targeting by index (0-based) -static HEADER_TARGET_WORKER: HeaderName = HeaderName::from_static("x-smg-target-worker"); -/// Header for consistent hash routing +/// Header for routing key based sticky sessions static HEADER_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-key"); -/// Execution branch for metrics #[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum Branch { +enum ExecutionBranch { NoHealthyWorkers, - TargetWorkerHit, - TargetWorkerMiss, - RoutingKeyHit, - RandomFallback, + FastPathHit, + SlowPathOccupiedHit, + SlowPathOccupiedMiss, + SlowPathVacant, + NoRoutingId, } -impl Branch { - #[inline] - const fn as_str(&self) -> &'static str { +impl ExecutionBranch { + fn as_str(&self) -> &'static str { match self { Self::NoHealthyWorkers => "no_healthy_workers", - Self::TargetWorkerHit => "target_worker_hit", - Self::TargetWorkerMiss => "target_worker_miss", - Self::RoutingKeyHit => "routing_key_hit", - Self::RandomFallback => "random_fallback", + Self::FastPathHit => "fast_path_hit", + Self::SlowPathOccupiedHit => "slow_path_occupied_hit", + Self::SlowPathOccupiedMiss => "slow_path_occupied_miss", + Self::SlowPathVacant => "slow_path_vacant", + Self::NoRoutingId => "no_routing_id", } } } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct RoutingId(String); + +impl RoutingId { + fn new(id: impl Into) -> Self { + Self(id.into()) + } +} + +const MAX_CANDIDATE_WORKERS: usize = 2; + +#[derive(Debug, Clone)] +struct RoutingInfo { + candi_worker_urls: Vec, +} + +impl RoutingInfo { + fn push_bounded(&mut self, url: String) { + while self.candi_worker_urls.len() >= MAX_CANDIDATE_WORKERS { + self.candi_worker_urls.remove(0); + } + self.candi_worker_urls.push(url); + } +} + +// TODO may optimize performance +// TODO evict old data periodically #[derive(Debug, Default)] -pub struct ManualPolicy; +pub struct ManualPolicy { + routing_map: DashMap, +} impl ManualPolicy { pub fn new() -> Self { - Self + Self { + routing_map: DashMap::new(), + } + } + + fn select_by_routing_id( + &self, + workers: &[Arc], + routing_id: &str, + healthy_indices: &[usize], + ) -> (usize, ExecutionBranch) { + let routing_id = RoutingId::new(routing_id); + + // Fast path + if let Some(info) = self.routing_map.get(&routing_id) { + if let Some(idx) = + find_healthy_worker(&info.candi_worker_urls, workers, healthy_indices) + { + return (idx, ExecutionBranch::FastPathHit); + } + } + + // Slow path + match self.routing_map.entry(routing_id) { + Entry::Occupied(mut entry) => { + if let Some(idx) = + find_healthy_worker(&entry.get().candi_worker_urls, workers, healthy_indices) + { + return (idx, ExecutionBranch::SlowPathOccupiedHit); + } + let selected_idx = random_select(healthy_indices); + entry + .get_mut() + .push_bounded(workers[selected_idx].url().to_string()); + (selected_idx, ExecutionBranch::SlowPathOccupiedMiss) + } + Entry::Vacant(entry) => { + let selected_idx = random_select(healthy_indices); + entry.insert(RoutingInfo { + candi_worker_urls: vec![workers[selected_idx].url().to_string()], + }); + (selected_idx, ExecutionBranch::SlowPathVacant) + } + } } fn select_worker_impl( &self, workers: &[Arc], info: &SelectWorkerInfo, - ) -> (Option, Branch) { - // O(n) - unavoidable, need to know which workers are healthy + ) -> (Option, ExecutionBranch) { let healthy_indices = get_healthy_worker_indices(workers); if healthy_indices.is_empty() { - return (None, Branch::NoHealthyWorkers); + return (None, ExecutionBranch::NoHealthyWorkers); } - // Extract routing headers - to_str() is O(1), just validates ASCII, no allocation - let target_worker = info - .headers - .and_then(|h| h.get(&HEADER_TARGET_WORKER)) - .and_then(|v| v.to_str().ok()) - .filter(|s| !s.is_empty()); - - let routing_key = info + // Extract routing key from header + let routing_id = info .headers .and_then(|h| h.get(&HEADER_ROUTING_KEY)) .and_then(|v| v.to_str().ok()) .filter(|s| !s.is_empty()); - // Priority 1: X-SMG-Target-Worker - direct routing by worker index - // O(1) parse + O(1) bounds check + O(1) health check - if let Some(idx_str) = target_worker { - if let Ok(idx) = idx_str.parse::() { - if idx < workers.len() && workers[idx].is_healthy() { - return (Some(idx), Branch::TargetWorkerHit); - } - } - return (None, Branch::TargetWorkerMiss); + if let Some(routing_id) = routing_id { + let (idx, branch) = self.select_by_routing_id(workers, routing_id, &healthy_indices); + return (Some(idx), branch); } - // Priority 2: X-SMG-Routing-Key - consistent hash routing - // O(key_len) hash + O(1) modulo + O(1) index - if let Some(key) = routing_key { - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - key.hash(&mut hasher); - let idx = (hasher.finish() as usize) % healthy_indices.len(); - return (Some(healthy_indices[idx]), Branch::RoutingKeyHit); - } - - // Fallback: random selection using thread-local RNG (fast, no allocation) - let idx = rand::rng().random_range(0..healthy_indices.len()); - (Some(healthy_indices[idx]), Branch::RandomFallback) + ( + Some(random_select(&healthy_indices)), + ExecutionBranch::NoRoutingId, + ) } } @@ -119,6 +173,32 @@ impl LoadBalancingPolicy for ManualPolicy { } } +fn find_healthy_worker( + urls: &[String], + workers: &[Arc], + healthy_indices: &[usize], +) -> Option { + for url in urls { + if let Some(idx) = find_worker_index_by_url(workers, url) { + if healthy_indices.contains(&idx) { + return Some(idx); + } + } + } + None +} + +fn find_worker_index_by_url(workers: &[Arc], url: &str) -> Option { + workers.iter().position(|w| w.url() == url) +} + +// TODO: use load-aware selection later +fn random_select(healthy_indices: &[usize]) -> usize { + let mut rng = rand::rng(); + let random_idx = rng.random_range(0..healthy_indices.len()); + healthy_indices[random_idx] +} + #[cfg(test)] mod tests { use std::collections::HashMap; @@ -126,18 +206,6 @@ mod tests { use super::*; use crate::core::{BasicWorkerBuilder, WorkerType}; - fn headers_with_routing_key(key: &str) -> http::HeaderMap { - let mut headers = http::HeaderMap::new(); - headers.insert("x-smg-routing-key", key.parse().unwrap()); - headers - } - - fn headers_with_target_worker(idx: usize) -> http::HeaderMap { - let mut headers = http::HeaderMap::new(); - headers.insert("x-smg-target-worker", idx.to_string().parse().unwrap()); - headers - } - fn create_workers(urls: &[&str]) -> Vec> { urls.iter() .map(|url| { @@ -150,8 +218,14 @@ mod tests { .collect() } + fn headers_with_routing_key(key: &str) -> http::HeaderMap { + let mut headers = http::HeaderMap::new(); + headers.insert("x-smg-routing-key", key.parse().unwrap()); + headers + } + #[test] - fn test_consistent_routing() { + fn test_manual_consistent_routing() { let policy = ManualPolicy::new(); let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); @@ -161,19 +235,23 @@ mod tests { ..Default::default() }; - let (first_result, _) = policy.select_worker_impl(&workers, &info); + let (first_result, branch) = policy.select_worker_impl(&workers, &info); let first_idx = first_result.unwrap(); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); - // Same key should always route to same worker for _ in 0..10 { let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, Some(first_idx)); - assert_eq!(branch, Branch::RoutingKeyHit); + assert_eq!( + result, + Some(first_idx), + "Same routing_id should route to same worker" + ); + assert_eq!(branch, ExecutionBranch::FastPathHit); } } #[test] - fn test_different_keys_distribute() { + fn test_manual_different_routing_ids() { let policy = ManualPolicy::new(); let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); @@ -184,133 +262,103 @@ mod tests { headers: Some(&headers), ..Default::default() }; - let (result, _) = policy.select_worker_impl(&workers, &info); - *distribution.entry(result.unwrap()).or_insert(0) += 1; - } - - assert!(distribution.len() > 1, "Should distribute across workers"); - } - - #[test] - fn test_target_worker_hit() { - let policy = ManualPolicy::new(); - let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - - let headers = headers_with_target_worker(1); - let info = SelectWorkerInfo { - headers: Some(&headers), - ..Default::default() - }; - - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, Some(1)); - assert_eq!(branch, Branch::TargetWorkerHit); - } - - #[test] - fn test_target_worker_miss_out_of_bounds() { - let policy = ManualPolicy::new(); - let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - - let headers = headers_with_target_worker(5); // Out of bounds - let info = SelectWorkerInfo { - headers: Some(&headers), - ..Default::default() - }; - - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, None); - assert_eq!(branch, Branch::TargetWorkerMiss); - } - - #[test] - fn test_target_worker_miss_unhealthy() { - let policy = ManualPolicy::new(); - let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - workers[1].set_healthy(false); - - let headers = headers_with_target_worker(1); - let info = SelectWorkerInfo { - headers: Some(&headers), - ..Default::default() - }; - - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, None); - assert_eq!(branch, Branch::TargetWorkerMiss); - } - - #[test] - fn test_target_worker_priority_over_routing_key() { - let policy = ManualPolicy::new(); - let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - - let mut headers = http::HeaderMap::new(); - headers.insert("x-smg-target-worker", "1".parse().unwrap()); - headers.insert("x-smg-routing-key", "some-key".parse().unwrap()); - - let info = SelectWorkerInfo { - headers: Some(&headers), - ..Default::default() - }; - - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, Some(1)); - assert_eq!(branch, Branch::TargetWorkerHit); - } - - #[test] - fn test_fallback_random_distribution() { - let policy = ManualPolicy::new(); - let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); - - // Without routing headers, should distribute randomly across workers - let mut distribution = HashMap::new(); - for _ in 0..100 { - let info = SelectWorkerInfo::default(); let (result, branch) = policy.select_worker_impl(&workers, &info); - assert!(result.is_some()); - assert_eq!(branch, Branch::RandomFallback); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); *distribution.entry(result.unwrap()).or_insert(0) += 1; } - // Should distribute across multiple workers (not always same one) assert!( distribution.len() > 1, - "Random fallback should distribute across workers" + "Should distribute across multiple workers" ); } #[test] - fn test_no_healthy_workers() { + fn test_manual_fallback_random() { let policy = ManualPolicy::new(); - let workers = create_workers(&["http://w1:8000"]); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + let mut counts = HashMap::new(); + for _ in 0..100 { + let info = SelectWorkerInfo::default(); + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(branch, ExecutionBranch::NoRoutingId); + if let Some(idx) = result { + *counts.entry(idx).or_insert(0) += 1; + } + } + + assert_eq!(counts.len(), 2, "Random fallback should use all workers"); + } + + #[test] + fn test_manual_with_unhealthy_workers() { + let policy = ManualPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + workers[0].set_healthy(false); + let headers = headers_with_routing_key("test-routing-id"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(1), "Should only select healthy worker"); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); + + for _ in 0..10 { + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(1), "Should only select healthy worker"); + assert_eq!(branch, ExecutionBranch::FastPathHit); + } + } + + #[test] + fn test_manual_no_healthy_workers() { + let policy = ManualPolicy::new(); + let workers = create_workers(&["http://w1:8000"]); + + workers[0].set_healthy(false); let headers = headers_with_routing_key("test"); let info = SelectWorkerInfo { headers: Some(&headers), ..Default::default() }; - let (result, branch) = policy.select_worker_impl(&workers, &info); assert_eq!(result, None); - assert_eq!(branch, Branch::NoHealthyWorkers); + assert_eq!(branch, ExecutionBranch::NoHealthyWorkers); } #[test] - fn test_empty_workers() { + fn test_manual_empty_routing_id() { let policy = ManualPolicy::new(); - let workers: Vec> = vec![]; + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - let info = SelectWorkerInfo::default(); - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, None); - assert_eq!(branch, Branch::NoHealthyWorkers); + let mut counts = HashMap::new(); + for _ in 0..100 { + let headers = headers_with_routing_key(""); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(branch, ExecutionBranch::NoRoutingId); + if let Some(idx) = result { + *counts.entry(idx).or_insert(0) += 1; + } + } + + assert_eq!( + counts.len(), + 2, + "Empty routing_id should use random fallback" + ); } #[test] - fn test_routing_key_remaps_when_worker_unhealthy() { + fn test_manual_remaps_when_worker_becomes_unhealthy() { let policy = ManualPolicy::new(); let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); @@ -320,37 +368,217 @@ mod tests { ..Default::default() }; - let (first_result, _) = policy.select_worker_impl(&workers, &info); + let (first_result, branch) = policy.select_worker_impl(&workers, &info); let first_idx = first_result.unwrap(); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); - // Mark that worker unhealthy workers[first_idx].set_healthy(false); - // Should now route to the other worker - let (new_result, _) = policy.select_worker_impl(&workers, &info); + let (new_result, branch) = policy.select_worker_impl(&workers, &info); let new_idx = new_result.unwrap(); - assert_ne!(new_idx, first_idx); + assert_ne!(new_idx, first_idx, "Should remap to healthy worker"); + assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + + for _ in 0..10 { + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!( + result, + Some(new_idx), + "Should consistently route to new worker" + ); + assert_eq!(branch, ExecutionBranch::FastPathHit); + } } #[test] - fn test_empty_routing_key_uses_fallback() { + fn test_manual_empty_workers() { let policy = ManualPolicy::new(); - let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + let workers: Vec> = vec![]; + let headers = headers_with_routing_key("test"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, ExecutionBranch::NoHealthyWorkers); + } - let headers = headers_with_routing_key(""); + #[test] + fn test_manual_single_worker() { + let policy = ManualPolicy::new(); + let workers = create_workers(&["http://w1:8000"]); + + let headers = headers_with_routing_key("single-test"); let info = SelectWorkerInfo { headers: Some(&headers), ..Default::default() }; let (result, branch) = policy.select_worker_impl(&workers, &info); - assert!(result.is_some()); - assert_eq!(branch, Branch::RandomFallback); + assert_eq!(result, Some(0)); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); + + for _ in 0..10 { + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(0)); + assert_eq!(branch, ExecutionBranch::FastPathHit); + } } #[test] - fn test_policy_name() { + fn test_manual_worker_recovery() { + let policy = ManualPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + let headers = headers_with_routing_key("recovery-test"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (first_result, branch) = policy.select_worker_impl(&workers, &info); + let first_idx = first_result.unwrap(); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); + + workers[first_idx].set_healthy(false); + + let (second_result, branch) = policy.select_worker_impl(&workers, &info); + let second_idx = second_result.unwrap(); + assert_ne!(second_idx, first_idx); + assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + + workers[first_idx].set_healthy(true); + + let (after_recovery, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!( + after_recovery, + Some(first_idx), + "Should return to original worker after recovery since it's first in candidate list" + ); + assert_eq!(branch, ExecutionBranch::FastPathHit); + } + + #[test] + fn test_manual_max_candidate_workers_eviction() { + let policy = ManualPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + + let headers = headers_with_routing_key("eviction-test"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + + let (first_result, branch) = policy.select_worker_impl(&workers, &info); + let first_idx = first_result.unwrap(); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); + + workers[first_idx].set_healthy(false); + + let (second_result, branch) = policy.select_worker_impl(&workers, &info); + let second_idx = second_result.unwrap(); + assert_ne!(second_idx, first_idx); + assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + + workers[second_idx].set_healthy(false); + + let remaining_idx = (0..3).find(|&i| i != first_idx && i != second_idx).unwrap(); + let (third_result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!( + third_result, + Some(remaining_idx), + "Should select the only remaining healthy worker" + ); + assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + + workers[first_idx].set_healthy(true); + + let (idx_after_restore, branch) = policy.select_worker_impl(&workers, &info); + assert_ne!( + idx_after_restore, + Some(first_idx), + "First worker should be evicted from candidates due to MAX_CANDIDATE_WORKERS=2" + ); + assert_eq!(branch, ExecutionBranch::FastPathHit); + } + + #[test] + fn test_manual_policy_name() { let policy = ManualPolicy::new(); assert_eq!(policy.name(), "manual"); } + + #[test] + fn test_manual_routing_info_push_bounded() { + let mut info = RoutingInfo { + candi_worker_urls: vec!["http://w1:8000".to_string()], + }; + + info.push_bounded("http://w2:8000".to_string()); + assert_eq!(info.candi_worker_urls.len(), 2); + assert_eq!(info.candi_worker_urls[0], "http://w1:8000"); + assert_eq!(info.candi_worker_urls[1], "http://w2:8000"); + + info.push_bounded("http://w3:8000".to_string()); + assert_eq!(info.candi_worker_urls.len(), 2); + assert_eq!( + info.candi_worker_urls[0], "http://w2:8000", + "Oldest entry should be removed" + ); + assert_eq!(info.candi_worker_urls[1], "http://w3:8000"); + } + + #[test] + fn test_manual_find_healthy_worker_priority() { + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + + let urls = vec![ + "http://w1:8000".to_string(), + "http://w2:8000".to_string(), + "http://w3:8000".to_string(), + ]; + let healthy_indices = vec![0, 1, 2]; + + let result = find_healthy_worker(&urls, &workers, &healthy_indices); + assert_eq!( + result, + Some(0), + "Should return first healthy worker in urls" + ); + + workers[0].set_healthy(false); + let healthy_indices = vec![1, 2]; + let result = find_healthy_worker(&urls, &workers, &healthy_indices); + assert_eq!(result, Some(1), "Should skip unhealthy and return next"); + + workers[1].set_healthy(false); + let healthy_indices = vec![2]; + let result = find_healthy_worker(&urls, &workers, &healthy_indices); + assert_eq!(result, Some(2), "Should return last healthy worker"); + + workers[2].set_healthy(false); + let healthy_indices: Vec = vec![]; + let result = find_healthy_worker(&urls, &workers, &healthy_indices); + assert_eq!(result, None, "Should return None when no healthy workers"); + } + + #[test] + fn test_manual_find_worker_index_by_url() { + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + assert_eq!( + find_worker_index_by_url(&workers, "http://w1:8000"), + Some(0) + ); + assert_eq!( + find_worker_index_by_url(&workers, "http://w2:8000"), + Some(1) + ); + assert_eq!( + find_worker_index_by_url(&workers, "http://w3:8000"), + None, + "Should return None for unknown URL" + ); + } } diff --git a/sgl-model-gateway/src/policies/mod.rs b/sgl-model-gateway/src/policies/mod.rs index 1f276d512..5b8902da9 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -5,10 +5,11 @@ use std::{fmt::Debug, sync::Arc}; -use crate::core::Worker; +use crate::core::{HashRing, Worker}; mod bucket; mod cache_aware; +mod consistent_hashing; mod factory; mod manual; mod power_of_two; @@ -19,6 +20,7 @@ pub mod tree; pub use bucket::BucketPolicy; pub use cache_aware::CacheAwarePolicy; +pub use consistent_hashing::ConsistentHashingPolicy; pub use factory::PolicyFactory; pub use manual::ManualPolicy; pub use power_of_two::PowerOfTwoPolicy; @@ -138,15 +140,18 @@ pub(crate) fn normalize_model_key(model_id: &str) -> &str { } /// Information passed to policy for worker selection -#[derive(Debug, Default, Clone)] +#[derive(Debug, Clone, Default)] pub struct SelectWorkerInfo<'a> { /// Request text for cache-aware routing pub request_text: Option<&'a str>, /// HTTP headers for header-based routing policies /// Policies can extract routing information from headers like: - /// - X-Target-Worker: Direct routing to a specific worker by URL - /// - X-Routing-Key: Consistent hash routing for session affinity + /// - X-SMG-Target-Worker: Direct routing to a specific worker by index + /// - X-SMG-Routing-Key: Consistent hash routing for session affinity pub headers: Option<&'a http::HeaderMap>, + /// Pre-computed hash ring for O(log n) consistent hashing + /// Built and cached by WorkerRegistry, passed through to avoid per-request rebuilds + pub hash_ring: Option>, } #[cfg(test)] diff --git a/sgl-model-gateway/src/policies/registry.rs b/sgl-model-gateway/src/policies/registry.rs index 02cd4cf3d..9e063282a 100644 --- a/sgl-model-gateway/src/policies/registry.rs +++ b/sgl-model-gateway/src/policies/registry.rs @@ -10,8 +10,8 @@ use tracing::{debug, info, warn}; /// All subsequent workers of the same model use the established policy. /// When the last worker of a model is removed, the policy mapping is cleaned up. use super::{ - BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, - ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, + BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy, + LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, }; use crate::{config::types::PolicyConfig, core::Worker}; @@ -168,6 +168,8 @@ impl PolicyRegistry { "cache_aware" => Arc::new(CacheAwarePolicy::new()), "power_of_two" => Arc::new(PowerOfTwoPolicy::new()), "bucket" => Arc::new(BucketPolicy::new()), + "manual" => Arc::new(ManualPolicy::new()), + "consistent_hashing" => Arc::new(ConsistentHashingPolicy::new()), _ => { warn!("Unknown policy type '{}', using default", policy_type); Arc::clone(&self.default_policy) @@ -210,6 +212,7 @@ impl PolicyRegistry { Arc::new(BucketPolicy::with_config(config)) } PolicyConfig::Manual => Arc::new(ManualPolicy::new()), + PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()), } } diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs index 0d136fa04..f19531543 100644 --- a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -8,7 +8,7 @@ use tracing::{error, warn}; use super::PipelineStage; use crate::{ - core::{ConnectionMode, Worker, WorkerRegistry, WorkerType}, + core::{ConnectionMode, Worker, WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID}, observability::metrics::{metrics_labels, Metrics}, policies::{PolicyRegistry, SelectWorkerInfo}, routers::{ @@ -148,12 +148,18 @@ impl WorkerSelectionStage { None => self.policy_registry.get_default_policy(), }; + // Get cached hash ring for consistent hashing (O(log n) lookup) + let hash_ring = self + .worker_registry + .get_hash_ring(model_id.unwrap_or(UNKNOWN_MODEL_ID)); + // Select worker using the policy let idx = policy.select_worker( &available, &SelectWorkerInfo { request_text: text, headers, + hash_ring, }, )?; let selected = available[idx].clone(); @@ -213,9 +219,15 @@ impl WorkerSelectionStage { None => self.policy_registry.get_default_policy(), }; + // Get cached hash ring for consistent hashing (O(log n) lookup) + let hash_ring = self + .worker_registry + .get_hash_ring(model_id.unwrap_or(UNKNOWN_MODEL_ID)); + let info = SelectWorkerInfo { request_text: text, headers, + hash_ring, }; let prefill_idx = policy.select_worker(&available_prefill, &info)?; let decode_idx = policy.select_worker(&available_decode, &info)?; diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 81b83bb37..8bceb261b 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -19,7 +19,8 @@ use super::pd_types::api_path; use crate::{ config::types::RetryConfig, core::{ - is_retryable_status, RetryExecutor, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType, + is_retryable_status, HashRing, RetryExecutor, Worker, WorkerLoadGuard, WorkerRegistry, + WorkerType, UNKNOWN_MODEL_ID, }, observability::{ events::{self, Event}, @@ -731,11 +732,17 @@ impl PDRouter { let prefill_policy = self.policy_registry.get_prefill_policy(); let decode_policy = self.policy_registry.get_decode_policy(); + // Get cached hash ring for consistent hashing + let hash_ring = self + .worker_registry + .get_hash_ring(effective_model_id.unwrap_or(UNKNOWN_MODEL_ID)); + let prefill = Self::pick_worker_by_policy_arc( &prefill_workers, &*prefill_policy, request_text, headers, + hash_ring.clone(), "prefill", )?; @@ -744,6 +751,7 @@ impl PDRouter { &*decode_policy, request_text, headers, + hash_ring, "decode", )?; @@ -770,6 +778,7 @@ impl PDRouter { policy: &dyn LoadBalancingPolicy, request_text: Option<&str>, headers: Option<&HeaderMap>, + hash_ring: Option>, worker_type: &str, ) -> Result, String> { if workers.is_empty() { @@ -798,6 +807,7 @@ impl PDRouter { &SelectWorkerInfo { request_text, headers, + hash_ring, }, ) .ok_or_else(|| { diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 27820e376..8e61459ac 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -20,7 +20,7 @@ use crate::{ config::types::RetryConfig, core::{ is_retryable_status, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuard, - WorkerRegistry, WorkerType, + WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID, }, observability::{ events::{self, Event}, @@ -169,11 +169,17 @@ impl Router { None => self.policy_registry.get_default_policy(), }; + // Get cached hash ring for consistent hashing (O(log n) lookup) + let hash_ring = self + .worker_registry + .get_hash_ring(effective_model_id.unwrap_or(UNKNOWN_MODEL_ID)); + let idx = policy.select_worker( &available, &SelectWorkerInfo { request_text: text, headers, + hash_ring, }, )?;