diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 9c463d163..9692ae9f7 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -211,6 +211,22 @@ pub struct Cli { /// `--dp-size` like the limit. #[arg(long)] pub saturation_queue_floor: Option, + /// Number of random candidates sampled for the min-load fallback; the + /// least-pressured of the sample wins. The default 2 keeps today's + /// power-of-2 behavior unchanged. `k >= pool` skips the shuffle and + /// returns the exact minimum, with ties broken randomly — an idle + /// fleet ties on every comparison, so a fixed order would pin every + /// fallback dispatch to one worker; `k = 1` is a uniform draw within + /// the tier, and because a one-member sample has no runner-up the + /// proposal carries no backup, which disables the backup-admission + /// and pressure-guard paths. Note + /// the division of labor with `--cache-candidate-min-workers`, + /// `--cache-candidate-ratio`, and `--cache-candidate-max-workers`: + /// those bound the cache-affinity OWNER candidate set; this flag + /// bounds the min-load FALLBACK sample used when no owner is usable. + /// Requires `--policy cache_aware`. + #[arg(long)] + pub min_load_choices: Option, // ---- score composition ---- /// Policies to sum, spelled exactly as `--policy` spells them and each @@ -399,7 +415,8 @@ impl Cli { || self.cache_candidate_max_workers.is_some() || self.cache_switch_margin_tokens.is_some() || self.worker_queue_limit.is_some() - || self.saturation_queue_floor.is_some(); + || self.saturation_queue_floor.is_some() + || self.min_load_choices.is_some(); // Value checks before the policy check: a value that is wrong under // every policy should say so, rather than pointing at --policy. if self.worker_queue_limit == Some(0) { @@ -427,6 +444,9 @@ impl Cli { )); } } + if self.min_load_choices == Some(0) { + return Err(anyhow!("--min-load-choices must be at least 1")); + } if tuned_cache_candidates && self.policy != PolicyKind::CacheAware { return Err(anyhow!( "cache candidate tuning flags require --policy cache_aware" @@ -652,6 +672,7 @@ impl Cli { .unwrap_or(d.cache_switch_margin_tokens), worker_queue_limit: self.worker_queue_limit.or(d.worker_queue_limit), saturation_queue_floor: self.saturation_queue_floor.or(d.saturation_queue_floor), + min_load_choices: self.min_load_choices.unwrap_or(d.min_load_choices), }) } else { None @@ -2096,6 +2117,43 @@ mod tests { ); } + #[test] + fn min_load_choices_is_plumbed_and_validated() { + let config = cfg_of("--policy cache_aware --min-load-choices 5").unwrap(); + assert_eq!( + config + .model + .affinity + .expect("cache-aware needs affinity config") + .min_load_choices, + 5 + ); + + // Unset keeps the pre-existing power-of-2 behavior. + let defaults = cfg_of("--policy cache_aware").unwrap(); + assert_eq!( + defaults + .model + .affinity + .expect("default affinity config") + .min_load_choices, + 2 + ); + + let err = cfg_of("--policy cache_aware --min-load-choices 0") + .expect_err("a zero sample size would select nothing") + .to_string(); + assert!(err.contains("--min-load-choices"), "got: {err}"); + + let err = cfg_of("--policy power_of_two --min-load-choices 3") + .expect_err("the knob only tunes the cache-aware fallback") + .to_string(); + assert!( + err.contains("cache candidate tuning flags require --policy cache_aware"), + "got: {err}" + ); + } + #[test] fn cache_candidate_cli_rejects_invalid_bounds() { for (args, expected) in [ diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index dfa02f28d..6aeebc0b4 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -469,6 +469,11 @@ pub const DEFAULT_SESSION_ID_HEADER: &str = "x-session-id"; /// Default external-indexer request limits. pub const DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT: usize = 32; +/// Default min-load sample size: the pre-existing power-of-2 behavior. +/// Every code path that has no `AffinityConfig` to read must fall back to +/// this, so the no-affinity path never drifts from the configured default. +pub const DEFAULT_MIN_LOAD_CHOICES: usize = 2; + /// Controls whether admission may select a session-affinity backup. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)] pub enum AffinityMode { @@ -562,6 +567,20 @@ pub struct AffinityConfig { /// The CLI enforces `floor <= worker_queue_limit` and requires the /// gate; like the limit, scale the floor with `dp_size`. pub saturation_queue_floor: Option, + /// Number of random candidates sampled for the min-load fallback + /// (`--min-load-choices`); the least-pressured of the sample wins. + /// [`DEFAULT_MIN_LOAD_CHOICES`] is the pre-existing power-of-2 + /// behavior, so upgrading changes nothing. `k >= pool` skips the + /// shuffle and returns the exact minimum, with ties broken randomly + /// (an idle fleet ties on every comparison, so a fixed order would pin + /// every fallback dispatch to one worker); `k = 1` is a uniform draw + /// within the tier, and its sample has no second member, so the + /// proposal carries no backup and admission loses its backup-admission + /// and pressure-guard paths. The + /// `--cache-candidate-*` knobs bound the cache-affinity OWNER candidate + /// set; this bounds the min-load FALLBACK sample used when no owner is + /// usable. + pub min_load_choices: usize, } impl Default for AffinityConfig { @@ -586,6 +605,7 @@ impl Default for AffinityConfig { cache_switch_margin_tokens: 1_024, worker_queue_limit: None, saturation_queue_floor: None, + min_load_choices: DEFAULT_MIN_LOAD_CHOICES, } } } diff --git a/experimental/sgl-router/src/policies/admission.rs b/experimental/sgl-router/src/policies/admission.rs index 76fb09a6d..75dce9554 100644 --- a/experimental/sgl-router/src/policies/admission.rs +++ b/experimental/sgl-router/src/policies/admission.rs @@ -18,7 +18,7 @@ //! prefix owner instead of cold-prefilling on a non-owner. use crate::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad, NativeCacheWorkerLoad}; -use crate::policies::power_of_two::select_with_snapshot; +use crate::policies::power_of_two::select_k_with_snapshot; use crate::policies::{CacheCandidate, CacheCandidateProposal, GuardHints, SelectionProposal}; use crate::workers::Worker; use std::cmp::Ordering; @@ -399,6 +399,7 @@ pub fn resolve_prefill( request_input_tokens: u64, snapshot: &EngineLoadSnapshot, queue_limit: Option, + min_load_choices: usize, ) -> Option { resolve_prefill_admitted(range, proposal, request_input_tokens, snapshot, queue_limit).or_else( || { @@ -411,7 +412,8 @@ pub fn resolve_prefill( .filter(|worker| contains_worker(range, worker)) .cloned(); let legal = legal_prefill_candidates(range, proposal); - let selected = select_with_snapshot(&legal, Some(snapshot))?; + let selected = + select_k_with_snapshot(&legal, Some(snapshot), min_load_choices, queue_limit)?; Some(FinalDecision { selected, primary: Arc::clone(&proposal.primary), @@ -1044,6 +1046,7 @@ fn pressure_guard_prefers_backup( mod tests { use super::*; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::power_of_two::select_k_with_snapshot; use std::time::Instant; fn worker(id: &str) -> Arc { @@ -1096,13 +1099,21 @@ mod tests { 20, &loads, None, + 2, ) .is_some()); assert_eq!( - resolve_prefill(&range, &SelectionProposal::primary(full), 20, &loads, None) - .expect("fallback selects the admitted worker") - .selected - .id, + resolve_prefill( + &range, + &SelectionProposal::primary(full), + 20, + &loads, + None, + 2 + ) + .expect("fallback selects the admitted worker") + .selected + .id, unknown.id ); } @@ -1131,6 +1142,7 @@ mod tests { 32, &loads, None, + 2, ) .expect("capacity exhaustion must degrade within the legal domain"); @@ -1149,7 +1161,7 @@ mod tests { let proposal = SelectionProposal::with_backup(Arc::clone(&primary), Arc::clone(&backup)); let explicit = snapshot(&[(&primary, 0, 0, 100, 100), (&backup, 0, 10, 100, 100)]); let opposite = snapshot(&[(&primary, 0, 10, 100, 100), (&backup, 0, 0, 100, 100)]); - let opposite_decision = select_with_snapshot(&workers, Some(&opposite)) + let opposite_decision = select_k_with_snapshot(&workers, Some(&opposite), 2, None) .expect("the opposite snapshot has the same legal workers"); assert_eq!(opposite_decision.id, backup.id); @@ -1159,6 +1171,7 @@ mod tests { 32, &explicit, None, + 2, ) .expect("capacity exhaustion must degrade to Power-of-Two"); @@ -1714,6 +1727,7 @@ mod tests { 32, &loads, Some(4), + 2, ) .expect("an admitted worker exists"); @@ -1808,6 +1822,7 @@ mod tests { 32, &loads, Some(4), + 2, ) .expect("an all-queueing fleet must still route"); diff --git a/experimental/sgl-router/src/policies/cache_aware.rs b/experimental/sgl-router/src/policies/cache_aware.rs index e434426d1..6ad349b3f 100644 --- a/experimental/sgl-router/src/policies/cache_aware.rs +++ b/experimental/sgl-router/src/policies/cache_aware.rs @@ -181,6 +181,7 @@ impl Policy for CacheAwarePolicy { } } PowerOfTwoChoicesPolicy::new() + .with_load_control(self.config.min_load_choices, self.config.worker_queue_limit) .propose(workers, ctx) .map(PrefillProposal::Pair) } diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index 3237e64d8..35095ee30 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -810,6 +810,7 @@ mod tests { 32, &loads, None, + 2, ) .expect("the admitted backup must become Final P"); assert_eq!(decision.selected.id, backup.id); @@ -1604,7 +1605,7 @@ mod tests { let range = CandidateRange::global(&workers); let proposal = SelectionProposal::with_backup(Arc::clone(&primary), Arc::clone(&backup)); - let decision = resolve_prefill(&range, &proposal, 32, &snapshot, None) + let decision = resolve_prefill(&range, &proposal, 32, &snapshot, None, 2) .expect("an admitted backup must be selected"); assert_eq!(decision.selected.id, backup.id); @@ -1623,6 +1624,7 @@ mod tests { 1_000_000, &snapshot, None, + 2, ) .expect("disabled reporting must preserve the healthy registry candidate"); @@ -1661,6 +1663,7 @@ mod tests { 80, &snapshot, None, + 2, ) .expect("both candidates fit capacity"); @@ -1711,6 +1714,7 @@ mod tests { 32, &snapshot, None, + 2, ) .expect("an admitted range fallback must be selected"); diff --git a/experimental/sgl-router/src/policies/power_of_two.rs b/experimental/sgl-router/src/policies/power_of_two.rs index b840157ca..35c44d841 100644 --- a/experimental/sgl-router/src/policies/power_of_two.rs +++ b/experimental/sgl-router/src/policies/power_of_two.rs @@ -1,48 +1,76 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 -use crate::policies::admission::compare_prefill_pressure; +//! Power-of-k-choices load balancing. Each selection samples `choices` +//! random distinct candidates and dispatches to the least-pressured one, +//! so N router replicas reading the same load snapshot do not converge on +//! one shared fleet minimum. When a queue limit is configured the sample +//! is drawn from the tier of workers the queue gate admits, and only from +//! the whole pool when every worker is queueing (the second tier keeps an +//! all-queueing fleet routable). +//! +//! `select` and `propose` share one scan so they cannot disagree on the +//! same pool: see [`best_two_of_sample`] for why it is a linear scan and +//! never a sort. + +use crate::config::DEFAULT_MIN_LOAD_CHOICES; +use crate::policies::admission::{compare_prefill_pressure, queue_gate_admits}; use crate::policies::engine_load::EngineLoadSnapshot; use crate::policies::{Policy, ProposalKind, SelectionContext, SelectionProposal}; use crate::workers::Worker; +use rand::seq::index::sample; use rand::Rng; +use std::borrow::Cow; use std::sync::Arc; -#[derive(Debug, Default)] -pub struct PowerOfTwoChoicesPolicy; +#[derive(Debug)] +pub struct PowerOfTwoChoicesPolicy { + choices: usize, + queue_limit: Option, +} + +impl Default for PowerOfTwoChoicesPolicy { + fn default() -> Self { + Self::new() + } +} impl PowerOfTwoChoicesPolicy { pub fn new() -> Self { - Self + Self { + choices: DEFAULT_MIN_LOAD_CHOICES, + queue_limit: None, + } + } + + /// Sets the sample size and the queue gate for the fallback path. + /// `choices` is clamped to at least 1. + pub fn with_load_control(mut self, choices: usize, queue_limit: Option) -> Self { + self.choices = choices.max(1); + self.queue_limit = queue_limit; + self } } impl Policy for PowerOfTwoChoicesPolicy { fn select(&self, workers: &[Arc], ctx: &SelectionContext<'_>) -> Option> { - select_with_snapshot(workers, ctx.load_snapshot()) + select_k_with_snapshot(workers, ctx.load_snapshot(), self.choices, self.queue_limit) } - /// Returns the primary and backup from one sample. + /// Returns the primary and backup from one sample. A sample of one + /// (`--min-load-choices 1`) has no second member, so the proposal + /// carries no backup and admission loses its backup paths. fn propose( &self, workers: &[Arc], ctx: &SelectionContext<'_>, ) -> Option { - match workers.len() { - 0 => None, - 1 => Some( - SelectionProposal::primary(workers[0].clone()).with_kind(ProposalKind::PowerOfTwo), - ), - len => { - let mut rng = rand::thread_rng(); - let i = rng.gen_range(0..len); - let mut j = rng.gen_range(0..len - 1); - if j >= i { - j += 1; - } - let (primary, backup) = ordered_pair(&workers[i], &workers[j], ctx); - Some(SelectionProposal::with_backup(primary, backup)) - } + let snapshot = ctx.load_snapshot(); + let pool = sample_pool(workers, snapshot, self.queue_limit); + let (primary, backup) = best_two_of_sample(&pool, snapshot, self.choices)?; + match backup { + Some(backup) => Some(SelectionProposal::with_backup(primary, backup)), + None => Some(SelectionProposal::primary(primary).with_kind(ProposalKind::PowerOfTwo)), } } @@ -51,49 +79,477 @@ impl Policy for PowerOfTwoChoicesPolicy { } } -pub(crate) fn select_with_snapshot( +pub(crate) fn select_k_with_snapshot( workers: &[Arc], snapshot: Option<&EngineLoadSnapshot>, + choices: usize, + queue_limit: Option, ) -> Option> { - match workers.len() { - 0 => None, - 1 => Some(workers[0].clone()), - len => { - let mut rng = rand::thread_rng(); - let i = rng.gen_range(0..len); - let mut j = rng.gen_range(0..len - 1); - if j >= i { - j += 1; + let pool = sample_pool(workers, snapshot, queue_limit); + best_two_of_sample(&pool, snapshot, choices).map(|(primary, _)| primary) +} + +/// The tier the sample is drawn from: workers the queue gate admits, +/// or the whole pool when every worker is queueing (the second tier +/// keeps an all-queueing fleet routable). A sample must never land on +/// a queueing worker while an unqueued one exists. +fn sample_pool<'w>( + workers: &'w [Arc], + snapshot: Option<&EngineLoadSnapshot>, + queue_limit: Option, +) -> Cow<'w, [Arc]> { + // Without a limit there is nothing to gate on, and without a snapshot + // the gate cannot be evaluated per-worker: both mean the whole pool. + let (Some(snapshot), Some(limit)) = (snapshot, queue_limit) else { + return Cow::Borrowed(workers); + }; + let unqueued: Vec> = workers + .iter() + .filter(|worker| queue_gate_admits(snapshot, worker.as_ref(), Some(limit))) + .cloned() + .collect(); + if unqueued.is_empty() { + Cow::Borrowed(workers) + } else { + Cow::Owned(unqueued) + } +} + +/// The sample's two least-pressured members, best first, as pool indices +/// resolved to workers. `choices >= pool` skips the shuffle and scans the +/// whole pool from a random offset, so the winner is the exact minimum +/// whenever pressures differ and ties resolve randomly rather than always +/// at `pool[0]`; a smaller `choices` draws that many distinct indices, +/// which `rand`'s `sample` returns fully shuffled, so ties inside a sample +/// resolve randomly too. +/// +/// Tie-breaking is not a detail of the large-`k` path: an idle fleet ties +/// on every comparison, and the default `k = 2` is already `>= pool` on a +/// two-worker fleet, so a fixed scan order would pin every fallback +/// dispatch to the first worker. +/// +/// Both tiers scan linearly and never sort. `compare_prefill_pressure` +/// is only a pairwise comparison: two workers that both publish +/// `estimated_prefill_queue_ms` are ordered on that estimate, and any +/// other pair on the waiting-token tuple, so it is not a total order +/// across a mixed set (an idle worker publishes no estimate). Handing it +/// to `sort_by` makes the standard library panic with "user-provided +/// comparison function does not correctly implement a total order", +/// unwinding whichever request task was selecting at the time. +fn best_two_of_sample( + pool: &[Arc], + snapshot: Option<&EngineLoadSnapshot>, + choices: usize, +) -> Option<(Arc, Option>)> { + let len = pool.len(); + if len == 0 { + return None; + } + let choices = choices.max(1); + let drawn: Vec = if choices >= len { + // The whole pool, presented from a random offset. The scan below + // keeps the incumbent on a tie, so a fixed start would hand every + // tie to `pool[0]`; one rotation is O(1) randomness and leaves the + // exact minimum intact whenever the pressures actually differ. + let start = rand::thread_rng().gen_range(0..len); + (0..len).map(|offset| (start + offset) % len).collect() + } else if choices == 1 { + vec![rand::thread_rng().gen_range(0..len)] + } else { + sample(&mut rand::thread_rng(), len, choices).into_vec() + }; + let mut best: Option = None; + let mut runner_up: Option = None; + for index in drawn { + // Only a strict improvement displaces the incumbent, so a tie keeps + // whichever member the draw presented first. + if best.is_none_or(|current| { + compare_prefill_pressure(&pool[index], &pool[current], snapshot).is_lt() + }) { + runner_up = best; + best = Some(index); + } else if runner_up.is_none_or(|current| { + compare_prefill_pressure(&pool[index], &pool[current], snapshot).is_lt() + }) { + runner_up = Some(index); + } + } + best.map(|index| { + ( + Arc::clone(&pool[index]), + runner_up.map(|index| Arc::clone(&pool[index])), + ) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + use crate::policies::engine_load::NativeCacheWorkerLoad; + use std::time::Instant; + + fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}:30000"), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("model".into())], + bootstrap_port: None, + })) + } + + /// Snapshot keyed on waiting depth; `waiting` sets both the queue-gate + /// reading (`num_waiting_reqs`) and the pressure ordering + /// (`num_waiting_uncached_tokens`), so one knob drives both. + fn snapshot(entries: &[(&Arc, u64)]) -> EngineLoadSnapshot { + EngineLoadSnapshot::from_native_cache_workers( + 7, + entries + .iter() + .map(|(worker, waiting)| { + ( + worker.url.clone(), + NativeCacheWorkerLoad { + num_running_reqs: 0, + num_waiting_reqs: *waiting, + num_waiting_uncached_tokens: *waiting, + num_used_tokens: 10, + num_total_tokens: 10, + max_total_num_tokens: 10_000, + max_running_requests: 64, + prefill_throughput_tokens_per_s: None, + estimated_prefill_queue_ms: None, + captured_at: Instant::now(), + }, + ) + }) + .collect(), + ) + } + + /// A fleet where only some workers publish `estimated_prefill_queue_ms`. + /// An idle worker has no throughput delta to derive one from, so this is + /// the steady state, not an edge case - and it makes + /// `compare_prefill_pressure` intransitive: a slow worker with a shallow + /// queue loses to a fast worker with a deep one on the estimate, while + /// both are ordered against an estimate-less worker on waiting tokens. + fn mixed_estimate_snapshot(workers: &[Arc]) -> EngineLoadSnapshot { + EngineLoadSnapshot::from_native_cache_workers( + 11, + workers + .iter() + .enumerate() + .map(|(index, worker)| { + let waiting = (index as u64 * 7) % 13; + ( + worker.url.clone(), + NativeCacheWorkerLoad { + num_running_reqs: 0, + num_waiting_reqs: 0, + num_waiting_uncached_tokens: waiting, + num_used_tokens: 10, + num_total_tokens: 10, + max_total_num_tokens: 10_000, + max_running_requests: 64, + prefill_throughput_tokens_per_s: None, + // Every third worker is idle and publishes no + // estimate; the rest rank inversely to `waiting`. + estimated_prefill_queue_ms: (index % 3 != 0) + .then(|| (13 - waiting) as f64), + captured_at: Instant::now(), + }, + ) + }) + .collect(), + ) + } + + /// `compare_prefill_pressure` is a pairwise comparison, not a total + /// order, so the k-way minimum must be a linear scan. Sorting a sample + /// this size panics with "user-provided comparison function does not + /// correctly implement a total order" - every request that samples a + /// mixed idle-and-busy fleet under a large `--min-load-choices` dies + /// with its task. + #[test] + fn a_large_sample_over_mixed_estimates_never_panics() { + let model = ModelId("model".into()); + let workers: Vec> = (0..64).map(|i| worker(&format!("w{i}"))).collect(); + let loads = mixed_estimate_snapshot(&workers); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&loads); + + for choices in [2, 3, 21, 32, 64, 128] { + let policy = PowerOfTwoChoicesPolicy::new().with_load_control(choices, None); + for _ in 0..16 { + let proposal = policy + .propose(&workers, &ctx) + .expect("a non-empty fleet must produce a proposal"); + assert!(workers.iter().any(|w| w.id == proposal.primary.id)); + let backup = proposal.backup.expect("k >= 2 must carry a runner-up"); + assert_ne!( + backup.id, proposal.primary.id, + "the sample draws distinct indices" + ); + select_k_with_snapshot(&workers, Some(&loads), choices, None) + .expect("select must agree that the fleet is routable"); } - Some(select_lower_pressure(&workers[i], &workers[j], snapshot)) + } + } + + /// `propose` and `select` must rank the same pool the same way. + #[test] + fn propose_and_select_agree_on_the_sample_minimum() { + let model = ModelId("model".into()); + let deep = worker("deep"); + let middle = worker("middle"); + let shallow = worker("shallow"); + let workers = vec![Arc::clone(&deep), Arc::clone(&middle), Arc::clone(&shallow)]; + let loads = snapshot(&[(&deep, 100), (&middle, 50), (&shallow, 1)]); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&loads); + + // `choices >= pool` on both paths: the exact minimum, and the + // runner-up is the second-lowest rather than a shuffle artifact. + let proposal = PowerOfTwoChoicesPolicy::new() + .with_load_control(3, None) + .propose(&workers, &ctx) + .expect("three candidates must produce a proposal"); + assert_eq!(proposal.primary.id, shallow.id); + assert_eq!( + proposal.backup.expect("a three-member sample has one").id, + middle.id + ); + assert_eq!( + select_k_with_snapshot(&workers, Some(&loads), 3, None) + .expect("the pool is non-empty") + .id, + proposal.primary.id + ); + } + + /// A one-member sample has no runner-up, so admission loses its backup + /// paths entirely. Documented on `--min-load-choices`; pinned here. + #[test] + fn one_choice_proposes_no_backup() { + let model = ModelId("model".into()); + let left = worker("left"); + let right = worker("right"); + let workers = vec![Arc::clone(&left), Arc::clone(&right)]; + let loads = snapshot(&[(&left, 1), (&right, 2)]); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&loads); + + let proposal = PowerOfTwoChoicesPolicy::new() + .with_load_control(1, None) + .propose(&workers, &ctx) + .expect("a non-empty pool must produce a proposal"); + assert!(proposal.backup.is_none()); + assert_eq!(proposal.kind, ProposalKind::PowerOfTwo); + } + + #[test] + fn sample_never_lands_on_a_queueing_worker_while_an_unqueued_one_exists() { + let queued_a = worker("queued_a"); + let queued_b = worker("queued_b"); + let queued_c = worker("queued_c"); + let unqueued = worker("unqueued"); + let workers = vec![ + Arc::clone(&queued_a), + Arc::clone(&queued_b), + Arc::clone(&queued_c), + Arc::clone(&unqueued), + ]; + let loads = snapshot(&[ + (&queued_a, 9), + (&queued_b, 5), + (&queued_c, 12), + (&unqueued, 3), + ]); + + for _ in 0..64 { + let selected = select_k_with_snapshot(&workers, Some(&loads), 2, Some(4)) + .expect("an unqueued worker exists"); + assert_eq!(selected.id, unqueued.id); + } + } + + #[test] + fn all_queueing_fleet_still_selects_from_the_whole_pool() { + let left = worker("left"); + let right = worker("right"); + let third = worker("third"); + let workers = vec![Arc::clone(&left), Arc::clone(&right), Arc::clone(&third)]; + let loads = snapshot(&[(&left, 9), (&right, 5), (&third, 12)]); + + // `choices >= pool` on the second tier, so the winner is the exact + // pressure minimum of the whole fleet rather than any pool member. + for _ in 0..64 { + let selected = select_k_with_snapshot(&workers, Some(&loads), 3, Some(4)) + .expect("an all-queueing fleet must still route"); + assert_eq!( + selected.id, right.id, + "the whole-pool tier must rank the all-queueing fleet by pressure" + ); + } + // With a two-member sample the winner is still never the worker the + // other two both beat on pressure. + for _ in 0..64 { + let selected = select_k_with_snapshot(&workers, Some(&loads), 2, Some(4)) + .expect("an all-queueing fleet must still route"); + assert!( + workers.iter().any(|worker| worker.id == selected.id), + "the whole-pool tier must return a pool member" + ); + } + } + + #[test] + fn choices_at_or_above_the_pool_size_returns_the_exact_minimum() { + let deep = worker("deep"); + let shallow = worker("shallow"); + let middle = worker("middle"); + // Pool order deliberately disagrees with the pressure ordering. + let workers = vec![Arc::clone(&deep), Arc::clone(&shallow), Arc::clone(&middle)]; + let loads = snapshot(&[(&deep, 100), (&shallow, 1), (&middle, 50)]); + + for choices in [3, 8] { + let selected = select_k_with_snapshot(&workers, Some(&loads), choices, None) + .expect("the pool is non-empty"); + assert_eq!( + selected.id, shallow.id, + "choices >= pool must return the exact minimum whatever the scan offset" + ); + } + } + + /// A fleet with nothing in flight ties on every pressure comparison, + /// and the scan keeps its incumbent on a tie, so the scan order alone + /// decides where the request goes. Resolving that in pool order sends + /// every tied dispatch to `pool[0]` — and with the default `k = 2` on + /// a two-worker fleet the `choices >= pool` path takes every dispatch, + /// so the whole fallback pins to one worker. Caught by the hicache + /// storage-tier e2e test: all of its filler traffic landed on a single + /// engine, turning that engine's host tier over before the probe could + /// read the primed prefix back from it. + #[test] + fn a_tied_pool_spreads_instead_of_pinning_the_first_worker() { + let model = ModelId("model".into()); + let first = worker("first"); + let second = worker("second"); + let workers = vec![Arc::clone(&first), Arc::clone(&second)]; + // Equal pressure on both: an idle fleet's steady state, not an + // edge case. + let loads = snapshot(&[(&first, 0), (&second, 0)]); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&loads); + let policy = PowerOfTwoChoicesPolicy::new().with_load_control(2, None); + + let mut selected_second = false; + let mut proposed_second = false; + for _ in 0..256 { + selected_second |= select_k_with_snapshot(&workers, Some(&loads), 2, None) + .expect("the pool is non-empty") + .id + == second.id; + proposed_second |= policy + .propose(&workers, &ctx) + .expect("the pool is non-empty") + .primary + .id + == second.id; + } + assert!( + selected_second, + "a tie must not always resolve to the first worker in pool order" + ); + assert!( + proposed_second, + "the proposal path shares the scan, so it must spread the same way" + ); + } + + #[test] + fn one_choice_draws_within_the_tier_and_ignores_pressure() { + let deep = worker("deep"); + let shallow = worker("shallow"); + let workers = vec![Arc::clone(&deep), Arc::clone(&shallow)]; + // Both sit under the gate, so the tier is the whole pool and a + // single draw must reach the deeper worker too - that is exactly + // what stops N replicas converging on one shared minimum. + let loads = snapshot(&[(&deep, 3), (&shallow, 0)]); + + let mut saw_deep = false; + for _ in 0..256 { + let selected = select_k_with_snapshot(&workers, Some(&loads), 1, Some(4)) + .expect("both workers are unqueued"); + saw_deep |= selected.id == deep.id; + } + assert!( + saw_deep, + "a one-member sample must be a draw, not the pressure minimum" + ); + } + + #[test] + fn one_choice_stays_inside_the_queue_gate_tier() { + let queued = worker("queued"); + let also_queued = worker("also_queued"); + let unqueued = worker("unqueued"); + let workers = vec![ + Arc::clone(&queued), + Arc::clone(&also_queued), + Arc::clone(&unqueued), + ]; + let loads = snapshot(&[(&queued, 9), (&also_queued, 20), (&unqueued, 0)]); + + for _ in 0..64 { + let selected = select_k_with_snapshot(&workers, Some(&loads), 1, Some(4)) + .expect("an unqueued worker exists"); + assert_eq!( + selected.id, unqueued.id, + "a single draw must still respect the queue-gate tier" + ); + } + } + + #[test] + fn two_workers_with_defaults_propose_both_ordered_by_pressure() { + let model = ModelId("model".into()); + let busy = worker("busy"); + let idle = worker("idle"); + let workers = vec![Arc::clone(&busy), Arc::clone(&idle)]; + let loads = snapshot(&[(&busy, 512), (&idle, 16)]); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&loads); + + let proposal = PowerOfTwoChoicesPolicy::new() + .with_load_control(2, None) + .propose(&workers, &ctx) + .expect("two candidates must produce a proposal"); + + assert_eq!(proposal.primary.id, idle.id); + assert_eq!( + proposal.backup.expect("P2 keeps its other sample").id, + busy.id + ); + } + + #[test] + fn proposal_primary_respects_the_queue_gate_tier() { + let model = ModelId("model".into()); + let queued_a = worker("queued_a"); + let queued_b = worker("queued_b"); + let unqueued = worker("unqueued"); + let workers = vec![ + Arc::clone(&queued_a), + Arc::clone(&queued_b), + Arc::clone(&unqueued), + ]; + let loads = snapshot(&[(&queued_a, 9), (&queued_b, 5), (&unqueued, 3)]); + let ctx = SelectionContext::new(&model, None).with_load_snapshot(&loads); + let policy = PowerOfTwoChoicesPolicy::new().with_load_control(2, Some(4)); + + for _ in 0..64 { + let proposal = policy + .propose(&workers, &ctx) + .expect("an unqueued worker exists"); + assert_eq!(proposal.primary.id, unqueued.id); } } } - -fn select_lower_pressure( - left: &Arc, - right: &Arc, - snapshot: Option<&EngineLoadSnapshot>, -) -> Arc { - ordered_pair_with_snapshot(left, right, snapshot).0 -} - -fn ordered_pair( - left: &Arc, - right: &Arc, - ctx: &SelectionContext<'_>, -) -> (Arc, Arc) { - ordered_pair_with_snapshot(left, right, ctx.load_snapshot()) -} - -fn ordered_pair_with_snapshot( - left: &Arc, - right: &Arc, - snapshot: Option<&EngineLoadSnapshot>, -) -> (Arc, Arc) { - if compare_prefill_pressure(left, right, snapshot).is_gt() { - (Arc::clone(right), Arc::clone(left)) - } else { - (Arc::clone(left), Arc::clone(right)) - } -} diff --git a/experimental/sgl-router/src/policies/scoring/mod.rs b/experimental/sgl-router/src/policies/scoring/mod.rs index d5a17df1a..6332d1a9d 100644 --- a/experimental/sgl-router/src/policies/scoring/mod.rs +++ b/experimental/sgl-router/src/policies/scoring/mod.rs @@ -747,9 +747,15 @@ mod tests { (&ws[2], 0, 0, 0, 4_096), ]); - let decision = - resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot, None) - .expect("capacity exhaustion must degrade inside the filtered domain"); + let decision = resolve_prefill( + &CandidateRange::global(&ws), + &proposal, + 32, + &snapshot, + None, + 2, + ) + .expect("capacity exhaustion must degrade inside the filtered domain"); assert!(matches!(decision.selected.id.0.as_str(), "a" | "b")); } @@ -784,9 +790,15 @@ mod tests { assert_eq!(proposal.primary.id, ws[2].id); let snapshot = EngineLoadSnapshot::default(); - let decision = - resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot, None) - .expect("an eligible escape worker exists"); + let decision = resolve_prefill( + &CandidateRange::global(&ws), + &proposal, + 32, + &snapshot, + None, + 2, + ) + .expect("an eligible escape worker exists"); assert_ne!(decision.selected.id, ws[2].id); assert!(matches!(decision.selected.id.0.as_str(), "a" | "b")); diff --git a/experimental/sgl-router/src/policies/selection.rs b/experimental/sgl-router/src/policies/selection.rs index 40e7e799b..92fc274b3 100644 --- a/experimental/sgl-router/src/policies/selection.rs +++ b/experimental/sgl-router/src/policies/selection.rs @@ -72,6 +72,8 @@ pub(crate) struct PrefillSelectionInputs<'a> { pub worker_queue_limit: Option, /// `--saturation-queue-floor`. `None` disables the saturation pin. pub saturation_queue_floor: Option, + /// `--min-load-choices`: sample size for the min-load capacity fallback. + pub min_load_choices: usize, } /// The queue-gate blind warn is sampled: it fires on a per-request path, and @@ -481,6 +483,7 @@ impl<'a> Selector<'a> { inputs.request_input_tokens, snapshot, inputs.worker_queue_limit, + inputs.min_load_choices, ) } else { resolve_prefill_admitted( @@ -827,6 +830,7 @@ mod tests { workers, worker_queue_limit: None, saturation_queue_floor: None, + min_load_choices: 2, ttft_slo_ms: None, tps_slo: None, session_affinity_mode: SessionAffinityMode::Bucket, diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index 65fe5806f..2f8711689 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -3,6 +3,7 @@ use crate::config::{ ConflictPolicy, ParamSpec, SamplingField, SamplingOverrides, SessionAffinityMode, + DEFAULT_MIN_LOAD_CHOICES, }; use crate::discovery::{ModelId, WorkerMode}; use crate::policies::kv_events::{compute_block_hashes, compute_block_hashes_bigram}; @@ -646,6 +647,16 @@ pub async fn chat_completions( .affinity .as_ref() .and_then(|config| config.saturation_queue_floor); + // Sample size for the min-load fallback beneath admission + // (`--min-load-choices`). A policy with no affinity config never reaches + // the cache-aware paths, so it keeps the pre-existing power-of-2 default. + let min_load_choices = ctx + .config + .model + .affinity + .as_ref() + .map(|config| config.min_load_choices) + .unwrap_or(DEFAULT_MIN_LOAD_CHOICES); // Each Bucket retry rebuilds the proposal and reruns Admission/Guard. let worker = select_prefill_worker(&PrefillSelectionInputs { policy: policy.as_ref(), @@ -666,6 +677,7 @@ pub async fn chat_completions( session_affinity_mode, worker_queue_limit, saturation_queue_floor, + min_load_choices, }) .map_err(|reason| policy_selection_failed(&ctx, &model_str, reason))?;