feat(global-router): least-in-flight places pools by affinity's rule

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 \

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,18 +71,23 @@ 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<String>,
     /// A pool's PoolRelay, which streams the pool's load, as
     /// `<cluster>/<namespace>/<DGD name>=<URL>`. Repeat the flag for each pool;
     /// each pool needs a Relay of its own, since a Relay reports the load of
     /// 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<dyn PoolSelector> = 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<SelectionCorePolicy> {
-    match policy {
+/// `signals` starts the load signals the policy weighs, if any.
+fn selection_core_policy(
+    args: &Args,
+    signals: impl FnOnce() -> Arc<PoolSignals>,
+) -> Result<SelectionCorePolicy> {
+    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<Arc<PoolSnapshot>>,
+    signals: &Arc<PoolSignals>,
+    shutdown: &CancellationToken,
+) -> Arc<PoolSignals> {
+    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<AffinitySelector>, 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);
     }

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,90 +97,53 @@ impl PoolSelector for AffinitySelector {
         candidates: &[&'a Pool],
     ) -> Option<PoolTarget<'a>> {
         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 {
                 return None;
             }
             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<f64>,
-    /// Round trip to the pool's region: zero for the region the request
-    /// entered, otherwise as last measured. `None` until measured.
-    distance: Option<Duration>,
-}
-
 /// 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<Candidate<'a>> {
+/// 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<Placed<&'a Pool>> {
     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::<Vec<_>>();
-    // 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<Pool>, Arc<PoolSignals>) {
+    pub(crate) fn fleet() -> (Vec<Pool>, Arc<PoolSignals>) {
         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<String> {
+    fn names(placed: &[Placed<&Pool>]) -> Vec<String> {
         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()
     }
 

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,
 };