chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[router] bucket policy (#11719)

Browse Source
This commit is contained in:
syy-hw
2025-11-10 18:02:53 +08:00
committed by GitHub
parent 9ea2c686c7
commit 611a4fd08b
12 changed files with 1435 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sgl-router/py_src/sglang_router/router.py
View File

@@ -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]

9
sgl-router/py_src/sglang_router/router_args.py
View File

@@ -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,

33
sgl-router/src/config/types.rs
View File

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

96
sgl-router/src/config/validation.rs
View File

@@ -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(

7
sgl-router/src/core/workflow/steps/worker_registration.rs
View File

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

11
sgl-router/src/lib.rs
View File

@@ -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,

1167
sgl-router/src/policies/bucket.rs Normal file
View File

File diff suppressed because it is too large Load Diff

34
sgl-router/src/policies/factory.rs
View File

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

19
sgl-router/src/policies/mod.rs
View File

@@ -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

34
sgl-router/src/policies/registry.rs
View File

@@ -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<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 {

5
sgl-router/src/routers/http/pd_types.rs
View File

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

28
sgl-router/tests/test_pd_routing.rs
View File

@@ -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 {