diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 7e3af70b02..669f407473 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -368,6 +368,7 @@ struct Router { host: String, port: u16, health_check_port: Option, + routing_key_override: bool, worker_urls: Vec, policy: PolicyType, worker_startup_timeout_secs: u64, @@ -504,6 +505,19 @@ impl Router { }) } + fn parse_assignment_mode(&self) -> Result { + match self.assignment_mode.as_str() { + "random" => Ok(config::ManualAssignmentMode::Random), + "min_load" => Ok(config::ManualAssignmentMode::MinLoad), + "min_group" => Ok(config::ManualAssignmentMode::MinGroup), + other => Err(config::ConfigError::InvalidValue { + field: "assignment_mode".to_string(), + value: other.to_string(), + reason: "expected 'random', 'min_load', or 'min_group'".to_string(), + }), + } + } + pub fn to_router_config(&self) -> config::ConfigResult { use config::{ DiscoveryConfig, MetricsConfig, PolicyConfig as ConfigPolicyConfig, RoutingMode, @@ -541,18 +555,7 @@ impl Router { PolicyType::Manual => ConfigPolicyConfig::Manual { eviction_interval_secs: self.eviction_interval_secs, max_idle_secs: self.max_idle_secs, - assignment_mode: match self.assignment_mode.as_str() { - "random" => config::ManualAssignmentMode::Random, - "min_load" => config::ManualAssignmentMode::MinLoad, - "min_group" => config::ManualAssignmentMode::MinGroup, - other => { - return Err(config::ConfigError::InvalidValue { - field: "assignment_mode".to_string(), - value: other.to_string(), - reason: "expected 'random', 'min_load', or 'min_group'".to_string(), - }); - } - }, + assignment_mode: self.parse_assignment_mode()?, }, PolicyType::ConsistentHashing => ConfigPolicyConfig::ConsistentHashing, PolicyType::PrefixHash => ConfigPolicyConfig::PrefixHash { @@ -756,6 +759,12 @@ impl Router { .maybe_storage_hook_wasm_path(self.storage_hook_wasm_path.as_deref()) .enable_wasm(self.enable_wasm) .dp_aware(self.dp_aware) + .routing_key_override(config::RoutingKeyOverrideConfig { + enabled: self.routing_key_override, + eviction_interval_secs: self.eviction_interval_secs, + max_idle_secs: self.max_idle_secs, + assignment_mode: self.parse_assignment_mode()?, + }) .retries(!self.disable_retries) .circuit_breaker(!self.disable_circuit_breaker) .igw(self.enable_igw) @@ -890,6 +899,7 @@ impl Router { // positional argument keeps its index for callers that construct // `_Router(...)` positionally. See the struct-field note above. health_check_port = None, + routing_key_override = false, ))] #[expect(clippy::too_many_arguments)] #[expect( @@ -1009,6 +1019,7 @@ impl Router { // Appended last to match the `#[pyo3(signature)]` order above and // preserve positional-argument compatibility. health_check_port: Option, + routing_key_override: bool, ) -> PyResult { let mut all_urls = worker_urls.clone(); @@ -1028,6 +1039,7 @@ impl Router { host, port, health_check_port, + routing_key_override, worker_urls, policy, worker_startup_timeout_secs, diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index c374c9697a..8b7ba74903 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -66,6 +66,7 @@ class RouterArgs: max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches bucket_adjust_interval_secs: int = 5 dp_aware: bool = False + routing_key_override: bool = False dp_minimum_tokens_scheduler: bool = False enable_igw: bool = False # Enable IGW (Inter-Gateway) mode for multi-model support api_key: str | None = None @@ -473,6 +474,11 @@ def add_cli_args( action="store_true", help="Enable data parallelism aware schedule", ) + routing_group.add_argument( + f"--{prefix}routing-key-override", + action="store_true", + help="Honor X-SMG-Routing-Key for sticky routing on any policy", + ) routing_group.add_argument( f"--{prefix}dp-minimum-tokens-scheduler", action="store_true", diff --git a/model_gateway/src/app_context.rs b/model_gateway/src/app_context.rs index 2873848039..a9fc289880 100644 --- a/model_gateway/src/app_context.rs +++ b/model_gateway/src/app_context.rs @@ -526,7 +526,10 @@ impl AppContextBuilder { /// Create policy registry fn with_policy_registry(mut self, config: &RouterConfig) -> Self { - self.policy_registry = Some(Arc::new(PolicyRegistry::new(config.policy.clone()))); + self.policy_registry = Some(Arc::new(PolicyRegistry::with_override( + config.policy.clone(), + config.routing_key_override.clone(), + ))); self } diff --git a/model_gateway/src/config/builder.rs b/model_gateway/src/config/builder.rs index b9bf025852..e019b54bbe 100644 --- a/model_gateway/src/config/builder.rs +++ b/model_gateway/src/config/builder.rs @@ -5,7 +5,8 @@ use smg_mcp::McpConfig; use super::{ CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig, HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RedisConfig, - RetryConfig, RouterConfig, RoutingMode, TokenizerCacheConfig, TraceConfig, + RetryConfig, RouterConfig, RoutingKeyOverrideConfig, RoutingMode, TokenizerCacheConfig, + TraceConfig, }; use crate::worker::ConnectionMode; @@ -526,6 +527,11 @@ impl RouterConfigBuilder { self } + pub fn routing_key_override(mut self, config: RoutingKeyOverrideConfig) -> Self { + self.config.routing_key_override = config; + self + } + /// Inverse of disable_retries field pub fn retries(mut self, enable: bool) -> Self { self.config.disable_retries = !enable; diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index 7a08ca8d0a..fab2e4836c 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -17,6 +17,9 @@ pub struct RouterConfig { #[serde(default)] pub connection_mode: ConnectionMode, pub policy: PolicyConfig, + /// Per-request sticky-routing override (honors `X-SMG-Routing-Key`). + #[serde(default)] + pub routing_key_override: RoutingKeyOverrideConfig, pub host: String, pub port: u16, /// Dedicated port for the isolated Kubernetes liveness/readiness/health @@ -303,6 +306,34 @@ pub enum ManualAssignmentMode { MinGroup, } +/// Per-request sticky-routing override: when `X-SMG-Routing-Key` is present, any +/// eligible policy routes via manual sticky-map semantics. Reuses the manual +/// policy knobs for the sticky map; eviction defaults match the manual policy so +/// config-file users with only `enabled: true` still get TTL eviction (no leak). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RoutingKeyOverrideConfig { + /// When false, policies are used unchanged. + #[serde(default)] + pub enabled: bool, + #[serde(default = "default_manual_eviction_interval_secs")] + pub eviction_interval_secs: u64, + #[serde(default = "default_manual_max_idle_secs")] + pub max_idle_secs: u64, + #[serde(default)] + pub assignment_mode: ManualAssignmentMode, +} + +impl Default for RoutingKeyOverrideConfig { + fn default() -> Self { + Self { + enabled: false, + eviction_interval_secs: default_manual_eviction_interval_secs(), + max_idle_secs: default_manual_max_idle_secs(), + assignment_mode: ManualAssignmentMode::default(), + } + } +} + /// Policy configuration for routing #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type")] @@ -659,6 +690,7 @@ impl Default for RouterConfig { worker_urls: vec![], }, policy: PolicyConfig::Random, + routing_key_override: RoutingKeyOverrideConfig::default(), host: "0.0.0.0".to_string(), port: 3001, health_check_port: None, diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 8a4c42e4d5..3271e9fd19 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -7,7 +7,7 @@ use smg::{ validate_mesh_server_name, CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig, HistoryBackend, ManualAssignmentMode, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RedisConfig, RetryConfig, RouterConfig, - RoutingMode, SchemaConfig, TokenizerCacheConfig, TraceConfig, + RoutingKeyOverrideConfig, RoutingMode, SchemaConfig, TokenizerCacheConfig, TraceConfig, }, observability::{ metrics::PrometheusConfig, @@ -239,6 +239,11 @@ struct CliArgs { #[arg(long, default_value_t = false, help_heading = "Routing Policy")] dp_aware: bool, + /// Honor X-SMG-Routing-Key for sticky routing on any policy (reuses the + /// manual eviction/idle/assignment knobs for the sticky map) + #[arg(long, default_value_t = false, help_heading = "Routing Policy")] + routing_key_override: bool, + /// Enable IGW (Inference Gateway) mode for multi-model support #[arg(long, default_value_t = false, help_heading = "Routing Policy")] enable_igw: bool, @@ -965,10 +970,6 @@ impl CliArgs { })) } - #[expect( - clippy::panic, - reason = "unreachable: clap value_parser restricts valid assignment modes" - )] fn parse_policy(&self, policy_str: &str) -> PolicyConfig { match policy_str { "random" => PolicyConfig::Random, @@ -1000,17 +1001,25 @@ impl CliArgs { "manual" => PolicyConfig::Manual { eviction_interval_secs: self.eviction_interval, max_idle_secs: self.max_idle_secs, - assignment_mode: match self.assignment_mode.as_str() { - "random" => ManualAssignmentMode::Random, - "min_load" => ManualAssignmentMode::MinLoad, - "min_group" => ManualAssignmentMode::MinGroup, - other => panic!("Unknown assignment mode: {other}"), - }, + assignment_mode: Self::parse_assignment_mode(&self.assignment_mode), }, _ => PolicyConfig::RoundRobin, } } + #[expect( + clippy::panic, + reason = "unreachable: clap value_parser restricts valid assignment modes" + )] + fn parse_assignment_mode(mode: &str) -> ManualAssignmentMode { + match mode { + "random" => ManualAssignmentMode::Random, + "min_load" => ManualAssignmentMode::MinLoad, + "min_group" => ManualAssignmentMode::MinGroup, + other => panic!("Unknown assignment mode: {other}"), + } + } + fn load_schema_config(&self) -> ConfigResult> { match &self.schema_config { Some(path) => { @@ -1340,6 +1349,12 @@ impl CliArgs { .maybe_tool_call_parser(self.tool_call_parser.as_ref()) .maybe_mcp_config_path(self.mcp_config_path.as_ref()) .dp_aware(self.dp_aware) + .routing_key_override(RoutingKeyOverrideConfig { + enabled: self.routing_key_override, + eviction_interval_secs: self.eviction_interval, + max_idle_secs: self.max_idle_secs, + assignment_mode: Self::parse_assignment_mode(&self.assignment_mode), + }) .retries(!self.disable_retries) .circuit_breaker(!self.disable_circuit_breaker) .enable_wasm(self.enable_wasm) diff --git a/model_gateway/src/policies/manual.rs b/model_gateway/src/policies/manual.rs index c94f265844..e694e0d886 100644 --- a/model_gateway/src/policies/manual.rs +++ b/model_gateway/src/policies/manual.rs @@ -21,6 +21,7 @@ use tracing::info; use super::{ get_healthy_worker_indices, utils::PeriodicTask, LoadBalancingPolicy, SelectWorkerInfo, + WorkerLeg, }; use crate::{ config::ManualAssignmentMode, observability::metrics::Metrics, @@ -202,7 +203,15 @@ impl ManualPolicy { } if let Some(routing_id) = extract_routing_key(info.headers) { - let (idx, branch) = self.select_by_routing_id(workers, routing_id, &healthy_indices); + // Single is the common leg; route on the bare key to skip the + // per-request allocation. PD legs namespace so prefill and decode + // stick independently. + let (idx, branch) = if info.leg == WorkerLeg::Single { + self.select_by_routing_id(workers, routing_id, &healthy_indices) + } else { + let namespaced = format!("{}{}", info.leg.routing_id_prefix(), routing_id); + self.select_by_routing_id(workers, &namespaced, &healthy_indices) + }; return (Some(idx), branch); } @@ -330,6 +339,34 @@ mod tests { headers } + #[test] + fn test_manual_leg_namespaces_sticky_entries() { + 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 prefill = SelectWorkerInfo { + headers: Some(&headers), + leg: WorkerLeg::Prefill, + ..Default::default() + }; + let decode = SelectWorkerInfo { + headers: Some(&headers), + leg: WorkerLeg::Decode, + ..Default::default() + }; + + let (p1, _) = policy.select_worker_impl(&workers, &prefill); + let (d1, _) = policy.select_worker_impl(&workers, &decode); + + // Same key under two legs -> two independent sticky entries. + assert_eq!(policy.routing_map.len(), 2); + for _ in 0..5 { + assert_eq!(policy.select_worker_impl(&workers, &prefill).0, p1); + assert_eq!(policy.select_worker_impl(&workers, &decode).0, d1); + } + } + #[test] fn test_manual_consistent_routing() { let policy = ManualPolicy::new(); diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index 8ad0776c08..d1f37c9b09 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -187,6 +187,28 @@ pub(crate) fn normalize_model_key(model_id: &str) -> &str { } } +/// Which PD leg a selection is for. `Single` is non-PD (the default) and keeps +/// routing-key stickiness byte-identical to pre-leg behavior. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum WorkerLeg { + #[default] + Single, + Prefill, + Decode, +} + +impl WorkerLeg { + /// Prefix used to namespace sticky routing IDs per leg. `Single` is empty so + /// non-PD entries are unchanged. + pub fn routing_id_prefix(self) -> &'static str { + match self { + WorkerLeg::Single => "", + WorkerLeg::Prefill => "prefill:", + WorkerLeg::Decode => "decode:", + } + } +} + /// Information passed to policy for worker selection #[derive(Debug, Clone, Default)] pub struct SelectWorkerInfo<'a> { @@ -203,6 +225,9 @@ pub struct SelectWorkerInfo<'a> { /// Pre-computed hash ring for O(log n) consistent hashing /// Built and cached by WorkerRegistry, passed through to avoid per-request rebuilds pub hash_ring: Option>, + /// Which PD leg this selection is for (default `Single`); namespaces + /// header-based sticky routing so prefill and decode stick independently. + pub leg: WorkerLeg, } #[cfg(test)] @@ -291,4 +316,13 @@ mod tests { ); } } + + #[test] + fn test_select_worker_info_leg_defaults_to_single() { + let info = SelectWorkerInfo::default(); + assert_eq!(info.leg, WorkerLeg::Single); + assert_eq!(WorkerLeg::Single.routing_id_prefix(), ""); + assert_eq!(WorkerLeg::Prefill.routing_id_prefix(), "prefill:"); + assert_eq!(WorkerLeg::Decode.routing_id_prefix(), "decode:"); + } } diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 6cf2224b07..872c58e8a0 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -13,10 +13,14 @@ use tracing::{debug, info, warn}; /// When the first worker of a new model is added, it determines the policy for that model. /// All subsequent workers of the same model use the established policy. /// When the last worker of a model is removed, the policy mapping is cleaned up. -use super::{BucketPolicy, CacheAwarePolicy, DPRankLoadPolicy, LoadBalancingPolicy, PolicyFactory}; +use super::{ + BucketPolicy, CacheAwarePolicy, DPRankLoadPolicy, LoadBalancingPolicy, ManualConfig, + ManualPolicy, PolicyFactory, SelectWorkerInfo, +}; use crate::{ - config::types::PolicyConfig, + config::types::{PolicyConfig, RoutingKeyOverrideConfig}, policies::cache_aware::LoadReceiver, + routers::common::header_utils::extract_routing_key, worker::{KvEventMonitor, Worker}, }; @@ -49,12 +53,33 @@ pub struct PolicyRegistry { // DP-rank policy: Supports the selection of dp-rank outside the engine. dp_rank_policy: Arc>>, + + /// Shared sticky selector for the `X-SMG-Routing-Key` override. `Some` when the + /// override is enabled; consulted (instead of the configured policy) for keyed + /// requests via [`PolicyRegistry::select_worker`]. + routing_key_sticky: Option>, } impl PolicyRegistry { - /// Create a new PolicyRegistry with a default policy + /// Create a new PolicyRegistry with a default policy (no routing-key override). pub fn new(default_policy_config: PolicyConfig) -> Self { + Self::with_override(default_policy_config, RoutingKeyOverrideConfig::default()) + } + + /// Create a PolicyRegistry. When `routing_key_override.enabled`, builds a shared + /// sticky selector consulted for keyed requests in [`Self::select_worker`]. + pub fn with_override( + default_policy_config: PolicyConfig, + routing_key_override: RoutingKeyOverrideConfig, + ) -> Self { let default_policy = Self::create_policy_from_config(&default_policy_config); + let routing_key_sticky = routing_key_override.enabled.then(|| { + Arc::new(ManualPolicy::with_config(ManualConfig { + eviction_interval_secs: routing_key_override.eviction_interval_secs, + max_idle_secs: routing_key_override.max_idle_secs, + assignment_mode: routing_key_override.assignment_mode, + })) + }); Self { model_policies: Arc::new(DashMap::new()), @@ -65,7 +90,34 @@ impl PolicyRegistry { kv_event_monitor: Arc::new(RwLock::new(None)), load_rx: Arc::new(RwLock::new(None)), dp_rank_policy: Arc::new(OnceLock::new()), + routing_key_sticky, + } + } + + /// Select a worker, applying the `X-SMG-Routing-Key` sticky override when it is + /// enabled, the request carries the header, and the configured policy does not + /// already honor the key (`manual` / `consistent_hashing`). Otherwise delegates + /// to `policy`. `policy.name()` stays the real policy (for metrics). + pub fn select_worker( + &self, + policy: &Arc, + workers: &[Arc], + info: &SelectWorkerInfo, + ) -> Option { + if let Some(sticky) = self.routing_key_sticky.as_ref() { + if Self::routing_key_override_applies(policy.name()) + && extract_routing_key(info.headers).is_some() + { + return sticky.select_worker(workers, info); + } } + policy.select_worker(workers, info) + } + + /// Policies that already honor `X-SMG-Routing-Key` keep their own handling; all + /// others (cache_aware, least_load, prefix_hash, ...) get the sticky override. + fn routing_key_override_applies(name: &str) -> bool { + !matches!(name, "manual" | "consistent_hashing") } /// Set KV event monitor (thread-safe, can be called after initialization). @@ -549,6 +601,91 @@ mod tests { })) } + fn headers_with_key(key: &str) -> http::HeaderMap { + let mut h = http::HeaderMap::new(); + h.insert("x-smg-routing-key", key.parse().unwrap()); + h + } + + #[test] + fn override_eligibility_skips_key_native_policies() { + // Policies that already honor X-SMG-Routing-Key are skipped; others (incl. + // prefix_hash, which routes by tokens) get the sticky override. + assert!(PolicyRegistry::routing_key_override_applies("cache_aware")); + assert!(PolicyRegistry::routing_key_override_applies("prefix_hash")); + assert!(PolicyRegistry::routing_key_override_applies("least_load")); + assert!(!PolicyRegistry::routing_key_override_applies("manual")); + assert!(!PolicyRegistry::routing_key_override_applies( + "consistent_hashing" + )); + } + + #[test] + fn override_routes_keyed_request_stickily() { + let reg = PolicyRegistry::with_override( + PolicyConfig::RoundRobin, + RoutingKeyOverrideConfig { + enabled: true, + ..Default::default() + }, + ); + let policy = reg.get_default_policy(); + let workers = vec![ + worker("http://w1", WorkerType::Regular), + worker("http://w2", WorkerType::Regular), + worker("http://w3", WorkerType::Regular), + ]; + let headers = headers_with_key("session-A"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + let first = reg.select_worker(&policy, &workers, &info).unwrap(); + for _ in 0..5 { + assert_eq!(reg.select_worker(&policy, &workers, &info), Some(first)); + } + } + + #[test] + fn override_without_key_uses_configured_policy() { + let reg = PolicyRegistry::with_override( + PolicyConfig::RoundRobin, + RoutingKeyOverrideConfig { + enabled: true, + ..Default::default() + }, + ); + let policy = reg.get_default_policy(); + let workers = vec![ + worker("http://w1", WorkerType::Regular), + worker("http://w2", WorkerType::Regular), + ]; + let info = SelectWorkerInfo::default(); // no key header + // RoundRobin alternates -> proves the configured policy is used, not sticky. + let a = reg.select_worker(&policy, &workers, &info).unwrap(); + let b = reg.select_worker(&policy, &workers, &info).unwrap(); + assert_ne!(a, b); + } + + #[test] + fn override_disabled_ignores_key() { + let reg = PolicyRegistry::new(PolicyConfig::RoundRobin); // override off + let policy = reg.get_default_policy(); + let workers = vec![ + worker("http://w1", WorkerType::Regular), + worker("http://w2", WorkerType::Regular), + ]; + let headers = headers_with_key("session-A"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + // Override off -> the key is ignored, RoundRobin alternates. + let a = reg.select_worker(&policy, &workers, &info).unwrap(); + let b = reg.select_worker(&policy, &workers, &info).unwrap(); + assert_ne!(a, b); + } + #[test] fn test_policy_registry_basic() { let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); diff --git a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs index a1b270de2e..0abca6f345 100644 --- a/model_gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/model_gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -153,14 +153,17 @@ impl WorkerSelectionStage { // Get cached hash ring for consistent hashing (O(log n) lookup) let hash_ring = self.worker_registry.get_hash_ring(model_id); - // Select worker using the policy - let idx = policy.select_worker( + // Select worker via the registry (applies the routing-key sticky override + // when enabled; otherwise delegates to the configured policy). + let idx = self.policy_registry.select_worker( + &policy, &available, &SelectWorkerInfo { request_text: text, tokens, headers, hash_ring, + leg: crate::policies::WorkerLeg::Single, }, )?; let selected = available[idx].clone(); @@ -267,14 +270,22 @@ impl WorkerSelectionStage { // Get cached hash ring for consistent hashing (O(log n) lookup) let hash_ring = self.worker_registry.get_hash_ring(model_id); - let info = SelectWorkerInfo { + // Prefill and decode are separate pools; tag each leg so the routing-key + // override keys its sticky map per leg (a key sticks independently). + let mut info = SelectWorkerInfo { request_text: text, tokens, headers, hash_ring, + leg: crate::policies::WorkerLeg::Prefill, }; - let prefill_idx = policy.select_worker(&available_prefill, &info)?; - let decode_idx = policy.select_worker(&available_decode, &info)?; + let prefill_idx = self + .policy_registry + .select_worker(&policy, &available_prefill, &info)?; + info.leg = crate::policies::WorkerLeg::Decode; + let decode_idx = self + .policy_registry + .select_worker(&policy, &available_decode, &info)?; let model = model_id; let policy_name = policy.name(); diff --git a/model_gateway/src/routers/http/pd_router.rs b/model_gateway/src/routers/http/pd_router.rs index ea9e72ce91..3628c839fd 100644 --- a/model_gateway/src/routers/http/pd_router.rs +++ b/model_gateway/src/routers/http/pd_router.rs @@ -843,22 +843,24 @@ impl PDRouter { // Get cached hash ring for consistent hashing let hash_ring = self.worker_registry.get_hash_ring(model_id); - let prefill = Self::pick_worker_by_policy_arc( + let prefill = self.pick_worker_by_policy_arc( &prefill_workers, - &*prefill_policy, + &prefill_policy, request_text, headers, hash_ring.clone(), "prefill", + crate::policies::WorkerLeg::Prefill, )?; - let decode = Self::pick_worker_by_policy_arc( + let decode = self.pick_worker_by_policy_arc( &decode_workers, - &*decode_policy, + &decode_policy, request_text, headers, hash_ring, "decode", + crate::policies::WorkerLeg::Decode, )?; // Record worker selection metrics (Layer 3) @@ -879,13 +881,19 @@ impl PDRouter { Ok((prefill, decode)) } + #[expect( + clippy::too_many_arguments, + reason = "HTTP PD worker pick threads policy + request context + leg" + )] fn pick_worker_by_policy_arc( + &self, workers: &[Arc], - policy: &dyn LoadBalancingPolicy, + policy: &Arc, request_text: Option<&str>, headers: Option<&HeaderMap>, hash_ring: Option>, worker_type: &str, + leg: crate::policies::WorkerLeg, ) -> Result, String> { if workers.is_empty() { return Err(format!( @@ -905,14 +913,17 @@ impl PDRouter { )); } - let selected_idx = policy + let selected_idx = self + .policy_registry .select_worker( + policy, &available_workers, &SelectWorkerInfo { request_text, tokens: None, // HTTP doesn't have tokens, use gRPC for PrefixHash headers, hash_ring, + leg, }, ) .ok_or_else(|| { diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 0799467f1e..15718138e9 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -173,13 +173,15 @@ impl Router { // Get cached hash ring for consistent hashing (O(log n) lookup) let hash_ring = self.worker_registry.get_hash_ring(model_id); - let idx = policy.select_worker( + let idx = self.policy_registry.select_worker( + &policy, &available, &SelectWorkerInfo { request_text: text, tokens: None, // HTTP doesn't have tokens, use gRPC for PrefixHash headers, hash_ring, + leg: crate::policies::WorkerLeg::Single, }, )?; @@ -566,13 +568,15 @@ impl Router { let policy = self.policy_registry.get_policy_or_default(model_id); let hash_ring = self.worker_registry.get_hash_ring(model_id); - let idx = match policy.select_worker( + let idx = match self.policy_registry.select_worker( + &policy, &available, &SelectWorkerInfo { request_text: Some(&text), tokens: None, headers, hash_ring, + leg: crate::policies::WorkerLeg::Single, }, ) { Some(i) => i, diff --git a/model_gateway/src/service_discovery.rs b/model_gateway/src/service_discovery.rs index f92274d8a7..d9ab9951fc 100644 --- a/model_gateway/src/service_discovery.rs +++ b/model_gateway/src/service_discovery.rs @@ -1286,8 +1286,9 @@ mod tests { router_config: router_config.clone(), rate_limiter: Some(Arc::new(TokenBucket::new(1000, 1000))), worker_registry: worker_registry.clone(), - policy_registry: Arc::new(crate::policies::PolicyRegistry::new( + policy_registry: Arc::new(crate::policies::PolicyRegistry::with_override( router_config.policy.clone(), + router_config.routing_key_override.clone(), )), reasoning_parser_factory: None, tool_parser_factory: None,