Skip to content

feat(policies): add least_load load-balancing policy #1629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
slin1237 merged 1 commit into from
Jun 10, 2026
Merged

feat(policies): add least_load load-balancing policy #1629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions bindings/python/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ pub enum PolicyType {
RoundRobin,
CacheAware,
PowerOfTwo,
LeastLoad,
Bucket,
Manual,
ConsistentHashing,
Expand Down Expand Up @@ -515,6 +516,10 @@ impl Router {
PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo {
load_check_interval_secs: 5,
},
PolicyType::LeastLoad => ConfigPolicyConfig::LeastLoad {
Comment thread
slin1237 marked this conversation as resolved.
load_check_interval_secs: 5,
lambda: 1.5,
},
PolicyType::Bucket => ConfigPolicyConfig::Bucket {
balance_abs_threshold: self.balance_abs_threshold,
balance_rel_threshold: self.balance_rel_threshold,
Expand Down
3 changes: 3 additions & 0 deletions bindings/python/src/smg/router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ def policy_from_str(policy_str: str | None) -> PolicyType:
"round_robin": PolicyType.RoundRobin,
"cache_aware": PolicyType.CacheAware,
"power_of_two": PolicyType.PowerOfTwo,
"least_load": PolicyType.LeastLoad,
"bucket": PolicyType.Bucket,
"manual": PolicyType.Manual,
"consistent_hashing": PolicyType.ConsistentHashing,
Expand Down Expand Up @@ -143,6 +144,8 @@ class Router:
balance
- PolicyType.PowerOfTwo: Select best of two random workers based on load
(PD mode only)
- PolicyType.LeastLoad: Route to the worker with the lowest load score
(in-flight requests plus KV-cache pressure)
host: Host address to bind the router server. Supports IPv4, IPv6 (e.g., ::,
::1), or 0.0.0.0 for all interfaces. Default: '0.0.0.0'
port: Port number to bind the router server. Default: 3001
Expand Down
1 change: 1 addition & 0 deletions bindings/python/src/smg/router_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
"round_robin",
"cache_aware",
"power_of_two",
"least_load",
"manual",
"consistent_hashing",
"prefix_hash",
Expand Down
1 change: 1 addition & 0 deletions bindings/python/tests/test_router_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -339,6 +339,7 @@ def test_policy_enum_conversion(self):
assert policy_from_str("round_robin") == PolicyType.RoundRobin
assert policy_from_str("cache_aware") == PolicyType.CacheAware
assert policy_from_str("power_of_two") == PolicyType.PowerOfTwo
assert policy_from_str("least_load") == PolicyType.LeastLoad

def test_invalid_policy_enum_conversion(self):
"""Test invalid policy string to enum conversion."""
Expand Down
3 changes: 2 additions & 1 deletion bindings/python/tests/test_startup_sequence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,12 +75,13 @@ class TestPolicyFromStr:
def test_policy_conversion_in_startup(self):
"""Test policy conversion during startup sequence."""
# Test all valid policies
policies = ["random", "round_robin", "cache_aware", "power_of_two"]
policies = ["random", "round_robin", "cache_aware", "power_of_two", "least_load"]
expected_enums = [
PolicyType.Random,
PolicyType.RoundRobin,
PolicyType.CacheAware,
PolicyType.PowerOfTwo,
PolicyType.LeastLoad,
]

for policy_str, expected_enum in zip(policies, expected_enums):
Expand Down
4 changes: 2 additions & 2 deletions bindings/python/tests/test_validation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,7 +318,7 @@ def test_pd_service_discovery_validation(self):
def test_policy_validation(self):
"""Test policy configuration validation."""
# Valid policies
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two"]
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two", "least_load"]

for policy in valid_policies:
args = RouterArgs(policy=policy)
Expand All @@ -327,7 +327,7 @@ def test_policy_validation(self):
def test_pd_policy_validation(self):
"""Test PD policy configuration validation."""
# Valid PD policies
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two"]
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two", "least_load"]

for prefill_policy in valid_policies:
for decode_policy in valid_policies:
Expand Down
22 changes: 22 additions & 0 deletions model_gateway/src/config/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,6 +373,19 @@ pub enum PolicyConfig {
#[serde(rename = "power_of_two")]
PowerOfTwo { load_check_interval_secs: u64 },

/// Least-load policy: routes to the worker minimizing
/// `in_flight + lambda * k/(1-k)` — the real-time in-flight request count plus
/// a convex KV-cache pressure term from the load monitor.
/// See `policies/least_load.rs`.
#[serde(rename = "least_load")]
LeastLoad {
Comment thread
slin1237 marked this conversation as resolved.
#[serde(default = "default_least_load_interval")]
load_check_interval_secs: u64,
/// KV-pressure weight (request-equivalents per unit of M/M/1 congestion).
#[serde(default = "default_least_load_lambda")]
lambda: f64,
},
Comment thread
slin1237 marked this conversation as resolved.

#[serde(rename = "bucket")]
Bucket {
/// Absolute load difference threshold for load balancing
Expand Down Expand Up @@ -445,13 +458,22 @@ fn default_manual_max_idle_secs() -> u64 {
4 * 3600
}

fn default_least_load_interval() -> u64 {
10
}

fn default_least_load_lambda() -> f64 {
1.5
}

impl PolicyConfig {
pub fn name(&self) -> &'static str {
match self {
PolicyConfig::Random => "random",
PolicyConfig::RoundRobin => "round_robin",
PolicyConfig::CacheAware { .. } => "cache_aware",
PolicyConfig::PowerOfTwo { .. } => "power_of_two",
PolicyConfig::LeastLoad { .. } => "least_load",
Comment thread
slin1237 marked this conversation as resolved.
PolicyConfig::Bucket { .. } => "bucket",
PolicyConfig::Manual { .. } => "manual",
PolicyConfig::ConsistentHashing => "consistent_hashing",
Expand Down
20 changes: 20 additions & 0 deletions model_gateway/src/config/validation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,6 +323,26 @@ impl ConfigValidator {
});
}
}
PolicyConfig::LeastLoad {
load_check_interval_secs,
lambda,
} => {
if *load_check_interval_secs == 0 {
return Err(ConfigError::InvalidValue {
field: "load_check_interval_secs".to_string(),
value: load_check_interval_secs.to_string(),
reason: "Must be > 0".to_string(),
});
}

if !lambda.is_finite() || *lambda < 0.0 {
return Err(ConfigError::InvalidValue {
field: "lambda".to_string(),
value: lambda.to_string(),
reason: "Must be finite and >= 0.0".to_string(),
});
}
}
Comment thread
slin1237 marked this conversation as resolved.
PolicyConfig::Bucket {
balance_abs_threshold: _,
balance_rel_threshold,
Expand Down
10 changes: 7 additions & 3 deletions model_gateway/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ struct CliArgs {

// ==================== Routing Policy ====================
/// Load balancing policy to use
#[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "Routing Policy")]
#[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "Routing Policy")]

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Add least_load to Helm values schema

When this policy is selected via the Helm chart, chart values validation still rejects it because deploy/helm/smg/values.schema.json line 10 only enumerates cache_aware, round_robin, power_of_two, manual, random, and prefix_hash (I checked the Helm chart policy schema and values comments). This means Kubernetes/Helm users cannot enable the new least_load policy even though the gateway CLI and Python args now accept it; update the chart schema (and related values comment) alongside this new accepted policy.

Useful? React with 👍 / 👎.

@coderabbitai coderabbitai Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

bucket is accepted by CLI but silently parsed as round_robin.

value_parser allows "bucket" (Lines 152/217/221), but parse_policy has no "bucket" match arm, so Line 949 fallback applies. This causes incorrect routing policy selection without user-visible error.

Suggested fix (fail closed until bucket parsing is explicit) 
-    #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "Routing Policy")]
+    #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual"], help_heading = "Routing Policy")]
@@
-    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")]
+    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual"], help_heading = "PD Disaggregation")]
@@
-    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")]
+    #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual"], help_heading = "PD Disaggregation")]

Also applies to: 217-222, 916-950

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model_gateway/src/main.rs` at line 152, The CLI accepts "bucket" in the
value_parser but parse_policy lacks a "bucket" match arm so inputs silently fall
back to "round_robin"; update the parse_policy function to explicitly handle
"bucket" (either map it to the correct RoutingPolicy variant or return a parse
error to fail closed), ensuring the match in parse_policy covers the "bucket"
string and returns a Result/Err instead of defaulting to the fallback; reference
the value_parser and parse_policy identifiers and add the "bucket" branch (or
explicit Err) to keep CLI parsing consistent and visible to users.

policy: String,

/// Cache threshold (0.0-1.0) for cache-aware routing
Expand Down Expand Up @@ -214,11 +214,11 @@ struct CliArgs {
decode: Vec<String>,

/// Specific policy for prefill nodes in PD mode
#[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")]
#[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")]
prefill_policy: Option<String>,

/// Specific policy for decode nodes in PD mode
#[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")]
#[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "least_load", "prefix_hash", "consistent_hashing", "manual", "bucket"], help_heading = "PD Disaggregation")]
decode_policy: Option<String>,

/// Timeout in seconds for worker startup and registration
Expand Down Expand Up @@ -928,6 +928,10 @@ impl CliArgs {
"power_of_two" => PolicyConfig::PowerOfTwo {
load_check_interval_secs: 5,
},
"least_load" => PolicyConfig::LeastLoad {
load_check_interval_secs: 5,
lambda: 1.5,
},
Comment thread
slin1237 marked this conversation as resolved.
"prefix_hash" => PolicyConfig::PrefixHash {
prefix_token_count: self.prefix_token_count,
load_factor: self.prefix_hash_load_factor,
Expand Down
8 changes: 6 additions & 2 deletions model_gateway/src/policies/factory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ use std::sync::Arc;

use super::{
BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy,
LoadBalancingPolicy, ManualConfig, ManualPolicy, PowerOfTwoPolicy, PrefixHashConfig,
PrefixHashPolicy, RandomPolicy, RoundRobinPolicy,
LeastLoadPolicy, LoadBalancingPolicy, ManualConfig, ManualPolicy, PowerOfTwoPolicy,
PrefixHashConfig, PrefixHashPolicy, RandomPolicy, RoundRobinPolicy,
};
use crate::config::PolicyConfig;

Expand All @@ -19,6 +19,9 @@ impl PolicyFactory {
PolicyConfig::Random => Arc::new(RandomPolicy::new()),
PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()),
PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()),
PolicyConfig::LeastLoad { lambda, .. } => {
Arc::new(LeastLoadPolicy::with_lambda(*lambda))
}
Comment thread
slin1237 marked this conversation as resolved.
PolicyConfig::CacheAware {
cache_threshold,
balance_abs_threshold,
Expand Down Expand Up @@ -81,6 +84,7 @@ impl PolicyFactory {
"random" => Some(Arc::new(RandomPolicy::new())),
"round_robin" | "roundrobin" => Some(Arc::new(RoundRobinPolicy::new())),
"power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())),
"least_load" | "leastload" => Some(Arc::new(LeastLoadPolicy::new())),
"cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())),
"bucket" => Some(Arc::new(BucketPolicy::new())),
"manual" => Some(Arc::new(ManualPolicy::new())),
Expand Down
Loading
Loading