From caa95c7eb4465046c565ee1a16ad3cd626857dfa Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Thu, 25 Dec 2025 16:04:43 -0500 Subject: [PATCH] [model-gateway] update ManualPolicy with header-based routing (#15847) --- sgl-model-gateway/src/config/types.rs | 3 + sgl-model-gateway/src/config/validation.rs | 3 +- sgl-model-gateway/src/main.rs | 7 +- .../src/observability/metrics.rs | 22 +- sgl-model-gateway/src/policies/cache_aware.rs | 14 +- sgl-model-gateway/src/policies/factory.rs | 5 - sgl-model-gateway/src/policies/manual.rs | 917 ++++-------------- sgl-model-gateway/src/policies/mod.rs | 12 +- .../grpc/common/stages/worker_selection.rs | 39 +- sgl-model-gateway/src/routers/grpc/context.rs | 3 - .../grpc/harmony/stages/preparation.rs | 3 - .../grpc/regular/stages/chat/preparation.rs | 2 - .../regular/stages/embedding/preparation.rs | 5 +- .../regular/stages/generate/preparation.rs | 2 - sgl-model-gateway/src/routers/header_utils.rs | 39 - .../src/routers/http/pd_router.rs | 49 +- sgl-model-gateway/src/routers/http/router.rs | 28 +- .../tests/cache_aware_backward_compat_test.rs | 28 +- 18 files changed, 320 insertions(+), 861 deletions(-) diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index d96167d5f..c318114f1 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -337,6 +337,9 @@ 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 #[serde(rename = "manual")] Manual, } diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index 47b86e927..363bf6723 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/sgl-model-gateway/src/config/validation.rs @@ -147,7 +147,7 @@ impl ConfigValidator { fn validate_policy(policy: &PolicyConfig) -> ConfigResult<()> { match policy { - PolicyConfig::Random | PolicyConfig::RoundRobin => {} + PolicyConfig::Random | PolicyConfig::RoundRobin | PolicyConfig::Manual => {} PolicyConfig::CacheAware { cache_threshold, balance_abs_threshold: _, @@ -226,7 +226,6 @@ impl ConfigValidator { }); } } - PolicyConfig::Manual => {} } Ok(()) } diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index 468d870a9..263e22b03 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -136,7 +136,7 @@ struct CliArgs { #[arg(long, num_args = 0..)] worker_urls: Vec, - #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two"])] + #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "manual"])] policy: String, #[arg(long, default_value_t = false)] @@ -145,10 +145,10 @@ struct CliArgs { #[arg(long, action = ArgAction::Append)] decode: Vec, - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two"])] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "manual"])] prefill_policy: Option, - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two"])] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "manual"])] decode_policy: Option, #[arg(long, default_value_t = 1800)] @@ -415,6 +415,7 @@ impl CliArgs { "power_of_two" => PolicyConfig::PowerOfTwo { load_check_interval_secs: 5, }, + "manual" => PolicyConfig::Manual, _ => PolicyConfig::RoundRobin, } } diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index 10266559b..265b69f9d 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -190,10 +190,6 @@ pub fn init_metrics() { "smg_worker_errors_total", "Worker-level errors by worker_type, connection_mode, error_type" ); - describe_counter!( - "smg_worker_manual_policy_branch_total", - "Manual policy execution branch by branch type" - ); // Layer 3: Worker resilience metrics (circuit breaker) describe_gauge!( @@ -805,6 +801,15 @@ impl Metrics { .increment(1); } + /// Record manual policy execution branch for routing decisions + pub fn record_worker_manual_policy_branch(branch: &'static str) { + counter!( + "smg_manual_policy_branch_total", + "branch" => branch + ) + .increment(1); + } + /// Set running requests per worker pub fn set_worker_requests_active(worker: &str, count: usize) { gauge!( @@ -814,15 +819,6 @@ impl Metrics { .set(count as f64); } - /// Record manual policy execution branch - pub fn record_worker_manual_policy_branch(branch: &'static str) { - counter!( - "smg_worker_manual_policy_branch_total", - "branch" => branch - ) - .increment(1); - } - /// Set worker health status pub fn set_worker_health(worker_url: &str, healthy: bool) { gauge!( diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index 9ddefc08b..f1fa05470 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -517,16 +517,12 @@ mod tests { policy.init_workers(&workers); // Should select worker2 (lower load) despite cache affinity + let info = SelectWorkerInfo { + request_text: Some("test"), + ..Default::default() + }; for _ in 0..5 { - let idx = policy - .select_worker( - &workers, - &SelectWorkerInfo { - request_text: Some("test"), - ..Default::default() - }, - ) - .unwrap(); + let idx = policy.select_worker(&workers, &info).unwrap(); assert_eq!(idx, 1); // Should always pick worker2 } } diff --git a/sgl-model-gateway/src/policies/factory.rs b/sgl-model-gateway/src/policies/factory.rs index b4de69925..3b536709b 100644 --- a/sgl-model-gateway/src/policies/factory.rs +++ b/sgl-model-gateway/src/policies/factory.rs @@ -96,9 +96,6 @@ 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"); } #[tokio::test] @@ -113,8 +110,6 @@ 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("unknown").is_none()); } } diff --git a/sgl-model-gateway/src/policies/manual.rs b/sgl-model-gateway/src/policies/manual.rs index 9663d735a..d55cf3eb5 100644 --- a/sgl-model-gateway/src/policies/manual.rs +++ b/sgl-model-gateway/src/policies/manual.rs @@ -1,139 +1,105 @@ -//! Manual routing policy based on routing_id +//! Manual 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 +//! +//! Complexity: O(n) for get_healthy_worker_indices (unavoidable), O(1) for routing decisions. -use std::sync::Arc; +use std::{ + hash::{Hash, Hasher}, + sync::Arc, +}; -use dashmap::{mapref::entry::Entry, DashMap}; -use rand::Rng; +use http::header::HeaderName; +use rand::Rng as _; 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 +static HEADER_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-key"); + +/// Execution branch for metrics #[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ExecutionBranch { +enum Branch { NoHealthyWorkers, - FastPathHit, - SlowPathOccupiedHit, - SlowPathOccupiedMiss, - SlowPathVacant, - NoRoutingId, + TargetWorkerHit, + TargetWorkerMiss, + RoutingKeyHit, + RandomFallback, } -impl ExecutionBranch { - // TODO auto generate - fn as_str(&self) -> &'static str { +impl Branch { + #[inline] + const fn as_str(&self) -> &'static str { match self { Self::NoHealthyWorkers => "no_healthy_workers", - 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", + Self::TargetWorkerHit => "target_worker_hit", + Self::TargetWorkerMiss => "target_worker_miss", + Self::RoutingKeyHit => "routing_key_hit", + Self::RandomFallback => "random_fallback", } } } -#[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 { - routing_map: DashMap, -} +pub struct ManualPolicy; impl ManualPolicy { pub fn new() -> 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) - } - } + Self } fn select_worker_impl( &self, workers: &[Arc], info: &SelectWorkerInfo, - ) -> (Option, ExecutionBranch) { + ) -> (Option, Branch) { + // O(n) - unavoidable, need to know which workers are healthy let healthy_indices = get_healthy_worker_indices(workers); if healthy_indices.is_empty() { - return (None, ExecutionBranch::NoHealthyWorkers); + return (None, Branch::NoHealthyWorkers); } - if let Some(routing_id) = info.routing_id { - if !routing_id.is_empty() { - let (idx, branch) = - self.select_by_routing_id(workers, routing_id, &healthy_indices); - return (Some(idx), branch); + // 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); } - ( - Some(random_select(&healthy_indices)), - ExecutionBranch::NoRoutingId, - ) + // 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) } } @@ -148,41 +114,11 @@ impl LoadBalancingPolicy for ManualPolicy { "manual" } - fn needs_routing_id(&self) -> bool { - true - } - fn as_any(&self) -> &dyn std::any::Any { self } } -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; @@ -190,654 +126,231 @@ mod tests { use super::*; use crate::core::{BasicWorkerBuilder, WorkerType}; - #[test] - fn test_manual_consistent_routing() { - let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w3:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; + 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 = ManualPolicy::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 { - routing_id: Some("user-123"), + headers: Some(&headers), ..Default::default() }; - let (first_result, branch) = policy.select_worker_impl(&workers, &info); + let (first_result, _) = 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), - "Same routing_id should route to same worker" - ); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(result, Some(first_idx)); + assert_eq!(branch, Branch::RoutingKeyHit); } } #[test] - fn test_manual_different_routing_ids() { + fn test_different_keys_distribute() { let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w3:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); let mut distribution = HashMap::new(); for i in 0..100 { - let routing_id = format!("user-{}", i); + let headers = headers_with_routing_key(&format!("user-{}", i)); let info = SelectWorkerInfo { - routing_id: Some(&routing_id), + headers: Some(&headers), ..Default::default() }; - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + let (result, _) = policy.select_worker_impl(&workers, &info); *distribution.entry(result.unwrap()).or_insert(0) += 1; } - assert!( - distribution.len() > 1, - "Should distribute across multiple workers" - ); + assert!(distribution.len() > 1, "Should distribute across workers"); } #[test] - fn test_manual_fallback_random() { + fn test_target_worker_hit() { let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); - let mut counts = HashMap::new(); + 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_eq!(branch, ExecutionBranch::NoRoutingId); - if let Some(idx) = result { - *counts.entry(idx).or_insert(0) += 1; - } + assert!(result.is_some()); + assert_eq!(branch, Branch::RandomFallback); + *distribution.entry(result.unwrap()).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: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - workers[0].set_healthy(false); - - let info = SelectWorkerInfo { - routing_id: Some("test-routing-id"), - ..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: Vec> = vec![Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - )]; - - workers[0].set_healthy(false); - let info = SelectWorkerInfo { - routing_id: Some("test"), - ..Default::default() - }; - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, None); - assert_eq!(branch, ExecutionBranch::NoHealthyWorkers); - } - - #[test] - fn test_manual_empty_routing_id() { - let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - let mut counts = HashMap::new(); - for _ in 0..100 { - let info = SelectWorkerInfo { - routing_id: Some(""), - ..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" + // Should distribute across multiple workers (not always same one) + assert!( + distribution.len() > 1, + "Random fallback should distribute across workers" ); } #[test] - fn test_manual_remaps_when_worker_becomes_unhealthy() { + fn test_no_healthy_workers() { let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; + let workers = create_workers(&["http://w1:8000"]); + workers[0].set_healthy(false); + let headers = headers_with_routing_key("test"); let info = SelectWorkerInfo { - routing_id: Some("sticky-user"), + 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 (new_result, branch) = policy.select_worker_impl(&workers, &info); - let new_idx = new_result.unwrap(); - 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); - } + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, Branch::NoHealthyWorkers); } #[test] - fn test_manual_empty_workers() { + fn test_empty_workers() { let policy = ManualPolicy::new(); let workers: Vec> = vec![]; - let info = SelectWorkerInfo { - routing_id: Some("test"), - ..Default::default() - }; - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(result, None); - assert_eq!(branch, ExecutionBranch::NoHealthyWorkers); - } - - #[test] - fn test_manual_single_worker() { - let policy = ManualPolicy::new(); - let workers: Vec> = vec![Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - )]; - - let info = SelectWorkerInfo { - routing_id: Some("single-test"), - ..Default::default() - }; - - let (result, branch) = policy.select_worker_impl(&workers, &info); - 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_manual_worker_recovery() { - let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - let info = SelectWorkerInfo { - routing_id: Some("recovery-test"), - ..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: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w3:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - let info = SelectWorkerInfo { - routing_id: Some("eviction-test"), - ..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_execution_branch_fast_path_hit() { - let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - let info = SelectWorkerInfo { - routing_id: Some("fast-path-test"), - ..Default::default() - }; - - let _ = policy.select_worker_impl(&workers, &info); - - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert!(result.is_some()); - assert_eq!(branch, ExecutionBranch::FastPathHit); - } - - #[test] - fn test_manual_execution_branch_no_routing_id() { - let policy = ManualPolicy::new(); - let workers: Vec> = vec![Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - )]; let info = SelectWorkerInfo::default(); let (result, branch) = policy.select_worker_impl(&workers, &info); - assert!(result.is_some()); - assert_eq!(branch, ExecutionBranch::NoRoutingId); + assert_eq!(result, None); + assert_eq!(branch, Branch::NoHealthyWorkers); } #[test] - fn test_manual_execution_branch_slow_path_occupied_miss() { + fn test_routing_key_remaps_when_worker_unhealthy() { let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + let headers = headers_with_routing_key("sticky-user"); let info = SelectWorkerInfo { - routing_id: Some("occupied-miss-test"), + headers: Some(&headers), ..Default::default() }; - let (first_result, branch) = policy.select_worker_impl(&workers, &info); + let (first_result, _) = 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); - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert!(result.is_some()); - assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + // Should now route to the other worker + let (new_result, _) = policy.select_worker_impl(&workers, &info); + let new_idx = new_result.unwrap(); + assert_ne!(new_idx, first_idx); } #[test] - fn test_manual_execution_branch_slow_path_occupied_hit() { + fn test_empty_routing_key_uses_fallback() { let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + let headers = headers_with_routing_key(""); let info = SelectWorkerInfo { - routing_id: Some("occupied-hit-test"), + headers: Some(&headers), ..Default::default() }; - let _ = policy.select_worker_impl(&workers, &info); - - policy.routing_map.clear(); - - policy.routing_map.insert( - RoutingId::new("occupied-hit-test"), - RoutingInfo { - candi_worker_urls: vec!["http://w1:8000".to_string()], - }, - ); - let (result, branch) = policy.select_worker_impl(&workers, &info); assert!(result.is_some()); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(branch, Branch::RandomFallback); } #[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: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w3:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - 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: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - 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" - ); - } - - #[test] - fn test_manual_policy_name() { + fn test_policy_name() { let policy = ManualPolicy::new(); assert_eq!(policy.name(), "manual"); } - - #[test] - fn test_manual_policy_needs_routing_id() { - let policy = ManualPolicy::new(); - assert!(policy.needs_routing_id()); - } - - #[test] - fn test_manual_all_workers_become_unhealthy_then_recover() { - let policy = ManualPolicy::new(); - let workers: Vec> = vec![ - Arc::new( - BasicWorkerBuilder::new("http://w1:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - Arc::new( - BasicWorkerBuilder::new("http://w2:8000") - .worker_type(WorkerType::Regular) - .build(), - ), - ]; - - let info = SelectWorkerInfo { - routing_id: Some("all-unhealthy-test"), - ..Default::default() - }; - - let (first_result, branch) = policy.select_worker_impl(&workers, &info); - let first_idx = first_result.unwrap(); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); - - workers[0].set_healthy(false); - workers[1].set_healthy(false); - - let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!( - result, None, - "Should return None when all workers are unhealthy" - ); - assert_eq!(branch, ExecutionBranch::NoHealthyWorkers); - - workers[first_idx].set_healthy(true); - - let (after_recovery, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!( - after_recovery, - Some(first_idx), - "Should route to recovered worker in candidate list" - ); - assert_eq!(branch, ExecutionBranch::FastPathHit); - } } diff --git a/sgl-model-gateway/src/policies/mod.rs b/sgl-model-gateway/src/policies/mod.rs index 43e32616a..1f276d512 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -57,11 +57,6 @@ pub trait LoadBalancingPolicy: Send + Sync + Debug { false // Default: most policies don't need request text } - /// Check if this policy needs routing_id for routing decisions - fn needs_routing_id(&self) -> bool { - false // Default: most policies don't need routing_id - } - /// Update worker load information /// /// This is called periodically with current load information for load-aware policies. @@ -147,8 +142,11 @@ pub(crate) fn normalize_model_key(model_id: &str) -> &str { pub struct SelectWorkerInfo<'a> { /// Request text for cache-aware routing pub request_text: Option<&'a str>, - /// Routing ID for manual routing policy (consistent hashing) - pub routing_id: 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 + pub headers: Option<&'a http::HeaderMap>, } #[cfg(test)] 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 aa94c4c4f..0d136fa04 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 @@ -61,18 +61,17 @@ impl PipelineStage for WorkerSelectionStage { // For Harmony, use selection_text produced during Harmony encoding // Otherwise, use original_text from regular preparation - let info = SelectWorkerInfo { - request_text: if prep.harmony_mode { - prep.selection_text.as_deref() - } else { - prep.original_text.as_deref() - }, - routing_id: prep.routing_id.as_deref(), + let text = if prep.harmony_mode { + prep.selection_text.as_deref() + } else { + prep.original_text.as_deref() }; + let headers = ctx.input.headers.as_ref(); + let workers = match self.mode { WorkerSelectionMode::Regular => { - match self.select_single_worker(ctx.input.model_id.as_deref(), &info) { + match self.select_single_worker(ctx.input.model_id.as_deref(), text, headers) { Some(w) => WorkerSelection::Single { worker: w }, None => { error!( @@ -89,7 +88,7 @@ impl PipelineStage for WorkerSelectionStage { } } WorkerSelectionMode::PrefillDecode => { - match self.select_pd_pair(ctx.input.model_id.as_deref(), &info) { + match self.select_pd_pair(ctx.input.model_id.as_deref(), text, headers) { Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode }, None => { error!( @@ -123,7 +122,8 @@ impl WorkerSelectionStage { fn select_single_worker( &self, model_id: Option<&str>, - info: &SelectWorkerInfo, + text: Option<&str>, + headers: Option<&http::HeaderMap>, ) -> Option> { // Get workers for the specified model, filtered by connection mode let workers = self.worker_registry.get_workers_filtered( @@ -149,7 +149,13 @@ impl WorkerSelectionStage { }; // Select worker using the policy - let idx = policy.select_worker(&available, info)?; + let idx = policy.select_worker( + &available, + &SelectWorkerInfo { + request_text: text, + headers, + }, + )?; let selected = available[idx].clone(); // Record worker selection metric @@ -166,7 +172,8 @@ impl WorkerSelectionStage { fn select_pd_pair( &self, model_id: Option<&str>, - info: &SelectWorkerInfo, + text: Option<&str>, + headers: Option<&http::HeaderMap>, ) -> Option<(Arc, Arc)> { let all_workers = self.worker_registry.get_workers_filtered( model_id, @@ -206,8 +213,12 @@ impl WorkerSelectionStage { None => self.policy_registry.get_default_policy(), }; - let prefill_idx = policy.select_worker(&available_prefill, info)?; - let decode_idx = policy.select_worker(&available_decode, info)?; + let info = SelectWorkerInfo { + request_text: text, + headers, + }; + let prefill_idx = policy.select_worker(&available_prefill, &info)?; + let decode_idx = policy.select_worker(&available_decode, &info)?; let model = model_id.unwrap_or("default"); let policy_name = policy.name(); diff --git a/sgl-model-gateway/src/routers/grpc/context.rs b/sgl-model-gateway/src/routers/grpc/context.rs index cb0469b2c..a5da72f01 100644 --- a/sgl-model-gateway/src/routers/grpc/context.rs +++ b/sgl-model-gateway/src/routers/grpc/context.rs @@ -94,9 +94,6 @@ pub struct PreparationOutput { /// Original text (for chat) or resolved text (for generate) pub original_text: Option, - /// Routing ID for manual routing policy - pub routing_id: Option, - /// Tokenized input pub token_ids: Vec, diff --git a/sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs b/sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs index 8089f197d..32b309319 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs @@ -19,7 +19,6 @@ use crate::{ context::{PreparationOutput, RequestContext, RequestType}, utils, }, - header_utils, }, }; @@ -124,7 +123,6 @@ impl HarmonyPreparationStage { // Step 4: Store results ctx.state.preparation = Some(PreparationOutput { original_text: None, - routing_id: header_utils::extract_routing_id(ctx.input.headers.as_ref()), token_ids: build_output.input_ids, processed_messages: None, tool_constraints, @@ -205,7 +203,6 @@ impl HarmonyPreparationStage { // Step 4: Store results with constraint ctx.state.preparation = Some(PreparationOutput { original_text: None, - routing_id: header_utils::extract_routing_id(ctx.input.headers.as_ref()), token_ids: build_output.input_ids, processed_messages: None, tool_constraints: constraint, diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/chat/preparation.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/chat/preparation.rs index 5c87d64ff..ee951ab64 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/chat/preparation.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/chat/preparation.rs @@ -15,7 +15,6 @@ use crate::{ context::{PreparationOutput, RequestContext}, utils, }, - header_utils, }, }; @@ -97,7 +96,6 @@ impl ChatPreparationStage { // Store results in context ctx.state.preparation = Some(PreparationOutput { original_text: Some(processed_messages.text.clone()), - routing_id: header_utils::extract_routing_id(ctx.input.headers.as_ref()), token_ids, processed_messages: Some(processed_messages), tool_constraints: tool_call_constraint, diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/embedding/preparation.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/embedding/preparation.rs index a49327b10..fae235a3d 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/embedding/preparation.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/embedding/preparation.rs @@ -13,7 +13,6 @@ use crate::{ context::{PreparationOutput, RequestContext, RequestType}, utils, }, - header_utils, }, }; @@ -48,9 +47,8 @@ impl PipelineStage for EmbeddingPreparationStage { )); }; - // Extract text from request before borrowing ctx mutably + // Extract text from request let text = request.extract_text_for_routing(); - let routing_id = header_utils::extract_routing_id(ctx.input.headers.as_ref()); if text.is_empty() { return Err(error::bad_request( "empty_input", @@ -79,7 +77,6 @@ impl PipelineStage for EmbeddingPreparationStage { // Store preparation output ctx.state.preparation = Some(PreparationOutput { original_text: Some(text), - routing_id, token_ids, processed_messages: None, tool_constraints: None, diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs index 4f6ca3f5a..d1aea48cb 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs @@ -15,7 +15,6 @@ use crate::{ context::{PreparationOutput, RequestContext}, utils, }, - header_utils, }, tokenizer::traits::Tokenizer, }; @@ -69,7 +68,6 @@ impl GeneratePreparationStage { ctx.state.preparation = Some(PreparationOutput { original_text, - routing_id: header_utils::extract_routing_id(ctx.input.headers.as_ref()), token_ids, processed_messages: None, tool_constraints: None, diff --git a/sgl-model-gateway/src/routers/header_utils.rs b/sgl-model-gateway/src/routers/header_utils.rs index d15655891..b7ca861f6 100644 --- a/sgl-model-gateway/src/routers/header_utils.rs +++ b/sgl-model-gateway/src/routers/header_utils.rs @@ -157,17 +157,6 @@ pub fn apply_provider_headers( req } -/// Header name for routing key used by manual routing policy -pub const ROUTING_KEY_HEADER: &str = "X-SMG-Routing-Key"; - -/// Extract routing ID from HTTP headers for manual routing policy -pub fn extract_routing_id(headers: Option<&HeaderMap>) -> Option { - headers - .and_then(|h| h.get(ROUTING_KEY_HEADER)) - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_string()) -} - /// Extract auth header with passthrough semantics. /// /// Passthrough mode: User's Authorization header takes priority. @@ -194,31 +183,3 @@ pub fn extract_auth_header( .and_then(|k| HeaderValue::from_str(&format!("Bearer {}", k)).ok()) }) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_extract_routing_id_with_header() { - let mut headers = HeaderMap::new(); - headers.insert(ROUTING_KEY_HEADER, HeaderValue::from_static("user-123")); - let result = extract_routing_id(Some(&headers)); - assert_eq!(result, Some("user-123".to_string())); - } - - #[test] - fn test_extract_routing_id_without_header() { - let headers = HeaderMap::new(); - let result = extract_routing_id(Some(&headers)); - assert_eq!(result, None); - } - - #[test] - fn test_extract_routing_id_empty_value() { - let mut headers = HeaderMap::new(); - headers.insert(ROUTING_KEY_HEADER, HeaderValue::from_static("")); - let result = extract_routing_id(Some(&headers)); - assert_eq!(result, Some("".to_string())); - } -} diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 3f348ea7a..81b83bb37 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -26,7 +26,7 @@ use crate::{ metrics::{bool_to_static_str, metrics_labels, Metrics}, otel_trace::inject_trace_context_http, }, - policies::{LoadBalancingPolicy, PolicyRegistry}, + policies::{LoadBalancingPolicy, PolicyRegistry, SelectWorkerInfo}, protocols::{ chat::{ChatCompletionRequest, ChatMessage, MessageContent}, common::{InputIds, StringOrArray}, @@ -58,8 +58,8 @@ struct PDRequestContext<'a> { is_stream: bool, return_logprob: bool, request_text: Option, - routing_id: Option, model_id: Option<&'a str>, + headers: Option, } impl PDRouter { @@ -306,8 +306,8 @@ impl PDRouter { let (prefill, decode) = match self .select_pd_pair( context.request_text.as_deref(), - context.routing_id.as_deref(), context.model_id, + context.headers.as_ref(), ) .await { @@ -696,8 +696,8 @@ impl PDRouter { async fn select_pd_pair( &self, request_text: Option<&str>, - routing_id: Option<&str>, model_id: Option<&str>, + headers: Option<&HeaderMap>, ) -> Result<(Arc, Arc), String> { let effective_model_id = if !self.enable_igw { None } else { model_id }; @@ -731,16 +731,21 @@ impl PDRouter { let prefill_policy = self.policy_registry.get_prefill_policy(); let decode_policy = self.policy_registry.get_decode_policy(); - let info = crate::policies::SelectWorkerInfo { + let prefill = Self::pick_worker_by_policy_arc( + &prefill_workers, + &*prefill_policy, request_text, - routing_id, - }; + headers, + "prefill", + )?; - let prefill = - Self::pick_worker_by_policy_arc(&prefill_workers, &*prefill_policy, &info, "prefill")?; - - let decode = - Self::pick_worker_by_policy_arc(&decode_workers, &*decode_policy, &info, "decode")?; + let decode = Self::pick_worker_by_policy_arc( + &decode_workers, + &*decode_policy, + request_text, + headers, + "decode", + )?; // Record worker selection metrics (Layer 3) let model = model_id.unwrap_or("default"); @@ -763,7 +768,8 @@ impl PDRouter { fn pick_worker_by_policy_arc( workers: &[Arc], policy: &dyn LoadBalancingPolicy, - info: &crate::policies::SelectWorkerInfo, + request_text: Option<&str>, + headers: Option<&HeaderMap>, worker_type: &str, ) -> Result, String> { if workers.is_empty() { @@ -787,7 +793,13 @@ impl PDRouter { } let selected_idx = policy - .select_worker(&available_workers, info) + .select_worker( + &available_workers, + &SelectWorkerInfo { + request_text, + headers, + }, + ) .ok_or_else(|| { format!( "Policy {} failed to select a {} worker", @@ -1245,8 +1257,8 @@ impl RouterTrait for PDRouter { is_stream, return_logprob, request_text, - routing_id: header_utils::extract_routing_id(headers), model_id, + headers: headers.cloned(), }; self.execute_dual_dispatch(headers, body, context).await @@ -1287,8 +1299,8 @@ impl RouterTrait for PDRouter { is_stream, return_logprob, request_text, - routing_id: header_utils::extract_routing_id(headers), model_id, + headers: headers.cloned(), }; self.execute_dual_dispatch(headers, body, context).await @@ -1321,8 +1333,8 @@ impl RouterTrait for PDRouter { is_stream, return_logprob, request_text, - routing_id: header_utils::extract_routing_id(headers), model_id, + headers: headers.cloned(), }; self.execute_dual_dispatch(headers, body, context).await @@ -1334,6 +1346,7 @@ impl RouterTrait for PDRouter { body: &RerankRequest, model_id: Option<&str>, ) -> Response { + // Extract text for cache-aware routing let req_text = if self.policies_need_request_text() { Some(body.query.clone()) } else { @@ -1346,8 +1359,8 @@ impl RouterTrait for PDRouter { is_stream: false, return_logprob: false, request_text: req_text, - routing_id: header_utils::extract_routing_id(headers), model_id, + headers: headers.cloned(), }; self.execute_dual_dispatch(headers, body, context).await diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 06e1b8951..27820e376 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -27,7 +27,7 @@ use crate::{ metrics::{bool_to_static_str, metrics_labels, Metrics}, otel_trace::inject_trace_context_http, }, - policies::PolicyRegistry, + policies::{PolicyRegistry, SelectWorkerInfo}, protocols::{ chat::ChatCompletionRequest, classify::ClassifyRequest, @@ -39,7 +39,7 @@ use crate::{ responses::{ResponsesGetParams, ResponsesRequest}, }, routers::{ - error, + error::{self, extract_error_code_from_response}, grpc::utils::{error_type_from_status, route_to_endpoint}, header_utils, RouterTrait, }, @@ -140,7 +140,8 @@ impl Router { fn select_worker_for_model( &self, model_id: Option<&str>, - info: &crate::policies::SelectWorkerInfo, + text: Option<&str>, + headers: Option<&HeaderMap>, ) -> Option> { let effective_model_id = if !self.enable_igw { None } else { model_id }; @@ -168,7 +169,13 @@ impl Router { None => self.policy_registry.get_default_policy(), }; - let idx = policy.select_worker(&available, info)?; + let idx = policy.select_worker( + &available, + &SelectWorkerInfo { + request_text: text, + headers, + }, + )?; // Record worker selection metric (Layer 3) Metrics::record_worker_selection( @@ -191,11 +198,6 @@ impl Router { let start = Instant::now(); let is_stream = typed_req.is_stream(); let text = typed_req.extract_text_for_routing(); - let routing_id = header_utils::extract_routing_id(headers); - let info = crate::policies::SelectWorkerInfo { - request_text: Some(&text), - routing_id: routing_id.as_deref(), - }; let model = model_id.unwrap_or("default"); let endpoint = route_to_endpoint(route); @@ -214,7 +216,7 @@ impl Router { // operation per attempt |_: u32| async { let res = self - .route_typed_request_once(headers, typed_req, route, model_id, is_stream, &info) + .route_typed_request_once(headers, typed_req, route, model_id, is_stream, &text) .await; // Need to be outside `route_typed_request_once` because that function has multiple return paths @@ -272,9 +274,9 @@ impl Router { route: &'static str, model_id: Option<&str>, is_stream: bool, - info: &crate::policies::SelectWorkerInfo<'_>, + text: &str, ) -> Response { - let worker = match self.select_worker_for_model(model_id, info) { + let worker = match self.select_worker_for_model(model_id, Some(text), headers) { Some(w) => w, None => { return error::service_unavailable( @@ -696,8 +698,6 @@ fn convert_reqwest_error(e: reqwest::Error) -> Response { use async_trait::async_trait; -use crate::routers::error::extract_error_code_from_response; - #[async_trait] impl RouterTrait for Router { fn as_any(&self) -> &dyn std::any::Any { diff --git a/sgl-model-gateway/tests/cache_aware_backward_compat_test.rs b/sgl-model-gateway/tests/cache_aware_backward_compat_test.rs index 938ae3432..f1d94c464 100644 --- a/sgl-model-gateway/tests/cache_aware_backward_compat_test.rs +++ b/sgl-model-gateway/tests/cache_aware_backward_compat_test.rs @@ -103,24 +103,16 @@ fn test_mixed_model_ids() { let default_workers: Vec> = vec![Arc::new(worker1.clone()), Arc::new(worker3.clone())]; - let selected = policy.select_worker( - &default_workers, - &SelectWorkerInfo { - request_text: Some("test request"), - ..Default::default() - }, - ); + let info = SelectWorkerInfo { + request_text: Some("test request"), + ..Default::default() + }; + let selected = policy.select_worker(&default_workers, &info); assert!(selected.is_some(), "Should select from default workers"); let llama_workers: Vec> = vec![Arc::new(worker2.clone()), Arc::new(worker4.clone())]; - let selected = policy.select_worker( - &llama_workers, - &SelectWorkerInfo { - request_text: Some("test request"), - ..Default::default() - }, - ); + let selected = policy.select_worker(&llama_workers, &info); assert!(selected.is_some(), "Should select from llama-3 workers"); let all_workers: Vec> = vec![ @@ -129,13 +121,7 @@ fn test_mixed_model_ids() { Arc::new(worker3.clone()), Arc::new(worker4.clone()), ]; - let selected = policy.select_worker( - &all_workers, - &SelectWorkerInfo { - request_text: Some("test request"), - ..Default::default() - }, - ); + let selected = policy.select_worker(&all_workers, &info); assert!(selected.is_some(), "Should select from all workers"); }