From e3374fbd7941988c0e1147fb568d4026c40cecba Mon Sep 17 00:00:00 2001 From: nnshah1 Date: Tue, 3 Nov 2026 13:05:18 -0800 Subject: [PATCH] feat(global-router): least-in-flight places pools by affinity's rule Least-in-flight ignored regions and each pool's capacity, so a request could leave its region for a pool with one fewer request in flight. It now uses affinity's placement rule on both paths: nearest region first, then the lower load (requests in flight over the pool's capacity, from its Relay), and it leaves the nearest pool only for one with more than --headroom-factor times its spare room. The rule moves to placement.rs, shared by affinity and both least-in-flight implementations, and --router-policy least-in-flight now needs --region. This intentionally changes the current least-in-flight selector. Co-Authored-By: Claude Opus 5.5 Signed-off-by: nnshah1 --- lib/global-frontend/README.md | 19 ++- lib/global-frontend/src/main.rs | 101 +++++++++---- lib/global-router/src/affinity.rs | 90 ++++-------- lib/global-router/src/lib.rs | 3 +- lib/global-router/src/placement.rs | 87 ++++++++++++ .../src/selection_core/catalog.rs | 7 +- .../src/selection_core/directory.rs | 20 ++- lib/global-router/src/selection_core/mod.rs | 69 +++++++-- .../src/selection_core/policies.rs | 81 +++++++---- lib/global-router/src/selector.rs | 134 ++++++++++++++++-- lib/global-router/src/signals.rs | 40 ++++-- 11 files changed, 481 insertions(+), 170 deletions(-) create mode 100644 lib/global-router/src/placement.rs diff --git a/lib/global-frontend/README.md b/lib/global-frontend/README.md index ade53a841ff6..bd5df67824f2 100644 --- a/lib/global-frontend/README.md +++ b/lib/global-frontend/README.md @@ -42,7 +42,7 @@ model; every policy then rotates across the pool's frontends. | Policy | Picks | |---|---| | `round-robin` (default) | Each pool in turn. | -| `least-in-flight` | The pool with the fewest of this process's requests in flight. | +| `least-in-flight` | The nearest pool with the fewest requests in flight, moving only for a pool with much more room. | | `affinity` | The session's pool, nearest region first, moving only for a pool with much more room. | `affinity` keeps each session's turns on one pool, so the pool's KV cache for @@ -57,10 +57,19 @@ nearest such pool, or when that pool is full, choosing the nearest pool with any room. The session's pin moves with it. Pins are per process and expire after `--session-idle-secs` (default 1800) without a request. -`affinity` needs `--region`, matching the pool's region (the cluster name for -explicit pools). It reads each pool's load from the pool's PoolRelay, given per pool -over plain HTTP. Each pool needs a Relay of its own, since a Relay reports the -load of everything it watches as one: +`least-in-flight` places requests by the same rule, with requests in flight +over each pool's capacity as the load: this process's requests in flight to the +pool over the requests its workers can run at once, as its Relay reports. Pools +are ordered nearest region first, then less loaded, then fewer in flight; the +request goes to the first, rotating among pools that tie with it, unless +another pool has more than `--headroom-factor` times its spare room. A pool +without a Relay report counts as full, so without `--pool-relay` requests stay +on the nearest pools, fewest in flight first. + +`affinity` and `least-in-flight` need `--region`, matching the pool's region +(the cluster name for explicit pools). They read each pool's load from the +pool's PoolRelay, given per pool over plain HTTP. Each pool needs a Relay of +its own, since a Relay reports the load of everything it watches as one: ```bash cargo run -p dynamo-global-frontend -- --port 8000 \ diff --git a/lib/global-frontend/src/main.rs b/lib/global-frontend/src/main.rs index 73c5e5a32c02..57f036531c4d 100644 --- a/lib/global-frontend/src/main.rs +++ b/lib/global-frontend/src/main.rs @@ -15,8 +15,8 @@ use dynamo_global_frontend::{ follow_relays, measure_round_trips, serve, serve_relay_views, }; use dynamo_global_router::{ - AffinitySelector, InFlightTracker, LeastInFlightSelector, Pool, PoolSelector, PoolSignals, - RoundRobinSelector, SelectionCorePolicy, SelectionCoreSelector, + AffinitySelector, InFlightTracker, LeastInFlightConfig, LeastInFlightSelector, Pool, + PoolSelector, PoolSignals, RoundRobinSelector, SelectionCorePolicy, SelectionCoreSelector, }; use tokio::sync::watch; use tokio_util::sync::CancellationToken; @@ -32,7 +32,8 @@ const PIN_SWEEP_INTERVAL: Duration = Duration::from_secs(60); enum RouterPolicy { /// Every pool serving the model in turn. RoundRobin, - /// The pool with the fewest of this process's requests in flight. + /// The nearest pool with the fewest requests in flight over its capacity + /// (from `--pool-relay`), moving only for much more room. LeastInFlight, /// Keep each session on one pool, nearest region first, and move it only /// for a pool with much more room (see `--headroom-factor`). Reads each @@ -70,9 +71,13 @@ struct Args { /// Select pools through the embedded router core (not yet with affinity). #[arg(long)] use_selection_core: bool, - /// The region this process runs in. Affinity prefers it, and metrics - /// distinguish local requests from requests sent to other regions. - #[arg(long, required_if_eq("router_policy", "affinity"))] + /// The region this process runs in. `affinity` and `least-in-flight` + /// prefer it, and metrics distinguish local requests from requests sent + /// to other regions. + #[arg( + long, + required_if_eq_any([("router_policy", "affinity"), ("router_policy", "least-in-flight")]) + )] region: Option, /// A pool's PoolRelay, which streams the pool's load, as /// `//=`. Repeat the flag for each pool; @@ -80,8 +85,9 @@ struct Args { /// everything it watches as one. #[arg(long = "pool-relay", value_parser = parse_pool_relay)] pool_relays: Vec<(PoolKey, Url)>, - /// `affinity`: a request leaves its pool only for one with more than this - /// many times its spare capacity, or when its pool is full. + /// `affinity` and `least-in-flight`: a request leaves its pool only for + /// one with more than this many times its spare capacity, or when its + /// pool is full. #[arg(long, default_value_t = 2.0)] headroom_factor: f64, /// `affinity`: how long a session keeps its pool without a request. @@ -172,9 +178,6 @@ async fn main() -> Result<()> { if !(args.headroom_factor >= 1.0 && args.headroom_factor.is_finite()) { bail!("--headroom-factor must be at least 1"); } - let selection_core_policy = (args.use_selection_core) - .then(|| selection_core_policy(args.router_policy)) - .transpose()?; let shutdown = CancellationToken::new(); let signal = shutdown.clone(); @@ -186,8 +189,9 @@ async fn main() -> Result<()> { }); // Keep the explicit catalog sender alive for the static-pool mode. - let (_catalog, pools) = watch::channel(Arc::new(PoolSnapshot::new(args.pools))); - let pools = if let Some(url) = args.global_discovery_url { + let (_catalog, pools) = + watch::channel(Arc::new(PoolSnapshot::new(std::mem::take(&mut args.pools)))); + let pools = if let Some(url) = args.global_discovery_url.take() { let client = dynamo_global_discovery_client::GlobalDiscoveryClient::new( url, Duration::from_secs(5), @@ -239,7 +243,9 @@ async fn main() -> Result<()> { .build() .context("build the pool HTTP client")?; let metrics = Metrics::new(args.region.clone(), in_flight.clone()); - let app = if let Some(policy) = selection_core_policy { + let load_signals = || load_signals(&args, &pools, &signals, &shutdown); + let app = if args.use_selection_core { + let policy = selection_core_policy(&args, load_signals)?; let selection_core = Arc::new(SelectionCoreSelector::new(policy, in_flight.clone()).await?); // Wait for the first pool snapshot, or early requests get a 503. let (first_snapshot, reconcile) = @@ -252,21 +258,15 @@ async fn main() -> Result<()> { } else { let selector: Arc = match args.router_policy { RouterPolicy::RoundRobin => Arc::new(RoundRobinSelector::new()), - RouterPolicy::LeastInFlight => Arc::new(LeastInFlightSelector::new(in_flight.clone())), + RouterPolicy::LeastInFlight => Arc::new(LeastInFlightSelector::new( + in_flight.clone(), + load_signals(), + args.region.clone().context("--region is required")?, + args.headroom_factor, + )), RouterPolicy::Affinity => { let region = args.region.clone().context("--region is required")?; - if args.pool_relays.is_empty() && args.relay_view_port.is_none() { - tracing::warn!( - "no --pool-relay is given, so affinity knows no pool's load \ - and keeps each session on its nearest pool" - ); - } - tokio::spawn(measure_round_trips( - pools.clone(), - signals.clone(), - ROUND_TRIP_INTERVAL, - shutdown.child_token(), - )); + let signals = load_signals(); let idle = Duration::from_secs(args.session_idle_secs); let selector = Arc::new(AffinitySelector::new( signals, @@ -304,16 +304,48 @@ async fn main() -> Result<()> { } } -fn selection_core_policy(policy: RouterPolicy) -> Result { - match policy { +/// `signals` starts the load signals the policy weighs, if any. +fn selection_core_policy( + args: &Args, + signals: impl FnOnce() -> Arc, +) -> Result { + match args.router_policy { RouterPolicy::RoundRobin => Ok(SelectionCorePolicy::RoundRobin), - RouterPolicy::LeastInFlight => Ok(SelectionCorePolicy::LeastInFlight), + RouterPolicy::LeastInFlight => { + Ok(SelectionCorePolicy::LeastInFlight(LeastInFlightConfig { + signals: signals(), + region: args.region.clone().context("--region is required")?, + headroom_factor: args.headroom_factor, + })) + } RouterPolicy::Affinity => { bail!("--use-selection-core does not support --router-policy affinity yet") } } } +/// Starts the round-trip probes; the Relay catalog fills the rest. +fn load_signals( + args: &Args, + pools: &watch::Receiver>, + signals: &Arc, + shutdown: &CancellationToken, +) -> Arc { + if args.pool_relays.is_empty() && args.relay_view_port.is_none() { + tracing::warn!( + "no --pool-relay is given, so no pool's load is known: requests stay on \ + their nearest pool" + ); + } + tokio::spawn(measure_round_trips( + pools.clone(), + signals.clone(), + ROUND_TRIP_INTERVAL, + shutdown.child_token(), + )); + signals.clone() +} + /// Drops expired session pins once a `PIN_SWEEP_INTERVAL`, off the request /// path, until `cancel` fires. async fn sweep_pins(selector: Arc, cancel: CancellationToken) { @@ -366,7 +398,7 @@ mod tests { } #[test] - fn parses_pool_relays_and_requires_a_region_for_affinity() { + fn parses_pool_relays_and_requires_a_region_for_load_aware_policies() { let (key, relay) = parse_pool_relay("ohio/dynamo/blue=http://relay:50051").unwrap(); assert_eq!(key.dgd_name(), "blue"); assert_eq!(relay.as_str(), "http://relay:50051/"); @@ -386,6 +418,13 @@ mod tests { assert!(Args::try_parse_from(affinity).is_err()); let args = Args::try_parse_from([&affinity[..], &["--region", "ohio"]].concat()).unwrap(); assert_eq!(args.router_policy, RouterPolicy::Affinity); + let least = [ + "global-frontend", + pool, + "--router-policy", + "least-in-flight", + ]; + assert!(Args::try_parse_from(least).is_err()); let args = Args::try_parse_from(["global-frontend", pool]).unwrap(); assert_eq!(args.router_policy, RouterPolicy::RoundRobin); } diff --git a/lib/global-router/src/affinity.rs b/lib/global-router/src/affinity.rs index 0e5eb85226e7..f5b85c84efec 100644 --- a/lib/global-router/src/affinity.rs +++ b/lib/global-router/src/affinity.rs @@ -8,6 +8,7 @@ use dashmap::DashMap; use dynamo_global_common::PoolId; use crate::Pool; +use crate::placement::{self, Placed}; use crate::selector::{Cursors, PoolSelector, PoolTarget, RoutingRequest, reachable, target}; use crate::signals::PoolSignals; @@ -96,6 +97,9 @@ impl PoolSelector for AffinitySelector { candidates: &[&'a Pool], ) -> Option> { let placed = place(candidates, &self.region, &self.signals); + if placed.is_empty() { + return None; + } let pinned = request.session.and_then(|session| { let pin = self.pins.get(session)?; if pin.used.elapsed() > self.idle { @@ -103,83 +107,43 @@ impl PoolSelector for AffinitySelector { } placed .iter() - .find(|candidate| candidate.pool.id == pin.pool) + .position(|candidate| candidate.item.id == pin.pool) }); - let choice = headroom(&placed, pinned.or(placed.first()), self.headroom_factor)?; + let choice = + placed[placement::headroom(&placed, pinned.unwrap_or(0), self.headroom_factor)].item; if let Some(session) = request.session { - self.pin(session, &choice.pool.id); + self.pin(session, &choice.id); } - target(&self.frontends, choice.pool) + target(&self.frontends, choice) } } -/// A reachable pool, with what the affinity selector weighs. -#[derive(Debug)] -struct Candidate<'a> { - pool: &'a Pool, - /// How full the pool is against its own capacity, 1 being full (see - /// [`PoolSignals::occupancy`]). `None` without a fresh report. - load: Option, - /// Round trip to the pool's region: zero for the region the request - /// entered, otherwise as last measured. `None` until measured. - distance: Option, -} - /// The candidates that have a frontend, for a request entering `region`, -/// from least to most expensive: nearest region first, then less full, then -/// catalog order. Pools whose distance or load is unknown come after those -/// whose is known. -fn place<'a>(candidates: &[&'a Pool], region: &str, signals: &PoolSignals) -> Vec> { +/// ordered by [`placement::order`] on their occupancy: nearest region first, +/// then less full, then catalog order. +fn place<'a>( + candidates: &[&'a Pool], + region: &str, + signals: &PoolSignals, +) -> Vec> { let round_trips = signals.region_round_trips(); let mut placed = candidates .iter() .copied() .filter(|pool| reachable(pool)) - .map(|pool| Candidate { - pool, + .map(|pool| Placed { + item: pool, load: signals.occupancy(&pool.id).map(|occupancy| occupancy.load), - distance: if pool.location.region == region { - Some(Duration::ZERO) - } else { - round_trips.get(&pool.location.region).copied() - }, + distance: placement::distance(&pool.location.region, region, &round_trips), + in_flight: 0, }) .collect::>(); - // Stable, so ties keep catalog order. - placed.sort_by(|left, right| { - let distance = |candidate: &Candidate<'_>| candidate.distance.unwrap_or(Duration::MAX); - let load = |candidate: &Candidate<'_>| candidate.load.unwrap_or(f64::INFINITY); - distance(left) - .cmp(&distance(right)) - .then_with(|| load(left).total_cmp(&load(right))) - }); + placement::order(&mut placed); placed } -/// `default`, unless another pool has more than `factor` times its spare -/// capacity (1 minus its load); then the first such pool in `candidates`' -/// order. A default whose load is unknown counts as full. Without a default, -/// the first candidate. -fn headroom<'a, 'b>( - candidates: &'b [Candidate<'a>], - default: Option<&'b Candidate<'a>>, - factor: f64, -) -> Option<&'b Candidate<'a>> { - let Some(default) = default else { - return candidates.first(); - }; - let need = default - .load - .map_or(0.0, |load| (factor * (1.0 - load)).max(0.0)); - candidates - .iter() - .filter(|candidate| candidate.pool.id != default.pool.id) - .find(|candidate| candidate.load.is_some_and(|load| 1.0 - load > need)) - .or(Some(default)) -} - #[cfg(test)] -mod tests { +pub(crate) mod tests { use std::collections::HashMap; use crate::load::InFlightTracker; @@ -188,13 +152,13 @@ mod tests { use super::*; - const OHIO: &str = "us-east-2"; + pub(crate) const OHIO: &str = "us-east-2"; const OREGON: &str = "us-west-2"; const TOKYO: &str = "ap-northeast-1"; /// Two pools in Ohio, one in Oregon and one in Tokyo, with round trips /// measured from Ohio. - fn fleet() -> (Vec, Arc) { + pub(crate) fn fleet() -> (Vec, Arc) { let pools = vec![ regional_pool(OHIO, "ohio-a", "qwen", &["http://ohio-a:8000"]), regional_pool(OHIO, "ohio-b", "qwen", &["http://ohio-b:8000"]), @@ -212,7 +176,7 @@ mod tests { /// Sets each pool's occupancy, in reports fresh for a minute: requests in /// flight, all generating, over 128 slots, with too little KV in use to /// bind. - fn set_occupancy(pools: &[Pool], signals: &PoolSignals, loads: [f64; 4]) { + pub(crate) fn set_occupancy(pools: &[Pool], signals: &PoolSignals, loads: [f64; 4]) { let window = Duration::from_secs(20); let source = RelaySource { runtime_id: 1, @@ -231,10 +195,10 @@ mod tests { } } - fn names(placed: &[Candidate<'_>]) -> Vec { + fn names(placed: &[Placed<&Pool>]) -> Vec { placed .iter() - .map(|candidate| candidate.pool.frontends[0].host_str().unwrap().to_owned()) + .map(|candidate| candidate.item.frontends[0].host_str().unwrap().to_owned()) .collect() } diff --git a/lib/global-router/src/lib.rs b/lib/global-router/src/lib.rs index a8c76041f967..3ab9b54b5af5 100644 --- a/lib/global-router/src/lib.rs +++ b/lib/global-router/src/lib.rs @@ -14,6 +14,7 @@ mod affinity; mod frontend_signals; mod load; +mod placement; mod pool; mod selection_core; mod selector; @@ -22,7 +23,7 @@ mod signals; pub use affinity::AffinitySelector; pub use load::{InFlightGuard, InFlightTracker}; pub use pool::Pool; -pub use selection_core::{SelectionCorePolicy, SelectionCoreSelector}; +pub use selection_core::{LeastInFlightConfig, SelectionCorePolicy, SelectionCoreSelector}; pub use selector::{ LeastInFlightSelector, PoolSelector, PoolTarget, RoundRobinSelector, RoutingRequest, }; diff --git a/lib/global-router/src/placement.rs b/lib/global-router/src/placement.rs new file mode 100644 index 000000000000..41f5cf673967 --- /dev/null +++ b/lib/global-router/src/placement.rs @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The placement rule shared by affinity and least-in-flight: nearest region +//! first, then less loaded; leave the default only for much more room. + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::time::Duration; + +#[derive(Debug)] +pub(crate) struct Placed { + pub(crate) item: T, + /// 1 is full; `None` (unknown) counts as full. + pub(crate) load: Option, + /// Round trip to the pool's region; zero for the local region. + pub(crate) distance: Option, + /// The last tie-break. + pub(crate) in_flight: u64, +} + +pub(crate) fn distance( + pool_region: &str, + region: &str, + round_trips: &HashMap, +) -> Option { + if pool_region == region { + Some(Duration::ZERO) + } else { + round_trips.get(pool_region).copied() + } +} + +/// Cheapest first; unknowns last; stable. +pub(crate) fn order(placed: &mut [Placed]) { + placed.sort_by(|left, right| cost(left, right)); +} + +/// How many leading candidates tie with the first. +pub(crate) fn leading_ties(placed: &[Placed]) -> usize { + let Some(first) = placed.first() else { + return 0; + }; + placed + .iter() + .take_while(|candidate| cost(first, candidate) == Ordering::Equal) + .count() +} + +/// `default`, unless a candidate has more than `factor` times its spare room; +/// then the first such candidate. +pub(crate) fn headroom(placed: &[Placed], default: usize, factor: f64) -> usize { + let need = placed[default] + .load + .map_or(0.0, |load| (factor * (1.0 - load)).max(0.0)); + placed + .iter() + .enumerate() + .filter(|(index, _)| *index != default) + .find(|(_, candidate)| candidate.load.is_some_and(|load| 1.0 - load > need)) + .map_or(default, |(index, _)| index) +} + +fn cost(left: &Placed, right: &Placed) -> Ordering { + known_first(left.distance, right.distance) + .then_with(|| known_first_f64(left.load, right.load)) + .then_with(|| left.in_flight.cmp(&right.in_flight)) +} + +/// Smaller first; unknown last. +fn known_first(left: Option, right: Option) -> Ordering { + match (left, right) { + (Some(left), Some(right)) => left.cmp(&right), + (Some(_), None) => Ordering::Less, + (None, Some(_)) => Ordering::Greater, + (None, None) => Ordering::Equal, + } +} + +fn known_first_f64(left: Option, right: Option) -> Ordering { + match (left, right) { + (Some(left), Some(right)) => left.total_cmp(&right), + (Some(_), None) => Ordering::Less, + (None, Some(_)) => Ordering::Greater, + (None, None) => Ordering::Equal, + } +} diff --git a/lib/global-router/src/selection_core/catalog.rs b/lib/global-router/src/selection_core/catalog.rs index f42370bb222f..b477cdc41864 100644 --- a/lib/global-router/src/selection_core/catalog.rs +++ b/lib/global-router/src/selection_core/catalog.rs @@ -150,7 +150,10 @@ mod tests { let first = source.next_snapshot().await.unwrap(); assert_eq!(ids(&first), [worker_id(&blue.id)]); assert_eq!( - directory.snapshot().get(&worker_id(&blue.id)), + directory + .snapshot() + .get(&worker_id(&blue.id)) + .map(|entry| &entry.id), Some(&blue.id) ); @@ -168,7 +171,7 @@ mod tests { assert_eq!(ids(&third), [worker_id(&green.id)]); catalog.send_replace(Arc::new(vec![green.clone()])); source.next_snapshot().await.unwrap(); - assert_eq!(directory.snapshot().get(&worker_id(&blue.id)), None); + assert!(!directory.snapshot().contains_key(&worker_id(&blue.id))); drop(catalog); assert!(source.next_snapshot().await.is_none()); diff --git a/lib/global-router/src/selection_core/directory.rs b/lib/global-router/src/selection_core/directory.rs index 82bedf682d97..a5a40f135973 100644 --- a/lib/global-router/src/selection_core/directory.rs +++ b/lib/global-router/src/selection_core/directory.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Worker id to pool, for pickers that weigh per-pool signals. +//! Worker id to pool and region, for pickers that weigh per-pool signals. use std::collections::HashMap; use std::sync::Arc; @@ -13,22 +13,34 @@ use parking_lot::RwLock; use super::catalog::worker_id; use crate::Pool; +#[derive(Debug)] +pub(super) struct PoolEntry { + pub(super) id: PoolId, + pub(super) region: String, +} + /// Replaced whole on each snapshot; read on every pick. #[derive(Debug, Default)] pub(super) struct PoolDirectory { - pools: RwLock>>, + pools: RwLock>>, } impl PoolDirectory { pub(super) fn replace(&self, pools: &[Pool]) { let pools = pools .iter() - .map(|pool| (worker_id(&pool.id), pool.id.clone())) + .map(|pool| { + let entry = PoolEntry { + id: pool.id.clone(), + region: pool.location.region.clone(), + }; + (worker_id(&pool.id), entry) + }) .collect(); *self.pools.write() = Arc::new(pools); } - pub(super) fn snapshot(&self) -> Arc> { + pub(super) fn snapshot(&self) -> Arc> { self.pools.read().clone() } } diff --git a/lib/global-router/src/selection_core/mod.rs b/lib/global-router/src/selection_core/mod.rs index 8acb794dace0..b71b8d952de0 100644 --- a/lib/global-router/src/selection_core/mod.rs +++ b/lib/global-router/src/selection_core/mod.rs @@ -33,6 +33,7 @@ use tokio_util::sync::CancellationToken; use crate::frontend_signals::FrontendSignals; use crate::load::{InFlightGuard, InFlightTracker}; use crate::selector::{Cursors, PoolTarget, reachable, target}; +use crate::signals::PoolSignals; use crate::{Pool, RoutingRequest}; use catalog::{PoolCatalogSource, ROUTING_GROUP, worker_id}; use directory::PoolDirectory; @@ -43,7 +44,17 @@ pub enum SelectionCorePolicy { /// Every pool in turn, as [`crate::RoundRobinSelector`]. RoundRobin, /// As [`crate::LeastInFlightSelector`]. - LeastInFlight, + LeastInFlight(LeastInFlightConfig), +} + +#[derive(Clone, Debug)] +pub struct LeastInFlightConfig { + /// Pool capacities and region round trips. + pub signals: Arc, + /// This process's region. + pub region: String, + /// Leave the nearest pool only for one with this many times its spare room. + pub headroom_factor: f64, } /// Selects a pool through the router core, then rotates over its frontends. @@ -127,7 +138,7 @@ impl SelectionCoreSelector { }; let in_flight = match self.policy { SelectionCorePolicy::RoundRobin => self.counts.start(&pool.id), - SelectionCorePolicy::LeastInFlight => { + SelectionCorePolicy::LeastInFlight(_) => { self.counts.take_pending(&pool.id, selected.worker_id) } }; @@ -142,7 +153,7 @@ impl SelectionCoreSelector { let request = advisory(model, allowed); let selected = match self.policy { SelectionCorePolicy::RoundRobin => self.service.select(request).await, - SelectionCorePolicy::LeastInFlight => { + SelectionCorePolicy::LeastInFlight(_) => { let service = self.service.clone(); let select = Box::pin(async move { service.select(request).await }); PendingSelect::new(select, self.counts.clone()) @@ -156,7 +167,7 @@ impl SelectionCoreSelector { } fn release(&self, worker: WorkerId) { - if let SelectionCorePolicy::LeastInFlight = self.policy { + if let SelectionCorePolicy::LeastInFlight(_) = self.policy { self.counts.release_pending(worker); } } @@ -253,8 +264,8 @@ mod tests { use url::Url; - use crate::selector::tests::pool; - use crate::{LeastInFlightSelector, PoolSelector, RoundRobinSelector}; + use crate::selector::{self, tests::pool}; + use crate::{PoolSelector, RoundRobinSelector}; use super::*; @@ -287,6 +298,19 @@ mod tests { (selection_core, catalog) } + fn least_in_flight_weighing(signals: &Arc) -> SelectionCorePolicy { + SelectionCorePolicy::LeastInFlight(LeastInFlightConfig { + signals: signals.clone(), + region: "us-east-2".into(), + headroom_factor: 2.0, + }) + } + + /// No capacities or round trips known: fewest in flight. + fn least_in_flight(in_flight: &Arc) -> SelectionCorePolicy { + least_in_flight_weighing(&Arc::new(PoolSignals::new(in_flight.clone()))) + } + /// Same picks per round as the current selector; order within a round may differ. #[tokio::test] async fn round_robin_picks_as_the_current_selector() { @@ -344,8 +368,9 @@ mod tests { let in_flight = Arc::new(InFlightTracker::new()); let pools = [blue.clone(), green.clone(), red.clone()]; let (selection_core, _catalog) = - selection_core(SelectionCorePolicy::LeastInFlight, &pools, &in_flight).await; - let current = LeastInFlightSelector::new(in_flight.clone()); + selection_core(least_in_flight(&in_flight), &pools, &in_flight).await; + let signals = Arc::new(PoolSignals::new(in_flight.clone())); + let current = selector::tests::least_in_flight(&in_flight, &signals); let mut held = Vec::new(); for turn in 0..12 { @@ -373,6 +398,30 @@ mod tests { } } + /// Same pool as the current selector: nearest until another region has + /// much more room. + #[tokio::test] + async fn least_in_flight_places_by_region_as_the_current_selector() { + let in_flight = Arc::new(InFlightTracker::new()); + let (ohio, tokyo, signals) = selector::tests::two_regions(&in_flight); + let pools = [ohio.clone(), tokyo.clone()]; + let (selection_core, _catalog) = + selection_core(least_in_flight_weighing(&signals), &pools, &in_flight).await; + let current = selector::tests::least_in_flight(&in_flight, &signals); + let candidates = [&tokyo, &ohio]; + + let mut held = Vec::new(); + for (busy, expected) in [(5, &ohio), (50, &tokyo)] { + held.extend((0..busy).map(|_| in_flight.start(&ohio.id))); + let request = RoutingRequest::new("qwen"); + let (target, counted) = selection_core.select(&request, &candidates).await.unwrap(); + assert_eq!(target.pool.id, expected.id); + drop(counted); + let from_current = current.select(&request, &candidates).unwrap(); + assert_eq!(from_current.pool.id, expected.id); + } + } + /// Concurrent picks see each other's pending picks. #[tokio::test] async fn concurrent_least_in_flight_picks_see_each_other() { @@ -381,7 +430,7 @@ mod tests { let in_flight = Arc::new(InFlightTracker::new()); let pools = [blue.clone(), green.clone()]; let (selection_core, _catalog) = - selection_core(SelectionCorePolicy::LeastInFlight, &pools, &in_flight).await; + selection_core(least_in_flight(&in_flight), &pools, &in_flight).await; let _busy = in_flight.start(&green.id); // Without pending picks, all three would go to blue. @@ -407,7 +456,7 @@ mod tests { let in_flight = Arc::new(InFlightTracker::new()); let pools = [blue.clone(), green.clone()]; let (selection_core, _catalog) = - selection_core(SelectionCorePolicy::LeastInFlight, &pools, &in_flight).await; + selection_core(least_in_flight(&in_flight), &pools, &in_flight).await; let request = RoutingRequest::new("qwen"); let candidates = [&blue, &green]; diff --git a/lib/global-router/src/selection_core/policies.rs b/lib/global-router/src/selection_core/policies.rs index e9a86a36ef53..008b0f79ac03 100644 --- a/lib/global-router/src/selection_core/policies.rs +++ b/lib/global-router/src/selection_core/policies.rs @@ -19,15 +19,17 @@ use dynamo_kv_router::services::selection::{ use dynamo_kv_router::{RoutingPartitionRef, WorkerType}; use dynamo_router_policies::router_modes; -use super::SelectionCorePolicy; use super::directory::PoolDirectory; +use super::{LeastInFlightConfig, SelectionCorePolicy}; use crate::frontend_signals::FrontendSignals; +use crate::placement::{self, Placed}; +use crate::selector::in_flight_load; impl SelectionCorePolicy { fn name(&self) -> &'static str { match self { Self::RoundRobin => router_modes::ROUND_ROBIN, - Self::LeastInFlight => "global-least-in-flight", + Self::LeastInFlight(_) => "global-least-in-flight", } } } @@ -42,11 +44,15 @@ pub(super) fn registry( let mut registry = dynamo_router_policies::default_registry(); let provider = match policy { SelectionCorePolicy::RoundRobin => return Ok(registry), - SelectionCorePolicy::LeastInFlight => provider(move || LeastInFlightPicker { - counts: counts.clone(), - directory: directory.clone(), - ties: AtomicUsize::new(0), - }), + SelectionCorePolicy::LeastInFlight(config) => { + let config = config.clone(); + provider(move || LeastInFlightPicker { + config: config.clone(), + counts: counts.clone(), + directory: directory.clone(), + ties: AtomicUsize::new(0), + }) + } }; registry.register_worker_selection(policy.name(), provider)?; Ok(registry) @@ -87,10 +93,11 @@ where ) } -/// Fewest in flight, as [`crate::LeastInFlightSelector`]; ties rotate by -/// starting the scan one candidate later each pick. The pick stays pending -/// on its pool until the request takes it over. A preview advances nothing. +/// [`crate::LeastInFlightSelector`]'s rule over the frontend's counts. The +/// pick stays pending on its pool until the request takes it over. A preview +/// advances nothing. struct LeastInFlightPicker { + config: LeastInFlightConfig, counts: Arc, directory: Arc, ties: AtomicUsize, @@ -103,27 +110,49 @@ impl WorkerPicker for LeastInFlightPicker { input: WorkerInputView<'_>, ) -> Result { let candidates = input.candidates(); - let count = candidates.len(); - if count == 0 { + let pools = self.directory.snapshot(); + let signals = &self.config.signals; + let round_trips = signals.region_round_trips(); + // Read each count once: requests start and finish while this runs. + let mut placed = candidates + .iter() + .enumerate() + .map(|(row, candidate)| { + let worker = candidate.worker().worker_id; + match pools.get(&worker) { + Some(pool) => { + let in_flight = self.counts.in_flight(&pool.id, worker); + Placed { + item: row, + load: in_flight_load(in_flight, signals.capacity(&pool.id)), + distance: placement::distance( + &pool.region, + &self.config.region, + &round_trips, + ), + in_flight, + } + } + None => Placed { + item: row, + load: None, + distance: None, + in_flight: u64::MAX, + }, + } + }) + .collect::>(); + placement::order(&mut placed); + let tied = placement::leading_ties(&placed); + if tied == 0 { return Err(WorkerSelectionPolicyError::failed("no candidates")); } - let start = if context.is_preview() { + let default = if context.is_preview() { self.ties.load(Ordering::Relaxed) } else { self.ties.fetch_add(1, Ordering::Relaxed) - } % count; - let pools = self.directory.snapshot(); - let (row, _) = (0..count) - .map(|offset| (start + offset) % count) - .map(|row| { - let worker = candidates[row].worker().worker_id; - let in_flight = pools - .get(&worker) - .map_or(u64::MAX, |pool| self.counts.in_flight(pool, worker)); - (row, in_flight) - }) - .min_by_key(|(_, in_flight)| *in_flight) - .ok_or_else(|| WorkerSelectionPolicyError::failed("no candidates"))?; + } % tied; + let row = placed[placement::headroom(&placed, default, self.config.headroom_factor)].item; if !context.is_preview() { self.counts.add_pending(candidates[row].worker().worker_id); } diff --git a/lib/global-router/src/selector.rs b/lib/global-router/src/selector.rs index 32031320cbf6..d71beefc4b51 100644 --- a/lib/global-router/src/selector.rs +++ b/lib/global-router/src/selector.rs @@ -12,6 +12,8 @@ use url::Url; use crate::Pool; use crate::load::InFlightTracker; +use crate::placement::{self, Placed}; +use crate::signals::PoolSignals; /// What a selector may know about a request. Routing inputs such as the /// client's region become fields here. @@ -86,26 +88,38 @@ impl PoolSelector for RoundRobinSelector { } } -/// Sends each request to the pool with the fewest of this process's requests -/// in flight. +/// Sends each request to the nearest pool with the fewest requests in +/// flight, by affinity's placement rule with load = requests in flight over +/// the pool's capacity ([`PoolSignals::capacity`]). Ties rotate. /// /// A slower pool holds its requests longer, so it shows more in flight and /// gets a smaller share. Only this process's requests count: replicas do not -/// share counts, and the pools' own load reports are not read. Ties rotate -/// among the tied pools. +/// share counts. #[derive(Debug)] pub struct LeastInFlightSelector { in_flight: Arc, + signals: Arc, + region: String, + headroom_factor: f64, ties: AtomicUsize, frontends: Cursors, } impl LeastInFlightSelector { /// `in_flight` must be the tracker the Global Frontend counts its - /// requests in. - pub fn new(in_flight: Arc) -> Self { + /// requests in; `region` is the region of this process, which every + /// request enters. + pub fn new( + in_flight: Arc, + signals: Arc, + region: String, + headroom_factor: f64, + ) -> Self { Self { in_flight, + signals, + region, + headroom_factor, ties: AtomicUsize::new(0), frontends: Cursors::default(), } @@ -118,21 +132,44 @@ impl PoolSelector for LeastInFlightSelector { _request: &RoutingRequest<'_>, candidates: &[&'a Pool], ) -> Option> { + let round_trips = self.signals.region_round_trips(); // Read each count once: requests start and finish while this runs. - let counts = candidates + let mut placed = candidates .iter() .copied() .filter(|pool| reachable(pool)) - .map(|pool| (pool, self.in_flight.get(&pool.id))) + .map(|pool| { + let in_flight = self.in_flight.get(&pool.id); + Placed { + item: pool, + load: in_flight_load(in_flight, self.signals.capacity(&pool.id)), + distance: placement::distance( + &pool.location.region, + &self.region, + &round_trips, + ), + in_flight, + } + }) .collect::>(); - let fewest = counts.iter().map(|(_, count)| *count).min()?; - let tied = counts.iter().filter(|(_, count)| *count == fewest); - let turn = self.ties.fetch_add(1, Ordering::Relaxed) % tied.clone().count(); - let (pool, _) = tied.clone().nth(turn)?; - target(&self.frontends, pool) + placement::order(&mut placed); + let tied = placement::leading_ties(&placed); + if tied == 0 { + return None; + } + let default = self.ties.fetch_add(1, Ordering::Relaxed) % tied; + let choice = placed[placement::headroom(&placed, default, self.headroom_factor)].item; + target(&self.frontends, choice) } } +/// `in_flight` over capacity; `None` without a capacity. +pub(crate) fn in_flight_load(in_flight: u64, capacity: Option) -> Option { + capacity + .filter(|capacity| *capacity > 0) + .map(|capacity| in_flight as f64 / capacity as f64) +} + /// Whether the pool has a frontend to send to. pub(crate) fn reachable(pool: &Pool) -> bool { !pool.frontends.is_empty() @@ -179,14 +216,51 @@ impl Cursors { #[cfg(test)] pub(crate) mod tests { + use std::collections::HashMap; + use std::time::Duration; + use dynamo_global_common::PoolKey; + use crate::signals::{LoadReport, RelaySource}; + use super::*; pub(crate) fn pool(dgd: &str, model: &str, frontends: &[&str]) -> Pool { regional_pool("us-east-2", dgd, model, frontends) } + pub(crate) fn least_in_flight( + in_flight: &Arc, + signals: &Arc, + ) -> LeastInFlightSelector { + LeastInFlightSelector::new(in_flight.clone(), signals.clone(), "us-east-2".into(), 2.0) + } + + /// Ohio and Tokyo, 135 ms apart, each able to run 100 requests at once. + pub(crate) fn two_regions(in_flight: &Arc) -> (Pool, Pool, Arc) { + let ohio = regional_pool("us-east-2", "ohio", "qwen", &["http://ohio:8000"]); + let tokyo = regional_pool("ap-northeast-1", "tokyo", "qwen", &["http://tokyo:8000"]); + let signals = Arc::new(PoolSignals::new(in_flight.clone())); + signals.set_region_round_trips(HashMap::from([( + "ap-northeast-1".to_owned(), + Duration::from_millis(135), + )])); + for pool in [&ohio, &tokyo] { + let source = RelaySource { + runtime_id: 1, + incarnation: 1, + }; + let report = LoadReport { + in_flight: 0, + generating: 0, + max_concurrency: Some(100), + window: Duration::from_secs(60), + }; + signals.record_load(&pool.id, source, report); + } + (ohio, tokyo, signals) + } + pub(crate) fn regional_pool(region: &str, dgd: &str, model: &str, frontends: &[&str]) -> Pool { let key = PoolKey::new(region, "dynamo", dgd).unwrap(); let frontends = frontends.iter().map(|url| url.parse().unwrap()).collect(); @@ -261,7 +335,7 @@ pub(crate) mod tests { let blue = pool("blue", "qwen", &["http://blue:8000"]); let green = pool("green", "qwen", &["http://green:8000"]); let in_flight = Arc::new(InFlightTracker::new()); - let selector = LeastInFlightSelector::new(in_flight.clone()); + let selector = least_in_flight(&in_flight, &Arc::new(PoolSignals::new(in_flight.clone()))); let request = RoutingRequest::new("qwen"); let _busy = [in_flight.start(&blue.id), in_flight.start(&blue.id)]; @@ -277,7 +351,8 @@ pub(crate) mod tests { let empty = pool("draining", "qwen", &[]); let blue = pool("blue", "qwen", &["http://blue:8000"]); let green = pool("green", "qwen", &["http://green:8000"]); - let selector = LeastInFlightSelector::new(Arc::new(InFlightTracker::new())); + let in_flight = Arc::new(InFlightTracker::new()); + let selector = least_in_flight(&in_flight, &Arc::new(PoolSignals::new(in_flight.clone()))); let request = RoutingRequest::new("qwen"); let chosen = (0..4) @@ -294,4 +369,33 @@ pub(crate) mod tests { assert_eq!(chosen, [blue.clone(), green.clone(), blue, green]); assert_eq!(selector.select(&request, &[&empty]), None); } + + /// Stays local until another region has twice the room; unknown capacity counts as full. + #[test] + fn least_in_flight_stays_in_its_region_until_another_has_much_more_room() { + let in_flight = Arc::new(InFlightTracker::new()); + let (ohio, tokyo, signals) = two_regions(&in_flight); + let selector = least_in_flight(&in_flight, &signals); + let picked = || { + let request = RoutingRequest::new("qwen"); + selector + .select(&request, &[&tokyo, &ohio]) + .unwrap() + .pool + .id + .clone() + }; + + // Ohio 5 of 100, Tokyo idle: 0.95 spare against 1.0. + let mut held = (0..5) + .map(|_| in_flight.start(&ohio.id)) + .collect::>(); + assert_eq!(picked(), ohio.id); + // Ohio 55 of 100: 0.45 spare, and Tokyo's 1.0 is more than twice it. + held.extend((0..50).map(|_| in_flight.start(&ohio.id))); + assert_eq!(picked(), tokyo.id); + // Tokyo's capacity unknown: it counts as full, so Ohio keeps it. + signals.forget(&tokyo.id); + assert_eq!(picked(), ohio.id); + } } diff --git a/lib/global-router/src/signals.rs b/lib/global-router/src/signals.rs index b81801be6703..4a4f4f2709df 100644 --- a/lib/global-router/src/signals.rs +++ b/lib/global-router/src/signals.rs @@ -196,19 +196,7 @@ impl PoolSignals { pub fn occupancy(&self, pool: &PoolId) -> Option { let now = Instant::now(); let sources = self.reports.get(pool)?; - let (source, current) = sources - .iter() - .filter_map(|(source, source_report)| { - source_report - .load - .as_ref() - .filter(|current| { - now <= current.expires - && current.report.max_concurrency.is_some_and(|n| n > 0) - }) - .map(|current| (source, current)) - }) - .max_by_key(|(source, current)| (current.received, **source))?; + let (source, current) = freshest(&sources, now)?; let capacity = current .report .max_concurrency @@ -236,6 +224,13 @@ impl PoolSignals { }) } + /// Requests the pool's workers can run at once, from the freshest report. + pub fn capacity(&self, pool: &PoolId) -> Option { + let sources = self.reports.get(pool)?; + let (_, current) = freshest(&sources, Instant::now())?; + current.report.max_concurrency + } + /// Replaces the measured round trip to each region. pub fn set_region_round_trips(&self, round_trips: HashMap) { *self.regions.write() = Arc::new(round_trips); @@ -254,6 +249,25 @@ fn prune(sources: &mut HashMap, now: Instant) { }); } +/// The freshest unexpired report that carries a capacity. +fn freshest( + sources: &HashMap, + now: Instant, +) -> Option<(&RelaySource, &Report)> { + sources + .iter() + .filter_map(|(source, source_report)| { + source_report + .load + .as_ref() + .filter(|current| { + now <= current.expires && current.report.max_concurrency.is_some_and(|n| n > 0) + }) + .map(|current| (source, current)) + }) + .max_by_key(|(source, current)| (current.received, **source)) +} + #[cfg(test)] mod tests { use crate::selector::tests::pool;