diff --git a/bindings/golang/src/policy.rs b/bindings/golang/src/policy.rs index b26d31989..0930b555e 100644 --- a/bindings/golang/src/policy.rs +++ b/bindings/golang/src/policy.rs @@ -24,8 +24,8 @@ use openai_protocol::{ }; use smg::{ policies::{ - BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy, - RoundRobinPolicy, SelectWorkerInfo, + BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, + RandomPolicy, RoundRobinPolicy, SelectWorkerInfo, }, routers::grpc::{backend_client::BackendClient, utils::process_chat_messages}, worker::{ diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 8407dd1c8..97571b9e6 100755 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -20,6 +20,7 @@ pub enum PolicyType { RoundRobin, Passthrough, CacheAware, + CacheAwareLength, PowerOfTwo, LeastLoad, Bucket, @@ -526,6 +527,11 @@ struct Router { worker_overload_token_usage: Option, worker_overload_protection: bool, disable_load_monitoring: bool, + chars_per_token: usize, + long_prefill_threshold: usize, + long_pool_max_load: usize, + short_pool_max_load: usize, + long_prefill_indices: Vec, } impl Router { @@ -629,6 +635,25 @@ impl Router { cache_ttl_secs: self.cache_ttl_secs, cache_boundaries: self.cache_boundaries.clone(), }, + PolicyType::CacheAwareLength => ConfigPolicyConfig::CacheAwareLength { + cache_threshold: self.cache_threshold, + balance_abs_threshold: self.balance_abs_threshold, + balance_rel_threshold: self.balance_rel_threshold, + eviction_interval_secs: self.eviction_interval_secs, + max_tree_size: self.max_tree_size, + block_size: self.block_size, + balance_token_usage_threshold: self.balance_token_usage_threshold, + overload_token_usage_threshold: self.overload_token_usage_threshold, + overlap_decay: self.overlap_decay, + selection_temperature: self.selection_temperature, + cache_index: self.parse_cache_index()?, + cache_ttl_secs: self.cache_ttl_secs, + cache_boundaries: self.cache_boundaries.clone(), + chars_per_token: self.chars_per_token, + long_prefill_threshold: self.long_prefill_threshold, + long_pool_max_load: self.long_pool_max_load, + short_pool_max_load: self.short_pool_max_load, + }, PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo { load_check_interval_secs: self.load_monitor_interval, }, @@ -931,6 +956,7 @@ impl Router { self.server_key_path.as_ref(), ) .dp_minimum_tokens_scheduler(self.dp_minimum_tokens_scheduler) + .long_prefill_indices(self.long_prefill_indices.clone()) .build() } } @@ -1083,6 +1109,11 @@ impl Router { worker_overload_token_usage = None, worker_overload_protection = false, disable_load_monitoring = false, + chars_per_token = 4, + long_prefill_threshold = 100_000, + long_pool_max_load = 4, + short_pool_max_load = 32, + long_prefill_indices = vec![], ))] #[expect(clippy::too_many_arguments)] #[expect( @@ -1233,6 +1264,11 @@ impl Router { worker_overload_token_usage: Option, worker_overload_protection: bool, disable_load_monitoring: bool, + chars_per_token: usize, + long_prefill_threshold: usize, + long_pool_max_load: usize, + short_pool_max_load: usize, + long_prefill_indices: Vec, ) -> PyResult { let mut all_urls = worker_urls.clone(); @@ -1397,6 +1433,11 @@ impl Router { worker_overload_token_usage, worker_overload_protection, disable_load_monitoring, + chars_per_token, + long_prefill_threshold, + long_pool_max_load, + short_pool_max_load, + long_prefill_indices, }) } diff --git a/bindings/python/src/smg/router.py b/bindings/python/src/smg/router.py index d4d69ea5e..cfc60627e 100644 --- a/bindings/python/src/smg/router.py +++ b/bindings/python/src/smg/router.py @@ -25,6 +25,7 @@ def policy_from_str(policy_str: str | None) -> PolicyType: "round_robin": PolicyType.RoundRobin, "passthrough": PolicyType.Passthrough, "cache_aware": PolicyType.CacheAware, + "cache_aware_length": PolicyType.CacheAwareLength, "power_of_two": PolicyType.PowerOfTwo, "least_load": PolicyType.LeastLoad, "bucket": PolicyType.Bucket, diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index 94f9c4190..18d6ec28b 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -15,6 +15,7 @@ "round_robin", "passthrough", "cache_aware", + "cache_aware_length", "power_of_two", "least_load", "manual", @@ -242,6 +243,12 @@ class RouterArgs: worker_overload_protection: bool = False # Restore the conditional load-monitor poll gate (default: poll always) disable_load_monitoring: bool = False + # cache_aware_length policy: long/short pool split + chars_per_token: int = 4 + long_prefill_threshold: int = 100_000 + long_pool_max_load: int = 4 + short_pool_max_load: int = 32 + long_prefill_indices: list[int] = dataclasses.field(default_factory=list) @staticmethod def add_cli_args( @@ -627,6 +634,7 @@ def add_cli_args( ) routing_group.add_argument( f"--{prefix}eviction-interval-secs", + f"--{prefix}eviction-interval", type=int, default=RouterArgs.eviction_interval_secs, help="Interval in seconds between cache eviction operations", @@ -677,6 +685,42 @@ def add_cli_args( " approximate serving-engine cache retention. Defaults to 180." ), ) + # cache_aware_length policy parameters + routing_group.add_argument( + f"--{prefix}chars-per-token", + type=int, + default=RouterArgs.chars_per_token, + help="Divisor for char-level token estimation when X-Prompt-Tokens" + " is absent (cache_aware_length policy). Default 4.", + ) + routing_group.add_argument( + f"--{prefix}long-prefill-threshold", + type=int, + default=RouterArgs.long_prefill_threshold, + help="Uncached-prefill-token boundary between long and short" + " requests (cache_aware_length policy). Default 100000.", + ) + routing_group.add_argument( + f"--{prefix}long-pool-max-load", + type=int, + default=RouterArgs.long_pool_max_load, + help="Load ceiling for the long pool (pool=long workers)" + " (cache_aware_length policy). Default 4.", + ) + routing_group.add_argument( + f"--{prefix}short-pool-max-load", + type=int, + default=RouterArgs.short_pool_max_load, + help="Load ceiling for the short pool (remaining workers)" + " (cache_aware_length policy). Default 32.", + ) + routing_group.add_argument( + f"--{prefix}long-prefill-indices", + type=_parse_int_csv, + default=[], + help="Comma-separated 0-based indices of --prefill URLs that belong" + " to the long pool (get pool=long label for cache_aware_length).", + ) routing_group.add_argument( f"--{prefix}max-idle-secs", f"--{prefix}sticky-key-idle-secs", @@ -1681,6 +1725,25 @@ def from_cli_args(cls, args: argparse.Namespace, use_router_prefix: bool = False return cls(**args_dict) def _validate_router_args(self): + if (self.prefill_urls or self.decode_urls) and not ( + self.pd_disaggregation or self.epd_disaggregation + ): + raise ValueError( + "--prefill/--decode require --pd-disaggregation or --epd-disaggregation" + ) + + if len(set(self.long_prefill_indices)) != len(self.long_prefill_indices): + raise ValueError("--long-prefill-indices must not contain duplicate values") + + if self.long_prefill_indices: + if min(self.long_prefill_indices) < 0: + raise ValueError("--long-prefill-indices values must be non-negative") + if max(self.long_prefill_indices) >= len(self.prefill_urls): + raise ValueError( + "--long-prefill-indices value out of range for " + f"{len(self.prefill_urls)} configured prefill workers" + ) + # Validate configuration based on mode if self.epd_disaggregation: if self.encode_policy: diff --git a/bindings/python/tests/test_arg_parser.py b/bindings/python/tests/test_arg_parser.py index 5459ac617..784d28fdd 100644 --- a/bindings/python/tests/test_arg_parser.py +++ b/bindings/python/tests/test_arg_parser.py @@ -10,7 +10,7 @@ import pytest from smg.launch_router import RouterArgs, parse_router_args -from smg.router import policy_from_str +from smg.router import Router, policy_from_str class TestRouterArgs: @@ -32,6 +32,11 @@ def test_default_values(self): # Test PD-specific defaults assert args.prefill_policy is None assert args.decode_policy is None + assert args.chars_per_token == 4 + assert args.long_prefill_threshold == 100_000 + assert args.long_pool_max_load == 4 + assert args.short_pool_max_load == 32 + assert args.long_prefill_indices == [] # Test service discovery defaults assert args.service_discovery is False @@ -176,6 +181,135 @@ def test_parse_decode_urls_valid(self): result = RouterArgs._parse_decode_urls(None) assert result == [] + def test_parse_cache_aware_length_prefill_pool_command(self): + """The Python entrypoint accepts a static long/short prefill pool.""" + router_args = parse_router_args( + [ + "--pd-disaggregation", + "--prefill-policy", + "cache_aware_length", + "--long-prefill-threshold", + "100000", + "--balance-abs-threshold", + "16", + "--balance-rel-threshold", + "2.0", + "--eviction-interval", + "120", + "--long-pool-max-load", + "4", + "--short-pool-max-load", + "32", + "--cache-threshold", + "0.25", + "--prefill", + "http://p1:8000", + "--prefill", + "http://p2:8000", + "--prefill", + "http://p3:8000", + "--prefill", + "http://p4:8000", + "--prefill", + "http://p5:8000", + "--decode", + "http://d1:8000", + "--long-prefill-indices", + "3,4", + ] + ) + router_args._validate_router_args() + + assert router_args.pd_disaggregation is True + assert router_args.prefill_policy == "cache_aware_length" + assert router_args.long_prefill_threshold == 100_000 + assert router_args.balance_abs_threshold == 16 + assert router_args.balance_rel_threshold == 2.0 + assert router_args.eviction_interval_secs == 120 + assert router_args.long_pool_max_load == 4 + assert router_args.short_pool_max_load == 32 + assert router_args.cache_threshold == 0.25 + assert router_args.long_prefill_indices == [3, 4] + + def test_cache_aware_length_pool_arguments_reach_rust_router(self, monkeypatch): + router_args = parse_router_args( + [ + "--pd-disaggregation", + "--prefill-policy", + "cache_aware_length", + "--prefill", + "http://p1:8000", + "--prefill", + "http://p2:8000", + "--long-prefill-indices", + "1", + "--long-prefill-threshold", + "100000", + "--long-pool-max-load", + "4", + "--short-pool-max-load", + "32", + "--decode", + "http://d1:8000", + ] + ) + captured = {} + + def fake_rust_router(**kwargs): + captured.update(kwargs) + return object() + + monkeypatch.setattr("smg.router._Router", fake_rust_router) + Router.from_args(router_args) + + assert captured["prefill_policy"] == policy_from_str("cache_aware_length") + assert captured["long_prefill_indices"] == [1] + assert captured["long_prefill_threshold"] == 100_000 + assert captured["long_pool_max_load"] == 4 + assert captured["short_pool_max_load"] == 32 + + @pytest.mark.parametrize( + ("indices", "message"), + [ + ([-1], "non-negative"), + ([3, 3], "duplicate"), + ([5], "out of range"), + ], + ) + def test_long_prefill_indices_reject_invalid_values(self, indices, message): + router_args = RouterArgs( + pd_disaggregation=True, + prefill_urls=[(f"http://p{i}:8000", None) for i in range(1, 6)], + decode_urls=["http://d1:8000"], + prefill_policy="cache_aware_length", + long_prefill_indices=indices, + ) + + with pytest.raises(ValueError, match=message): + router_args._validate_router_args() + + def test_prefill_decode_urls_require_disaggregation_mode(self): + router_args = parse_router_args( + [ + "--prefill", + "http://p1:8000", + "--decode", + "http://d1:8000", + ] + ) + + with pytest.raises(ValueError, match="--pd-disaggregation"): + router_args._validate_router_args() + + def test_prefixed_eviction_interval_alias_maps_to_router_args(self): + parser = argparse.ArgumentParser() + RouterArgs.add_cli_args(parser, use_router_prefix=True) + + namespace = parser.parse_args(["--router-eviction-interval", "120"]) + router_args = RouterArgs.from_cli_args(namespace, use_router_prefix=True) + + assert router_args.eviction_interval_secs == 120 + def test_from_cli_args_basic(self): """Test creating RouterArgs from basic CLI arguments.""" args = SimpleNamespace( @@ -1401,6 +1535,11 @@ class TestRouterArgsFieldOrder: "worker_overload_token_usage", "worker_overload_protection", "disable_load_monitoring", + "chars_per_token", + "long_prefill_threshold", + "long_pool_max_load", + "short_pool_max_load", + "long_prefill_indices", ] def test_complete_field_sequence_is_frozen(self): @@ -1433,8 +1572,37 @@ def test_new_fields_appended_after_positional_reserve(self): "worker_overload_token_usage", "worker_overload_protection", "disable_load_monitoring", + "chars_per_token", + "long_prefill_threshold", + "long_pool_max_load", + "short_pool_max_load", + "long_prefill_indices", ): assert names.index(appended) > marker, ( f"{appended} must be appended after worker_startup_delay to " "preserve positional callers" ) + + def test_cache_aware_length_cli_args(self): + """Test --chars-per-token, --long-prefill-* CLI args are parsed.""" + parser = argparse.ArgumentParser() + RouterArgs.add_cli_args(parser) + args = parser.parse_args( + [ + "--chars-per-token", + "8", + "--long-prefill-threshold", + "200000", + "--long-pool-max-load", + "10", + "--short-pool-max-load", + "64", + "--long-prefill-indices", + "0,2", + ] + ) + assert args.chars_per_token == 8 + assert args.long_prefill_threshold == 200000 + assert args.long_pool_max_load == 10 + assert args.short_pool_max_load == 64 + assert args.long_prefill_indices == [0, 2] diff --git a/bindings/python/tests/test_router_config.py b/bindings/python/tests/test_router_config.py index 554a8f3c2..d35b3978b 100644 --- a/bindings/python/tests/test_router_config.py +++ b/bindings/python/tests/test_router_config.py @@ -433,3 +433,32 @@ def test_config_with_empty_dicts(self): assert args.prefill_selector == {} assert args.decode_selector == {} assert args.storage_context_headers == {} + + def test_cache_aware_length_params_defaults(self): + """Test cache_aware_length parameters have correct defaults.""" + args = RouterArgs() + assert args.chars_per_token == 4 + assert args.long_prefill_threshold == 100_000 + assert args.long_pool_max_load == 4 + assert args.short_pool_max_load == 32 + assert args.long_prefill_indices == [] + + def test_cache_aware_length_params_custom(self): + """Test cache_aware_length parameters accept custom values.""" + args = RouterArgs( + chars_per_token=8, + long_prefill_threshold=200_000, + long_pool_max_load=10, + short_pool_max_load=64, + long_prefill_indices=[0, 2], + ) + assert args.chars_per_token == 8 + assert args.long_prefill_threshold == 200_000 + assert args.long_pool_max_load == 10 + assert args.short_pool_max_load == 64 + assert args.long_prefill_indices == [0, 2] + + def test_policy_from_str_cache_aware_length(self): + """Test policy_from_str maps cache_aware_length correctly.""" + result = policy_from_str("cache_aware_length") + assert result == PolicyType.CacheAwareLength diff --git a/model_gateway/src/app_context.rs b/model_gateway/src/app_context.rs index 0e991d94c..2b3c0ad97 100644 --- a/model_gateway/src/app_context.rs +++ b/model_gateway/src/app_context.rs @@ -712,10 +712,16 @@ impl AppContextBuilder { fn with_kv_event_monitor(mut self, config: &RouterConfig) -> Self { use crate::config::types::{PolicyConfig, RoutingMode}; - let role_is_cache_aware = - |policy: &Option| matches!(policy, Some(PolicyConfig::CacheAware { .. })); - let is_cache_aware = matches!(config.policy, PolicyConfig::CacheAware { .. }) - || match &config.mode { + let role_is_cache_aware = |policy: &Option| { + matches!( + policy, + Some(PolicyConfig::CacheAware { .. } | PolicyConfig::CacheAwareLength { .. }) + ) + }; + let is_cache_aware = matches!( + config.policy, + PolicyConfig::CacheAware { .. } | PolicyConfig::CacheAwareLength { .. } + ) || match &config.mode { RoutingMode::PrefillDecode { prefill_policy, decode_policy, diff --git a/model_gateway/src/config/builder.rs b/model_gateway/src/config/builder.rs index b7321eede..eb2cda340 100644 --- a/model_gateway/src/config/builder.rs +++ b/model_gateway/src/config/builder.rs @@ -112,6 +112,11 @@ impl RouterConfigBuilder { self } + pub fn long_prefill_indices(mut self, indices: Vec) -> Self { + self.config.long_prefill_indices = indices; + self + } + pub fn random_policy(mut self) -> Self { self.config.policy = PolicyConfig::Random; self diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index fc17c72b2..3b96d3e6e 100755 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -20,6 +20,10 @@ pub struct RouterConfig { pub mode: RoutingMode, #[serde(default)] pub connection_mode: ConnectionMode, + /// 0-based indices of `--prefill` URLs that belong to the long pool + /// (get `pool=long` label for cache_aware_length). Empty = no long pool. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub long_prefill_indices: Vec, /// Explicit runtime for the startup workers (`--worker-urls`), set from /// `--backend` when the connection mode is ZMQ. The ZMQ handshake is shared /// across engine runtimes, so the wire protocol cannot be probed and must @@ -644,6 +648,55 @@ pub enum PolicyConfig { cache_boundaries: Vec, }, + /// Cache-aware length policy: a full superset of `cache_aware` that adds + /// a long/short pool split on the no-cache branch, driven by the `pool` + /// worker label (`pool=long` → long pool, otherwise short pool). Inherits + /// all cache_aware features (string tree, token tree, event-driven + /// routing, hash index, mesh sync, KV pressure). See + /// `policies/cache_aware_length.rs`. + #[serde(rename = "cache_aware_length")] + CacheAwareLength { + // --- Inherited from cache_aware --- + #[serde(alias = "cache_match_threshold")] + #[serde(default = "default_cal_cache_threshold")] + cache_threshold: f32, + #[serde(alias = "spill_abs_threshold")] + #[serde(default = "default_cal_balance_abs_threshold")] + balance_abs_threshold: usize, + #[serde(alias = "spill_rel_threshold")] + #[serde(default = "default_cal_balance_rel_threshold")] + balance_rel_threshold: f32, + #[serde(default = "default_cal_eviction_interval_secs")] + eviction_interval_secs: u64, + #[serde(default = "default_cal_max_tree_size")] + max_tree_size: usize, + #[serde(default = "default_block_size")] + block_size: usize, + #[serde(default = "default_balance_token_usage_threshold")] + balance_token_usage_threshold: f32, + #[serde(default = "default_balance_token_usage_threshold")] + overload_token_usage_threshold: f32, + #[serde(default = "default_overlap_decay")] + overlap_decay: f32, + #[serde(default = "default_selection_temperature")] + selection_temperature: f32, + #[serde(default)] + cache_index: CacheIndexKind, + #[serde(default = "default_cache_ttl_secs")] + cache_ttl_secs: u64, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + cache_boundaries: Vec, + // --- Length-specific --- + #[serde(default = "default_cal_chars_per_token")] + chars_per_token: usize, + #[serde(default = "default_cal_long_prefill_threshold")] + long_prefill_threshold: usize, + #[serde(default = "default_cal_long_pool_max_load")] + long_pool_max_load: usize, + #[serde(default = "default_cal_short_pool_max_load")] + short_pool_max_load: usize, + }, + /// Power-of-two choices policy: samples two workers and routes to the one /// with the lower expected wait, scored like `least_load` /// (`(queued_tokens + inflight_tokens) / throughput + kv_pressure_weight * k/(1-k)`). @@ -777,6 +830,35 @@ fn default_cache_ttl_secs() -> u64 { 180 } +// cache_aware_length defaults (kept aligned with CacheAwareLengthConfig::default). +fn default_cal_cache_threshold() -> f32 { + 0.3 +} +fn default_cal_balance_abs_threshold() -> usize { + 32 +} +fn default_cal_balance_rel_threshold() -> f32 { + 1.1 +} +fn default_cal_eviction_interval_secs() -> u64 { + 30 +} +fn default_cal_max_tree_size() -> usize { + 10000 +} +fn default_cal_chars_per_token() -> usize { + 4 +} +fn default_cal_long_prefill_threshold() -> usize { + 100_000 +} +fn default_cal_long_pool_max_load() -> usize { + 4 +} +fn default_cal_short_pool_max_load() -> usize { + 32 +} + fn default_prefix_token_count() -> usize { 256 } @@ -828,6 +910,7 @@ impl PolicyConfig { PolicyConfig::RoundRobin => "round_robin", PolicyConfig::Passthrough => "passthrough", PolicyConfig::CacheAware { .. } => "cache_aware", + PolicyConfig::CacheAwareLength { .. } => "cache_aware_length", PolicyConfig::PowerOfTwo { .. } => "power_of_two", PolicyConfig::LeastLoad { .. } => "least_load", PolicyConfig::Bucket { .. } => "bucket", @@ -1040,6 +1123,7 @@ impl Default for RouterConfig { }, policy: PolicyConfig::Random, cache_boundaries: Vec::new(), + long_prefill_indices: Vec::new(), routing_key_override: RoutingKeyOverrideConfig::default(), host: "0.0.0.0".to_string(), port: 3001, diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index ad70151f6..907949109 100644 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -93,6 +93,7 @@ impl ConfigValidator { Self::validate_mode(&config.mode)?; Self::validate_policy(&config.policy)?; Self::validate_cache_boundaries(&config.cache_boundaries)?; + Self::validate_long_prefill_indices(config)?; Self::validate_server_settings(config)?; Self::validate_storage_context_headers(config)?; Self::validate_routing_key_headers(config)?; @@ -453,6 +454,40 @@ impl ConfigValidator { Ok(()) } + fn validate_long_prefill_indices(config: &RouterConfig) -> ConfigResult<()> { + let indices = &config.long_prefill_indices; + if indices.is_empty() { + return Ok(()); + } + let mut seen = std::collections::HashSet::new(); + for &i in indices { + if !seen.insert(i) { + return Err(ConfigError::InvalidValue { + field: "long_prefill_indices".to_string(), + value: i.to_string(), + reason: "must not contain duplicate values".to_string(), + }); + } + } + let prefill_count = match &config.mode { + RoutingMode::PrefillDecode { prefill_urls, .. } + | RoutingMode::EncodePrefillDecode { prefill_urls, .. } => prefill_urls.len(), + _ => 0, + }; + if let Some(&max) = indices.iter().max() { + if max >= prefill_count { + return Err(ConfigError::InvalidValue { + field: "long_prefill_indices".to_string(), + value: max.to_string(), + reason: format!( + "out of range for {prefill_count} configured prefill workers" + ), + }); + } + } + Ok(()) + } + fn validate_policy(policy: &PolicyConfig) -> ConfigResult<()> { match policy { PolicyConfig::Random @@ -475,92 +510,81 @@ impl ConfigValidator { cache_ttl_secs, cache_boundaries, } => { - Self::validate_cache_boundaries(cache_boundaries)?; + Self::validate_cache_aware_shared( + cache_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + block_size, + balance_token_usage_threshold, + overload_token_usage_threshold, + overlap_decay, + selection_temperature, + cache_index, + cache_ttl_secs, + cache_boundaries, + )?; + } + PolicyConfig::CacheAwareLength { + cache_threshold, + balance_abs_threshold: _, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + block_size, + balance_token_usage_threshold, + overload_token_usage_threshold, + overlap_decay, + selection_temperature, + cache_index, + cache_ttl_secs, + cache_boundaries, + chars_per_token, + long_prefill_threshold, + long_pool_max_load, + short_pool_max_load, + } => { + Self::validate_cache_aware_shared( + cache_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + block_size, + balance_token_usage_threshold, + overload_token_usage_threshold, + overlap_decay, + selection_temperature, + cache_index, + cache_ttl_secs, + cache_boundaries, + )?; - if *cache_ttl_secs == 0 { + // ---- Length-specific checks ---- + if *chars_per_token == 0 { return Err(ConfigError::InvalidValue { - field: "cache_ttl_secs".to_string(), - value: cache_ttl_secs.to_string(), + field: "chars_per_token".to_string(), + value: chars_per_token.to_string(), reason: "Must be > 0".to_string(), }); } - - if *cache_index == CacheIndexKind::Hash && cache_boundaries.is_empty() { - return Err(ConfigError::InvalidValue { - field: "cache_index".to_string(), - value: "hash".to_string(), - reason: "cache_index=hash requires non-empty cache_boundaries".to_string(), - }); - } - - if !overlap_decay.is_finite() || *overlap_decay < 0.0 { - return Err(ConfigError::InvalidValue { - field: "overlap_decay".to_string(), - value: overlap_decay.to_string(), - reason: "Must be finite and >= 0.0 (0.0 disables)".to_string(), - }); - } - - if !selection_temperature.is_finite() || *selection_temperature < 0.0 { - return Err(ConfigError::InvalidValue { - field: "selection_temperature".to_string(), - value: selection_temperature.to_string(), - reason: "Must be finite and >= 0.0 (0.0 is argmax)".to_string(), - }); - } - - if *block_size == 0 { + if *long_prefill_threshold == 0 { return Err(ConfigError::InvalidValue { - field: "block_size".to_string(), - value: block_size.to_string(), + field: "long_prefill_threshold".to_string(), + value: long_prefill_threshold.to_string(), reason: "Must be > 0".to_string(), }); } - - if *balance_token_usage_threshold <= 0.0 { + if *long_pool_max_load == 0 { return Err(ConfigError::InvalidValue { - field: "balance_token_usage_threshold".to_string(), - value: balance_token_usage_threshold.to_string(), - reason: "Must be > 0.0 (use >= 1.0 to disable)".to_string(), - }); - } - - if *overload_token_usage_threshold <= 0.0 { - return Err(ConfigError::InvalidValue { - field: "overload_token_usage_threshold".to_string(), - value: overload_token_usage_threshold.to_string(), - reason: "Must be > 0.0 (use >= 1.0 to disable)".to_string(), - }); - } - - if !(0.0..=1.0).contains(cache_threshold) { - return Err(ConfigError::InvalidValue { - field: "cache_threshold".to_string(), - value: cache_threshold.to_string(), - reason: "Must be between 0.0 and 1.0".to_string(), - }); - } - - if *balance_rel_threshold < 1.0 { - return Err(ConfigError::InvalidValue { - field: "balance_rel_threshold".to_string(), - value: balance_rel_threshold.to_string(), - reason: "Must be >= 1.0".to_string(), - }); - } - - if *eviction_interval_secs == 0 { - return Err(ConfigError::InvalidValue { - field: "eviction_interval_secs".to_string(), - value: eviction_interval_secs.to_string(), + field: "long_pool_max_load".to_string(), + value: long_pool_max_load.to_string(), reason: "Must be > 0".to_string(), }); } - - if *max_tree_size == 0 { + if *short_pool_max_load == 0 { return Err(ConfigError::InvalidValue { - field: "max_tree_size".to_string(), - value: max_tree_size.to_string(), + field: "short_pool_max_load".to_string(), + value: short_pool_max_load.to_string(), reason: "Must be > 0".to_string(), }); } @@ -671,6 +695,120 @@ impl ConfigValidator { Ok(()) } + /// Shared validation for the cache-aware fields common to both + /// `CacheAware` and `CacheAwareLength` policy variants. + #[expect(clippy::too_many_arguments, reason = "mirrors the PolicyConfig fields")] + fn validate_cache_aware_shared( + cache_threshold: &f32, + balance_rel_threshold: &f32, + eviction_interval_secs: &u64, + max_tree_size: &usize, + block_size: &usize, + balance_token_usage_threshold: &f32, + overload_token_usage_threshold: &f32, + overlap_decay: &f32, + selection_temperature: &f32, + cache_index: &CacheIndexKind, + cache_ttl_secs: &u64, + cache_boundaries: &[usize], + ) -> ConfigResult<()> { + Self::validate_cache_boundaries(cache_boundaries)?; + + if *cache_ttl_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "cache_ttl_secs".to_string(), + value: cache_ttl_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *cache_index == CacheIndexKind::Hash && cache_boundaries.is_empty() { + return Err(ConfigError::InvalidValue { + field: "cache_index".to_string(), + value: "hash".to_string(), + reason: "cache_index=hash requires non-empty cache_boundaries".to_string(), + }); + } + + if !overlap_decay.is_finite() || *overlap_decay < 0.0 { + return Err(ConfigError::InvalidValue { + field: "overlap_decay".to_string(), + value: overlap_decay.to_string(), + reason: "Must be finite and >= 0.0 (0.0 disables)".to_string(), + }); + } + + if !selection_temperature.is_finite() || *selection_temperature < 0.0 { + return Err(ConfigError::InvalidValue { + field: "selection_temperature".to_string(), + value: selection_temperature.to_string(), + reason: "Must be finite and >= 0.0 (0.0 is argmax)".to_string(), + }); + } + + if *block_size == 0 { + return Err(ConfigError::InvalidValue { + field: "block_size".to_string(), + value: block_size.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if !balance_token_usage_threshold.is_finite() + || *balance_token_usage_threshold <= 0.0 + { + return Err(ConfigError::InvalidValue { + field: "balance_token_usage_threshold".to_string(), + value: balance_token_usage_threshold.to_string(), + reason: "Must be finite and > 0.0 (use >= 1.0 to disable)".to_string(), + }); + } + + if !overload_token_usage_threshold.is_finite() + || *overload_token_usage_threshold <= 0.0 + { + return Err(ConfigError::InvalidValue { + field: "overload_token_usage_threshold".to_string(), + value: overload_token_usage_threshold.to_string(), + reason: "Must be finite and > 0.0 (use >= 1.0 to disable)".to_string(), + }); + } + + if !(0.0..=1.0).contains(cache_threshold) { + return Err(ConfigError::InvalidValue { + field: "cache_threshold".to_string(), + value: cache_threshold.to_string(), + reason: "Must be between 0.0 and 1.0".to_string(), + }); + } + + if !balance_rel_threshold.is_finite() || *balance_rel_threshold < 1.0 { + return Err(ConfigError::InvalidValue { + field: "balance_rel_threshold".to_string(), + value: balance_rel_threshold.to_string(), + reason: "Must be finite and >= 1.0".to_string(), + }); + } + + if *eviction_interval_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "eviction_interval_secs".to_string(), + value: eviction_interval_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *max_tree_size == 0 { + return Err(ConfigError::InvalidValue { + field: "max_tree_size".to_string(), + value: max_tree_size.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + Ok(()) + } + fn validate_cache_boundaries(boundaries: &[usize]) -> ConfigResult<()> { if boundaries.first() == Some(&0) { return Err(ConfigError::InvalidValue { @@ -1624,6 +1762,94 @@ mod tests { assert!(ConfigValidator::validate(&make(0.0, f32::INFINITY)).is_err()); } + #[test] + fn test_validate_cache_aware_length_rejects_nan_thresholds() { + let make = |balance_token_usage_threshold: f32, + overload_token_usage_threshold: f32| { + RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec![ + "http://worker1:8000".to_string(), + "http://worker2:8000".to_string(), + ], + }, + PolicyConfig::CacheAwareLength { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 60, + max_tree_size: 1000, + block_size: 16, + balance_token_usage_threshold, + overload_token_usage_threshold, + overlap_decay: 0.0, + selection_temperature: 0.0, + cache_index: Default::default(), + cache_ttl_secs: 180, + cache_boundaries: Vec::new(), + chars_per_token: 4, + long_prefill_threshold: 100_000, + long_pool_max_load: 4, + short_pool_max_load: 32, + }, + ) + }; + + // Valid defaults pass. + assert!(ConfigValidator::validate(&make(1.0, 1.0)).is_ok()); + // NaN rejected for both fields. + assert!(ConfigValidator::validate(&make(f32::NAN, 1.0)).is_err()); + assert!(ConfigValidator::validate(&make(1.0, f32::NAN)).is_err()); + // Infinity rejected for both fields. + assert!(ConfigValidator::validate(&make(f32::INFINITY, 1.0)).is_err()); + assert!(ConfigValidator::validate(&make(1.0, f32::INFINITY)).is_err()); + // Zero rejected. + assert!(ConfigValidator::validate(&make(0.0, 1.0)).is_err()); + assert!(ConfigValidator::validate(&make(1.0, 0.0)).is_err()); + } + + #[test] + fn test_validate_cache_aware_length_rejects_non_finite_balance_rel_threshold() { + let make = |balance_rel_threshold: f32| { + RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec![ + "http://worker1:8000".to_string(), + "http://worker2:8000".to_string(), + ], + }, + PolicyConfig::CacheAwareLength { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold, + eviction_interval_secs: 60, + max_tree_size: 1000, + block_size: 16, + balance_token_usage_threshold: 1.0, + overload_token_usage_threshold: 1.0, + overlap_decay: 0.0, + selection_temperature: 0.0, + cache_index: Default::default(), + cache_ttl_secs: 180, + cache_boundaries: Vec::new(), + chars_per_token: 4, + long_prefill_threshold: 100_000, + long_pool_max_load: 4, + short_pool_max_load: 32, + }, + ) + }; + + // Valid value passes. + assert!(ConfigValidator::validate(&make(1.1)).is_ok()); + // NaN passes the < 1.0 check (NaN < 1.0 is false) but must be rejected. + assert!(ConfigValidator::validate(&make(f32::NAN)).is_err()); + // Infinity must be rejected. + assert!(ConfigValidator::validate(&make(f32::INFINITY)).is_err()); + // Below 1.0 rejected. + assert!(ConfigValidator::validate(&make(0.9)).is_err()); + } + #[test] fn test_validate_cache_index_fields() { let make = |cache_index: CacheIndexKind, cache_ttl_secs: u64, boundaries: Vec| { diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 29aea9f88..374ea56d0 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -230,7 +230,7 @@ struct CliArgs { // ==================== Routing Policy ==================== /// Load balancing policy to use - #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "passthrough", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "Routing Policy")] + #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "passthrough", "cache_aware", "cache_aware_length", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "Routing Policy")] policy: String, /// Minimum matched-prefix share (0.0-1.0) before cache-aware routing @@ -369,6 +369,33 @@ struct CliArgs { #[arg(long, default_value_t = 180, value_parser = clap::value_parser!(u64).range(1..), help_heading = "Routing Policy")] cache_ttl_secs: u64, + // ---- cache_aware_length policy ---- + /// Divisor for char-level token estimation when X-Prompt-Tokens is absent + /// (cache_aware_length policy). Default 4. + #[arg(long, default_value_t = 4, value_parser = parse_positive_usize, help_heading = "Routing Policy")] + chars_per_token: usize, + + /// Uncached-prefill-token boundary between long and short requests + /// (cache_aware_length policy). Default 100000. + #[arg(long, default_value_t = 100_000, value_parser = parse_positive_usize, help_heading = "Routing Policy")] + long_prefill_threshold: usize, + + /// Load ceiling for the long pool (pool=long workers) in the + /// cache_aware_length policy. Default 4. + #[arg(long, default_value_t = 4, value_parser = parse_positive_usize, help_heading = "Routing Policy")] + long_pool_max_load: usize, + + /// Load ceiling for the short pool (remaining workers) in the + /// cache_aware_length policy. Default 32. + #[arg(long, default_value_t = 32, value_parser = parse_positive_usize, help_heading = "Routing Policy")] + short_pool_max_load: usize, + + /// Comma-separated 0-based indices of --prefill URLs that belong to the + /// long pool (get pool=long label for cache_aware_length). E.g. "3,4" + /// marks the 4th and 5th prefill workers as long pool. + #[arg(long, value_delimiter = ',', help_heading = "PD Disaggregation")] + long_prefill_indices: Vec, + /// How long an unused sticky routing key stays pinned: keys idle beyond /// this many seconds are evicted from the manual-policy / sticky-session /// map @@ -470,11 +497,11 @@ struct CliArgs { decode: Vec, /// Specific policy for prefill nodes in PD mode - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "cache_aware_length", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] prefill_policy: Option, /// Specific policy for decode nodes in PD mode - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "cache_aware_length", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")] decode_policy: Option, /// Specific policy for encode nodes in EPD mode. Defaults to consistent_hashing. @@ -1382,6 +1409,25 @@ impl CliArgs { cache_ttl_secs: self.cache_ttl_secs, cache_boundaries: self.cache_boundaries.clone(), }, + "cache_aware_length" => PolicyConfig::CacheAwareLength { + cache_threshold: self.cache_threshold, + balance_abs_threshold: self.balance_abs_threshold, + balance_rel_threshold: self.balance_rel_threshold, + eviction_interval_secs: self.eviction_interval, + max_tree_size: self.max_tree_size, + block_size: self.block_size, + balance_token_usage_threshold: self.balance_token_usage_threshold, + overload_token_usage_threshold: self.overload_token_usage_threshold, + overlap_decay: self.overlap_decay, + selection_temperature: self.selection_temperature, + cache_index: Self::parse_cache_index(&self.cache_index), + cache_ttl_secs: self.cache_ttl_secs, + cache_boundaries: self.cache_boundaries.clone(), + chars_per_token: self.chars_per_token, + long_prefill_threshold: self.long_prefill_threshold, + long_pool_max_load: self.long_pool_max_load, + short_pool_max_load: self.short_pool_max_load, + }, "power_of_two" => PolicyConfig::PowerOfTwo { load_check_interval_secs: 5, }, @@ -1766,6 +1812,7 @@ impl CliArgs { .mode(mode) .policy(policy) .cache_boundaries(self.cache_boundaries.clone()) + .long_prefill_indices(self.long_prefill_indices.clone()) .connection_mode(connection_mode) .startup_worker_runtime_type(startup_worker_runtime_type) .zmq_engine_count(self.zmq_engine_count) diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index 942ee1e1f..12f45313d 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -85,6 +85,46 @@ use crate::{ /// Latest per-worker backend load snapshot stream, keyed by worker URL. pub(crate) type LoadReceiver = watch::Receiver>; +/// Hint about the uncached prefill portion from a partial cache match that +/// fell below `cache_threshold`. Lets the no-cache strategy classify by +/// actual prefill work instead of the full request size. +#[derive(Debug, Clone, Copy)] +pub(crate) enum UncachedHint { + /// Uncached token count (gRPC / token-tree path). + Tokens(usize), + /// Uncached character count (HTTP / string-tree path); the strategy + /// converts to tokens via its `chars_per_token` setting. + Chars(usize), +} + +/// Strategy for the no-cache branch: when a request does not hit the +/// cache tree (or KV events / hash index), this trait selects which worker +/// receives the request. The default behavior (no strategy set) routes to +/// the least-loaded healthy worker (`min_load_idx`). `CacheAwareLengthPolicy` +/// injects a strategy that splits workers into long/short pools by +/// uncached prefill tokens. +pub(crate) trait NoCacheStrategy: Send + Sync + std::fmt::Debug { + /// Select a worker for the no-cache (miss) branch. `min_load_idx` is + /// the pre-computed least-loaded healthy worker the caller would use + /// by default; the strategy may return it or a pool-selected alternative. + /// + /// `uncached_hint` carries the estimated uncached-prefill size when the + /// call originates from a partial cache match that fell below + /// `cache_threshold`. It is `None` when no matching was attempted (e.g. + /// imbalanced fallback, event-driven no-overlap, hash-path fallback), in + /// which case the strategy should estimate from `info` as before. + fn select_no_cache( + &self, + workers: &[Arc], + info: &SelectWorkerInfo, + healthy_indices: &[usize], + min_load_idx: Option, + avg_load: f64, + model_id: &str, + uncached_hint: Option, + ) -> Option; +} + /// Cache-aware routing policy /// /// Routes requests based on cache affinity when load is balanced, @@ -152,6 +192,11 @@ pub struct CacheAwarePolicy { /// cloned inner handle stays canonical and walking it never holds /// an outer-shard guard. placement_index: Arc>>, + /// Optional no-cache strategy. When `None`, the no-cache fallback + /// routes to `min_load_idx` (the default behavior). When `Some`, + /// the strategy selects the worker instead — used by + /// `CacheAwareLengthPolicy` for long/short pool split. + no_cache_strategy: Option>, } /// Hash-mode per-model placement map: (boundary position, xxh3 of the token @@ -327,6 +372,47 @@ impl CacheAwarePolicy { populate_hash_index: AtomicBool::new(false), mesh_tree_sync: RwLock::new(None), placement_index, + no_cache_strategy: None, + } + } + + /// Attach a no-cache strategy, enabling custom worker selection on the + /// cache-miss branch. Used by `CacheAwareLengthPolicy` for long/short + /// pool split. Returns `self` for chaining. + pub(crate) fn with_no_cache_strategy(mut self, strategy: Arc) -> Self { + self.no_cache_strategy = Some(strategy); + self + } + + /// Resolve the no-cache branch: if a strategy is attached, delegate to + /// it; otherwise fall back to `min_load_idx` (the default behavior). + /// Callers still own tree update + `increment_processed` for the + /// returned index. + /// + /// `uncached_hint` is the estimated uncached-prefill token count from a + /// partial cache match (below `cache_threshold`); `None` when no matching + /// was attempted. + fn resolve_no_cache( + &self, + workers: &[Arc], + info: &SelectWorkerInfo, + healthy_indices: &[usize], + min_load_idx: Option, + avg_load: f64, + model_id: &str, + uncached_hint: Option, + ) -> Option { + match &self.no_cache_strategy { + Some(strategy) => strategy.select_no_cache( + workers, + info, + healthy_indices, + min_load_idx, + avg_load, + model_id, + uncached_hint, + ), + None => min_load_idx, } } @@ -687,15 +773,26 @@ impl CacheAwarePolicy { &self, workers: &[Arc], info: &SelectWorkerInfo, + healthy_indices: &[usize], min_load_idx: Option, + avg_load: f64, model_id: &str, ) -> Option { // Shortest queue when imbalanced. The min-load index is gathered upstream // in select_worker with the (load, processed_requests, idx) tie-break - // from #1714 (spreads load when decode outpaces prefill). - let min_load_idx = min_load_idx?; + // from #1714 (spreads load when decode outpaces prefill). When a + // no-cache strategy is attached, it may override the selection. + let min_load_idx = self.resolve_no_cache( + workers, + info, + healthy_indices, + min_load_idx, + avg_load, + model_id, + None, + )?; - let worker_url = workers[min_load_idx].url(); + let worker_url = workers[min_load_idx].url().to_string(); // Even in imbalanced mode, update the appropriate tree to maintain cache state // Prefer token tree for gRPC requests, fall back to string tree for HTTP @@ -715,7 +812,7 @@ impl CacheAwarePolicy { // prefix length the standalone match returned. When we don't // populate the index, a plain insert (no match) suffices. if self.should_populate_hash_index() { - let result = tree.match_and_insert(tokens, worker_url); + let result = tree.match_and_insert(tokens, &worker_url); let matched_prefix: Vec = tokens[..result.matched_token_count].to_vec(); self.hash_index .entry(model_id.to_string()) @@ -723,7 +820,7 @@ impl CacheAwarePolicy { .token_tree .insert(kv_index::hash_token_path(tokens), matched_prefix); } else { - tree.insert_tokens(tokens, worker_url); + tree.insert_tokens(tokens, &worker_url); } } } else if let Some(text) = info.request_text { @@ -741,7 +838,7 @@ impl CacheAwarePolicy { // prefix length the standalone match returned. When we don't // populate the index, a plain insert (no match) suffices. if self.should_populate_hash_index() { - let result = tree.match_and_insert(text, worker_url); + let result = tree.match_and_insert(text, &worker_url); let matched_prefix: String = text.chars().take(result.matched_char_count).collect(); let path_hash = kv_index::hash_node_path(text); @@ -751,7 +848,7 @@ impl CacheAwarePolicy { .string_tree .insert(path_hash, matched_prefix); } else { - tree.insert_text(text, worker_url); + tree.insert_text(text, &worker_url); } } else { debug!( @@ -766,7 +863,7 @@ impl CacheAwarePolicy { debug!( branch = "kv_pressure_min_load", - worker = worker_url, + worker = %worker_url, model_id, "Cache-aware selection" ); @@ -1057,7 +1154,14 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // request-count pressure is applied per request to the selected // candidate inside each affinity path. if self.is_kv_imbalanced(workers, &healthy_indices) { - return self.select_worker_min_load(workers, info, min_load_idx, model_id); + return self.select_worker_min_load( + workers, + info, + &healthy_indices, + min_load_idx, + avg_load, + model_id, + ); } // Cache-aware routing when balanced — three types (mutually exclusive): @@ -1068,6 +1172,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { if self.has_event_indexer(model_id) { self.select_worker_event_driven( workers, + info, tokens, &healthy_indices, min_load_idx, @@ -1077,6 +1182,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { } else { self.select_worker_with_tokens( workers, + info, tokens, &healthy_indices, min_load_idx, @@ -1088,6 +1194,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { let text = request_text.unwrap_or(""); self.select_worker_with_text( workers, + info, text, &healthy_indices, min_load_idx, @@ -1246,6 +1353,7 @@ impl CacheAwarePolicy { fn select_worker_event_driven( &self, workers: &[Arc], + info: &SelectWorkerInfo, tokens: &[u32], healthy_indices: &[usize], min_load_idx: Option, @@ -1279,8 +1387,16 @@ impl CacheAwarePolicy { return self.gate_selected_candidate(workers, idx, avg_load, min_load_idx); } - // No cache overlap — min-load fallback (min-load index gathered upstream) - let min_idx = min_load_idx?; + // No cache overlap — delegate to no-cache strategy (default: min-load). + let min_idx = self.resolve_no_cache( + workers, + info, + healthy_indices, + min_load_idx, + avg_load, + model_id, + None, + )?; debug!( worker = workers[min_idx].url(), model_id, "Event-driven routing: no overlap, min-load fallback" @@ -1538,7 +1654,10 @@ impl CacheAwarePolicy { let Some(tokens) = info.tokens.filter(|t| !t.is_empty()) else { return self.hash_min_load( workers, + info, + healthy_indices, min_load_idx, + avg_load, model_id, "min_load_fallback", &[], @@ -1557,7 +1676,10 @@ impl CacheAwarePolicy { if applicable.is_empty() { return self.hash_min_load( workers, + info, + healthy_indices, min_load_idx, + avg_load, model_id, "short_request", tokens, @@ -1569,7 +1691,10 @@ impl CacheAwarePolicy { if self.is_kv_imbalanced(workers, healthy_indices) { return self.hash_min_load( workers, + info, + healthy_indices, min_load_idx, + avg_load, model_id, "kv_pressure_min_load", tokens, @@ -1601,7 +1726,10 @@ impl CacheAwarePolicy { self.hash_min_load( workers, + info, + healthy_indices, min_load_idx, + avg_load, model_id, "min_load_fallback", tokens, @@ -1616,14 +1744,25 @@ impl CacheAwarePolicy { fn hash_min_load( &self, workers: &[Arc], + info: &SelectWorkerInfo, + healthy_indices: &[usize], min_load_idx: Option, + avg_load: f64, model_id: &str, branch: &'static str, tokens: &[u32], applicable: &[usize], now: Instant, ) -> Option { - let idx = min_load_idx?; + let idx = self.resolve_no_cache( + workers, + info, + healthy_indices, + min_load_idx, + avg_load, + model_id, + None, + )?; if !applicable.is_empty() { self.record_placement(model_id, tokens, applicable, workers[idx].url(), now); } @@ -1747,6 +1886,7 @@ impl CacheAwarePolicy { fn select_worker_with_tokens( &self, workers: &[Arc], + info: &SelectWorkerInfo, tokens: &[u32], healthy_indices: &[usize], min_load_idx: Option, @@ -1797,7 +1937,28 @@ impl CacheAwarePolicy { min_load_idx, ) } else { - min_load_idx + // Partial match below threshold: pass the uncached + // portion (input - matched) so the length strategy can + // classify by actual prefill work, not full size. + // When matched == 0 (no match at all), pass None so the + // strategy falls through to its normal priority chain + // (header → tokens → char estimate). + let uncached_hint = (result.matched_token_count > 0).then(|| { + UncachedHint::Tokens( + result + .input_token_count + .saturating_sub(result.matched_token_count), + ) + }); + self.resolve_no_cache( + workers, + info, + healthy_indices, + min_load_idx, + avg_load, + model_id, + uncached_hint, + ) }; // Insert for the selected worker (None => no insert, exactly @@ -1874,6 +2035,7 @@ impl CacheAwarePolicy { fn select_worker_with_text( &self, workers: &[Arc], + info: &SelectWorkerInfo, text: &str, healthy_indices: &[usize], min_load_idx: Option, @@ -1914,7 +2076,27 @@ impl CacheAwarePolicy { min_load_idx, ) } else { - min_load_idx + // Partial match below threshold: pass the uncached + // char count (input - matched) so the length strategy + // can classify by actual prefill work, not full size. + // When matched == 0 (no match at all), pass None so the + // strategy falls through to its normal priority chain. + let uncached_hint = (result.matched_char_count > 0).then(|| { + UncachedHint::Chars( + result + .input_char_count + .saturating_sub(result.matched_char_count), + ) + }); + self.resolve_no_cache( + workers, + info, + healthy_indices, + min_load_idx, + avg_load, + model_id, + uncached_hint, + ) }; // Insert for the selected worker (None => no insert, exactly diff --git a/model_gateway/src/policies/cache_aware_length.rs b/model_gateway/src/policies/cache_aware_length.rs new file mode 100644 index 000000000..3c6f7b2cb --- /dev/null +++ b/model_gateway/src/policies/cache_aware_length.rs @@ -0,0 +1,798 @@ +/* + Cache-Aware Length Load Balancing Router (cache_aware_length) + + A full superset of `cache_aware` that adds a long/short pool split on the + no-cache (miss) branch. Uses a `NoCacheStrategy` trait injection: the + inner `CacheAwarePolicy` handles all cache-affinity routing (string tree, + token tree, event-driven, hash index, mesh sync, KV pressure), and when a + request misses the cache, the `LengthStrategy` selects a worker based on + uncached prefill tokens and the `pool` worker label. + + Pool membership: + long pool = healthy workers with labels["pool"] == "long" + short pool = remaining healthy workers + + No-cache branch (Step 4): + token source (priority): + 1. X-Prompt-Tokens header (exact, supplied by an upstream gateway). + 2. info.tokens length (token-only gRPC requests). + 3. (input_chars - matched_chars) / chars_per_token (char estimate). + 4. None computable → all-healthy min-load. + uncached >= long_prefill_threshold (long request): + long pool has free worker (load < long_pool_max_load) + → long pool min-load + else short pool has an idle worker (load == 0) + → that worker (long→short overflow) + else long pool has a healthy worker + → long pool min-load (queue) + else → all-healthy min-load + uncached < long_prefill_threshold (short request): + short pool has free worker (load < short_pool_max_load) + → short pool min-load + else long pool has free worker + → long pool min-load (short→long overflow) + else short pool has a worker + → short pool min-load (fallback queue) + else long pool has a worker + → long pool min-load + else (both pools empty) → all-healthy min-load + + This policy does NOT duplicate cache_aware's routing code — it holds an + inner `CacheAwarePolicy` and delegates `select_worker` to it. The only + customisation is the `NoCacheStrategy` implementation injected via + `CacheAwarePolicy::with_no_cache_strategy`. +*/ + +use std::sync::Arc; + +use tracing::debug; + +use super::{ + normalize_model_key, CacheAwareConfig, CacheAwareLengthConfig, CacheAwarePolicy, + LoadBalancingPolicy, NoCacheStrategy, SelectWorkerInfo, UncachedHint, +}; +use crate::{ + mesh::adapters::tree_sync::TreeSyncAdapter, + worker::{KvEventMonitor, Worker}, +}; + +use super::cache_aware::LoadReceiver; + +/// HTTP header carrying the exact prompt token count, supplied by an +/// upstream gateway that has already tokenized the request. +const HEADER_PROMPT_TOKENS: &str = "x-prompt-tokens"; + +/// Cache-aware length routing policy — a full superset of `cache_aware` +/// that adds long/short pool split on the no-cache branch. +/// +/// Holds an inner `CacheAwarePolicy` that owns the trees, mesh, KV monitor, +/// and all cache-affinity routing. The `LengthStrategy` (implementing +/// `NoCacheStrategy`) is injected into the inner policy so it intercepts +/// only the cache-miss branch. +#[derive(Debug)] +pub struct CacheAwareLengthPolicy { + inner: CacheAwarePolicy, + #[allow(dead_code)] + config: CacheAwareLengthConfig, +} + +/// The no-cache strategy: selects a worker by splitting the healthy fleet +/// into long/short pools based on uncached prefill tokens and the `pool` +/// worker label. +#[derive(Debug)] +struct LengthStrategy { + chars_per_token: usize, + long_prefill_threshold: usize, + long_pool_max_load: usize, + short_pool_max_load: usize, +} + +impl Default for CacheAwareLengthPolicy { + fn default() -> Self { + Self::with_config(CacheAwareLengthConfig::default()) + } +} + +impl CacheAwareLengthPolicy { + pub fn new() -> Self { + Self::with_config(CacheAwareLengthConfig::default()) + } + + pub fn with_config(config: CacheAwareLengthConfig) -> Self { + let strategy = Arc::new(LengthStrategy { + chars_per_token: config.chars_per_token, + long_prefill_threshold: config.long_prefill_threshold, + long_pool_max_load: config.long_pool_max_load, + short_pool_max_load: config.short_pool_max_load, + }); + let inner = CacheAwarePolicy::with_config(config.base.clone()) + .with_no_cache_strategy(strategy as Arc); + Self { + inner, + config, + } + } + + // --- Delegated setters (forward to inner CacheAwarePolicy) --- + + pub fn set_kv_event_monitor(&self, monitor: Option>) { + self.inner.set_kv_event_monitor(monitor); + } + + pub fn set_load_receiver(&self, rx: Option) { + self.inner.set_load_receiver(rx); + } + + pub fn set_mesh_tree_sync(&self, adapter: Option>) { + self.inner.set_mesh_tree_sync(adapter); + } + + pub fn init_workers(&self, workers: &[Arc]) { + self.inner.init_workers(workers); + } + + pub fn add_worker(&self, worker: &dyn Worker) { + self.inner.add_worker(worker); + } + + pub fn remove_worker_by_url(&self, url: &str) { + self.inner.remove_worker_by_url(url); + } + + /// Test-only access to the config so factory tests can verify values. + #[cfg(test)] + pub(crate) fn config_for_test(&self) -> &CacheAwareLengthConfig { + &self.config + } +} + +impl LoadBalancingPolicy for CacheAwareLengthPolicy { + fn select_worker(&self, workers: &[Arc], info: &SelectWorkerInfo) -> Option { + // Delegate entirely to the inner CacheAwarePolicy. The LengthStrategy + // (injected via with_no_cache_strategy) intercepts only the no-cache + // branch; cache-hit, KV-pressure, hash-mode, and event-driven paths + // run unchanged inside the inner policy. + self.inner.select_worker(workers, info) + } + + fn name(&self) -> &'static str { + "cache_aware_length" + } + + fn needs_request_text(&self) -> bool { + true + } + + fn needs_backend_loads(&self) -> bool { + self.inner.needs_backend_loads() + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +impl NoCacheStrategy for LengthStrategy { + fn select_no_cache( + &self, + workers: &[Arc], + info: &SelectWorkerInfo, + healthy_indices: &[usize], + min_load_idx: Option, + _avg_load: f64, + _model_id: &str, + uncached_hint: Option, + ) -> Option { + if healthy_indices.is_empty() { + return min_load_idx; + } + + // Compute uncached prefill tokens by priority: + // 0. Partial-cache hint (input - matched) when a tree match fell below + // cache_threshold — classify by actual prefill work, not full size. + // 1. X-Prompt-Tokens header (exact). + // 2. info.tokens length (token-only gRPC). + // 3. input_chars / chars_per_token (char estimate). + // 4. None → all-healthy min-load. + let uncached_tokens = + self.compute_uncached_tokens(info, uncached_hint); + + let Some(uncached) = uncached_tokens else { + // Neither source computable → all-healthy min-load. + return min_load_idx; + }; + + // Split healthy workers into long/short pools by label, capturing + // (load, processed) in the same pass so pool helpers don't re-read + // routing_state() per call. + let long_pool: Vec<(usize, usize, usize)> = healthy_indices + .iter() + .copied() + .filter(|&i| is_long_pool(&*workers[i])) + .map(|i| { + let s = workers[i].routing_state(); + (i, s.load, s.processed) + }) + .collect(); + let short_pool: Vec<(usize, usize, usize)> = healthy_indices + .iter() + .copied() + .filter(|&i| !is_long_pool(&*workers[i])) + .map(|i| { + let s = workers[i].routing_state(); + (i, s.load, s.processed) + }) + .collect(); + + let selected = if uncached >= self.long_prefill_threshold { + self.select_long_request(&long_pool, &short_pool, min_load_idx) + } else { + self.select_short_request(&long_pool, &short_pool, min_load_idx) + }; + + // The inner policy's caller (select_worker_min_load, the tree + // closures, or hash_min_load) handles tree update + + // increment_processed for the returned index. + selected.or(min_load_idx) + } +} + +impl LengthStrategy { + /// Compute uncached prefill tokens by priority. + fn compute_uncached_tokens( + &self, + info: &SelectWorkerInfo, + hint: Option, + ) -> Option { + // 0. Partial-cache hint from a tree match below cache_threshold. + if let Some(h) = hint { + return match h { + UncachedHint::Tokens(n) => Some(n), + UncachedHint::Chars(n) => { + if self.chars_per_token > 0 { + Some(n.div_ceil(self.chars_per_token)) + } else { + Some(n) + } + } + }; + } + // 1. Exact header value. + if let Some(n) = parse_prompt_tokens_header(info.headers) { + return Some(n); + } + // 2. Token-only request (no text): use the token count directly. + if let Some(tokens) = info.tokens { + if !tokens.is_empty() { + return Some(tokens.len()); + } + } + // 3. Char-level estimate from the request text. + if let Some(text) = info.request_text { + let input_chars = text.chars().count(); + if input_chars > 0 && self.chars_per_token > 0 { + return Some(input_chars.div_ceil(self.chars_per_token)); + } + } + None + } + + /// Long request (uncached >= long_prefill_threshold). + fn select_long_request( + &self, + long_pool: &[(usize, usize, usize)], + short_pool: &[(usize, usize, usize)], + min_load_idx: Option, + ) -> Option { + if pool_has_free(long_pool, self.long_pool_max_load) { + return pool_min_load_worker(long_pool); + } + // Long pool full/unhealthy: overflow to an idle short-pool worker only. + if let Some(idx) = pool_idle_worker(short_pool) { + return Some(idx); + } + // Short pool all busy: queue on long pool if it still has a worker. + if let Some(idx) = pool_min_load_worker(long_pool) { + return Some(idx); + } + // Long pool fully unhealthy and short pool busy: all-healthy min-load. + min_load_idx + } + + /// Short request (uncached < long_prefill_threshold). + fn select_short_request( + &self, + long_pool: &[(usize, usize, usize)], + short_pool: &[(usize, usize, usize)], + min_load_idx: Option, + ) -> Option { + if pool_has_free(short_pool, self.short_pool_max_load) { + return pool_min_load_worker(short_pool); + } + // Short pool full: overflow to long pool if it has a free worker. + if pool_has_free(long_pool, self.long_pool_max_load) { + return pool_min_load_worker(long_pool); + } + // Both full: queue on short pool if it has a worker. + if let Some(idx) = pool_min_load_worker(short_pool) { + return Some(idx); + } + // Short pool empty: queue on long pool if it has a worker. + if let Some(idx) = pool_min_load_worker(long_pool) { + return Some(idx); + } + // Both pools empty: all-healthy min-load. + min_load_idx + } +} + +/// Whether a worker belongs to the long pool (`labels["pool"] == "long"`). +fn is_long_pool(worker: &dyn Worker) -> bool { + worker + .metadata() + .spec + .labels + .get("pool") + .is_some_and(|v| v == "long") +} + +/// Does any worker in `pool` have `load < max_load`? +fn pool_has_free(pool: &[(usize, usize, usize)], max_load: usize) -> bool { + pool.iter().any(|&(_, load, _)| load < max_load) +} + +/// Return the worker index of an idle (`load == 0`) entry, if any. +fn pool_idle_worker(pool: &[(usize, usize, usize)]) -> Option { + pool.iter().find(|(_, load, _)| *load == 0).map(|(idx, _, _)| *idx) +} + +/// Lowest-load worker in `pool` with the `(load, processed, idx)` tie-break. +/// Returns `None` when `pool` is empty. +fn pool_min_load_worker(pool: &[(usize, usize, usize)]) -> Option { + pool.iter() + .min_by_key(|(idx, load, processed)| (*load, *processed, *idx)) + .map(|(idx, _, _)| *idx) +} + +/// Parse the `X-Prompt-Tokens` header into a token count. Returns `None` on +/// missing/unparseable values. +fn parse_prompt_tokens_header(headers: Option<&http::HeaderMap>) -> Option { + let headers = headers?; + headers + .get(HEADER_PROMPT_TOKENS) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.trim().parse::().ok()) + .filter(|n| *n > 0) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use openai_protocol::worker::{HealthCheckConfig, WorkerStatus}; + + use super::*; + use crate::worker::{BasicWorkerBuilder, WorkerType}; + + fn no_health_check() -> HealthCheckConfig { + HealthCheckConfig { + disable_health_check: true, + ..Default::default() + } + } + + /// Build a worker with an optional `pool` label and a pre-set load. + fn make_worker(url: &str, pool: Option<&str>, load: usize) -> Arc { + let mut builder = BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .health_config(no_health_check()); + if let Some(p) = pool { + builder = builder.label("pool", p); + } + let worker: Arc = Arc::new(builder.build()); + for _ in 0..load { + std::mem::forget(crate::worker::WorkerLoadGuard::new( + Arc::clone(&worker), + None, + )); + } + worker + } + + fn info_with_text(text: &str) -> SelectWorkerInfo<'_> { + SelectWorkerInfo { + request_text: Some(text), + ..Default::default() + } + } + + fn info_with_header<'a>(headers: &'a http::HeaderMap, text: &'a str) -> SelectWorkerInfo<'a> { + SelectWorkerInfo { + request_text: Some(text), + headers: Some(headers), + ..Default::default() + } + } + + fn tokens_headers(tokens: usize) -> http::HeaderMap { + let mut headers = http::HeaderMap::new(); + headers.insert( + HEADER_PROMPT_TOKENS, + http::HeaderValue::from_str(&tokens.to_string()).unwrap(), + ); + headers + } + + fn test_config() -> CacheAwareLengthConfig { + CacheAwareLengthConfig { + base: CacheAwareConfig { + eviction_interval_secs: 0, + ..Default::default() + }, + long_prefill_threshold: 100_000, + long_pool_max_load: 2, + short_pool_max_load: 2, + chars_per_token: 4, + ..Default::default() + } + } + + #[test] + fn step1_returns_none_when_all_unhealthy() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .api_key("k") + .health_config(no_health_check()) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .api_key("k") + .health_config(no_health_check()) + .build(), + ), + ]; + for w in &workers { + w.set_status(WorkerStatus::NotReady); + } + assert!(policy + .select_worker(&workers, &info_with_text("hello")) + .is_none()); + } + + #[test] + fn step3_cache_hit_pins_to_same_worker() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", None, 0), + ]; + policy.init_workers(&workers); + + let prompt = "shared long prompt prefix that both workers could cache"; + let idx1 = policy + .select_worker(&workers, &info_with_text(prompt)) + .unwrap(); + let idx2 = policy + .select_worker(&workers, &info_with_text(prompt)) + .unwrap(); + assert_eq!(idx1, idx2); + } + + #[test] + fn step3_tree_missing_falls_back_random() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", None, 0), + ]; + let idx = policy + .select_worker(&workers, &info_with_text("novel prompt")) + .unwrap(); + assert!(idx < workers.len()); + } + + #[test] + fn step4_long_request_uses_long_pool_when_free() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn step4_long_request_overflows_to_idle_short() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 2), + ]; + policy.init_workers(&workers); + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn step4_long_request_queues_on_long_when_short_busy() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 1), + make_worker("http://w2:8000", Some("long"), 2), + ]; + policy.init_workers(&workers); + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn step4_short_request_uses_short_pool_when_free() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + let headers = tokens_headers(1_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn step4_short_request_overflows_to_long_when_short_full() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 2), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + let headers = tokens_headers(1_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn step4_short_request_falls_back_to_short_when_both_full() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 2), + make_worker("http://w2:8000", Some("long"), 2), + ]; + policy.init_workers(&workers); + let headers = tokens_headers(1_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn step4_short_request_uses_long_when_short_pool_empty() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![make_worker("http://w2:8000", Some("long"), 0)]; + policy.init_workers(&workers); + let headers = tokens_headers(1_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn char_estimate_falls_back_when_no_header() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + let prompt = "a".repeat(400); + let idx = policy + .select_worker(&workers, &info_with_text(&prompt)) + .unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn is_long_pool_reads_label() { + let w_long = make_worker("http://w:1", Some("long"), 0); + let w_short = make_worker("http://w:2", None, 0); + let w_other = make_worker("http://w:3", Some("short"), 0); + assert!(is_long_pool(&*w_long)); + assert!(!is_long_pool(&*w_short)); + assert!(!is_long_pool(&*w_other)); + } + + #[test] + fn parse_prompt_tokens_header_works() { + let mut headers = http::HeaderMap::new(); + headers.insert( + HEADER_PROMPT_TOKENS, + http::HeaderValue::from_static("12345"), + ); + assert_eq!(parse_prompt_tokens_header(Some(&headers)), Some(12345)); + assert_eq!(parse_prompt_tokens_header(None), None); + headers.insert( + HEADER_PROMPT_TOKENS, + http::HeaderValue::from_static("notanum"), + ); + assert_eq!(parse_prompt_tokens_header(Some(&headers)), None); + } + + #[test] + fn step3_hit_unhealthy_falls_back_to_first_healthy() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + + let prompt = "shared cache-building prompt for affinity"; + let first = policy + .select_worker(&workers, &info_with_text(prompt)) + .unwrap(); + workers[first].set_status(WorkerStatus::NotReady); + let second = policy + .select_worker(&workers, &info_with_text(prompt)) + .unwrap(); + assert_ne!(workers[second].url(), workers[first].url()); + } + + #[test] + fn step4_long_pool_unhealthy_overflows_to_idle_short() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + workers[1].set_status(WorkerStatus::NotReady); + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt no match yet"); + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w1:8000"); + } + + #[test] + fn step4_header_overrides_char_estimate() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + let mut headers = http::HeaderMap::new(); + headers.insert( + HEADER_PROMPT_TOKENS, + http::HeaderValue::from_static("200000"), + ); + let info = SelectWorkerInfo { + request_text: Some("short"), + headers: Some(&headers), + ..Default::default() + }; + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + #[test] + fn step4_token_only_request_uses_token_count() { + let policy = CacheAwareLengthPolicy::with_config(test_config()); + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + policy.init_workers(&workers); + let tokens: Vec = (0..200_000).collect(); + let info = SelectWorkerInfo { + request_text: None, + tokens: Some(&tokens), + ..Default::default() + }; + let idx = policy.select_worker(&workers, &info).unwrap(); + assert_eq!(workers[idx].url(), "http://w2:8000"); + } + + /// A partial cache hit below threshold must classify by the uncached + /// portion, not the full request size. A 200K-token request with 190K + /// matched → uncached = 10K < threshold → short pool, even though the + /// full request would have been "long." + #[test] + fn step4_partial_cache_hit_uses_uncached_not_full_size_tokens() { + let cfg = test_config(); + let strategy = LengthStrategy { + chars_per_token: cfg.chars_per_token, + long_prefill_threshold: cfg.long_prefill_threshold, + long_pool_max_load: cfg.long_pool_max_load, + short_pool_max_load: cfg.short_pool_max_load, + }; + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + let healthy_indices: Vec = vec![0, 1]; + // Hint: 200K total - 190K matched = 10K uncached < 100K threshold. + let hint = Some(UncachedHint::Tokens(10_000)); + let info = SelectWorkerInfo::default(); + let idx = strategy + .select_no_cache(&workers, &info, &healthy_indices, None, 0.0, "m", hint) + .unwrap(); + assert_eq!( + workers[idx].url(), + "http://w1:8000", + "partial match (10K uncached < 100K threshold) → short pool" + ); + } + + /// Same scenario but via char hint: 400K chars total - 390K matched = + /// 10K uncached chars → 2500 tokens < 100K threshold → short pool. + #[test] + fn step4_partial_cache_hit_uses_uncached_not_full_size_chars() { + let cfg = test_config(); + let strategy = LengthStrategy { + chars_per_token: cfg.chars_per_token, + long_prefill_threshold: cfg.long_prefill_threshold, + long_pool_max_load: cfg.long_pool_max_load, + short_pool_max_load: cfg.short_pool_max_load, + }; + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + let healthy_indices: Vec = vec![0, 1]; + // 10K uncached chars / 4 chars_per_token = 2500 tokens < 100K threshold. + let hint = Some(UncachedHint::Chars(10_000)); + let info = SelectWorkerInfo::default(); + let idx = strategy + .select_no_cache(&workers, &info, &healthy_indices, None, 0.0, "m", hint) + .unwrap(); + assert_eq!( + workers[idx].url(), + "http://w1:8000", + "partial match (10K uncached chars = 2500 tokens < 100K) → short pool" + ); + } + + /// Without a hint (no match at all), a 200K-token request stays "long" + /// via the header, proving the hint does not override when absent. + #[test] + fn step4_no_hint_falls_through_to_header() { + let cfg = test_config(); + let strategy = LengthStrategy { + chars_per_token: cfg.chars_per_token, + long_prefill_threshold: cfg.long_prefill_threshold, + long_pool_max_load: cfg.long_pool_max_load, + short_pool_max_load: cfg.short_pool_max_load, + }; + let workers: Vec> = vec![ + make_worker("http://w1:8000", None, 0), + make_worker("http://w2:8000", Some("long"), 0), + ]; + let healthy_indices: Vec = vec![0, 1]; + let headers = tokens_headers(200_000); + let info = info_with_header(&headers, "novel prompt"); + let idx = strategy + .select_no_cache(&workers, &info, &healthy_indices, None, 0.0, "m", None) + .unwrap(); + assert_eq!( + workers[idx].url(), + "http://w2:8000", + "no hint → header 200K > 100K threshold → long pool" + ); + } +} diff --git a/model_gateway/src/policies/factory.rs b/model_gateway/src/policies/factory.rs index e18213870..0ad2fdc41 100755 --- a/model_gateway/src/policies/factory.rs +++ b/model_gateway/src/policies/factory.rs @@ -3,9 +3,10 @@ use std::sync::Arc; use super::{ - BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy, - LeastLoadPolicy, LoadBalancingPolicy, ManualConfig, ManualPolicy, PassthroughPolicy, - PowerOfTwoPolicy, PrefixHashConfig, PrefixHashPolicy, RandomPolicy, RoundRobinPolicy, + BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwareLengthConfig, CacheAwareLengthPolicy, + CacheAwarePolicy, ConsistentHashingPolicy, LeastLoadPolicy, LoadBalancingPolicy, ManualConfig, + ManualPolicy, PassthroughPolicy, PowerOfTwoPolicy, PrefixHashConfig, PrefixHashPolicy, + RandomPolicy, RoundRobinPolicy, }; use crate::config::PolicyConfig; @@ -72,6 +73,49 @@ impl PolicyFactory { }; Arc::new(CacheAwarePolicy::with_config(config)) } + PolicyConfig::CacheAwareLength { + cache_threshold, + balance_abs_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + block_size, + balance_token_usage_threshold, + overload_token_usage_threshold, + overlap_decay, + selection_temperature, + cache_index, + cache_ttl_secs, + cache_boundaries, + chars_per_token, + long_prefill_threshold, + long_pool_max_load, + short_pool_max_load, + } => { + let base = CacheAwareConfig { + cache_threshold: *cache_threshold, + balance_abs_threshold: *balance_abs_threshold, + balance_rel_threshold: *balance_rel_threshold, + eviction_interval_secs: *eviction_interval_secs, + max_tree_size: *max_tree_size, + block_size: *block_size, + balance_token_usage_threshold: *balance_token_usage_threshold, + overload_token_usage_threshold: *overload_token_usage_threshold, + overlap_decay: *overlap_decay, + selection_temperature: *selection_temperature, + cache_index: *cache_index, + cache_ttl_secs: *cache_ttl_secs, + cache_boundaries: cache_boundaries.clone(), + }; + let config = CacheAwareLengthConfig { + base, + chars_per_token: *chars_per_token, + long_prefill_threshold: *long_prefill_threshold, + long_pool_max_load: *long_pool_max_load, + short_pool_max_load: *short_pool_max_load, + }; + Arc::new(CacheAwareLengthPolicy::with_config(config)) + } PolicyConfig::Bucket { balance_abs_threshold, balance_rel_threshold, @@ -123,6 +167,9 @@ impl PolicyFactory { "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())), "least_load" | "leastload" => Some(Arc::new(LeastLoadPolicy::new())), "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())), + "cache_aware_length" | "cacheawarelength" => { + Some(Arc::new(CacheAwareLengthPolicy::new())) + } "bucket" => Some(Arc::new(BucketPolicy::new())), "manual" => Some(Arc::new(ManualPolicy::new())), "consistent_hashing" | "consistenthashing" => { @@ -171,6 +218,39 @@ mod tests { }); assert_eq!(policy.name(), "cache_aware"); + let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAwareLength { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 30, + max_tree_size: 10000, + block_size: 16, + balance_token_usage_threshold: 1.0, + overload_token_usage_threshold: 1.0, + overlap_decay: 0.0, + selection_temperature: 0.0, + cache_index: Default::default(), + cache_ttl_secs: 180, + cache_boundaries: Vec::new(), + chars_per_token: 4, + long_prefill_threshold: 100_000, + long_pool_max_load: 4, + short_pool_max_load: 32, + }); + assert_eq!(policy.name(), "cache_aware_length"); + // Verify config values are preserved, not just the policy name. + let cal = policy + .as_any() + .downcast_ref::() + .unwrap(); + let cfg = cal.config_for_test(); + assert_eq!(cfg.chars_per_token, 4); + assert_eq!(cfg.long_prefill_threshold, 100_000); + assert_eq!(cfg.long_pool_max_load, 4); + assert_eq!(cfg.short_pool_max_load, 32); + assert_eq!(cfg.base.cache_threshold, 0.5); + assert_eq!(cfg.base.block_size, 16); + let policy = PolicyFactory::create_from_config(&PolicyConfig::Bucket { balance_abs_threshold: 10, balance_rel_threshold: 1.5, @@ -212,6 +292,18 @@ mod tests { assert!(PolicyFactory::create_by_name("PowerOfTwo").is_some()); assert!(PolicyFactory::create_by_name("cache_aware").is_some()); assert!(PolicyFactory::create_by_name("CacheAware").is_some()); + assert_eq!( + PolicyFactory::create_by_name("cache_aware_length") + .unwrap() + .name(), + "cache_aware_length" + ); + assert_eq!( + PolicyFactory::create_by_name("CacheAwareLength") + .unwrap() + .name(), + "cache_aware_length" + ); assert!(PolicyFactory::create_by_name("bucket").is_some()); assert!(PolicyFactory::create_by_name("Bucket").is_some()); assert!(PolicyFactory::create_by_name("manual").is_some()); diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index d6b192592..177a2d9c1 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -14,6 +14,7 @@ use crate::{ mod bucket; mod cache_aware; +mod cache_aware_length; mod consistent_hashing; mod dp_min_token; mod factory; @@ -28,7 +29,9 @@ mod round_robin; pub(crate) mod utils; pub use bucket::BucketPolicy; +pub(crate) use cache_aware::{NoCacheStrategy, UncachedHint}; pub use cache_aware::{CacheAwarePolicy, TreeHandle, TreeKind}; +pub use cache_aware_length::CacheAwareLengthPolicy; pub use consistent_hashing::ConsistentHashingPolicy; pub use dp_min_token::MinimumTokensPolicy; pub use factory::PolicyFactory; @@ -198,6 +201,44 @@ impl Default for CacheAwareConfig { } } +/// Configuration for the cache_aware_length policy (long/short pool split). +/// +/// Embeds [`CacheAwareConfig`] as `base` so the policy inherits all +/// cache_aware features (string tree, token tree, event-driven routing, +/// hash index, mesh sync, KV pressure) and adds 4 length-specific fields +/// for the no-cache long/short pool split. +#[derive(Debug, Clone)] +pub struct CacheAwareLengthConfig { + /// Full cache_aware configuration: cache_threshold, balance thresholds, + /// eviction, block_size, KV pressure knobs, overlap decay, selection + /// temperature, cache_index, cache_ttl_secs, cache_boundaries. + pub base: CacheAwareConfig, + /// Divisor for char-level token estimation when `X-Prompt-Tokens` is + /// absent: `uncached_tokens = (input_chars - matched_chars) / + /// chars_per_token`. + pub chars_per_token: usize, + /// Uncached-prefill-token boundary between long and short requests. + pub long_prefill_threshold: usize, + /// Load ceiling for the long pool (`pool=long` workers). + pub long_pool_max_load: usize, + /// Load ceiling for the short pool (remaining workers). + pub short_pool_max_load: usize, +} + +impl Default for CacheAwareLengthConfig { + fn default() -> Self { + Self { + base: CacheAwareConfig::default(), + chars_per_token: 4, + long_prefill_threshold: 100_000, + // 0 would reject every worker as non-free; pick generous defaults so + // the policy is usable out of the box before operators tune. + long_pool_max_load: 4, + short_pool_max_load: 32, + } + } +} + #[derive(Debug, Clone)] pub struct BucketConfig { pub balance_abs_threshold: usize, diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 42777507e..d8c4200eb 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -17,8 +17,8 @@ use tracing::{debug, info, warn}; use super::{ get_healthy_worker_indices, manual::{ExecutionBranch, PinState}, - BucketPolicy, CacheAwarePolicy, DPRankLoadPolicy, LoadBalancingPolicy, ManualConfig, - ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, + BucketPolicy, CacheAwareLengthPolicy, CacheAwarePolicy, DPRankLoadPolicy, LoadBalancingPolicy, + ManualConfig, ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, }; use crate::{ config::types::{ManualAssignmentMode, PolicyConfig, RoutingKeyOverrideConfig}, @@ -807,6 +807,15 @@ impl PolicyRegistry { ); cache_aware.init_workers(workers); } + } else if policy.name() == "cache_aware_length" { + if let Some(cal) = policy.as_any().downcast_ref::() { + debug!( + "Initializing cache_aware_length policy with {} workers for model {}", + workers.len(), + model_id + ); + cal.init_workers(workers); + } } } } @@ -824,6 +833,14 @@ impl PolicyRegistry { worker_url, model_id ); } + } else if policy.name() == "cache_aware_length" { + if let Some(cal) = policy.as_any().downcast_ref::() { + cal.remove_worker_by_url(worker_url); + debug!( + "Removed worker {} from cache_aware_length policy for model {}", + worker_url, model_id + ); + } } } } @@ -845,6 +862,14 @@ impl PolicyRegistry { worker_url, worker_type ); } + } else if policy.name() == "cache_aware_length" { + if let Some(cal) = policy.as_any().downcast_ref::() { + cal.remove_worker_by_url(worker_url); + debug!( + "Removed worker {} from {} cache_aware_length policy", + worker_url, worker_type + ); + } } } } @@ -882,6 +907,19 @@ impl PolicyRegistry { cache_aware.init_workers(prefill_workers); } } + } else if prefill_policy.name() == "cache_aware_length" { + if let Some(cal) = prefill_policy + .as_any() + .downcast_ref::() + { + if !prefill_workers.is_empty() { + debug!( + "Initializing prefill cache_aware_length policy with {} workers", + prefill_workers.len() + ); + cal.init_workers(prefill_workers); + } + } } } @@ -898,6 +936,19 @@ impl PolicyRegistry { cache_aware.init_workers(decode_workers); } } + } else if decode_policy.name() == "cache_aware_length" { + if let Some(cal) = decode_policy + .as_any() + .downcast_ref::() + { + if !decode_workers.is_empty() { + debug!( + "Initializing decode cache_aware_length policy with {} workers", + decode_workers.len() + ); + cal.init_workers(decode_workers); + } + } } } } diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 654958dc3..453a1541f 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -2123,7 +2123,7 @@ mod tests { use super::*; use crate::{ config::types::{PolicyConfig, RoutingKeyOverrideConfig}, - policies::CacheAwarePolicy, + policies::{CacheAwareLengthPolicy, CacheAwarePolicy}, routers::common::request_lease::test_probe::{spawn_release_gated_stub, DropProbeRequest}, worker::BasicWorkerBuilder, }; @@ -2476,6 +2476,40 @@ mod tests { .build() } + /// Like `plain_worker` but attaches a `pool=` label so the + /// cache_aware_length policy can split workers into long/short pools. + fn labeled_worker(url: &str, pool: Option<&str>) -> crate::worker::BasicWorker { + let mut builder = BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .health_config(no_health_check()); + if let Some(p) = pool { + builder = builder.label("pool", p); + } + builder.build() + } + + fn cache_aware_length_policy() -> PolicyConfig { + PolicyConfig::CacheAwareLength { + cache_threshold: 0.3, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 0, + max_tree_size: 4096, + block_size: 16, + balance_token_usage_threshold: 1.0, + overload_token_usage_threshold: 1.0, + overlap_decay: 0.0, + selection_temperature: 0.0, + cache_index: Default::default(), + cache_ttl_secs: 180, + cache_boundaries: Vec::new(), + chars_per_token: 4, + long_prefill_threshold: 100_000, + long_pool_max_load: 2, + short_pool_max_load: 2, + } + } + type CapturedUpstreamRequest = Arc>>; /// Loopback engine stub: captures the forwarded `/generate` request and @@ -2988,4 +3022,450 @@ mod tests { let router = streaming_router(least_load_policy(), 1024 * 1024, vec![]); assert_falls_back_with_body_intact(&router).await; } + + // ===== cache_aware_length E2E: all 16 decision-table scenarios ===== + // + // Each test drives the real routing decision point + // (`select_worker_for_model`), which builds the exact `SelectWorkerInfo` + // (request_text + headers) the policy receives in production — the same + // path the buffered HTTP request takes. Workers carry `pool=long` labels; + // pre-set loads use WorkerLoadGuard + mem::forget (same as the unit tests). + + /// Pin a worker's in-flight load at `load` for the lifetime of the test by + /// leaking the RAII guard (the process tears down on test exit). Accepts + /// the registered `Arc` directly — no downcast needed since + /// `WorkerLoadGuard::new` takes `Arc`. + fn pin_load(worker: &Arc, load: usize) { + for _ in 0..load { + std::mem::forget(WorkerLoadGuard::new(Arc::clone(worker), None)); + } + } + + /// Spawn a two-pool router: short workers (no label) + long workers + /// (`pool=long`), with cache_aware_length as the default policy and trees + /// seeded. Returns the router so tests can assert routed worker. + async fn length_router(short_urls: &[&str], long_urls: &[&str]) -> Router { + let mut workers: Vec = Vec::new(); + for u in short_urls { + workers.push(labeled_worker(u, None)); + } + for u in long_urls { + workers.push(labeled_worker(u, Some("long"))); + } + let router = streaming_router(cache_aware_length_policy(), 1024 * 1024, workers); + let live = router.worker_registry.get_all(); + router + .policy_registry + .get_default_policy() + .as_any() + .downcast_ref::() + .unwrap() + .init_workers(&live); + router + } + + /// Route `prompt` through the real `select_worker_for_model` (the exact + /// path a buffered HTTP request takes), returning the selected worker's + /// URL — the value the router writes to `x-smg-routed-worker-id`. + /// Returns `None` when the fleet rejects (503). Panics if the selection + /// unexpectedly fails (use `route_or_none` for the reject case). + fn route_to_url(router: &Router, prompt: &str, headers: Option<&HeaderMap>) -> String { + router + .select_worker_for_model( + crate::worker::UNKNOWN_MODEL_ID, + Some(prompt), + None, + headers, + None, + ) + .map(|w| w.url().to_string()) + .expect("selection should succeed with healthy workers") + } + + /// Like `route_to_url` but returns `None` when the fleet rejects (503). + fn route_or_none(router: &Router, prompt: &str, headers: Option<&HeaderMap>) -> Option { + router + .select_worker_for_model( + crate::worker::UNKNOWN_MODEL_ID, + Some(prompt), + None, + headers, + None, + ) + .map(|w| w.url().to_string()) + } + + /// Build a HeaderMap with `X-Prompt-Tokens: `. + fn tokens_header(tokens: usize) -> HeaderMap { + let mut h = HeaderMap::new(); + h.insert( + "x-prompt-tokens", + http::HeaderValue::from_str(&tokens.to_string()).unwrap(), + ); + h + } + + /// Find a live worker by URL in the router's registry and pin its load. + fn pin_worker(router: &Router, url: &str, load: usize) { + let worker = router + .worker_registry + .get_all() + .iter() + .find(|w| w.url() == url) + .cloned() + .unwrap(); + pin_load(&worker, load); + } + + /// Mark a worker unhealthy by URL. + fn mark_unhealthy(router: &Router, url: &str) { + router + .worker_registry + .get_all() + .iter() + .find(|w| w.url() == url) + .cloned() + .unwrap() + .set_status(openai_protocol::worker::WorkerStatus::NotReady); + } + + // --- Step 1: health filter --- + + #[tokio::test] + async fn cal_step1_all_unhealthy_returns_503() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + mark_unhealthy(&router, &url_s); + mark_unhealthy(&router, &url_l); + let h = tokens_header(200_000); + let routed = route_or_none(&router, "hello", Some(&h)); + assert!( + routed.is_none(), + "all unhealthy → 503 (None), got {routed:?}" + ); + } + + // --- Step 2: global imbalance --- + + #[tokio::test] + async fn cal_step2_global_imbalance_picks_min_load() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + // Pin the long-pool worker high so the fleet is imbalanced (100 vs 0). + pin_worker(&router, &url_l, 100); + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "imbalanced fleet → healthy min-load worker"); + } + + // --- Step 3: cache hit --- + + #[tokio::test] + async fn cal_step3_cache_hit_sticks_regardless_of_pool() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let prompt = "shared long prompt prefix that builds cache affinity"; + + // Seed the cache with a long-request header so the first request goes + // to the long-pool worker (url_l), not the default short-pool routing. + let h = tokens_header(200_000); + let first = route_to_url(&router, prompt, Some(&h)); + assert_eq!(first, url_l, "header-classified long request → long pool"); + + // Same prompt without the header: cache lookup (Step 3) takes + // precedence over pool selection (Step 4), so it stays on url_l. + let second = route_to_url(&router, prompt, None); + assert_eq!(second, url_l, "cache hit overrides short-pool routing"); + } + + #[tokio::test] + async fn cal_step3_no_tree_random_healthy_does_not_panic() { + // Do NOT call init_workers → no tree → random healthy fallback. + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = streaming_router( + cache_aware_length_policy(), + 1024 * 1024, + vec![ + labeled_worker(&url_s, None), + labeled_worker(&url_l, Some("long")), + ], + ); + let routed = route_to_url(&router, "novel", None); + assert!( + routed == url_s || routed == url_l, + "random fallback picks a healthy worker: {routed}" + ); + } + + // --- Step 4 long request (uncached ≥ 100K): 4 paths --- + + #[tokio::test] + async fn cal_step4_long_uses_long_pool_when_free() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "long request, free long pool → long pool"); + } + + #[tokio::test] + async fn cal_step4_long_overflows_to_idle_short() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_l, 2); // long pool full (load = max) + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "full long pool → idle (load 0) short worker"); + } + + #[tokio::test] + async fn cal_step4_long_queues_on_long_when_short_busy() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_l, 2); // long full + pin_worker(&router, &url_s, 1); // short busy (load > 0, not idle) + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "long full + short busy → queue on long pool"); + } + + // Step 4: long pool all unhealthy + short pool all load>0 → all-healthy min-load. + #[tokio::test] + async fn cal_step4_long_unhealthy_short_busy_all_healthy_min_load() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + mark_unhealthy(&router, &url_l); // long pool unhealthy + pin_worker(&router, &url_s, 1); // short busy (load > 0) + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!( + routed, url_s, + "long unhealthy + short busy → all-healthy min-load (short worker)" + ); + } + + // --- Step 4 short request (uncached < 100K): 5 paths --- + + #[tokio::test] + async fn cal_step4_short_uses_short_pool_when_free() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "short request, free short pool → short pool"); + } + + #[tokio::test] + async fn cal_step4_short_overflows_to_long_when_short_full() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_s, 2); // short pool full + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "short pool full → overflow to long pool"); + } + + #[tokio::test] + async fn cal_step4_short_falls_back_to_short_when_both_full() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_s, 2); // short full + pin_worker(&router, &url_l, 2); // long full + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "both full → queue on short pool min-load"); + } + + #[tokio::test] + async fn cal_step4_short_uses_long_when_short_pool_empty() { + // Only long-pool workers configured; short pool is empty. + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[], &[&url_l]).await; + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_l, "no short pool → long pool min-load"); + } + + #[tokio::test] + async fn cal_step4_short_both_pools_empty_all_healthy_min_load() { + // Only one healthy short worker (long pool empty); it is the sole + // all-healthy min-load candidate. + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[]).await; + let h = tokens_header(1_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!( + routed, url_s, + "single healthy worker → all-healthy min-load" + ); + } + + // --- Step 4: token source priority (header vs char estimate vs none) --- + + #[tokio::test] + async fn cal_step4_char_estimate_fallback_when_no_header() { + // No X-Prompt-Tokens header → char-level estimate. A 400-char novel + // prompt → 100 tokens < 100K threshold → short request → short pool. + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let prompt = "a".repeat(400); + let routed = route_to_url(&router, &prompt, None); + assert_eq!(routed, url_s, "char estimate < threshold → short pool"); + } + + #[tokio::test] + async fn cal_step4_uncached_unknown_all_healthy_min_load() { + // Empty prompt, no header, no tokens → uncached not computable → + // all-healthy min-load. Pin the long worker higher so the short worker + // is the unique minimum-load candidate. + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + pin_worker(&router, &url_l, 10); // long worker at load 10 + let routed = route_to_url(&router, "", None); + assert_eq!( + routed, url_s, + "uncached unknown → all-healthy min-load (short worker at load 0)" + ); + } + + // --- Step 5: tree recording (cache hit after pool routing) --- + + #[tokio::test] + async fn cal_step5_pool_routing_records_tree_for_future_hit() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let prompt = "novel pool-routed prompt that will be recorded"; + // First request: novel → pool split (short, since <100K by char estimate) + let first = route_to_url(&router, prompt, None); + // Second request: same prompt → now a cache hit → same worker. + let second = route_to_url(&router, prompt, None); + assert_eq!(first, second, "pool routing recorded tree → subsequent hit"); + } + + // --- Additional scenarios --- + + // Step 3: cache hit but matched worker unhealthy → clean stale + first healthy + #[tokio::test] + async fn cal_step3_hit_unhealthy_falls_back_to_first_healthy() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let prompt = "shared cache-building prompt for affinity test"; + // First request: seed the cache + let first = route_to_url(&router, prompt, None); + // Mark the selected worker unhealthy + mark_unhealthy(&router, &first); + // Second request: same prompt hits cache, but matched worker is + // unhealthy → fall back to another healthy worker + let second = route_to_url(&router, prompt, None); + assert_ne!( + second, first, + "unhealthy matched worker must not be selected" + ); + } + + // Step 4 ≥100K: long pool all unhealthy + short pool has load=0 worker → long→short overflow to idle worker + #[tokio::test] + async fn cal_step4_long_pool_unhealthy_overflows_to_idle_short() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + mark_unhealthy(&router, &url_l); // long pool all unhealthy + // short pool worker idle (load 0) + let h = tokens_header(200_000); + let routed = route_to_url(&router, "novel", Some(&h)); + assert_eq!(routed, url_s, "long pool all unhealthy → idle short worker"); + } + + // Step 4 token source priority: header exact value overrides char estimate + #[tokio::test] + async fn cal_step4_header_overrides_char_estimate() { + // Short prompt (char estimate → 1 token → short request → short pool), + // but header says 200000 (→ long request → long pool). + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let h = tokens_header(200_000); + let routed = route_to_url(&router, "short", Some(&h)); + assert_eq!( + routed, url_l, + "header (200K) must override char estimate (1 token) → long pool" + ); + } + + // --- Full HTTP path via route_typed_request --- + + /// Construct a minimal CompletionRequest for routing tests. + fn completion_request(prompt: &str) -> CompletionRequest { + serde_json::from_value(serde_json::json!({ + "model": "", + "prompt": prompt, + "stream": false, + })) + .unwrap() + } + + /// E2E via route_typed_request: a header-classified long request routes + /// to the long-pool worker, and the response carries x-smg-routed-worker-id. + #[tokio::test] + async fn cal_http_long_request_routes_to_long_pool() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + let headers = tokens_header(200_000); + let resp = router + .route_typed_request( + Some(&headers), + completion_request("novel prompt"), + "/generate", + crate::worker::UNKNOWN_MODEL_ID, + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + let routed = resp + .headers() + .get("x-smg-routed-worker-id") + .unwrap() + .to_str() + .unwrap() + .to_string(); + assert_eq!(routed, url_l, "header-classified long request → long pool"); + } + + /// E2E via route_typed_request: an all-unhealthy fleet returns 503. + #[tokio::test] + async fn cal_http_all_unhealthy_returns_503() { + let (url_s, _cap_s) = spawn_capture_stub("application/json", "{}").await; + let (url_l, _cap_l) = spawn_capture_stub("application/json", "{}").await; + let router = length_router(&[&url_s], &[&url_l]).await; + mark_unhealthy(&router, &url_s); + mark_unhealthy(&router, &url_l); + let resp = router + .route_typed_request( + None, + completion_request("hello"), + "/generate", + crate::worker::UNKNOWN_MODEL_ID, + ) + .await; + assert_eq!( + resp.status(), + StatusCode::SERVICE_UNAVAILABLE, + "all unhealthy → 503" + ); + } } diff --git a/model_gateway/src/workflow/job_queue.rs b/model_gateway/src/workflow/job_queue.rs index 14a74ade3..c4039167e 100644 --- a/model_gateway/src/workflow/job_queue.rs +++ b/model_gateway/src/workflow/job_queue.rs @@ -493,23 +493,29 @@ impl JobQueue { let api_key = router_config.api_key.clone(); let mut worker_count = 0; - // Create iterator of (url, worker_type, bootstrap_port) tuples based on mode - let workers: Vec<(String, &str, Option)> = match &router_config.mode { + // Create iterator of (url, worker_type, bootstrap_port, is_long_pool) + // tuples. The 4th element marks whether the worker gets pool=long + // label, determined by matching the prefill index against + // router_config.long_prefill_indices. + let long_indices = &router_config.long_prefill_indices; + let workers: Vec<(String, &str, Option, bool)> = match &router_config.mode { RoutingMode::Regular { worker_urls } => worker_urls .iter() - .map(|url| (url.clone(), "regular", None)) + .map(|url| (url.clone(), "regular", None, false)) .collect(), RoutingMode::PrefillDecode { prefill_urls, decode_urls, .. } => { - let prefill_workers = prefill_urls - .iter() - .map(|(url, port)| (url.clone(), "prefill", *port)); + let prefill_workers = + prefill_urls.iter().enumerate().map(|(i, (url, port))| { + (url.clone(), "prefill", *port, is_long_pool_index(i, long_indices)) + }); - let decode_workers = - decode_urls.iter().map(|url| (url.clone(), "decode", None)); + let decode_workers = decode_urls + .iter() + .map(|url| (url.clone(), "decode", None, false)); prefill_workers.chain(decode_workers).collect() } @@ -521,12 +527,14 @@ impl JobQueue { } => { let encode_workers = encode_urls .iter() - .map(|(url, port)| (url.clone(), "encode", *port)); - let prefill_workers = prefill_urls + .map(|(url, port)| (url.clone(), "encode", *port, false)); + let prefill_workers = + prefill_urls.iter().enumerate().map(|(i, (url, port))| { + (url.clone(), "prefill", *port, is_long_pool_index(i, long_indices)) + }); + let decode_workers = decode_urls .iter() - .map(|(url, port)| (url.clone(), "prefill", *port)); - let decode_workers = - decode_urls.iter().map(|url| (url.clone(), "decode", None)); + .map(|url| (url.clone(), "decode", None, false)); encode_workers .chain(prefill_workers) @@ -553,7 +561,7 @@ impl JobQueue { ); // Process all workers with unified loop - for (url, worker_type, bootstrap_port) in workers { + for (url, worker_type, bootstrap_port, is_long_pool) in workers { let url_for_error = url.clone(); // Clone for error message let proto_worker_type = match worker_type { "prefill" => WorkerType::Prefill, @@ -565,6 +573,9 @@ impl JobQueue { spec.worker_type = proto_worker_type; spec.api_key.clone_from(&api_key); spec.bootstrap_port = bootstrap_port; + if is_long_pool { + spec.labels.insert("pool".to_string(), "long".to_string()); + } apply_startup_worker_config(&mut spec, router_config); let config = spec; @@ -756,6 +767,12 @@ impl JobQueue { } } +/// Whether prefill worker at `index` belongs to the long pool, determined by +/// membership in `long_prefill_indices`. +fn is_long_pool_index(index: usize, long_indices: &[usize]) -> bool { + long_indices.contains(&index) +} + /// Stamp the router-config-derived fields onto a startup worker's spec: the /// pinned runtime, the grouped-ZMQ engine count, and the connection budget. /// Identity fields (url, worker type, api key, bootstrap port) are the @@ -967,4 +984,77 @@ mod tests { 30 ); } + + /// `--long-prefill-indices 3,4` marks the 4th and 5th prefill workers with + /// `pool=long`; other prefill workers and all decode workers stay unlabelled. + #[test] + fn long_prefill_indices_tags_correct_workers() { + let prefill_urls: Vec<(String, Option)> = vec![ + ("http://p1:8000".into(), None), + ("http://p2:8000".into(), None), + ("http://p3:8000".into(), None), + ("http://p4:8000".into(), None), + ("http://p5:8000".into(), None), + ]; + let long_indices = vec![3, 4]; + + // Simulate the worker-tuple construction from InitializeWorkersFromConfig + let workers: Vec<(String, &str, Option, bool)> = prefill_urls + .iter() + .enumerate() + .map(|(i, (url, port))| { + (url.clone(), "prefill", *port, is_long_pool_index(i, &long_indices)) + }) + .collect(); + + // P1 (index 0), P2 (index 1), P3 (index 2) → short pool (is_long_pool=false) + assert!(!workers[0].3, "P1 should be short pool"); + assert!(!workers[1].3, "P2 should be short pool"); + assert!(!workers[2].3, "P3 should be short pool"); + // P4 (index 3), P5 (index 4) → long pool (is_long_pool=true) + assert!(workers[3].3, "P4 should be long pool"); + assert!(workers[4].3, "P5 should be long pool"); + + // Verify the label is actually applied to the WorkerSpec + for (_, _, _, is_long_pool) in &workers { + let mut spec = WorkerSpec::new("http://test:8000"); + if *is_long_pool { + spec.labels.insert("pool".to_string(), "long".to_string()); + } + if *is_long_pool { + assert_eq!( + spec.labels.get("pool").map(|s| s.as_str()), + Some("long"), + "long-pool worker must have pool=long label" + ); + } else { + assert!( + spec.labels.get("pool").is_none(), + "short-pool worker must NOT have pool label" + ); + } + } + } + + /// Empty `long_prefill_indices` means no long pool — all prefill workers + /// are short pool. + #[test] + fn empty_long_prefill_indices_means_all_short() { + let prefill_urls: Vec<(String, Option)> = vec![ + ("http://p1:8000".into(), None), + ("http://p2:8000".into(), None), + ]; + let long_indices: Vec = vec![]; + + let workers: Vec<(String, &str, Option, bool)> = prefill_urls + .iter() + .enumerate() + .map(|(i, (url, port))| { + (url.clone(), "prefill", *port, is_long_pool_index(i, &long_indices)) + }) + .collect(); + + assert!(!workers[0].3, "no long indices → all short pool"); + assert!(!workers[1].3, "no long indices → all short pool"); + } } diff --git a/model_gateway/src/workflow/steps/local/update_policies_for_worker.rs b/model_gateway/src/workflow/steps/local/update_policies_for_worker.rs index 99beccc40..291ddf105 100644 --- a/model_gateway/src/workflow/steps/local/update_policies_for_worker.rs +++ b/model_gateway/src/workflow/steps/local/update_policies_for_worker.rs @@ -45,14 +45,18 @@ impl StepExecutor for UpdatePoliciesForWorkerStep { let workers = app_context.worker_registry.get_by_model(model_id); if let Some(policy) = app_context.policy_registry.get_policy(model_id) { - if policy.name() == "cache_aware" && !workers.is_empty() { - // Re-initialize cache-aware policy with updated workers + if (policy.name() == "cache_aware" || policy.name() == "cache_aware_length") + && !workers.is_empty() + { + // Re-initialize cache-aware (or cache_aware_length) policy + // with updated workers. app_context .policy_registry .init_cache_aware_policy(model_id, &workers); debug!( - "Updated cache-aware policy for model {} ({} workers)", + "Updated {} policy for model {} ({} workers)", + policy.name(), model_id, workers.len() ); diff --git a/model_gateway/src/workflow/steps/local/update_remaining_policies.rs b/model_gateway/src/workflow/steps/local/update_remaining_policies.rs index b6188afbe..ee1f50989 100644 --- a/model_gateway/src/workflow/steps/local/update_remaining_policies.rs +++ b/model_gateway/src/workflow/steps/local/update_remaining_policies.rs @@ -35,13 +35,16 @@ impl StepExecutor for UpdateRemainingPoliciesStep { let remaining_workers = app_context.worker_registry.get_by_model(model_id); if let Some(policy) = app_context.policy_registry.get_policy(model_id) { - if policy.name() == "cache_aware" && !remaining_workers.is_empty() { + if (policy.name() == "cache_aware" || policy.name() == "cache_aware_length") + && !remaining_workers.is_empty() + { app_context .policy_registry .init_cache_aware_policy(model_id, &remaining_workers); debug!( - "Updated cache-aware policy for model {} ({} remaining workers)", + "Updated {} policy for model {} ({} remaining workers)", + policy.name(), model_id, remaining_workers.len() ); diff --git a/model_gateway/src/workflow/steps/shared/update_policies.rs b/model_gateway/src/workflow/steps/shared/update_policies.rs index 4e8681db2..23f53c173 100644 --- a/model_gateway/src/workflow/steps/shared/update_policies.rs +++ b/model_gateway/src/workflow/steps/shared/update_policies.rs @@ -147,14 +147,19 @@ impl StepExecutor for UpdatePolicie .policy_registry .get_policy(&model_id) .is_some_and(|policy| policy.name() == "cache_aware"); - if cache_aware { + let cache_aware_length = app_context + .policy_registry + .get_policy(&model_id) + .is_some_and(|policy| policy.name() == "cache_aware_length"); + if cache_aware || cache_aware_length { app_context .policy_registry .init_cache_aware_policy(&model_id, &all_workers); } - // Start KV event subscription for gRPC workers with cache_aware policy - if cache_aware { + // Start KV event subscription for gRPC workers with cache_aware or + // cache_aware_length policy (both are KV-event-capable). + if cache_aware || cache_aware_length { if let Some(ref monitor) = app_context.kv_event_monitor { if *worker.connection_mode() == ConnectionMode::Grpc { monitor.on_worker_added(worker).await; @@ -162,7 +167,11 @@ impl StepExecutor for UpdatePolicie } } - Self::warn_on_cache_aware_without_kv_events(&model_id, worker, cache_aware); + Self::warn_on_cache_aware_without_kv_events( + &model_id, + worker, + cache_aware || cache_aware_length, + ); if !updated_models.contains(&model_id) { updated_models.push(model_id);