From 611a4fd08ba0435f9ab08863e71620140bafe969 Mon Sep 17 00:00:00 2001 From: syy-hw Date: Mon, 10 Nov 2025 18:02:53 +0800 Subject: [PATCH] [router] bucket policy (#11719) --- sgl-router/py_src/sglang_router/router.py | 1 + .../py_src/sglang_router/router_args.py | 9 +- sgl-router/src/config/types.rs | 33 + sgl-router/src/config/validation.rs | 96 ++ .../workflow/steps/worker_registration.rs | 7 + sgl-router/src/lib.rs | 11 + sgl-router/src/policies/bucket.rs | 1167 +++++++++++++++++ sgl-router/src/policies/factory.rs | 34 +- sgl-router/src/policies/mod.rs | 19 + sgl-router/src/policies/registry.rs | 34 +- sgl-router/src/routers/http/pd_types.rs | 5 + sgl-router/tests/test_pd_routing.rs | 28 + 12 files changed, 1435 insertions(+), 9 deletions(-) create mode 100644 sgl-router/src/policies/bucket.rs diff --git a/sgl-router/py_src/sglang_router/router.py b/sgl-router/py_src/sglang_router/router.py index 0c06f4a41..a71febfab 100644 --- a/sgl-router/py_src/sglang_router/router.py +++ b/sgl-router/py_src/sglang_router/router.py @@ -19,6 +19,7 @@ def policy_from_str(policy_str: Optional[str]) -> PolicyType: "round_robin": PolicyType.RoundRobin, "cache_aware": PolicyType.CacheAware, "power_of_two": PolicyType.PowerOfTwo, + "bucket": PolicyType.Bucket, } return policy_map[policy_str] diff --git a/sgl-router/py_src/sglang_router/router_args.py b/sgl-router/py_src/sglang_router/router_args.py index 813a0b2a1..bec64c8d9 100644 --- a/sgl-router/py_src/sglang_router/router_args.py +++ b/sgl-router/py_src/sglang_router/router_args.py @@ -34,6 +34,7 @@ class RouterArgs: eviction_interval_secs: int = 120 max_tree_size: int = 2**26 max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches + bucket_adjust_interval_secs: int = 5 dp_aware: bool = False enable_igw: bool = False # Enable IGW (Inter-Gateway) mode for multi-model support api_key: Optional[str] = None @@ -167,7 +168,7 @@ class RouterArgs: f"--{prefix}prefill-policy", type=str, default=None, - choices=["random", "round_robin", "cache_aware", "power_of_two"], + choices=["random", "round_robin", "cache_aware", "power_of_two", "bucket"], help="Specific policy for prefill nodes in PD mode. If not specified, uses the main policy", ) parser.add_argument( @@ -234,6 +235,12 @@ class RouterArgs: default=RouterArgs.balance_rel_threshold, help="Load balancing is triggered when (max_load - min_load) > abs_threshold AND max_load > min_load * rel_threshold. Otherwise, use cache aware", ) + parser.add_argument( + f"--{prefix}bucket-adjust-interval-secs", + type=int, + default=RouterArgs.bucket_adjust_interval_secs, + help="Interval in seconds between bucket boundary adjustment operations", + ) parser.add_argument( f"--{prefix}eviction-interval-secs", type=int, diff --git a/sgl-router/src/config/types.rs b/sgl-router/src/config/types.rs index d25c3106d..ae267d069 100644 --- a/sgl-router/src/config/types.rs +++ b/sgl-router/src/config/types.rs @@ -263,6 +263,16 @@ pub enum PolicyConfig { #[serde(rename = "power_of_two")] PowerOfTwo { load_check_interval_secs: u64 }, + + #[serde(rename = "bucket")] + Bucket { + /// Absolute load difference threshold for load balancing + balance_abs_threshold: usize, + /// Relative load ratio threshold for load balancing + balance_rel_threshold: f32, + /// Interval between bucket boundary adjustment cycles (seconds) + bucket_adjust_interval_secs: usize, + }, } impl PolicyConfig { @@ -272,6 +282,7 @@ impl PolicyConfig { PolicyConfig::RoundRobin => "round_robin", PolicyConfig::CacheAware { .. } => "cache_aware", PolicyConfig::PowerOfTwo { .. } => "power_of_two", + PolicyConfig::Bucket { .. } => "bucket", } } } @@ -728,6 +739,28 @@ mod tests { } } + #[test] + fn test_bucket_parameters() { + let bucket = PolicyConfig::Bucket { + balance_abs_threshold: 20, + balance_rel_threshold: 2.0, + bucket_adjust_interval_secs: 5, + }; + + match bucket { + PolicyConfig::Bucket { + balance_abs_threshold, + balance_rel_threshold, + bucket_adjust_interval_secs, + } => { + assert_eq!(balance_abs_threshold, 20); + assert!((balance_rel_threshold - 2.0).abs() < 0.0001); + assert_eq!(bucket_adjust_interval_secs, 5); + } + _ => panic!("Expected Bucket"), + } + } + #[test] fn test_discovery_config_default() { let config = DiscoveryConfig::default(); diff --git a/sgl-router/src/config/validation.rs b/sgl-router/src/config/validation.rs index 7e698e41f..b56a116c6 100644 --- a/sgl-router/src/config/validation.rs +++ b/sgl-router/src/config/validation.rs @@ -209,6 +209,34 @@ impl ConfigValidator { }); } } + PolicyConfig::Bucket { + balance_abs_threshold: _, + balance_rel_threshold, + bucket_adjust_interval_secs, + } => { + 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 *bucket_adjust_interval_secs < 1 { + return Err(ConfigError::InvalidValue { + field: "bucket_adjust_interval_secs".to_string(), + value: bucket_adjust_interval_secs.to_string(), + reason: "Must be >= 1s".to_string(), + }); + } + if *bucket_adjust_interval_secs >= 4294967296 { + return Err(ConfigError::InvalidValue { + field: "bucket_adjust_interval_secs".to_string(), + value: bucket_adjust_interval_secs.to_string(), + reason: "Must be < 4294967296s".to_string(), + }); + } + } } Ok(()) } @@ -505,6 +533,13 @@ impl ConfigValidator { }); } } + + // Check bucket for decode + if let Some(PolicyConfig::Bucket { .. }) = decode_policy { + return Err(ConfigError::IncompatibleConfig { + reason: "Decode policy should not be allowed to be bucket".to_string(), + }); + } } } @@ -792,6 +827,67 @@ mod tests { } } + #[test] + fn test_validate_pd_mode_bucket_policy_restrictions() { + let config = RouterConfig::new( + RoutingMode::PrefillDecode { + prefill_urls: vec![ + ("http://prefill1:8000".to_string(), None), + ("http://prefill2:8000".to_string(), None), + ], + decode_urls: vec![ + "http://decode1:8000".to_string(), + "http://decode2:8000".to_string(), + ], + prefill_policy: Some(PolicyConfig::Bucket { + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + bucket_adjust_interval_secs: 5, + }), + decode_policy: Some(PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }), + }, + PolicyConfig::Random, // Main policy as fallback + ); + + let result = ConfigValidator::validate(&config); + assert!( + result.is_ok(), + "Prefill policy should be allowed to be bucket" + ); + + let config = RouterConfig::new( + RoutingMode::PrefillDecode { + prefill_urls: vec![ + ("http://prefill1:8000".to_string(), None), + ("http://prefill2:8000".to_string(), None), + ], + decode_urls: vec![ + "http://decode1:8000".to_string(), + "http://decode2:8000".to_string(), + ], + prefill_policy: Some(PolicyConfig::Bucket { + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + bucket_adjust_interval_secs: 5, + }), + decode_policy: Some(PolicyConfig::Bucket { + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + bucket_adjust_interval_secs: 5, + }), + }, + PolicyConfig::Random, // Main policy as fallback + ); + + let result = ConfigValidator::validate(&config); + assert!( + result.is_err(), + "Decode policy should not be allowed to be bucket" + ); + } + #[test] fn test_validate_grpc_requires_tokenizer() { let mut config = RouterConfig::new( diff --git a/sgl-router/src/core/workflow/steps/worker_registration.rs b/sgl-router/src/core/workflow/steps/worker_registration.rs index 8c44ff798..f2b7a959e 100644 --- a/sgl-router/src/core/workflow/steps/worker_registration.rs +++ b/sgl-router/src/core/workflow/steps/worker_registration.rs @@ -756,6 +756,13 @@ impl StepExecutor for UpdatePoliciesStep { .init_cache_aware_policy(&model_id, &all_workers); } } + let prefill_workers = app_context.worker_registry.get_prefill_workers(); + let policy = app_context.policy_registry.get_prefill_policy(); + if policy.name() == "bucket" { + app_context + .policy_registry + .init_pd_bucket_policies(&prefill_workers); + } debug!( "Updated policies for worker {} (model: {})", diff --git a/sgl-router/src/lib.rs b/sgl-router/src/lib.rs index be6aac9db..b18e76749 100644 --- a/sgl-router/src/lib.rs +++ b/sgl-router/src/lib.rs @@ -29,6 +29,7 @@ pub enum PolicyType { RoundRobin, CacheAware, PowerOfTwo, + Bucket, } #[pyclass(eq)] @@ -169,6 +170,8 @@ struct Router { request_timeout_secs: u64, request_id_headers: Option>, pd_disaggregation: bool, + // Takes effect in PD mode and when policy = bucket + bucket_adjust_interval_secs: usize, prefill_urls: Option)>>, decode_urls: Option>, prefill_policy: Option, @@ -244,6 +247,11 @@ impl Router { PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo { load_check_interval_secs: 5, }, + PolicyType::Bucket => ConfigPolicyConfig::Bucket { + balance_abs_threshold: self.balance_abs_threshold, + balance_rel_threshold: self.balance_rel_threshold, + bucket_adjust_interval_secs: self.bucket_adjust_interval_secs, + }, } }; @@ -407,6 +415,7 @@ impl Router { request_timeout_secs = 1800, request_id_headers = None, pd_disaggregation = false, + bucket_adjust_interval_secs = 5, prefill_urls = None, decode_urls = None, prefill_policy = None, @@ -480,6 +489,7 @@ impl Router { request_timeout_secs: u64, request_id_headers: Option>, pd_disaggregation: bool, + bucket_adjust_interval_secs: usize, prefill_urls: Option)>>, decode_urls: Option>, prefill_policy: Option, @@ -566,6 +576,7 @@ impl Router { request_timeout_secs, request_id_headers, pd_disaggregation, + bucket_adjust_interval_secs, prefill_urls, decode_urls, prefill_policy, diff --git a/sgl-router/src/policies/bucket.rs b/sgl-router/src/policies/bucket.rs new file mode 100644 index 000000000..c491c5c95 --- /dev/null +++ b/sgl-router/src/policies/bucket.rs @@ -0,0 +1,1167 @@ +use std::{ + collections::{HashMap, HashSet, VecDeque}, + sync::{Arc, Mutex, RwLock}, + thread, + time::{Duration, SystemTime}, +}; + +use dashmap::DashMap; +use rand::Rng; +use tracing::{error, info, warn}; +use uuid::Uuid; + +use super::{get_healthy_worker_indices, BucketConfig, LoadBalancingPolicy}; +use crate::core::Worker; + +#[derive(Debug)] +pub struct BucketPolicy { + config: BucketConfig, + buckets: Arc>>>, + adjustment_handle: Option>, +} + +impl Default for BucketPolicy { + fn default() -> Self { + Self::new() + } +} + +impl Drop for BucketPolicy { + fn drop(&mut self) { + if let Some(handle) = self.adjustment_handle.take() { + drop(handle); + } + } +} + +impl BucketPolicy { + pub fn new() -> Self { + Self::with_config(BucketConfig::default()) + } + + pub fn with_config(config: BucketConfig) -> Self { + let buckets = Arc::new(DashMap::>>::new()); + + let adjustment_handle = { + let buckets_clone = Arc::clone(&buckets); + + let interval_secs = config.bucket_adjust_interval_secs; + + Some(thread::spawn(move || loop { + thread::sleep(Duration::from_secs(interval_secs as u64)); + + for bucket_ref in buckets_clone.iter() { + let model_id = bucket_ref.key(); + let bucket = bucket_ref.value(); + match bucket.write() { + Ok(mut bucket_guard) => { + bucket_guard.adjust_boundary(); + } + Err(e) => { + eprintln!( + "Failed to acquire write lock for bucket {}: {}", + model_id, e + ); + } + } + } + })) + }; + + Self { + config, + buckets, + adjustment_handle, + } + } + + pub fn init_prefill_worker_urls(&self, prefill_workers: &[Arc]) { + // Group workers by model + let mut model_workers: HashMap>> = HashMap::new(); + for worker in prefill_workers { + // Use "default" for unknown/empty model_ids for backward compatibility + let model_id = worker.model_id(); + let model_key = if model_id.is_empty() || model_id == "unknown" { + "default" + } else { + model_id + }; + model_workers + .entry(model_key.to_string()) + .or_default() + .push(worker); + } + // Initialize bucket for each model + for (model_key, model_workers) in model_workers { + let bucket = self + .buckets + .entry(model_key) + .or_insert_with(|| { + Arc::new(RwLock::new(Bucket::new( + self.config.bucket_adjust_interval_secs * 1000, + ))) + }) + .clone(); + + let worker_urls: Vec = model_workers + .iter() + .map(|worker| worker.url().to_string()) + .collect(); + + let lock_result = bucket.write(); + if let Ok(mut bucket_guard) = lock_result { + bucket_guard.init_prefill_worker_urls(worker_urls); + } else { + eprintln!("Failed to acquire write lock for bucket initialization"); + } + } + } + + pub fn add_prefill_url(&self, worker: &dyn Worker) { + let model_id = worker.model_id(); + let model_key = if model_id.is_empty() || model_id == "unknown" { + "default" + } else { + model_id + }; + let bucket = self + .buckets + .entry(model_key.to_string()) + .or_insert_with(|| { + Arc::new(RwLock::new(Bucket::new( + self.config.bucket_adjust_interval_secs * 1000, + ))) + }) + .clone(); + + let lock_result = bucket.write(); + if let Ok(mut bucket_guard) = lock_result { + let worker_url = worker.url().to_string(); + + let prefill_worker_urls_clone = { + let mut prefill_worker_urls = bucket_guard.prefill_worker_urls.lock().unwrap(); + if !prefill_worker_urls.contains(&worker_url) { + prefill_worker_urls.push(worker_url.clone()); + } + let cloned = prefill_worker_urls.clone(); + + let mut chars_per_url = bucket_guard.chars_per_url.lock().unwrap(); + chars_per_url.entry(worker_url.clone()).or_insert(0); + + cloned + }; + + bucket_guard.init_prefill_worker_urls(prefill_worker_urls_clone); + + info!( + "Added worker {} to bucket for model {}", + worker_url, model_key + ); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + + pub fn remove_prefill_url(&self, worker: &dyn Worker) { + let model_id = worker.model_id(); + let model_key = if model_id.is_empty() || model_id == "unknown" { + "default" + } else { + model_id + }; + + if let Some(bucket_entry) = self.buckets.get(model_key) { + let bucket = bucket_entry.value(); + let worker_url = worker.url().to_string(); + + let lock_result = bucket.write(); + if let Ok(mut bucket_guard) = lock_result { + let (updated_len, updated_urls) = { + let mut prefill_worker_urls = bucket_guard.prefill_worker_urls.lock().unwrap(); + prefill_worker_urls.retain(|u| u != &worker_url); + let len = prefill_worker_urls.len(); + let urls_clone = prefill_worker_urls.clone(); + + let mut chars_per_url = bucket_guard.chars_per_url.lock().unwrap(); + chars_per_url.remove(&worker_url); + + (len, urls_clone) + }; + + bucket_guard.bucket_cnt = updated_len; + + if updated_len > 0 { + bucket_guard.init_prefill_worker_urls(updated_urls); + } + + info!( + "Removed worker {} from bucket for model {} (remaining workers: {})", + worker_url, model_key, bucket_guard.bucket_cnt + ); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } else { + warn!( + "No bucket found for model {} when trying to remove worker", + model_key + ); + } + } +} + +impl LoadBalancingPolicy for BucketPolicy { + fn select_worker( + &self, + workers: &[Arc], + request_text: Option<&str>, + ) -> Option { + let healthy_indices = get_healthy_worker_indices(workers); + + if healthy_indices.is_empty() { + return None; + } + + let char_count = match request_text { + None => 0, + Some(text) => text.chars().count(), + }; + + // Determine the model for this set of workers (router pre-filters by model) + // All workers should be from the same model + let first_model = workers[healthy_indices[0]].model_id(); + let model_key = if first_model.is_empty() || first_model == "unknown" { + "default" + } else { + first_model + }; + + let bucket = self + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + let prefill_url = if let Some(bucket) = bucket { + let (choiced_url, chars_per_url_snapshot) = { + let buc = bucket.read().unwrap(); + let chars_per_url_snapshot = buc.chars_per_url.lock().unwrap().clone(); + let choiced_url = buc.find_boundary(char_count); + (choiced_url, chars_per_url_snapshot) + }; + let max_load = chars_per_url_snapshot.values().copied().max().unwrap_or(0); + let min_load = chars_per_url_snapshot.values().copied().min().unwrap_or(0); + let abs_diff = max_load.saturating_sub(min_load); + let rel_threshold = self.config.balance_rel_threshold * min_load as f32; + let is_imbalanced = + abs_diff > self.config.balance_abs_threshold && max_load as f32 > rel_threshold; + info!( + "Current PD instance status | is_imbalanced={}", + is_imbalanced + ); + + let mut rng = rand::rng(); + let prefill_url = if is_imbalanced { + info!("select prefill instance by Load Balance policy"); + let min_url = chars_per_url_snapshot + .iter() + .min_by_key(|(_, &chars)| chars) + .map(|(url, _)| url.clone()) + .unwrap_or_else(|| { + let idx = rng.random_range(0..healthy_indices.len()); + let url = workers[healthy_indices[idx]].url(); + warn!("No URL found, randomly selecting: {}", url); + url.to_string() + }); + min_url + } else { + info!("select prefill instance by Bucket policy"); + match choiced_url { + Some(url) if !url.is_empty() => url, + _ => { + let idx = rng.random_range(0..healthy_indices.len()); + let selected_url = workers[healthy_indices[idx]].url(); + warn!("Boundary not found, randomly selection: {}", selected_url); + selected_url.to_string() + } + } + }; + + { + let mut buc = bucket.write().unwrap(); + buc.post_process_request(char_count, prefill_url.clone()); + } + + prefill_url + } else { + warn!( + "No bucket found for model {}, randomly selecting healthy worker", + model_key + ); + let mut rng = rand::rng(); + let idx = rng.random_range(0..healthy_indices.len()); + let selected_worker = &workers[healthy_indices[idx]]; + let prefill_url = selected_worker.url().to_string(); + prefill_url + }; + + workers.iter().position(|w| w.url() == prefill_url) + } + + fn select_worker_pair( + &self, + prefill_workers: &[Arc], + decode_workers: &[Arc], + request_text: Option<&str>, + ) -> Option<(usize, usize)> { + let prefill_idx = self.select_worker(prefill_workers, request_text)?; + + let healthy_decode = get_healthy_worker_indices(decode_workers); + if healthy_decode.is_empty() { + return None; + } + + let mut rng = rand::rng(); + let decode_idx = rng.random_range(0..healthy_decode.len()); + + Some((prefill_idx, decode_idx)) + } + + fn name(&self) -> &'static str { + "bucket" + } + + fn needs_request_text(&self) -> bool { + true // Bucket policy needs request text + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[derive(Debug, Clone)] +pub struct Bucket { + l_max: usize, + bucket_cnt: usize, + pub prefill_worker_urls: Arc>>, + load_total: usize, + pub period: usize, + bucket_load: usize, + boundary: Vec, + request_list: VecDeque, + t_req_loads: HashMap, + pub chars_per_url: Arc>>, +} + +#[derive(Debug, Clone)] +pub struct SequencerRequest { + pub id: String, + pub char_cnt: usize, + pub timestamp: SystemTime, + pub prefill_worker_url: String, +} + +#[derive(Debug, Clone)] +pub struct Boundary { + pub url: String, + pub range: [usize; 2], +} + +impl Boundary { + pub fn new(url: String, range: [usize; 2]) -> Self { + Boundary { url, range } + } +} + +impl Bucket { + pub fn new(period: usize) -> Self { + let l_max = 4096; + + let bucket_cnt = 0; + + let load_total = 0; + let bucket_load = 0; + + let t_req_loads = HashMap::new(); + let request_list = VecDeque::new(); + + let initial_map = HashMap::new(); + + let boundary = Vec::new(); + + let prefill_worker_urls = Arc::new(Mutex::new(Vec::new())); + + Bucket { + l_max, + bucket_cnt, + prefill_worker_urls, + load_total, + period, + bucket_load, + boundary, + request_list, + t_req_loads, + chars_per_url: Arc::new(Mutex::new(initial_map)), + } + } + + pub fn init_prefill_worker_urls(&mut self, prefill_worker_urls: Vec) { + let bucket_cnt = prefill_worker_urls.len(); + self.bucket_cnt = bucket_cnt; + let mut urls_lock = self.prefill_worker_urls.lock().unwrap(); + *urls_lock = prefill_worker_urls.clone(); + + let mut chars_lock = self.chars_per_url.lock().unwrap(); + chars_lock.clear(); + + for url in prefill_worker_urls.iter() { + chars_lock.insert(url.clone(), 0); + } + + let worker_cnt = bucket_cnt; + let boundary = if worker_cnt == 0 { + Vec::new() + } else { + let gap = self.l_max / worker_cnt; + self.l_max = usize::MAX; + prefill_worker_urls + .iter() + .enumerate() + .map(|(i, url)| { + let min = i * gap; + let max = if i == worker_cnt - 1 { + self.l_max + } else { + (i + 1) * gap - 1 + }; + Boundary::new(url.clone(), [min, max]) + }) + .collect() + }; + + self.boundary = boundary; + info!("Init boundary:{:?}", self.boundary); + } + + pub fn post_process_request(&mut self, char_cnt: usize, prefill_url: String) { + { + let mut map = self.chars_per_url.lock().unwrap(); + *map.entry(prefill_url.clone()).or_insert(0) += char_cnt; + } + + let now = SystemTime::now(); + let time_window_duration = Duration::from_millis(self.period as u64); + let mut removed_load = 0; + + while let Some(req) = self.request_list.front() { + let expired = match now.duration_since(req.timestamp) { + Ok(duration) => duration > time_window_duration, + Err(_) => true, + }; + + if !expired { + break; + } + + if let Some(removed_req) = self.request_list.pop_front() { + self.t_req_loads.remove(&removed_req.id); + removed_load += removed_req.char_cnt; + + let mut map = self.chars_per_url.lock().unwrap(); + if let Some(count) = map.get_mut(&removed_req.prefill_worker_url) { + *count = count.saturating_sub(removed_req.char_cnt); + } + } + } + + self.load_total = self.load_total.saturating_sub(removed_load); + + let id = Uuid::new_v4().to_string(); + + self.t_req_loads.insert(id.clone(), char_cnt); + + self.request_list.push_back(SequencerRequest { + id, + char_cnt, + timestamp: now, + prefill_worker_url: prefill_url, + }); + + self.load_total = self.load_total.saturating_add(char_cnt); + } + + pub fn find_boundary(&self, char_count: usize) -> Option { + let mut left = 0; + let mut right = self.boundary.len(); + let mut _steps = 0; + + while left < right { + _steps += 1; + let mid = left + (right - left) / 2; + let range = self.boundary[mid].range; + + if char_count < range[0] { + right = mid; + } else if char_count > range[1] { + left = mid + 1; + } else { + return Some(self.boundary[mid].url.clone()); + } + } + None + } + + pub fn get_total_load(&self) -> usize { + self.load_total + } + + fn update_workers_cnt(&mut self) { + let pwu = self.prefill_worker_urls.lock().unwrap(); + self.bucket_cnt = pwu.len(); + + let mut char_map = self.chars_per_url.lock().unwrap(); + let current_urls: HashSet<_> = char_map.keys().cloned().collect(); + let new_urls: HashSet<_> = pwu.iter().cloned().collect(); + + for url in new_urls.difference(¤t_urls) { + char_map.insert(url.clone(), 0); + } + + for url in current_urls.difference(&new_urls) { + if char_map.get(url) == Some(&0) { + char_map.remove(url); + } + } + } + + pub fn adjust_boundary(&mut self) { + if self.t_req_loads.is_empty() { + return; + } + + self.update_workers_cnt(); + let worker_cnt = self.bucket_cnt; + if worker_cnt == 0 { + return; + } + let new_single_bucket_load = self.get_total_load() / worker_cnt; + let old_single_bucket_load = self.bucket_load; + + if new_single_bucket_load <= 2 * old_single_bucket_load + && (old_single_bucket_load <= 2 * new_single_bucket_load && old_single_bucket_load != 0) + { + info!("No need to adjust the bucket boundaries."); + return; + } + + info!("Before adjusting boundary | {:?}", self.boundary); + self.bucket_load = new_single_bucket_load; + let mut new_boundary = Vec::new(); + let mut hist_load: Vec = self.t_req_loads.values().cloned().collect(); + hist_load.sort(); + let mut upper_bound: usize = 0; + let mut last_load_index: usize = 0; + let max_value = usize::MAX; + + let worker_url = { + let guard = self.prefill_worker_urls.lock().unwrap(); + (*guard).clone() + }; + + let mut iter = worker_url.iter().peekable(); + // let mut curr_worker_id = 0; + while let Some(url) = iter.next() { + if last_load_index >= hist_load.len() && iter.peek().is_none() { + new_boundary.push(Boundary::new(url.clone(), [upper_bound, max_value])); + break; + } + let mut load_accumulator = 0; + let mut break_flag = false; + for &load in hist_load[last_load_index..].iter() { + load_accumulator += load; + if load_accumulator >= new_single_bucket_load { + if iter.peek().is_none() { + new_boundary.push(Boundary::new(url.clone(), [upper_bound, max_value])); + break_flag = true; + break; + } + let real_load = upper_bound + new_single_bucket_load; + if load <= upper_bound { + new_boundary.push(Boundary::new(url.clone(), [upper_bound, real_load])); + upper_bound = real_load + 1; + } else { + new_boundary.push(Boundary::new(url.clone(), [upper_bound, load])); + upper_bound = load + 1; + } + last_load_index += 1; + break_flag = true; + break; + } else { + last_load_index += 1; + } + } + if !break_flag { + let mut right_bound_value = upper_bound + new_single_bucket_load; + if iter.peek().is_none() { + right_bound_value = max_value; + new_boundary.push(Boundary::new(url.clone(), [upper_bound, right_bound_value])); + break; + } + new_boundary.push(Boundary::new(url.clone(), [upper_bound, right_bound_value])); + upper_bound = right_bound_value + 1; + } + } + self.boundary = new_boundary; + info!("After adjusting boundary | {:?}", self.boundary); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorkerBuilder, WorkerType}; + + #[tokio::test] + async fn test_load_balancing_conditions() { + // Test 1: Basic load balancing trigger + let config = BucketConfig { + balance_abs_threshold: 32, + balance_rel_threshold: 1.0001, + bucket_adjust_interval_secs: 10, + }; + let policy = BucketPolicy::with_config(config); + let prefill_workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w3:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + ]; + + // Initialize the policy with prefill_workers + policy.init_prefill_worker_urls(&prefill_workers); + + // === Phase S1: Construct bucket boundaries === + // Requests len =33 -> Bucket 1(expected range: 0-33) + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(33))) + .unwrap(); + // Two requests len =34 ->load balancing + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(34))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(34))) + .unwrap(); + + tokio::time::sleep(Duration::from_secs(10)).await; + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 33); + assert_eq!(bucket_guard.boundary[1].range[1], 67); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + // === Phase S2: Validate load balancing === + // Three consecutive len=33 requests (Should route to different buckets) + let idx_1 = policy + .select_worker(&prefill_workers, Some(&*"a".repeat(33))) + .unwrap(); + let idx_2 = policy + .select_worker(&prefill_workers, Some(&*"a".repeat(33))) + .unwrap(); + let idx_3 = policy + .select_worker(&prefill_workers, Some(&*"a".repeat(33))) + .unwrap(); + assert_eq!(idx_1, 0, "Should not trigger load balancing"); + assert_ne!(idx_2, idx_3, "Should trigger load balancing"); + assert_ne!(idx_2, 0, "Should trigger load balancing"); + assert_ne!(idx_3, 0, "Should trigger load balancing"); + + // Test 2: Not triggering when absolute threshold not met + let config = BucketConfig { + balance_abs_threshold: 30, + balance_rel_threshold: 2.0, + ..Default::default() + }; + let policy = BucketPolicy::with_config(config); + policy.init_prefill_worker_urls(&prefill_workers); + + // Create load difference below absolute threshold(20 + 8 = 28 < 30) + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(20))) + .unwrap(); // worker1: 20 + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(8))) + .unwrap(); // worker1: 8 + + // Next request should not use bucket scheduling (no load balancing) + let idx = policy + .select_worker(&prefill_workers, Some("request")) + .unwrap(); + assert_eq!( + idx, 0, + "Should not trigger load balancing when relative threshold not met" + ); + + // Test 3: Not triggering when relative threshold not met + let config = BucketConfig { + balance_abs_threshold: 5, + balance_rel_threshold: 3.0, + ..Default::default() + }; + let policy = BucketPolicy::with_config(config); + policy.init_prefill_worker_urls(&prefill_workers); + + // Create load difference (but relative threshold not met) + // Max/Min ratio = 15/5 = 3.0 + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(15))) + .unwrap(); // worker1: 15 + policy + .select_worker(&prefill_workers, Some("short")) + .unwrap(); // worker2: 5 + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(10))) + .unwrap(); // worker3: 10 + + // Next request should use bucket scheduling (load balancing) + let idx = policy + .select_worker(&prefill_workers, Some("request")) + .unwrap(); + assert_eq!( + idx, 0, + "Should not trigger load balancing when relative threshold not met" + ); + } + + #[tokio::test] + async fn test_adjust_boundary_1() { + // Test configuration: Set high threshold to prevent load balancing policy. + let config = BucketConfig { + balance_abs_threshold: 300, + balance_rel_threshold: 1.0001, + bucket_adjust_interval_secs: 3, + }; + let policy = BucketPolicy::with_config(config); + let prefill_workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w3:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + ]; + + // Initialize the policy with prefill_workers + policy.init_prefill_worker_urls(&prefill_workers); + + // Initial boundary + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 1364); + assert_eq!(bucket_guard.boundary[1].range[1], 2729); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + + // ===Phase S1: Initial requests to trigger boundary adjustment === + // Send requests with lengths: [5, 10, 15, 20, 24, 26] (total = 100) + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(5))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(10))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(15))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(20))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(24))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(26))) + .unwrap(); + + tokio::time::sleep(Duration::from_secs(3)).await; + // Verify boundaries adjusted to: [0, 20], [21, 26], [27, MAX] + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 20); + assert_eq!(bucket_guard.boundary[1].range[1], 26); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + + // ===Phase S2: Second set of requests to trigger boundary adjustment === + // Send requests with lengths: [10, 20, 30, 40, 45, 57] (total = 202) + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(10))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(20))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(30))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(40))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(45))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(57))) + .unwrap(); + + tokio::time::sleep(Duration::from_secs(3)).await; + // Verify boundaries adjusted to: [0, 40], [41, 57], [58, MAX] + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 40); + assert_eq!(bucket_guard.boundary[1].range[1], 57); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + } + + #[tokio::test] + async fn test_adjust_boundary_2() { + let config = BucketConfig { + balance_abs_threshold: 300, + balance_rel_threshold: 1.0001, + bucket_adjust_interval_secs: 3, + }; + let policy = BucketPolicy::with_config(config); + let prefill_workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w3:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + ]; + + // Initialize the policy with prefill_workers + policy.init_prefill_worker_urls(&prefill_workers); + + // Initial boundary + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 1364); + assert_eq!(bucket_guard.boundary[1].range[1], 2729); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + + // Send requests with char_count 20 + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(20))) + .unwrap(); + + tokio::time::sleep(Duration::from_secs(3)).await; + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 20); + assert_eq!(bucket_guard.boundary[1].range[1], 27); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(7))) + .unwrap(); + + tokio::time::sleep(Duration::from_secs(3)).await; + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 7); + assert_eq!(bucket_guard.boundary[1].range[1], 10); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + } + + #[tokio::test] + async fn test_not_adjust_boundary() { + let config = BucketConfig { + balance_abs_threshold: 300, + balance_rel_threshold: 1.0001, + bucket_adjust_interval_secs: 3, + }; + let policy = BucketPolicy::with_config(config); + let prefill_workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w3:8000") + .worker_type(WorkerType::Regular) + .api_key("test_api_key") + .build(), + ), + ]; + + // Initialize the policy with prefill_workers + policy.init_prefill_worker_urls(&prefill_workers); + + // Initial boundary + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 1364); + assert_eq!(bucket_guard.boundary[1].range[1], 2729); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(5))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(10))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(15))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(20))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(24))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(26))) + .unwrap(); + + tokio::time::sleep(Duration::from_secs(3)).await; + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 20); + assert_eq!(bucket_guard.boundary[1].range[1], 26); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(10))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(20))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(30))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(32))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(45))) + .unwrap(); + policy + .select_worker(&prefill_workers, Some(&*"a".repeat(55))) + .unwrap(); + + tokio::time::sleep(Duration::from_secs(3)).await; + { + let model_key = "default"; + + let bucket = policy + .buckets + .get(model_key) + .map(|entry| entry.value().clone()); + if let Some(bucket) = bucket { + let lock_result = bucket.write(); + if let Ok(bucket_guard) = lock_result { + // Expected Boundary: [0, 33] [34, 67] [68, MAX] + assert_eq!(bucket_guard.boundary[0].range[1], 20); + assert_eq!(bucket_guard.boundary[1].range[1], 26); + } else { + error!( + "Failed to acquire write lock for bucket of model {}", + model_key + ); + } + } + } + } +} diff --git a/sgl-router/src/policies/factory.rs b/sgl-router/src/policies/factory.rs index f03e8f1a0..a9db1842d 100644 --- a/sgl-router/src/policies/factory.rs +++ b/sgl-router/src/policies/factory.rs @@ -3,8 +3,8 @@ use std::sync::Arc; use super::{ - CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy, - RoundRobinPolicy, + BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, + PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, }; use crate::config::PolicyConfig; @@ -34,6 +34,18 @@ impl PolicyFactory { }; Arc::new(CacheAwarePolicy::with_config(config)) } + PolicyConfig::Bucket { + balance_abs_threshold, + balance_rel_threshold, + bucket_adjust_interval_secs, + } => { + let config = BucketConfig { + balance_abs_threshold: *balance_abs_threshold, + balance_rel_threshold: *balance_rel_threshold, + bucket_adjust_interval_secs: *bucket_adjust_interval_secs, + }; + Arc::new(BucketPolicy::with_config(config)) + } } } @@ -44,6 +56,7 @@ impl PolicyFactory { "round_robin" | "roundrobin" => Some(Arc::new(RoundRobinPolicy::new())), "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())), "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())), + "bucket" => Some(Arc::new(BucketPolicy::new())), _ => None, } } @@ -53,8 +66,8 @@ impl PolicyFactory { mod tests { use super::*; - #[test] - fn test_create_from_config() { + #[tokio::test] + async fn test_create_from_config() { let policy = PolicyFactory::create_from_config(&PolicyConfig::Random); assert_eq!(policy.name(), "random"); @@ -74,10 +87,17 @@ mod tests { max_tree_size: 1000, }); assert_eq!(policy.name(), "cache_aware"); + + let policy = PolicyFactory::create_from_config(&PolicyConfig::Bucket { + balance_abs_threshold: 10, + balance_rel_threshold: 1.5, + bucket_adjust_interval_secs: 5, + }); + assert_eq!(policy.name(), "bucket"); } - #[test] - fn test_create_by_name() { + #[tokio::test] + async fn test_create_by_name() { assert!(PolicyFactory::create_by_name("random").is_some()); assert!(PolicyFactory::create_by_name("RANDOM").is_some()); assert!(PolicyFactory::create_by_name("round_robin").is_some()); @@ -86,6 +106,8 @@ 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!(PolicyFactory::create_by_name("bucket").is_some()); + assert!(PolicyFactory::create_by_name("Bucket").is_some()); assert!(PolicyFactory::create_by_name("unknown").is_none()); } } diff --git a/sgl-router/src/policies/mod.rs b/sgl-router/src/policies/mod.rs index 564daa73e..67764eebb 100644 --- a/sgl-router/src/policies/mod.rs +++ b/sgl-router/src/policies/mod.rs @@ -7,6 +7,7 @@ use std::{fmt::Debug, sync::Arc}; use crate::core::Worker; +mod bucket; mod cache_aware; mod factory; mod power_of_two; @@ -14,6 +15,7 @@ mod random; mod registry; mod round_robin; +pub use bucket::BucketPolicy; pub use cache_aware::CacheAwarePolicy; pub use factory::PolicyFactory; pub use power_of_two::PowerOfTwoPolicy; @@ -108,6 +110,23 @@ impl Default for CacheAwareConfig { } } +#[derive(Debug, Clone)] +pub struct BucketConfig { + pub balance_abs_threshold: usize, + pub balance_rel_threshold: f32, + pub bucket_adjust_interval_secs: usize, +} + +impl Default for BucketConfig { + fn default() -> Self { + Self { + balance_abs_threshold: 32, + balance_rel_threshold: 1.0001, + bucket_adjust_interval_secs: 5, + } + } +} + /// Helper function to filter healthy workers and return their indices pub(crate) fn get_healthy_worker_indices(workers: &[Arc]) -> Vec { workers diff --git a/sgl-router/src/policies/registry.rs b/sgl-router/src/policies/registry.rs index 2904340ef..5fe5b24ae 100644 --- a/sgl-router/src/policies/registry.rs +++ b/sgl-router/src/policies/registry.rs @@ -12,8 +12,8 @@ use tracing::{debug, info, warn}; /// All subsequent workers of the same model use the established policy. /// When the last worker of a model is removed, the policy mapping is cleaned up. use super::{ - CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy, - RoundRobinPolicy, + BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, + PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, }; use crate::{config::types::PolicyConfig, core::Worker}; @@ -176,6 +176,7 @@ impl PolicyRegistry { "random" => Arc::new(RandomPolicy::new()), "cache_aware" => Arc::new(CacheAwarePolicy::new()), "power_of_two" => Arc::new(PowerOfTwoPolicy::new()), + "bucket" => Arc::new(BucketPolicy::new()), _ => { warn!("Unknown policy type '{}', using default", policy_type); Arc::clone(&self.default_policy) @@ -205,6 +206,18 @@ impl PolicyRegistry { Arc::new(CacheAwarePolicy::with_config(cache_config)) } PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()), + PolicyConfig::Bucket { + balance_abs_threshold, + balance_rel_threshold, + bucket_adjust_interval_secs, + } => { + let config = BucketConfig { + balance_abs_threshold: *balance_abs_threshold, + balance_rel_threshold: *balance_rel_threshold, + bucket_adjust_interval_secs: *bucket_adjust_interval_secs, + }; + Arc::new(BucketPolicy::with_config(config)) + } } } @@ -375,6 +388,23 @@ impl PolicyRegistry { } } } + + pub fn init_pd_bucket_policies(&self, prefill_workers: &[Arc]) { + // Initialize prefill policy if it's bucket + if let Some(prefill_policy) = self.prefill_policy.read().unwrap().as_ref() { + if prefill_policy.name() == "bucket" { + if let Some(bucket) = prefill_policy.as_any().downcast_ref::() { + if !prefill_workers.is_empty() { + debug!( + "Initializing prefill bucket policy with {} workers", + prefill_workers.len() + ); + bucket.init_prefill_worker_urls(prefill_workers); + } + } + } + } + } } impl std::fmt::Debug for PolicyRegistry { diff --git a/sgl-router/src/routers/http/pd_types.rs b/sgl-router/src/routers/http/pd_types.rs index 78c93d82e..9cb6ad9e0 100644 --- a/sgl-router/src/routers/http/pd_types.rs +++ b/sgl-router/src/routers/http/pd_types.rs @@ -70,4 +70,9 @@ pub enum PDSelectionPolicy { balance_abs_threshold: usize, balance_rel_threshold: f32, }, + Bucket { + balance_abs_threshold: usize, + balance_rel_threshold: f32, + bucket_adjust_interval_secs: usize, + }, } diff --git a/sgl-router/tests/test_pd_routing.rs b/sgl-router/tests/test_pd_routing.rs index 4a6ba7504..1151e144b 100644 --- a/sgl-router/tests/test_pd_routing.rs +++ b/sgl-router/tests/test_pd_routing.rs @@ -92,6 +92,11 @@ mod test_pd_routing { balance_abs_threshold: 32, balance_rel_threshold: 1.1, }, + PDSelectionPolicy::Bucket { + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + bucket_adjust_interval_secs: 5, + }, ]; for policy in policies { @@ -107,6 +112,12 @@ mod test_pd_routing { } => { assert!(*cache_threshold >= 0.0 && *cache_threshold <= 1.0); } + PDSelectionPolicy::Bucket { + balance_rel_threshold, + .. + } => { + assert!(*balance_rel_threshold >= 1.0); + } } } } @@ -160,6 +171,23 @@ mod test_pd_routing { max_tree_size: 1000000, }, ), + ( + RoutingMode::PrefillDecode { + prefill_urls: vec![ + ("http://p1:8080".to_string(), Some(9000)), + ("http://p2:8080".to_string(), Some(9001)), + ("http://p3:8080".to_string(), Some(9002)), + ], + decode_urls: vec!["http://d1:8080".to_string(), "http://d2:8080".to_string()], + prefill_policy: None, + decode_policy: None, + }, + PolicyConfig::Bucket { + balance_abs_threshold: 20, + balance_rel_threshold: 1.2, + bucket_adjust_interval_secs: 5, + }, + ), ]; for (mode, policy) in test_cases {