From 45adad37d0f59cd5a6b3cb0ecc2ee40f93ac8b66 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Thu, 25 Dec 2025 11:27:36 +0800 Subject: [PATCH] Add manual routing policy for router (#15586) --- .../benches/request_processing.rs | 2 + .../bindings/python/sglang_router/router.py | 1 + .../python/sglang_router/router_args.py | 13 +- sgl-model-gateway/bindings/python/src/lib.rs | 2 + sgl-model-gateway/src/config/types.rs | 4 + sgl-model-gateway/src/config/validation.rs | 1 + .../src/observability/metrics.rs | 13 + sgl-model-gateway/src/policies/bucket.rs | 39 + sgl-model-gateway/src/policies/cache_aware.rs | 19 +- sgl-model-gateway/src/policies/factory.rs | 9 +- sgl-model-gateway/src/policies/manual.rs | 843 ++++++++++++++++++ sgl-model-gateway/src/policies/mod.rs | 9 + sgl-model-gateway/src/policies/registry.rs | 3 +- sgl-model-gateway/src/protocols/chat.rs | 8 + sgl-model-gateway/src/protocols/classify.rs | 8 + sgl-model-gateway/src/protocols/common.rs | 3 + sgl-model-gateway/src/protocols/completion.rs | 8 + sgl-model-gateway/src/protocols/embedding.rs | 8 + sgl-model-gateway/src/protocols/generate.rs | 8 + sgl-model-gateway/src/protocols/rerank.rs | 9 + sgl-model-gateway/src/protocols/responses.rs | 9 + .../grpc/common/stages/worker_selection.rs | 26 +- sgl-model-gateway/src/routers/grpc/context.rs | 3 + .../grpc/harmony/stages/preparation.rs | 2 + .../grpc/regular/responses/tool_loop.rs | 2 + .../grpc/regular/stages/chat/preparation.rs | 1 + .../regular/stages/embedding/preparation.rs | 4 +- .../regular/stages/generate/preparation.rs | 1 + .../src/routers/http/pd_router.rs | 44 +- sgl-model-gateway/src/routers/http/router.rs | 31 +- .../tests/cache_aware_backward_compat_test.rs | 29 +- sgl-model-gateway/tests/responses_api_test.rs | 8 + sgl-model-gateway/tests/spec/embedding.rs | 6 + sgl-model-gateway/tests/spec/rerank.rs | 15 + .../tests/test_openai_routing.rs | 2 + 35 files changed, 1128 insertions(+), 65 deletions(-) create mode 100644 sgl-model-gateway/src/policies/manual.rs diff --git a/sgl-model-gateway/benches/request_processing.rs b/sgl-model-gateway/benches/request_processing.rs index 60d51faf0..5f967e5d4 100644 --- a/sgl-model-gateway/benches/request_processing.rs +++ b/sgl-model-gateway/benches/request_processing.rs @@ -67,6 +67,7 @@ fn default_generate_request() -> GenerateRequest { return_bytes: false, return_entropy: false, rid: None, + routing_id: None, } } @@ -122,6 +123,7 @@ fn default_completion_request() -> CompletionRequest { return_hidden_states: false, sampling_seed: None, other: serde_json::Map::new(), + routing_id: None, } } diff --git a/sgl-model-gateway/bindings/python/sglang_router/router.py b/sgl-model-gateway/bindings/python/sglang_router/router.py index 4abe19d29..ae0bc04aa 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router.py @@ -21,6 +21,7 @@ def policy_from_str(policy_str: Optional[str]) -> PolicyType: "cache_aware": PolicyType.CacheAware, "power_of_two": PolicyType.PowerOfTwo, "bucket": PolicyType.Bucket, + "manual": PolicyType.Manual, } return policy_map[policy_str] diff --git a/sgl-model-gateway/bindings/python/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/sglang_router/router_args.py index 9fa3125b4..03d206546 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router_args.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router_args.py @@ -171,21 +171,28 @@ class RouterArgs: f"--{prefix}policy", type=str, default=RouterArgs.policy, - choices=["random", "round_robin", "cache_aware", "power_of_two"], + choices=["random", "round_robin", "cache_aware", "power_of_two", "manual"], help="Load balancing policy to use. In PD mode, this is used for both prefill and decode unless overridden", ) parser.add_argument( f"--{prefix}prefill-policy", type=str, default=None, - choices=["random", "round_robin", "cache_aware", "power_of_two", "bucket"], + choices=[ + "random", + "round_robin", + "cache_aware", + "power_of_two", + "manual", + "bucket", + ], help="Specific policy for prefill nodes in PD mode. If not specified, uses the main policy", ) parser.add_argument( f"--{prefix}decode-policy", type=str, default=None, - choices=["random", "round_robin", "cache_aware", "power_of_two"], + choices=["random", "round_robin", "cache_aware", "power_of_two", "manual"], help="Specific policy for decode nodes in PD mode. If not specified, uses the main policy", ) diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index 79337d466..0a159b46e 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -12,6 +12,7 @@ pub enum PolicyType { CacheAware, PowerOfTwo, Bucket, + Manual, } #[pyclass(eq)] @@ -267,6 +268,7 @@ impl Router { balance_rel_threshold: self.balance_rel_threshold, bucket_adjust_interval_secs: self.bucket_adjust_interval_secs, }, + PolicyType::Manual => ConfigPolicyConfig::Manual, } }; diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index a933f1a41..d96167d5f 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -336,6 +336,9 @@ pub enum PolicyConfig { /// Interval between bucket boundary adjustment cycles (seconds) bucket_adjust_interval_secs: usize, }, + + #[serde(rename = "manual")] + Manual, } impl PolicyConfig { @@ -346,6 +349,7 @@ impl PolicyConfig { PolicyConfig::CacheAware { .. } => "cache_aware", PolicyConfig::PowerOfTwo { .. } => "power_of_two", PolicyConfig::Bucket { .. } => "bucket", + PolicyConfig::Manual => "manual", } } } diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index 44d46bd9a..47b86e927 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/sgl-model-gateway/src/config/validation.rs @@ -226,6 +226,7 @@ impl ConfigValidator { }); } } + PolicyConfig::Manual => {} } Ok(()) } diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index 6f3d8d4eb..10266559b 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -190,6 +190,10 @@ 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!( @@ -810,6 +814,15 @@ 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/bucket.rs b/sgl-model-gateway/src/policies/bucket.rs index 5db591e18..3665fb361 100644 --- a/sgl-model-gateway/src/policies/bucket.rs +++ b/sgl-model-gateway/src/policies/bucket.rs @@ -625,6 +625,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(33)), + ..Default::default() }, ) .unwrap(); @@ -634,6 +635,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(34)), + ..Default::default() }, ) .unwrap(); @@ -642,6 +644,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(34)), + ..Default::default() }, ) .unwrap(); @@ -675,6 +678,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(33)), + ..Default::default() }, ) .unwrap(); @@ -683,6 +687,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(33)), + ..Default::default() }, ) .unwrap(); @@ -691,6 +696,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(33)), + ..Default::default() }, ) .unwrap(); @@ -714,6 +720,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(20)), + ..Default::default() }, ) .unwrap(); // worker1: 20 @@ -722,6 +729,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(8)), + ..Default::default() }, ) .unwrap(); // worker1: 8 @@ -732,6 +740,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some("request"), + ..Default::default() }, ) .unwrap(); @@ -756,6 +765,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(15)), + ..Default::default() }, ) .unwrap(); // worker1: 15 @@ -764,6 +774,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some("short"), + ..Default::default() }, ) .unwrap(); // worker2: 5 @@ -772,6 +783,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(10)), + ..Default::default() }, ) .unwrap(); // worker3: 10 @@ -782,6 +794,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some("request"), + ..Default::default() }, ) .unwrap(); @@ -854,6 +867,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(5)), + ..Default::default() }, ) .unwrap(); @@ -862,6 +876,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(10)), + ..Default::default() }, ) .unwrap(); @@ -870,6 +885,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(15)), + ..Default::default() }, ) .unwrap(); @@ -878,6 +894,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(20)), + ..Default::default() }, ) .unwrap(); @@ -886,6 +903,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(24)), + ..Default::default() }, ) .unwrap(); @@ -894,6 +912,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(26)), + ..Default::default() }, ) .unwrap(); @@ -929,6 +948,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(10)), + ..Default::default() }, ) .unwrap(); @@ -937,6 +957,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(20)), + ..Default::default() }, ) .unwrap(); @@ -945,6 +966,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(30)), + ..Default::default() }, ) .unwrap(); @@ -953,6 +975,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(40)), + ..Default::default() }, ) .unwrap(); @@ -961,6 +984,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(45)), + ..Default::default() }, ) .unwrap(); @@ -969,6 +993,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(57)), + ..Default::default() }, ) .unwrap(); @@ -1059,6 +1084,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(20)), + ..Default::default() }, ) .unwrap(); @@ -1091,6 +1117,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(7)), + ..Default::default() }, ) .unwrap(); @@ -1179,6 +1206,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(5)), + ..Default::default() }, ) .unwrap(); @@ -1187,6 +1215,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(10)), + ..Default::default() }, ) .unwrap(); @@ -1195,6 +1224,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(15)), + ..Default::default() }, ) .unwrap(); @@ -1203,6 +1233,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(20)), + ..Default::default() }, ) .unwrap(); @@ -1211,6 +1242,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(24)), + ..Default::default() }, ) .unwrap(); @@ -1219,6 +1251,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(26)), + ..Default::default() }, ) .unwrap(); @@ -1251,6 +1284,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(10)), + ..Default::default() }, ) .unwrap(); @@ -1259,6 +1293,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(20)), + ..Default::default() }, ) .unwrap(); @@ -1267,6 +1302,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(30)), + ..Default::default() }, ) .unwrap(); @@ -1275,6 +1311,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(32)), + ..Default::default() }, ) .unwrap(); @@ -1283,6 +1320,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(45)), + ..Default::default() }, ) .unwrap(); @@ -1291,6 +1329,7 @@ mod tests { &prefill_workers, &SelectWorkerInfo { request_text: Some(&*"a".repeat(55)), + ..Default::default() }, ) .unwrap(); diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index 30cff7d8a..9ddefc08b 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -460,6 +460,7 @@ mod tests { &workers, &SelectWorkerInfo { request_text: Some("hello world"), + ..Default::default() }, ) .unwrap(); @@ -470,6 +471,7 @@ mod tests { &workers, &SelectWorkerInfo { request_text: Some("hello world"), + ..Default::default() }, ) .unwrap(); @@ -481,6 +483,7 @@ mod tests { &workers, &SelectWorkerInfo { request_text: Some("hello"), + ..Default::default() }, ) .unwrap(); @@ -514,11 +517,16 @@ mod tests { policy.init_workers(&workers); // Should select worker2 (lower load) despite cache affinity - let info = SelectWorkerInfo { - request_text: Some("test"), - }; for _ in 0..5 { - let idx = policy.select_worker(&workers, &info).unwrap(); + let idx = policy + .select_worker( + &workers, + &SelectWorkerInfo { + request_text: Some("test"), + ..Default::default() + }, + ) + .unwrap(); assert_eq!(idx, 1); // Should always pick worker2 } } @@ -550,12 +558,14 @@ mod tests { &workers, &SelectWorkerInfo { request_text: Some("test1"), + ..Default::default() }, ); policy.select_worker( &workers, &SelectWorkerInfo { request_text: Some("test2"), + ..Default::default() }, ); @@ -569,6 +579,7 @@ mod tests { &workers, &SelectWorkerInfo { request_text: Some("test1"), + ..Default::default() }, ) .unwrap(); diff --git a/sgl-model-gateway/src/policies/factory.rs b/sgl-model-gateway/src/policies/factory.rs index a9db1842d..b4de69925 100644 --- a/sgl-model-gateway/src/policies/factory.rs +++ b/sgl-model-gateway/src/policies/factory.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use super::{ BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, - PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, + ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, }; use crate::config::PolicyConfig; @@ -46,6 +46,7 @@ impl PolicyFactory { }; Arc::new(BucketPolicy::with_config(config)) } + PolicyConfig::Manual => Arc::new(ManualPolicy::new()), } } @@ -57,6 +58,7 @@ impl PolicyFactory { "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())), "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())), "bucket" => Some(Arc::new(BucketPolicy::new())), + "manual" => Some(Arc::new(ManualPolicy::new())), _ => None, } } @@ -94,6 +96,9 @@ 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] @@ -108,6 +113,8 @@ 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 new file mode 100644 index 000000000..9663d735a --- /dev/null +++ b/sgl-model-gateway/src/policies/manual.rs @@ -0,0 +1,843 @@ +//! Manual routing policy based on routing_id + +use std::sync::Arc; + +use dashmap::{mapref::entry::Entry, DashMap}; +use rand::Rng; + +use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo}; +use crate::{core::Worker, observability::metrics::Metrics}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ExecutionBranch { + NoHealthyWorkers, + FastPathHit, + SlowPathOccupiedHit, + SlowPathOccupiedMiss, + SlowPathVacant, + NoRoutingId, +} + +impl ExecutionBranch { + // TODO auto generate + 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", + } + } +} + +#[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, +} + +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) + } + } + } + + fn select_worker_impl( + &self, + workers: &[Arc], + info: &SelectWorkerInfo, + ) -> (Option, ExecutionBranch) { + let healthy_indices = get_healthy_worker_indices(workers); + if healthy_indices.is_empty() { + return (None, ExecutionBranch::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); + } + } + + ( + Some(random_select(&healthy_indices)), + ExecutionBranch::NoRoutingId, + ) + } +} + +impl LoadBalancingPolicy for ManualPolicy { + fn select_worker(&self, workers: &[Arc], info: &SelectWorkerInfo) -> Option { + let (result, branch) = self.select_worker_impl(workers, info); + Metrics::record_worker_manual_policy_branch(branch.as_str()); + result + } + + fn name(&self) -> &'static str { + "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; + + 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(), + ), + ]; + + let info = SelectWorkerInfo { + routing_id: Some("user-123"), + ..Default::default() + }; + + let (first_result, branch) = policy.select_worker_impl(&workers, &info); + let first_idx = first_result.unwrap(); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); + + 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); + } + } + + #[test] + fn test_manual_different_routing_ids() { + 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 mut distribution = HashMap::new(); + for i in 0..100 { + let routing_id = format!("user-{}", i); + let info = SelectWorkerInfo { + routing_id: Some(&routing_id), + ..Default::default() + }; + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(branch, ExecutionBranch::SlowPathVacant); + *distribution.entry(result.unwrap()).or_insert(0) += 1; + } + + assert!( + distribution.len() > 1, + "Should distribute across multiple workers" + ); + } + + #[test] + fn test_manual_fallback_random() { + 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::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: 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" + ); + } + + #[test] + fn test_manual_remaps_when_worker_becomes_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 info = SelectWorkerInfo { + routing_id: Some("sticky-user"), + ..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); + } + } + + #[test] + fn test_manual_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); + } + + #[test] + fn test_manual_execution_branch_slow_path_occupied_miss() { + 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("occupied-miss-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 (result, branch) = policy.select_worker_impl(&workers, &info); + assert!(result.is_some()); + assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + } + + #[test] + fn test_manual_execution_branch_slow_path_occupied_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("occupied-hit-test"), + ..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); + } + + #[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() { + 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 71482dd3b..43e32616a 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -10,6 +10,7 @@ use crate::core::Worker; mod bucket; mod cache_aware; mod factory; +mod manual; mod power_of_two; mod random; mod registry; @@ -19,6 +20,7 @@ pub mod tree; pub use bucket::BucketPolicy; pub use cache_aware::CacheAwarePolicy; pub use factory::PolicyFactory; +pub use manual::ManualPolicy; pub use power_of_two::PowerOfTwoPolicy; pub use random::RandomPolicy; pub use registry::PolicyRegistry; @@ -55,6 +57,11 @@ 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. @@ -140,6 +147,8 @@ 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>, } #[cfg(test)] diff --git a/sgl-model-gateway/src/policies/registry.rs b/sgl-model-gateway/src/policies/registry.rs index ee6f0135a..02cd4cf3d 100644 --- a/sgl-model-gateway/src/policies/registry.rs +++ b/sgl-model-gateway/src/policies/registry.rs @@ -11,7 +11,7 @@ use tracing::{debug, info, warn}; /// When the last worker of a model is removed, the policy mapping is cleaned up. use super::{ BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, - PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, + ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, }; use crate::{config::types::PolicyConfig, core::Worker}; @@ -209,6 +209,7 @@ impl PolicyRegistry { }; Arc::new(BucketPolicy::with_config(config)) } + PolicyConfig::Manual => Arc::new(ManualPolicy::new()), } } diff --git a/sgl-model-gateway/src/protocols/chat.rs b/sgl-model-gateway/src/protocols/chat.rs index 0a0900165..3ac0b91ff 100644 --- a/sgl-model-gateway/src/protocols/chat.rs +++ b/sgl-model-gateway/src/protocols/chat.rs @@ -359,6 +359,10 @@ pub struct ChatCompletionRequest { /// Random seed for sampling for deterministic outputs #[serde(skip_serializing_if = "Option::is_none")] pub sampling_seed: Option, + + /// Routing ID for manual routing policy + #[serde(skip_serializing_if = "Option::is_none")] + pub routing_id: Option, } // ============================================================================ @@ -696,6 +700,10 @@ impl GenerationRequest for ChatCompletionRequest { buffer } + + fn get_routing_id(&self) -> Option<&str> { + self.routing_id.as_deref() + } } // ============================================================================ diff --git a/sgl-model-gateway/src/protocols/classify.rs b/sgl-model-gateway/src/protocols/classify.rs index fc7e8b871..2be5a12d1 100644 --- a/sgl-model-gateway/src/protocols/classify.rs +++ b/sgl-model-gateway/src/protocols/classify.rs @@ -30,6 +30,10 @@ pub struct ClassifyRequest { /// SGLang extension: request id for tracking #[serde(skip_serializing_if = "Option::is_none")] pub rid: Option, + + /// Routing ID for manual routing policy + #[serde(skip_serializing_if = "Option::is_none")] + pub routing_id: Option, } impl GenerationRequest for ClassifyRequest { @@ -54,4 +58,8 @@ impl GenerationRequest for ClassifyRequest { _ => String::new(), } } + + fn get_routing_id(&self) -> Option<&str> { + self.routing_id.as_deref() + } } diff --git a/sgl-model-gateway/src/protocols/common.rs b/sgl-model-gateway/src/protocols/common.rs index 83e2d238a..c1e18a4e2 100644 --- a/sgl-model-gateway/src/protocols/common.rs +++ b/sgl-model-gateway/src/protocols/common.rs @@ -36,6 +36,9 @@ pub trait GenerationRequest: Send + Sync { /// Extract text content for routing decisions fn extract_text_for_routing(&self) -> String; + + /// Get routing ID for manual routing policy + fn get_routing_id(&self) -> Option<&str>; } // ============================================================================ diff --git a/sgl-model-gateway/src/protocols/completion.rs b/sgl-model-gateway/src/protocols/completion.rs index c6a4f638a..32411b4e5 100644 --- a/sgl-model-gateway/src/protocols/completion.rs +++ b/sgl-model-gateway/src/protocols/completion.rs @@ -145,6 +145,10 @@ pub struct CompletionRequest { /// Additional fields including bootstrap info for PD routing #[serde(flatten)] pub other: Map, + + /// Routing ID for manual routing policy + #[serde(skip_serializing_if = "Option::is_none")] + pub routing_id: Option, } impl GenerationRequest for CompletionRequest { @@ -162,6 +166,10 @@ impl GenerationRequest for CompletionRequest { StringOrArray::Array(v) => v.join(" "), } } + + fn get_routing_id(&self) -> Option<&str> { + self.routing_id.as_deref() + } } // ============================================================================ diff --git a/sgl-model-gateway/src/protocols/embedding.rs b/sgl-model-gateway/src/protocols/embedding.rs index 12e3daf19..22f44105e 100644 --- a/sgl-model-gateway/src/protocols/embedding.rs +++ b/sgl-model-gateway/src/protocols/embedding.rs @@ -31,6 +31,10 @@ pub struct EmbeddingRequest { #[serde(skip_serializing_if = "Option::is_none")] pub rid: Option, + /// Routing ID for manual routing policy + #[serde(skip_serializing_if = "Option::is_none")] + pub routing_id: Option, + /// SGLang extension: enable/disable logging of metrics for this request #[serde(skip_serializing_if = "Option::is_none")] pub log_metrics: Option, @@ -58,6 +62,10 @@ impl GenerationRequest for EmbeddingRequest { _ => String::new(), } } + + fn get_routing_id(&self) -> Option<&str> { + self.routing_id.as_deref() + } } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/sgl-model-gateway/src/protocols/generate.rs b/sgl-model-gateway/src/protocols/generate.rs index d5819095a..f1f50da15 100644 --- a/sgl-model-gateway/src/protocols/generate.rs +++ b/sgl-model-gateway/src/protocols/generate.rs @@ -167,6 +167,10 @@ pub struct GenerateRequest { /// Request ID for tracking (inherited from BaseReq in Python) #[serde(skip_serializing_if = "Option::is_none")] pub rid: Option, + + /// Routing ID for manual routing policy + #[serde(skip_serializing_if = "Option::is_none")] + pub routing_id: Option, } impl Normalizable for GenerateRequest { @@ -235,6 +239,10 @@ impl GenerationRequest for GenerateRequest { // No text input found String::new() } + + fn get_routing_id(&self) -> Option<&str> { + self.routing_id.as_deref() + } } // ============================================================================ diff --git a/sgl-model-gateway/src/protocols/rerank.rs b/sgl-model-gateway/src/protocols/rerank.rs index 6775f5d8a..0b2f91494 100644 --- a/sgl-model-gateway/src/protocols/rerank.rs +++ b/sgl-model-gateway/src/protocols/rerank.rs @@ -52,6 +52,10 @@ pub struct RerankRequest { /// User identifier pub user: Option, + + /// Routing ID for manual routing policy + #[serde(skip_serializing_if = "Option::is_none")] + pub routing_id: Option, } impl GenerationRequest for RerankRequest { @@ -66,6 +70,10 @@ impl GenerationRequest for RerankRequest { fn extract_text_for_routing(&self) -> String { self.query.clone() } + + fn get_routing_id(&self) -> Option<&str> { + self.routing_id.as_deref() + } } impl super::validated::Normalizable for RerankRequest { @@ -207,6 +215,7 @@ impl From for RerankRequest { return_documents: true, rid: None, user: None, + routing_id: None, } } } diff --git a/sgl-model-gateway/src/protocols/responses.rs b/sgl-model-gateway/src/protocols/responses.rs index f2f41d348..167f0e938 100644 --- a/sgl-model-gateway/src/protocols/responses.rs +++ b/sgl-model-gateway/src/protocols/responses.rs @@ -616,6 +616,10 @@ pub struct ResponsesRequest { #[serde(default = "default_repetition_penalty")] #[validate(range(min = 0.0, max = 2.0))] pub repetition_penalty: f32, + + /// Routing ID for manual routing policy + #[serde(skip_serializing_if = "Option::is_none")] + pub routing_id: Option, } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -659,6 +663,7 @@ impl Default for ResponsesRequest { top_k: default_top_k(), min_p: 0.0, repetition_penalty: default_repetition_penalty(), + routing_id: None, } } } @@ -770,6 +775,10 @@ impl GenerationRequest for ResponsesRequest { .join(" "), } } + + fn get_routing_id(&self) -> Option<&str> { + self.routing_id.as_deref() + } } /// Validate conversation ID format 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 740a3fb00..aa94c4c4f 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,15 +61,18 @@ impl PipelineStage for WorkerSelectionStage { // For Harmony, use selection_text produced during Harmony encoding // Otherwise, use original_text from regular preparation - let text = if prep.harmony_mode { - prep.selection_text.as_deref() - } else { - prep.original_text.as_deref() + 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 workers = match self.mode { WorkerSelectionMode::Regular => { - match self.select_single_worker(ctx.input.model_id.as_deref(), text) { + match self.select_single_worker(ctx.input.model_id.as_deref(), &info) { Some(w) => WorkerSelection::Single { worker: w }, None => { error!( @@ -86,7 +89,7 @@ impl PipelineStage for WorkerSelectionStage { } } WorkerSelectionMode::PrefillDecode => { - match self.select_pd_pair(ctx.input.model_id.as_deref(), text) { + match self.select_pd_pair(ctx.input.model_id.as_deref(), &info) { Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode }, None => { error!( @@ -120,7 +123,7 @@ impl WorkerSelectionStage { fn select_single_worker( &self, model_id: Option<&str>, - text: Option<&str>, + info: &SelectWorkerInfo, ) -> Option> { // Get workers for the specified model, filtered by connection mode let workers = self.worker_registry.get_workers_filtered( @@ -146,7 +149,7 @@ impl WorkerSelectionStage { }; // Select worker using the policy - let idx = policy.select_worker(&available, &SelectWorkerInfo { request_text: text })?; + let idx = policy.select_worker(&available, info)?; let selected = available[idx].clone(); // Record worker selection metric @@ -163,7 +166,7 @@ impl WorkerSelectionStage { fn select_pd_pair( &self, model_id: Option<&str>, - text: Option<&str>, + info: &SelectWorkerInfo, ) -> Option<(Arc, Arc)> { let all_workers = self.worker_registry.get_workers_filtered( model_id, @@ -203,9 +206,8 @@ impl WorkerSelectionStage { None => self.policy_registry.get_default_policy(), }; - let info = SelectWorkerInfo { request_text: text }; - let prefill_idx = policy.select_worker(&available_prefill, &info)?; - let decode_idx = policy.select_worker(&available_decode, &info)?; + 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 a5da72f01..cb0469b2c 100644 --- a/sgl-model-gateway/src/routers/grpc/context.rs +++ b/sgl-model-gateway/src/routers/grpc/context.rs @@ -94,6 +94,9 @@ 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 32b309319..7be967da6 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs @@ -123,6 +123,7 @@ impl HarmonyPreparationStage { // Step 4: Store results ctx.state.preparation = Some(PreparationOutput { original_text: None, + routing_id: request.routing_id.clone(), token_ids: build_output.input_ids, processed_messages: None, tool_constraints, @@ -203,6 +204,7 @@ impl HarmonyPreparationStage { // Step 4: Store results with constraint ctx.state.preparation = Some(PreparationOutput { original_text: None, + routing_id: request.routing_id.clone(), token_ids: build_output.input_ids, processed_messages: None, tool_constraints: constraint, diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs index 0f2df521f..10897f1fd 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs @@ -492,6 +492,7 @@ pub(super) async fn execute_tool_loop( top_k: current_request.top_k, min_p: current_request.min_p, repetition_penalty: current_request.repetition_penalty, + routing_id: current_request.routing_id.clone(), }; // Continue to next iteration @@ -1070,6 +1071,7 @@ async fn execute_tool_loop_streaming_internal( top_k: current_request.top_k, min_p: current_request.min_p, repetition_penalty: current_request.repetition_penalty, + routing_id: current_request.routing_id.clone(), }; continue; 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 ee951ab64..9f637cde3 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 @@ -96,6 +96,7 @@ impl ChatPreparationStage { // Store results in context ctx.state.preparation = Some(PreparationOutput { original_text: Some(processed_messages.text.clone()), + routing_id: request.routing_id.clone(), 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 fae235a3d..f4c3a0b8b 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 @@ -47,8 +47,9 @@ impl PipelineStage for EmbeddingPreparationStage { )); }; - // Extract text from request + // Extract text and routing_id from request before borrowing ctx mutably let text = request.extract_text_for_routing(); + let routing_id = request.routing_id.clone(); if text.is_empty() { return Err(error::bad_request( "empty_input", @@ -77,6 +78,7 @@ 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 d1aea48cb..86c6f53df 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 @@ -68,6 +68,7 @@ impl GeneratePreparationStage { ctx.state.preparation = Some(PreparationOutput { original_text, + routing_id: request.routing_id.clone(), token_ids, processed_messages: None, tool_constraints: None, diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 94f1d52de..bf9744eac 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, SelectWorkerInfo}, + policies::{LoadBalancingPolicy, PolicyRegistry}, protocols::{ chat::{ChatCompletionRequest, ChatMessage, MessageContent}, common::{InputIds, StringOrArray}, @@ -58,6 +58,7 @@ struct PDRequestContext<'a> { is_stream: bool, return_logprob: bool, request_text: Option, + routing_id: Option, model_id: Option<&'a str>, } @@ -303,7 +304,11 @@ impl PDRouter { let context = context.clone(); async move { let (prefill, decode) = match self - .select_pd_pair(context.request_text.as_deref(), context.model_id) + .select_pd_pair( + context.request_text.as_deref(), + context.routing_id.as_deref(), + context.model_id, + ) .await { Ok(pair) => pair, @@ -691,6 +696,7 @@ impl PDRouter { async fn select_pd_pair( &self, request_text: Option<&str>, + routing_id: Option<&str>, model_id: Option<&str>, ) -> Result<(Arc, Arc), String> { let effective_model_id = if !self.enable_igw { None } else { model_id }; @@ -725,19 +731,16 @@ impl PDRouter { let prefill_policy = self.policy_registry.get_prefill_policy(); let decode_policy = self.policy_registry.get_decode_policy(); - let prefill = Self::pick_worker_by_policy_arc( - &prefill_workers, - &*prefill_policy, + let info = crate::policies::SelectWorkerInfo { request_text, - "prefill", - )?; + routing_id, + }; - let decode = Self::pick_worker_by_policy_arc( - &decode_workers, - &*decode_policy, - request_text, - "decode", - )?; + 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")?; // Record worker selection metrics (Layer 3) let model = model_id.unwrap_or("default"); @@ -760,7 +763,7 @@ impl PDRouter { fn pick_worker_by_policy_arc( workers: &[Arc], policy: &dyn LoadBalancingPolicy, - request_text: Option<&str>, + info: &crate::policies::SelectWorkerInfo, worker_type: &str, ) -> Result, String> { if workers.is_empty() { @@ -784,7 +787,7 @@ impl PDRouter { } let selected_idx = policy - .select_worker(&available_workers, &SelectWorkerInfo { request_text }) + .select_worker(&available_workers, info) .ok_or_else(|| { format!( "Policy {} failed to select a {} worker", @@ -1120,7 +1123,7 @@ impl RouterTrait for PDRouter { // Note: This endpoint actually causes the model to generate tokens, so we only test one pair // Select a random worker pair using the policy - let (prefill, decode) = match self.select_pd_pair(None, None).await { + let (prefill, decode) = match self.select_pd_pair(None, None, None).await { Ok(pair) => pair, Err(e) => { return error::service_unavailable( @@ -1242,6 +1245,7 @@ impl RouterTrait for PDRouter { is_stream, return_logprob, request_text, + routing_id: body.routing_id.clone(), model_id, }; @@ -1283,6 +1287,7 @@ impl RouterTrait for PDRouter { is_stream, return_logprob, request_text, + routing_id: body.routing_id.clone(), model_id, }; @@ -1316,6 +1321,7 @@ impl RouterTrait for PDRouter { is_stream, return_logprob, request_text, + routing_id: body.routing_id.clone(), model_id, }; @@ -1328,7 +1334,6 @@ 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 { @@ -1341,6 +1346,7 @@ impl RouterTrait for PDRouter { is_stream: false, return_logprob: false, request_text: req_text, + routing_id: body.routing_id.clone(), model_id, }; @@ -1405,7 +1411,7 @@ mod tests { router.worker_registry.register(Arc::from(healthy_worker)); router.worker_registry.register(Arc::from(decode_worker)); - let result = router.select_pd_pair(None, None).await; + let result = router.select_pd_pair(None, None, None).await; assert!(result.is_ok()); let (prefill, _decode) = result.unwrap(); @@ -1418,7 +1424,7 @@ mod tests { async fn test_empty_worker_lists() { let router = create_test_pd_router(); - let result = router.select_pd_pair(None, None).await; + let result = router.select_pd_pair(None, None, None).await; assert!(result.is_err()); assert!(result.unwrap_err().contains("No prefill workers available")); diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 89e1c12aa..50e1071ca 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, SelectWorkerInfo}, + policies::PolicyRegistry, protocols::{ chat::ChatCompletionRequest, classify::ClassifyRequest, @@ -39,7 +39,7 @@ use crate::{ responses::{ResponsesGetParams, ResponsesRequest}, }, routers::{ - error::{self, extract_error_code_from_response}, + error, grpc::utils::{error_type_from_status, route_to_endpoint}, header_utils, RouterTrait, }, @@ -140,7 +140,7 @@ impl Router { fn select_worker_for_model( &self, model_id: Option<&str>, - text: Option<&str>, + info: &crate::policies::SelectWorkerInfo, ) -> Option> { let effective_model_id = if !self.enable_igw { None } else { model_id }; @@ -168,7 +168,7 @@ impl Router { None => self.policy_registry.get_default_policy(), }; - let idx = policy.select_worker(&available, &SelectWorkerInfo { request_text: text })?; + let idx = policy.select_worker(&available, info)?; // Record worker selection metric (Layer 3) Metrics::record_worker_selection( @@ -191,6 +191,11 @@ impl Router { let start = Instant::now(); let is_stream = typed_req.is_stream(); let text = typed_req.extract_text_for_routing(); + let routing_id = typed_req.get_routing_id().map(|s| s.to_string()); + 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); @@ -208,18 +213,8 @@ impl Router { &self.retry_config, // operation per attempt |_: u32| async { - let res = self - .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 - Metrics::record_router_upstream_response( - metrics_labels::ROUTER_HTTP, - res.status().as_u16(), - extract_error_code_from_response(&res), - ); - - res + self.route_typed_request_once(headers, typed_req, route, model_id, is_stream, &info) + .await }, // should_retry predicate |res, _attempt| is_retryable_status(res.status()), @@ -267,9 +262,9 @@ impl Router { route: &'static str, model_id: Option<&str>, is_stream: bool, - text: &str, + info: &crate::policies::SelectWorkerInfo<'_>, ) -> Response { - let worker = match self.select_worker_for_model(model_id, Some(text)) { + let worker = match self.select_worker_for_model(model_id, info) { Some(w) => w, None => { return error::service_unavailable( 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 6412c4129..938ae3432 100644 --- a/sgl-model-gateway/tests/cache_aware_backward_compat_test.rs +++ b/sgl-model-gateway/tests/cache_aware_backward_compat_test.rs @@ -44,6 +44,7 @@ fn test_backward_compatibility_with_empty_model_id() { &workers, &SelectWorkerInfo { request_text: Some("test request"), + ..Default::default() }, ); assert!(selected.is_some(), "Should select a worker"); @@ -102,15 +103,24 @@ fn test_mixed_model_ids() { let default_workers: Vec> = vec![Arc::new(worker1.clone()), Arc::new(worker3.clone())]; - let info = SelectWorkerInfo { - request_text: Some("test request"), - }; - let selected = policy.select_worker(&default_workers, &info); + let selected = policy.select_worker( + &default_workers, + &SelectWorkerInfo { + request_text: Some("test request"), + ..Default::default() + }, + ); 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, &info); + let selected = policy.select_worker( + &llama_workers, + &SelectWorkerInfo { + request_text: Some("test request"), + ..Default::default() + }, + ); assert!(selected.is_some(), "Should select from llama-3 workers"); let all_workers: Vec> = vec![ @@ -119,7 +129,13 @@ fn test_mixed_model_ids() { Arc::new(worker3.clone()), Arc::new(worker4.clone()), ]; - let selected = policy.select_worker(&all_workers, &info); + let selected = policy.select_worker( + &all_workers, + &SelectWorkerInfo { + request_text: Some("test request"), + ..Default::default() + }, + ); assert!(selected.is_some(), "Should select from all workers"); } @@ -156,6 +172,7 @@ fn test_remove_worker_by_url_backward_compat() { &workers, &SelectWorkerInfo { request_text: Some("test"), + ..Default::default() }, ); assert_eq!(selected, Some(0), "Should only have worker2 left"); diff --git a/sgl-model-gateway/tests/responses_api_test.rs b/sgl-model-gateway/tests/responses_api_test.rs index e58c124cb..fd82b00a1 100644 --- a/sgl-model-gateway/tests/responses_api_test.rs +++ b/sgl-model-gateway/tests/responses_api_test.rs @@ -105,6 +105,7 @@ async fn test_non_streaming_mcp_minimal_e2e_with_persistence() { min_p: 0.0, repetition_penalty: 1.0, conversation: None, + routing_id: None, }; let resp = router @@ -328,6 +329,7 @@ fn test_responses_request_creation() { min_p: 0.0, repetition_penalty: 1.0, conversation: None, + routing_id: None, }; assert!(!request.is_stream()); @@ -372,6 +374,7 @@ fn test_responses_request_sglang_extensions() { min_p: 0.05, repetition_penalty: 1.1, conversation: None, + routing_id: None, }; // Verify SGLang extensions are present @@ -487,6 +490,7 @@ fn test_json_serialization() { min_p: 0.1, repetition_penalty: 1.2, conversation: None, + routing_id: None, }; let json = serde_json::to_string(&request).expect("Serialization should work"); @@ -593,6 +597,7 @@ async fn test_multi_turn_loop_with_mcp() { min_p: 0.0, repetition_penalty: 1.0, conversation: None, + routing_id: None, }; // Execute the request (this should trigger the multi-turn loop) @@ -742,6 +747,7 @@ async fn test_max_tool_calls_limit() { min_p: 0.0, repetition_penalty: 1.0, conversation: None, + routing_id: None, }; let response = router.route_responses(None, &req, None).await; @@ -914,6 +920,7 @@ async fn test_streaming_with_mcp_tool_calls() { min_p: 0.0, repetition_penalty: 1.0, conversation: None, + routing_id: None, }; let response = router.route_responses(None, &req, None).await; @@ -1194,6 +1201,7 @@ async fn test_streaming_multi_turn_with_mcp() { min_p: 0.0, repetition_penalty: 1.0, conversation: None, + routing_id: None, }; let response = router.route_responses(None, &req, None).await; diff --git a/sgl-model-gateway/tests/spec/embedding.rs b/sgl-model-gateway/tests/spec/embedding.rs index 721c0a5ff..2925776ab 100644 --- a/sgl-model-gateway/tests/spec/embedding.rs +++ b/sgl-model-gateway/tests/spec/embedding.rs @@ -10,6 +10,7 @@ fn test_embedding_request_serialization_string_input() { user: Some("user-1".to_string()), dimensions: Some(128), rid: Some("rid-123".to_string()), + routing_id: None, log_metrics: None, }; @@ -33,6 +34,7 @@ fn test_embedding_request_serialization_array_input() { user: None, dimensions: None, rid: None, + routing_id: None, log_metrics: None, }; @@ -51,6 +53,7 @@ fn test_embedding_generation_request_trait_string() { user: None, dimensions: None, rid: None, + routing_id: None, log_metrics: None, }; assert!(!req.is_stream()); @@ -67,6 +70,7 @@ fn test_embedding_generation_request_trait_array() { user: None, dimensions: None, rid: None, + routing_id: None, log_metrics: None, }; assert_eq!(req.extract_text_for_routing(), "hello world"); @@ -81,6 +85,7 @@ fn test_embedding_generation_request_trait_non_text() { user: None, dimensions: None, rid: None, + routing_id: None, log_metrics: None, }; assert_eq!(req.extract_text_for_routing(), ""); @@ -95,6 +100,7 @@ fn test_embedding_generation_request_trait_mixed_array_ignores_nested() { user: None, dimensions: None, rid: None, + routing_id: None, log_metrics: None, }; // Only top-level string elements are extracted diff --git a/sgl-model-gateway/tests/spec/rerank.rs b/sgl-model-gateway/tests/spec/rerank.rs index 4a40990d3..b8e7cf4c1 100644 --- a/sgl-model-gateway/tests/spec/rerank.rs +++ b/sgl-model-gateway/tests/spec/rerank.rs @@ -17,6 +17,7 @@ fn test_rerank_request_serialization() { return_documents: true, rid: Some(StringOrArray::String("req-123".to_string())), user: Some("user-456".to_string()), + routing_id: None, }; let serialized = to_string(&request).unwrap(); @@ -59,6 +60,7 @@ fn test_rerank_request_validation_success() { return_documents: true, rid: None, user: None, + routing_id: None, }; assert!(request.validate().is_ok()); @@ -74,6 +76,7 @@ fn test_rerank_request_validation_empty_query() { return_documents: true, rid: None, user: None, + routing_id: None, }; let result = request.validate(); @@ -90,6 +93,7 @@ fn test_rerank_request_validation_whitespace_query() { return_documents: true, rid: None, user: None, + routing_id: None, }; let result = request.validate(); @@ -106,6 +110,7 @@ fn test_rerank_request_validation_empty_documents() { return_documents: true, rid: None, user: None, + routing_id: None, }; let result = request.validate(); @@ -122,6 +127,7 @@ fn test_rerank_request_validation_top_k_zero() { return_documents: true, rid: None, user: None, + routing_id: None, }; let result = request.validate(); @@ -138,6 +144,7 @@ fn test_rerank_request_validation_top_k_greater_than_docs() { return_documents: true, rid: None, user: None, + routing_id: None, }; // This should pass but log a warning @@ -154,6 +161,7 @@ fn test_rerank_request_effective_top_k() { return_documents: true, rid: None, user: None, + routing_id: None, }; assert_eq!(request.effective_top_k(), 2); @@ -169,6 +177,7 @@ fn test_rerank_request_effective_top_k_none() { return_documents: true, rid: None, user: None, + routing_id: None, }; assert_eq!(request.effective_top_k(), 3); @@ -390,6 +399,7 @@ fn test_rerank_request_generation_request_trait() { return_documents: true, rid: None, user: None, + routing_id: None, }; assert_eq!(request.get_model(), Some("test-model")); @@ -408,6 +418,7 @@ fn test_rerank_request_very_long_query() { return_documents: true, rid: None, user: None, + routing_id: None, }; assert!(request.validate().is_ok()); @@ -424,6 +435,7 @@ fn test_rerank_request_many_documents() { return_documents: true, rid: None, user: None, + routing_id: None, }; assert!(request.validate().is_ok()); @@ -443,6 +455,7 @@ fn test_rerank_request_special_characters() { return_documents: true, rid: Some(StringOrArray::String("req-🚀-123".to_string())), user: Some("user-🎉-456".to_string()), + routing_id: None, }; assert!(request.validate().is_ok()); @@ -461,6 +474,7 @@ fn test_rerank_request_rid_array() { "req2".to_string(), ])), user: None, + routing_id: None, }; assert!(request.validate().is_ok()); @@ -515,6 +529,7 @@ fn test_full_rerank_workflow() { return_documents: true, rid: Some(StringOrArray::String("req-123".to_string())), user: Some("user-456".to_string()), + routing_id: None, }; // Validate request diff --git a/sgl-model-gateway/tests/test_openai_routing.rs b/sgl-model-gateway/tests/test_openai_routing.rs index 282522d5e..aac49c041 100644 --- a/sgl-model-gateway/tests/test_openai_routing.rs +++ b/sgl-model-gateway/tests/test_openai_routing.rs @@ -89,6 +89,7 @@ fn create_minimal_completion_request() -> CompletionRequest { return_hidden_states: false, sampling_seed: None, other: serde_json::Map::new(), + routing_id: None, } } @@ -639,6 +640,7 @@ async fn test_unsupported_endpoints() { return_bytes: false, return_entropy: false, rid: None, + routing_id: None, }; let response = router.route_generate(None, &generate_request, None).await;