Skip to content
Open
Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
2bbb6cb
feat(router): add cache_aware_length policy with long/short pool split
jokerhaos Aug 21, 2026
9fa69e1
test: cover hit-unhealthy fallback, long-pool-unhealthy overflow, and…
jokerhaos Aug 21, 2026
8cec459
fix: address CodeRabbit review — token-only handling, perf, tests, CL…
jokerhaos Aug 21, 2026
4063ab6
fix: address remaining CodeRabbit comments — HTTP path tests, Python …
jokerhaos Aug 21, 2026
c685b4f
refactor: make cache_aware_length a full superset of cache_aware via …
jokerhaos Aug 21, 2026
039e3e9
feat: add --long-prefill-indices CLI param for long/short pool label …
jokerhaos Aug 21, 2026
522d613
fix: address CodeRabbit review — NaN validation, factory config asser…
jokerhaos Aug 21, 2026
ba3c416
fix: map cache_aware_length in Python policy_from_str + add NaN valid…
jokerhaos Aug 21, 2026
ae5843f
fix: add cache_aware_length CLI params to Python and Go bindings
jokerhaos Aug 21, 2026
54cda08
test: add Python binding tests for cache_aware_length CLI params
jokerhaos Aug 21, 2026
29c00d1
fix(python): complete cache aware length CLI contract
WyattJia Aug 21, 2026
669b894
fix: pass uncached hint through NoCacheStrategy and snapshot worker load
jokerhaos Aug 21, 2026
332651f
fix: address all CodeRabbit review comments on PR #2253
jokerhaos Aug 21, 2026
60aaac7
fix: reject non-finite balance_rel_threshold and add NaN regression test
jokerhaos Aug 21, 2026
6783cbd
fix: resolve CI pre-commit and clippy failures
jokerhaos Aug 24, 2026
6841326
test: add coverage for empty long_indices → short pool fallback
jokerhaos Aug 24, 2026
49af1a1
fix: address CodeRabbit nit comments on job_queue and router tests
jokerhaos Aug 24, 2026
f6bf400
test: pin worker in no-long-pool test for unique min-load assertion
jokerhaos Aug 24, 2026
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
4 changes: 2 additions & 2 deletions bindings/golang/src/policy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,8 @@ use openai_protocol::{
};
use smg::{
policies::{
BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy, RandomPolicy,
RoundRobinPolicy, SelectWorkerInfo,
BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PowerOfTwoPolicy,
RandomPolicy, RoundRobinPolicy, SelectWorkerInfo,
},
routers::grpc::{backend_client::BackendClient, utils::process_chat_messages},
worker::{
Expand Down
41 changes: 41 additions & 0 deletions bindings/python/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ pub enum PolicyType {
RoundRobin,
Passthrough,
CacheAware,
CacheAwareLength,
Comment thread
coderabbitai[bot] marked this conversation as resolved.
PowerOfTwo,
LeastLoad,
Bucket,
Expand Down Expand Up @@ -526,6 +527,11 @@ struct Router {
worker_overload_token_usage: Option<f64>,
worker_overload_protection: bool,
disable_load_monitoring: bool,
chars_per_token: usize,
long_prefill_threshold: usize,
long_pool_max_load: usize,
short_pool_max_load: usize,
long_prefill_indices: Vec<usize>,
}

impl Router {
Expand Down Expand Up @@ -629,6 +635,25 @@ impl Router {
cache_ttl_secs: self.cache_ttl_secs,
cache_boundaries: self.cache_boundaries.clone(),
},
PolicyType::CacheAwareLength => ConfigPolicyConfig::CacheAwareLength {
cache_threshold: self.cache_threshold,
balance_abs_threshold: self.balance_abs_threshold,
balance_rel_threshold: self.balance_rel_threshold,
eviction_interval_secs: self.eviction_interval_secs,
max_tree_size: self.max_tree_size,
block_size: self.block_size,
balance_token_usage_threshold: self.balance_token_usage_threshold,
overload_token_usage_threshold: self.overload_token_usage_threshold,
overlap_decay: self.overlap_decay,
selection_temperature: self.selection_temperature,
cache_index: self.parse_cache_index()?,
cache_ttl_secs: self.cache_ttl_secs,
cache_boundaries: self.cache_boundaries.clone(),
chars_per_token: self.chars_per_token,
long_prefill_threshold: self.long_prefill_threshold,
long_pool_max_load: self.long_pool_max_load,
short_pool_max_load: self.short_pool_max_load,
},
PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo {
load_check_interval_secs: self.load_monitor_interval,
},
Expand Down Expand Up @@ -931,6 +956,7 @@ impl Router {
self.server_key_path.as_ref(),
)
.dp_minimum_tokens_scheduler(self.dp_minimum_tokens_scheduler)
.long_prefill_indices(self.long_prefill_indices.clone())
.build()
}
}
Expand Down Expand Up @@ -1083,6 +1109,11 @@ impl Router {
worker_overload_token_usage = None,
worker_overload_protection = false,
disable_load_monitoring = false,
chars_per_token = 4,
long_prefill_threshold = 100_000,
long_pool_max_load = 4,
short_pool_max_load = 32,
long_prefill_indices = vec![],
))]
#[expect(clippy::too_many_arguments)]
#[expect(
Expand Down Expand Up @@ -1233,6 +1264,11 @@ impl Router {
worker_overload_token_usage: Option<f64>,
worker_overload_protection: bool,
disable_load_monitoring: bool,
chars_per_token: usize,
long_prefill_threshold: usize,
long_pool_max_load: usize,
short_pool_max_load: usize,
long_prefill_indices: Vec<usize>,
) -> PyResult<Self> {
let mut all_urls = worker_urls.clone();

Expand Down Expand Up @@ -1397,6 +1433,11 @@ impl Router {
worker_overload_token_usage,
worker_overload_protection,
disable_load_monitoring,
chars_per_token,
long_prefill_threshold,
long_pool_max_load,
short_pool_max_load,
long_prefill_indices,
})
}

