Support min load besides random routing key assignment policy in ManualPolicy (#16767)

This commit is contained in:
fzyzcjy
2026-01-12 18:31:18 -08:00
committed by GitHub
parent a83484275d
commit 9d3018f484
8 changed files with 163 additions and 15 deletions
@@ -348,6 +348,7 @@ struct Router {
eviction_interval_secs: u64,
max_tree_size: usize,
max_idle_secs: u64,
assignment_mode: String,
max_payload_size: usize,
dp_aware: bool,
api_key: Option<String>,
@@ -457,6 +458,11 @@ impl Router {
PolicyType::Manual => ConfigPolicyConfig::Manual {
eviction_interval_secs: self.eviction_interval_secs,
max_idle_secs: self.max_idle_secs,
assignment_mode: match self.assignment_mode.as_str() {
"random" => config::ManualAssignmentMode::Random,
"min_load" => config::ManualAssignmentMode::MinLoad,
other => panic!("Unknown assignment mode: {}", other),
},
},
PolicyType::ConsistentHashing => ConfigPolicyConfig::ConsistentHashing,
PolicyType::PrefixHash => ConfigPolicyConfig::PrefixHash {
@@ -640,6 +646,7 @@ impl Router {
eviction_interval_secs = 120,
max_tree_size = 2usize.pow(26),
max_idle_secs = 14400,
assignment_mode = String::from("random"),
max_payload_size = 512 * 1024 * 1024,
dp_aware = false,
api_key = None,
@@ -724,6 +731,7 @@ impl Router {
eviction_interval_secs: u64,
max_tree_size: usize,
max_idle_secs: u64,
assignment_mode: String,
max_payload_size: usize,
dp_aware: bool,
api_key: Option<String>,
@@ -821,6 +829,7 @@ impl Router {
eviction_interval_secs,
max_tree_size,
max_idle_secs,
assignment_mode,
max_payload_size,
dp_aware,
api_key,
@@ -36,6 +36,7 @@ class RouterArgs:
eviction_interval_secs: int = 60
max_tree_size: int = 2**26
max_idle_secs: int = 4 * 3600
assignment_mode: str = "random"
max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches
bucket_adjust_interval_secs: int = 5
dp_aware: bool = False
@@ -310,6 +311,13 @@ class RouterArgs:
default=RouterArgs.max_idle_secs,
help="Maximum idle time in seconds before eviction (for manual policy)",
)
routing_group.add_argument(
f"--{prefix}assignment-mode",
type=str,
default=RouterArgs.assignment_mode,
choices=["random", "min_load"],
help="Mode for assigning new routing keys in manual policy: random (default), min_load (worker with fewest requests)",
)
routing_group.add_argument(
f"--{prefix}max-payload-size",
type=int,