From 00562ee14ac125c5a9074c7cb739b1aff705adb2 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 2 Jan 2026 08:52:01 +0800 Subject: [PATCH] Support cache eviction for Manual Policy (#16263) --- .../python/sglang_router/router_args.py | 9 +- sgl-model-gateway/bindings/python/src/lib.rs | 9 +- sgl-model-gateway/src/config/types.rs | 20 +- sgl-model-gateway/src/config/validation.rs | 2 +- sgl-model-gateway/src/main.rs | 9 +- sgl-model-gateway/src/policies/factory.rs | 20 +- sgl-model-gateway/src/policies/manual.rs | 240 ++++++++++++++---- sgl-model-gateway/src/policies/mod.rs | 2 +- 8 files changed, 246 insertions(+), 65 deletions(-) diff --git a/sgl-model-gateway/bindings/python/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/sglang_router/router_args.py index dab8b50cd..d846d46c8 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router_args.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router_args.py @@ -33,8 +33,9 @@ class RouterArgs: cache_threshold: float = 0.3 balance_abs_threshold: int = 64 balance_rel_threshold: float = 1.5 - eviction_interval_secs: int = 120 + eviction_interval_secs: int = 60 max_tree_size: int = 2**26 + max_idle_secs: int = 4 * 3600 max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches bucket_adjust_interval_secs: int = 5 dp_aware: bool = False @@ -297,6 +298,12 @@ class RouterArgs: default=RouterArgs.max_tree_size, help="Maximum size of the approximation tree for cache-aware routing", ) + routing_group.add_argument( + f"--{prefix}max-idle-secs", + type=int, + default=RouterArgs.max_idle_secs, + help="Maximum idle time in seconds before eviction (for manual policy)", + ) routing_group.add_argument( f"--{prefix}max-payload-size", type=int, diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index e66138ab3..e1cfbc67f 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -312,6 +312,7 @@ struct Router { balance_rel_threshold: f32, eviction_interval_secs: u64, max_tree_size: usize, + max_idle_secs: u64, max_payload_size: usize, dp_aware: bool, api_key: Option, @@ -417,7 +418,10 @@ impl Router { balance_rel_threshold: self.balance_rel_threshold, bucket_adjust_interval_secs: self.bucket_adjust_interval_secs, }, - PolicyType::Manual => ConfigPolicyConfig::Manual, + PolicyType::Manual => ConfigPolicyConfig::Manual { + eviction_interval_secs: self.eviction_interval_secs, + max_idle_secs: self.max_idle_secs, + }, PolicyType::ConsistentHashing => ConfigPolicyConfig::ConsistentHashing, PolicyType::PrefixHash => ConfigPolicyConfig::PrefixHash { prefix_token_count: 256, @@ -589,6 +593,7 @@ impl Router { balance_rel_threshold = 1.5, eviction_interval_secs = 120, max_tree_size = 2usize.pow(26), + max_idle_secs = 14400, max_payload_size = 512 * 1024 * 1024, dp_aware = false, api_key = None, @@ -671,6 +676,7 @@ impl Router { balance_rel_threshold: f32, eviction_interval_secs: u64, max_tree_size: usize, + max_idle_secs: u64, max_payload_size: usize, dp_aware: bool, api_key: Option, @@ -766,6 +772,7 @@ impl Router { balance_rel_threshold, eviction_interval_secs, max_tree_size, + max_idle_secs, max_payload_size, dp_aware, api_key, diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index cbca5f614..25894b530 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -341,8 +341,16 @@ pub enum PolicyConfig { /// - X-SMG-Routing-Key: Routes to a cached worker or assigns a new one /// - Provides true sticky sessions with zero key redistribution on worker add /// - Falls back to random selection if no routing key is provided + /// - Supports LRU eviction when cache size exceeds max_entries #[serde(rename = "manual")] - Manual, + Manual { + /// Interval between TTL eviction cycles (seconds, default: 60) + #[serde(default = "default_manual_eviction_interval_secs")] + eviction_interval_secs: u64, + /// Maximum idle time before eviction (seconds, default: 14400 = 4 hours) + #[serde(default = "default_manual_max_idle_secs")] + max_idle_secs: u64, + }, /// Consistent hashing policy using hash ring for session affinity: /// - X-SMG-Target-Worker: Direct routing to a specific worker by URL @@ -376,6 +384,14 @@ fn default_load_factor() -> f64 { 1.25 } +fn default_manual_eviction_interval_secs() -> u64 { + 60 +} + +fn default_manual_max_idle_secs() -> u64 { + 4 * 3600 +} + impl PolicyConfig { pub fn name(&self) -> &'static str { match self { @@ -384,7 +400,7 @@ impl PolicyConfig { PolicyConfig::CacheAware { .. } => "cache_aware", PolicyConfig::PowerOfTwo { .. } => "power_of_two", PolicyConfig::Bucket { .. } => "bucket", - PolicyConfig::Manual => "manual", + PolicyConfig::Manual { .. } => "manual", PolicyConfig::ConsistentHashing => "consistent_hashing", PolicyConfig::PrefixHash { .. } => "prefix_hash", } diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index 53c04c4d1..36eee6209 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/sgl-model-gateway/src/config/validation.rs @@ -149,7 +149,7 @@ impl ConfigValidator { match policy { PolicyConfig::Random | PolicyConfig::RoundRobin - | PolicyConfig::Manual + | PolicyConfig::Manual { .. } | PolicyConfig::ConsistentHashing => {} PolicyConfig::CacheAware { cache_threshold, diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index 53e5b2ebd..f99907abe 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -166,6 +166,10 @@ struct CliArgs { #[arg(long, default_value_t = 67108864, help_heading = "Routing Policy")] max_tree_size: usize, + /// Maximum idle time in seconds before eviction (for manual policy) + #[arg(long, default_value_t = 14400, help_heading = "Routing Policy")] + max_idle_secs: u64, + /// Number of prefix tokens to use for prefix_hash policy #[arg(long, default_value_t = 256, help_heading = "Routing Policy")] prefix_token_count: usize, @@ -688,7 +692,10 @@ impl CliArgs { prefix_token_count: self.prefix_token_count, load_factor: self.prefix_hash_load_factor, }, - "manual" => PolicyConfig::Manual, + "manual" => PolicyConfig::Manual { + eviction_interval_secs: self.eviction_interval, + max_idle_secs: self.max_idle_secs, + }, _ => PolicyConfig::RoundRobin, } } diff --git a/sgl-model-gateway/src/policies/factory.rs b/sgl-model-gateway/src/policies/factory.rs index 2562303d0..fddb2bc95 100644 --- a/sgl-model-gateway/src/policies/factory.rs +++ b/sgl-model-gateway/src/policies/factory.rs @@ -4,8 +4,8 @@ use std::sync::Arc; use super::{ BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy, - LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, PrefixHashConfig, PrefixHashPolicy, - RandomPolicy, RoundRobinPolicy, + LoadBalancingPolicy, ManualConfig, ManualPolicy, PowerOfTwoPolicy, PrefixHashConfig, + PrefixHashPolicy, RandomPolicy, RoundRobinPolicy, }; use crate::config::PolicyConfig; @@ -47,7 +47,16 @@ impl PolicyFactory { }; Arc::new(BucketPolicy::with_config(config)) } - PolicyConfig::Manual => Arc::new(ManualPolicy::new()), + PolicyConfig::Manual { + eviction_interval_secs, + max_idle_secs, + } => { + let config = ManualConfig { + eviction_interval_secs: *eviction_interval_secs, + max_idle_secs: *max_idle_secs, + }; + Arc::new(ManualPolicy::with_config(config)) + } PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()), PolicyConfig::PrefixHash { prefix_token_count, @@ -113,7 +122,10 @@ mod tests { }); assert_eq!(policy.name(), "bucket"); - let policy = PolicyFactory::create_from_config(&PolicyConfig::Manual); + let policy = PolicyFactory::create_from_config(&PolicyConfig::Manual { + eviction_interval_secs: 60, + max_idle_secs: 4 * 3600, + }); assert_eq!(policy.name(), "manual"); let policy = PolicyFactory::create_from_config(&PolicyConfig::ConsistentHashing); diff --git a/sgl-model-gateway/src/policies/manual.rs b/sgl-model-gateway/src/policies/manual.rs index 55d3c6cb2..aa1b0cbb7 100644 --- a/sgl-model-gateway/src/policies/manual.rs +++ b/sgl-model-gateway/src/policies/manual.rs @@ -13,13 +13,16 @@ //! ## Header //! - `X-SMG-Routing-Key`: The routing key for sticky session routing -use std::sync::Arc; +use std::{sync::Arc, time::Instant}; use dashmap::{mapref::entry::Entry, DashMap}; use http::header::HeaderName; use rand::Rng; +use tracing::info; -use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo}; +use super::{ + get_healthy_worker_indices, utils::PeriodicTask, LoadBalancingPolicy, SelectWorkerInfo, +}; use crate::{core::Worker, observability::metrics::Metrics}; /// Header for routing key based sticky sessions @@ -28,10 +31,9 @@ static HEADER_ROUTING_KEY: HeaderName = HeaderName::from_static("x-smg-routing-k #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ExecutionBranch { NoHealthyWorkers, - FastPathHit, - SlowPathOccupiedHit, - SlowPathOccupiedMiss, - SlowPathVacant, + OccupiedHit, + OccupiedMiss, + Vacant, NoRoutingId, } @@ -39,10 +41,9 @@ impl ExecutionBranch { fn as_str(&self) -> &'static str { match self { Self::NoHealthyWorkers => "no_healthy_workers", - Self::FastPathHit => "fast_path_hit", - Self::SlowPathOccupiedHit => "slow_path_occupied_hit", - Self::SlowPathOccupiedMiss => "slow_path_occupied_miss", - Self::SlowPathVacant => "slow_path_vacant", + Self::OccupiedHit => "occupied_hit", + Self::OccupiedMiss => "occupied_miss", + Self::Vacant => "vacant", Self::NoRoutingId => "no_routing_id", } } @@ -60,11 +61,27 @@ impl RoutingId { const MAX_CANDIDATE_WORKERS: usize = 2; #[derive(Debug, Clone)] -struct RoutingInfo { - candi_worker_urls: Vec, +pub struct ManualConfig { + pub eviction_interval_secs: u64, + pub max_idle_secs: u64, } -impl RoutingInfo { +impl Default for ManualConfig { + fn default() -> Self { + Self { + eviction_interval_secs: 60, + max_idle_secs: 4 * 3600, + } + } +} + +#[derive(Debug, Clone)] +struct Node { + candi_worker_urls: Vec, + last_access: Instant, +} + +impl Node { fn push_bounded(&mut self, url: String) { while self.candi_worker_urls.len() >= MAX_CANDIDATE_WORKERS { self.candi_worker_urls.remove(0); @@ -74,16 +91,60 @@ impl RoutingInfo { } // TODO may optimize performance -// TODO evict old data periodically -#[derive(Debug, Default)] +#[derive(Debug)] pub struct ManualPolicy { - routing_map: DashMap, + routing_map: Arc>, + _eviction_task: Option, +} + +impl Default for ManualPolicy { + fn default() -> Self { + Self::new() + } } impl ManualPolicy { pub fn new() -> Self { + Self::with_config(ManualConfig::default()) + } + + pub fn with_config(config: ManualConfig) -> Self { + use std::time::Duration; + + let routing_map = Arc::new(DashMap::::new()); + + let eviction_task = if config.eviction_interval_secs > 0 && config.max_idle_secs > 0 { + let routing_map_clone = Arc::clone(&routing_map); + let max_idle = Duration::from_secs(config.max_idle_secs); + + Some(PeriodicTask::spawn( + config.eviction_interval_secs, + "ManualPolicyEviction", + move || { + let now = Instant::now(); + let before_size = routing_map_clone.len(); + + routing_map_clone + .retain(|_, node| now.duration_since(node.last_access) < max_idle); + + let evicted_count = before_size - routing_map_clone.len(); + if evicted_count > 0 { + info!( + "ManualPolicy TTL eviction: evicted {} entries, remaining {} (max_idle: {}s)", + evicted_count, + routing_map_clone.len(), + max_idle.as_secs() + ); + } + }, + )) + } else { + None + }; + Self { - routing_map: DashMap::new(), + routing_map, + _eviction_task: eviction_task, } } @@ -95,35 +156,27 @@ impl ManualPolicy { ) -> (usize, ExecutionBranch) { let routing_id = RoutingId::new(routing_id); - // Fast path - if let Some(info) = self.routing_map.get(&routing_id) { - if let Some(idx) = - find_healthy_worker(&info.candi_worker_urls, workers, healthy_indices) - { - return (idx, ExecutionBranch::FastPathHit); - } - } - - // Slow path match self.routing_map.entry(routing_id) { Entry::Occupied(mut entry) => { + let node = entry.get_mut(); + node.last_access = Instant::now(); if let Some(idx) = - find_healthy_worker(&entry.get().candi_worker_urls, workers, healthy_indices) + find_healthy_worker(&node.candi_worker_urls, workers, healthy_indices) { - return (idx, ExecutionBranch::SlowPathOccupiedHit); + (idx, ExecutionBranch::OccupiedHit) + } else { + let selected_idx = random_select(healthy_indices); + node.push_bounded(workers[selected_idx].url().to_string()); + (selected_idx, ExecutionBranch::OccupiedMiss) } - let selected_idx = random_select(healthy_indices); - entry - .get_mut() - .push_bounded(workers[selected_idx].url().to_string()); - (selected_idx, ExecutionBranch::SlowPathOccupiedMiss) } Entry::Vacant(entry) => { let selected_idx = random_select(healthy_indices); - entry.insert(RoutingInfo { + entry.insert(Node { candi_worker_urls: vec![workers[selected_idx].url().to_string()], + last_access: Instant::now(), }); - (selected_idx, ExecutionBranch::SlowPathVacant) + (selected_idx, ExecutionBranch::Vacant) } } } @@ -238,7 +291,7 @@ mod tests { let (first_result, branch) = policy.select_worker_impl(&workers, &info); let first_idx = first_result.unwrap(); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + assert_eq!(branch, ExecutionBranch::Vacant); for _ in 0..10 { let (result, branch) = policy.select_worker_impl(&workers, &info); @@ -247,7 +300,7 @@ mod tests { Some(first_idx), "Same routing_id should route to same worker" ); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(branch, ExecutionBranch::OccupiedHit); } } @@ -264,7 +317,7 @@ mod tests { ..Default::default() }; let (result, branch) = policy.select_worker_impl(&workers, &info); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + assert_eq!(branch, ExecutionBranch::Vacant); *distribution.entry(result.unwrap()).or_insert(0) += 1; } @@ -307,12 +360,12 @@ mod tests { let (result, branch) = policy.select_worker_impl(&workers, &info); assert_eq!(result, Some(1), "Should only select healthy worker"); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + assert_eq!(branch, ExecutionBranch::Vacant); for _ in 0..10 { let (result, branch) = policy.select_worker_impl(&workers, &info); assert_eq!(result, Some(1), "Should only select healthy worker"); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(branch, ExecutionBranch::OccupiedHit); } } @@ -371,14 +424,14 @@ mod tests { let (first_result, branch) = policy.select_worker_impl(&workers, &info); let first_idx = first_result.unwrap(); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + assert_eq!(branch, ExecutionBranch::Vacant); workers[first_idx].set_healthy(false); let (new_result, branch) = policy.select_worker_impl(&workers, &info); let new_idx = new_result.unwrap(); assert_ne!(new_idx, first_idx, "Should remap to healthy worker"); - assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + assert_eq!(branch, ExecutionBranch::OccupiedMiss); for _ in 0..10 { let (result, branch) = policy.select_worker_impl(&workers, &info); @@ -387,7 +440,7 @@ mod tests { Some(new_idx), "Should consistently route to new worker" ); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(branch, ExecutionBranch::OccupiedHit); } } @@ -418,12 +471,12 @@ mod tests { let (result, branch) = policy.select_worker_impl(&workers, &info); assert_eq!(result, Some(0)); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + assert_eq!(branch, ExecutionBranch::Vacant); for _ in 0..10 { let (result, branch) = policy.select_worker_impl(&workers, &info); assert_eq!(result, Some(0)); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(branch, ExecutionBranch::OccupiedHit); } } @@ -440,14 +493,14 @@ mod tests { let (first_result, branch) = policy.select_worker_impl(&workers, &info); let first_idx = first_result.unwrap(); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + assert_eq!(branch, ExecutionBranch::Vacant); workers[first_idx].set_healthy(false); let (second_result, branch) = policy.select_worker_impl(&workers, &info); let second_idx = second_result.unwrap(); assert_ne!(second_idx, first_idx); - assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + assert_eq!(branch, ExecutionBranch::OccupiedMiss); workers[first_idx].set_healthy(true); @@ -457,7 +510,7 @@ mod tests { Some(first_idx), "Should return to original worker after recovery since it's first in candidate list" ); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(branch, ExecutionBranch::OccupiedHit); } #[test] @@ -473,14 +526,14 @@ mod tests { let (first_result, branch) = policy.select_worker_impl(&workers, &info); let first_idx = first_result.unwrap(); - assert_eq!(branch, ExecutionBranch::SlowPathVacant); + assert_eq!(branch, ExecutionBranch::Vacant); workers[first_idx].set_healthy(false); let (second_result, branch) = policy.select_worker_impl(&workers, &info); let second_idx = second_result.unwrap(); assert_ne!(second_idx, first_idx); - assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + assert_eq!(branch, ExecutionBranch::OccupiedMiss); workers[second_idx].set_healthy(false); @@ -491,7 +544,7 @@ mod tests { Some(remaining_idx), "Should select the only remaining healthy worker" ); - assert_eq!(branch, ExecutionBranch::SlowPathOccupiedMiss); + assert_eq!(branch, ExecutionBranch::OccupiedMiss); workers[first_idx].set_healthy(true); @@ -501,7 +554,7 @@ mod tests { Some(first_idx), "First worker should be evicted from candidates due to MAX_CANDIDATE_WORKERS=2" ); - assert_eq!(branch, ExecutionBranch::FastPathHit); + assert_eq!(branch, ExecutionBranch::OccupiedHit); } #[test] @@ -512,8 +565,9 @@ mod tests { #[test] fn test_manual_routing_info_push_bounded() { - let mut info = RoutingInfo { + let mut info = Node { candi_worker_urls: vec!["http://w1:8000".to_string()], + last_access: Instant::now(), }; info.push_bounded("http://w2:8000".to_string()); @@ -582,4 +636,82 @@ mod tests { "Should return None for unknown URL" ); } + + #[test] + fn test_manual_config_default() { + let config = ManualConfig::default(); + assert_eq!(config.eviction_interval_secs, 60); + assert_eq!(config.max_idle_secs, 4 * 3600); + } + + #[test] + fn test_manual_with_disabled_eviction() { + let config = ManualConfig { + eviction_interval_secs: 0, + max_idle_secs: 3600, + }; + let policy = ManualPolicy::with_config(config); + assert!(policy._eviction_task.is_none()); + } + + #[test] + fn test_manual_last_access_updates() { + let policy = ManualPolicy::new(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + let headers = headers_with_routing_key("test-key"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + let routing_id = RoutingId::new("test-key"); + + // Vacant: first access + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(branch, ExecutionBranch::Vacant); + let first_idx = result.unwrap(); + let access_after_vacant = policy.routing_map.get(&routing_id).unwrap().last_access; + assert!(access_after_vacant.elapsed().as_millis() < 100); + + std::thread::sleep(std::time::Duration::from_millis(10)); + + // OccupiedHit: same worker still healthy + let (_, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(branch, ExecutionBranch::OccupiedHit); + let access_after_hit = policy.routing_map.get(&routing_id).unwrap().last_access; + assert!(access_after_hit > access_after_vacant); + + std::thread::sleep(std::time::Duration::from_millis(10)); + + // OccupiedMiss: worker becomes unhealthy + workers[first_idx].set_healthy(false); + let (_, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(branch, ExecutionBranch::OccupiedMiss); + let access_after_miss = policy.routing_map.get(&routing_id).unwrap().last_access; + assert!(access_after_miss > access_after_hit); + } + + #[test] + fn test_manual_ttl_eviction_logic() { + use std::time::Duration; + + let config = ManualConfig { + eviction_interval_secs: 2, + max_idle_secs: 2, + }; + let policy = ManualPolicy::with_config(config); + let workers = create_workers(&["http://w1:8000", "http://w2:8000"]); + + let headers = headers_with_routing_key("key-0"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + policy.select_worker_impl(&workers, &info); + + assert_eq!(policy.routing_map.len(), 1); + + std::thread::sleep(Duration::from_secs(4)); + + assert_eq!(policy.routing_map.len(), 0); + } } diff --git a/sgl-model-gateway/src/policies/mod.rs b/sgl-model-gateway/src/policies/mod.rs index 3ad4b47ae..3d4dd9a4f 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -23,7 +23,7 @@ pub use bucket::BucketPolicy; pub use cache_aware::CacheAwarePolicy; pub use consistent_hashing::ConsistentHashingPolicy; pub use factory::PolicyFactory; -pub use manual::ManualPolicy; +pub use manual::{ManualConfig, ManualPolicy}; pub use power_of_two::PowerOfTwoPolicy; pub use prefix_hash::{PrefixHashConfig, PrefixHashPolicy}; pub use random::RandomPolicy;