Expand Down
1 change: 1 addition & 0 deletions bindings/python/src/smg/router.py
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,
"passthrough": PolicyType.Passthrough,
"cache_aware": PolicyType.CacheAware,
"cache_aware_length": PolicyType.CacheAwareLength,
"power_of_two": PolicyType.PowerOfTwo,
"least_load": PolicyType.LeastLoad,
"bucket": PolicyType.Bucket,
Expand Down
63 changes: 63 additions & 0 deletions bindings/python/src/smg/router_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
"round_robin",
"passthrough",
"cache_aware",
"cache_aware_length",
"power_of_two",
"least_load",
"manual",
Expand Down Expand Up @@ -242,6 +243,12 @@ class RouterArgs:
worker_overload_protection: bool = False
# Restore the conditional load-monitor poll gate (default: poll always)
disable_load_monitoring: bool = False
# cache_aware_length policy: long/short pool split
chars_per_token: int = 4
long_prefill_threshold: int = 100_000
long_pool_max_load: int = 4
short_pool_max_load: int = 32
long_prefill_indices: list[int] = dataclasses.field(default_factory=list)

@staticmethod
def add_cli_args(
Expand Down Expand Up @@ -627,6 +634,7 @@ def add_cli_args(
)
routing_group.add_argument(
f"--{prefix}eviction-interval-secs",
f"--{prefix}eviction-interval",
type=int,
default=RouterArgs.eviction_interval_secs,
help="Interval in seconds between cache eviction operations",
Expand Down Expand Up @@ -677,6 +685,42 @@ def add_cli_args(
" approximate serving-engine cache retention. Defaults to 180."
),
)
# cache_aware_length policy parameters
routing_group.add_argument(
f"--{prefix}chars-per-token",
type=int,
default=RouterArgs.chars_per_token,
help="Divisor for char-level token estimation when X-Prompt-Tokens"
" is absent (cache_aware_length policy). Default 4.",
)
routing_group.add_argument(
f"--{prefix}long-prefill-threshold",
type=int,
default=RouterArgs.long_prefill_threshold,
help="Uncached-prefill-token boundary between long and short"
" requests (cache_aware_length policy). Default 100000.",
)
routing_group.add_argument(
f"--{prefix}long-pool-max-load",
type=int,
default=RouterArgs.long_pool_max_load,
help="Load ceiling for the long pool (pool=long workers)"
" (cache_aware_length policy). Default 4.",
)
routing_group.add_argument(
f"--{prefix}short-pool-max-load",
type=int,
default=RouterArgs.short_pool_max_load,
help="Load ceiling for the short pool (remaining workers)"
" (cache_aware_length policy). Default 32.",
)
routing_group.add_argument(
f"--{prefix}long-prefill-indices",
type=_parse_int_csv,
default=[],
help="Comma-separated 0-based indices of --prefill URLs that belong"
" to the long pool (get pool=long label for cache_aware_length).",
)
routing_group.add_argument(
f"--{prefix}max-idle-secs",
f"--{prefix}sticky-key-idle-secs",
Expand Down Expand Up @@ -1681,6 +1725,25 @@ def from_cli_args(cls, args: argparse.Namespace, use_router_prefix: bool = False
return cls(**args_dict)

def _validate_router_args(self):
if (self.prefill_urls or self.decode_urls) and not (
self.pd_disaggregation or self.epd_disaggregation
):
raise ValueError(
"--prefill/--decode require --pd-disaggregation or --epd-disaggregation"
)

if len(set(self.long_prefill_indices)) != len(self.long_prefill_indices):
raise ValueError("--long-prefill-indices must not contain duplicate values")

if self.long_prefill_indices:
if min(self.long_prefill_indices) < 0:
raise ValueError("--long-prefill-indices values must be non-negative")
if max(self.long_prefill_indices) >= len(self.prefill_urls):
raise ValueError(
"--long-prefill-indices value out of range for "
f"{len(self.prefill_urls)} configured prefill workers"
)

# Validate configuration based on mode
if self.epd_disaggregation:
if self.encode_policy:
Expand Down
Loading
Loading