[router] bucket policy (#11719)

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]
 

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,

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();

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(

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: {})",

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<Vec<String>>,
     pd_disaggregation: bool,
+    // Takes effect in PD mode and when policy = bucket
+    bucket_adjust_interval_secs: usize,
     prefill_urls: Option<Vec<(String, Option<u16>)>>,
     decode_urls: Option<Vec<String>>,
     prefill_policy: Option<PolicyType>,
@@ -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<Vec<String>>,
         pd_disaggregation: bool,
+        bucket_adjust_interval_secs: usize,
         prefill_urls: Option<Vec<(String, Option<u16>)>>,
         decode_urls: Option<Vec<String>>,
         prefill_policy: Option<PolicyType>,
@@ -566,6 +576,7 @@ impl Router {
             request_timeout_secs,
             request_id_headers,
             pd_disaggregation,
+            bucket_adjust_interval_secs,
             prefill_urls,
             decode_urls,
             prefill_policy,

sgl-router/src/policies/factory.rs

@@ -3,8 +3,8 @@
 use std::sync::Arc;
 
 use super::{
-    CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy,
+    BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy,
-    RoundRobinPolicy,
+    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]
+    #[tokio::test]
-    fn test_create_from_config() {
+    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]
+    #[tokio::test]
-    fn test_create_by_name() {
+    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());
     }
 }

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<dyn Worker>]) -> Vec<usize> {
     workers

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,
+    BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy,
-    RoundRobinPolicy,
+    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<dyn Worker>]) {
+        // 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::<BucketPolicy>() {
+                    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 {

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,
+    },
 }

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 {