From e0c4b9a6e50c3c192ca1ab0619b97406ac718e34 Mon Sep 17 00:00:00 2001 From: nnshah1 Date: Tue, 3 Nov 2026 13:35:08 -0800 Subject: [PATCH] feat(global-router): session affinity on the router core Affinity is the last policy to run on the router core. Its plugin applies the affinity selector's headroom rule over the same signals, and the core's session table (soft mode) holds each session's pin. A session books only to set, move or refresh its pin, and the booking is freed as soon as the pin is committed: the binding outlives the booking and expires after the idle window. Adds dynamo_global_frontend_core_pin_bookings_total. The current selectors do not change. Co-Authored-By: Claude Opus 5.5 Signed-off-by: nnshah1 --- .../observability/metrics-catalog.mdx | 6 + lib/global-frontend/README.md | 12 +- lib/global-frontend/src/main.rs | 18 +- lib/global-frontend/src/metrics.rs | 45 ++- lib/global-frontend/src/proxy.rs | 13 +- lib/global-frontend/tests/forwarding.rs | 59 +++- lib/global-router/src/lib.rs | 5 +- lib/global-router/src/selection_core/mod.rs | 283 ++++++++++++++++-- .../src/selection_core/policies.rs | 63 +++- 9 files changed, 454 insertions(+), 50 deletions(-) diff --git a/docs/fern/pages/reference/observability/metrics-catalog.mdx b/docs/fern/pages/reference/observability/metrics-catalog.mdx index e3eb3acea803..4d9e9276aa15 100644 --- a/docs/fern/pages/reference/observability/metrics-catalog.mdx +++ b/docs/fern/pages/reference/observability/metrics-catalog.mdx @@ -958,6 +958,12 @@ Emitted by the Global Frontend at `/metrics` on its serving port. It forwards ea **Experimental.** Time to select a pool for a request, in milliseconds, with either selection path (`--use-selection-core` or not) and any `--router-policy`. Counts every request that reaches selection, also one that no pool can take. No labels: a process runs one path and one policy, so compare paths across runs. Unit and buckets match the router's `dynamo_router_overhead_*_ms` histograms. +The following is emitted with `--use-selection-core`. It is read from the embedded router core at scrape time. + + + **Experimental.** Affinity pin bookings: a session's first request, a move to another pool, and a periodic refresh. Under `round-robin` and `least-in-flight` nothing is booked. A rate close to the request rate under `affinity` means sessions are moving or starting rather than staying on their pools. + + ## Mocker Metrics The mocker (`python -m dynamo.mocker`) simulates an inference engine. It emits imitation engine metrics with simulated values on `DYN_SYSTEM_PORT`, choosing the set by `--engine-type`. They carry the engine's usual labels plus `dynamo_namespace`, `dynamo_component`, `dynamo_endpoint`, `worker_id`, `model`, and `model_name`. These imitate engine metrics, so this table has no Labels, Description, or Status column; each metric means what it means in the engine it imitates. diff --git a/lib/global-frontend/README.md b/lib/global-frontend/README.md index bd5df67824f2..159407eda5ab 100644 --- a/lib/global-frontend/README.md +++ b/lib/global-frontend/README.md @@ -138,7 +138,8 @@ named `dynamo_global_frontend_*`: requests per pool by response status, `no_response` when the pool sends no response, or `cancelled` when the client leaves before the pool's response headers arrive; requests answered without a pool; time to first byte and duration per pool; requests in flight per -pool; and the time to select a pool. The `pool` label is the pool's +pool; the time to select a pool; and, with `--use-selection-core`, the +router core's affinity pin bookings. The `pool` label is the pool's `//` and `region` its region. On requests, `local` says whether the pool is in `--region` (`unknown` without it), which gives the share of requests kept in the region. Nothing is labelled from the request @@ -163,9 +164,12 @@ for each family. With `--use-selection-core`, the Global Frontend selects pools through an embedded router core: each pool is a core worker and the policy is a picker -plugin. `round-robin` and `least-in-flight` run on the core, picking as the -selectors do. A least-in-flight pick counts as pending on its pool until the -request is counted, so concurrent picks see each other. +plugin. Every policy runs on the core, picking as the selectors do. A +least-in-flight pick counts as pending on its pool until the request is +counted, so concurrent picks see each other. Under `affinity` the core's +session table holds each session's pin; a session books only to set, move or +refresh its pin, and the booking is freed at once. Session IDs longer than +256 bytes are refused with 400. ## Limitations diff --git a/lib/global-frontend/src/main.rs b/lib/global-frontend/src/main.rs index 57f036531c4d..2d0c38007715 100644 --- a/lib/global-frontend/src/main.rs +++ b/lib/global-frontend/src/main.rs @@ -15,8 +15,9 @@ use dynamo_global_frontend::{ follow_relays, measure_round_trips, serve, serve_relay_views, }; use dynamo_global_router::{ - AffinitySelector, InFlightTracker, LeastInFlightConfig, LeastInFlightSelector, Pool, - PoolSelector, PoolSignals, RoundRobinSelector, SelectionCorePolicy, SelectionCoreSelector, + AffinityConfig, AffinitySelector, InFlightTracker, LeastInFlightConfig, LeastInFlightSelector, + Pool, PoolSelector, PoolSignals, RoundRobinSelector, SelectionCorePolicy, + SelectionCoreSelector, }; use tokio::sync::watch; use tokio_util::sync::CancellationToken; @@ -68,7 +69,7 @@ struct Args { connect_timeout_secs: u64, #[arg(long, value_enum, default_value_t = RouterPolicy::RoundRobin)] router_policy: RouterPolicy, - /// Select pools through the embedded router core (not yet with affinity). + /// Select pools through the embedded router core. #[arg(long)] use_selection_core: bool, /// The region this process runs in. `affinity` and `least-in-flight` @@ -91,6 +92,7 @@ struct Args { #[arg(long, default_value_t = 2.0)] headroom_factor: f64, /// `affinity`: how long a session keeps its pool without a request. + /// On the router core, at most 65,536 sessions per model are pinned. #[arg(long, default_value_t = 1800)] session_idle_secs: u64, } @@ -247,6 +249,7 @@ async fn main() -> Result<()> { 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?); + metrics.watch_selection_core(selection_core.clone()); // Wait for the first pool snapshot, or early requests get a 503. let (first_snapshot, reconcile) = selection_core.reconcile_catalog(pools.clone(), shutdown.child_token()); @@ -318,9 +321,12 @@ fn selection_core_policy( headroom_factor: args.headroom_factor, })) } - RouterPolicy::Affinity => { - bail!("--use-selection-core does not support --router-policy affinity yet") - } + RouterPolicy::Affinity => Ok(SelectionCorePolicy::Affinity(AffinityConfig { + signals: signals(), + region: args.region.clone().context("--region is required")?, + headroom_factor: args.headroom_factor, + idle: Duration::from_secs(args.session_idle_secs), + })), } } diff --git a/lib/global-frontend/src/metrics.rs b/lib/global-frontend/src/metrics.rs index 37018baea06e..a2420400d178 100644 --- a/lib/global-frontend/src/metrics.rs +++ b/lib/global-frontend/src/metrics.rs @@ -8,11 +8,11 @@ //! and would mint a series per string. Pool frontends report what needs the //! stream parsed, such as inter-token latency and token counts. -use std::sync::{Arc, Mutex, PoisonError}; +use std::sync::{Arc, Mutex, OnceLock, PoisonError}; use std::time::Instant; use axum::http::StatusCode; -use dynamo_global_router::{InFlightTracker, Pool}; +use dynamo_global_router::{InFlightTracker, Pool, SelectionCoreSelector}; use dynamo_llm::http::service::metrics::generate_log_buckets; use dynamo_runtime::metrics::prometheus_names::clamp_u64_to_i64; use prometheus::{ @@ -44,6 +44,9 @@ pub struct Metrics { duration: HistogramVec, in_flight: IntGaugeVec, selection: Histogram, + /// Read at each scrape ([`Self::watch_core`]). + selection_core: OnceLock, + pin_bookings: IntCounter, /// A scrape resets and refills the gauges, so scrapes must not overlap. scrape: Mutex<()>, } @@ -134,6 +137,13 @@ impl Metrics { .buckets(prometheus::exponential_buckets(0.01, 3.0, 17).expect("valid buckets")), ), ); + let pin_bookings = register( + ®istry, + IntCounter::with_opts(Opts::new( + name("pin_bookings_total"), + "Affinity pin bookings: a session's first request, a move or a refresh", + )), + ); Self { registry, region, @@ -145,6 +155,8 @@ impl Metrics { duration, in_flight, selection, + selection_core: OnceLock::new(), + pin_bookings, scrape: Mutex::new(()), } } @@ -154,6 +166,20 @@ impl Metrics { .observe(started.elapsed().as_secs_f64() * 1000.0); } + /// Reads `selection_core`'s counts at each scrape. A second call is ignored. + pub fn watch_selection_core(&self, selection_core: Arc) { + let _ = self + .selection_core + .set(WatchedSelectionCore(selection_core)); + } + + fn read_selection_core(&self) { + let Some(WatchedSelectionCore(selection_core)) = self.selection_core.get() else { + return; + }; + advance(&self.pin_bookings, selection_core.pin_bookings()); + } + /// Counts a request answered without a pool. `error_type` is one of /// `prometheus_names::frontend_service::error_type`. pub(crate) fn rejected(&self, error_type: &str) { @@ -198,6 +224,7 @@ impl Metrics { /// read into the gauges now. pub(crate) fn render(&self, pools: &PoolSnapshot) -> prometheus::Result { let _scrape = self.scrape.lock().unwrap_or_else(PoisonError::into_inner); + self.read_selection_core(); // Pools that left the catalog lose their series. self.in_flight.reset(); for pool in pools.pools() { @@ -209,6 +236,20 @@ impl Metrics { } } +/// `SelectionCoreSelector` has no `Debug`. +struct WatchedSelectionCore(Arc); + +impl std::fmt::Debug for WatchedSelectionCore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("CoreSelector") + } +} + +/// Moves `counter` up to `total`, a monotonic count kept elsewhere. +fn advance(counter: &IntCounter, total: u64) { + counter.inc_by(total.saturating_sub(counter.get())); +} + fn name(suffix: &str) -> String { format!("{PREFIX}_{suffix}") } diff --git a/lib/global-frontend/src/proxy.rs b/lib/global-frontend/src/proxy.rs index d7800730e06e..2987036337f8 100644 --- a/lib/global-frontend/src/proxy.rs +++ b/lib/global-frontend/src/proxy.rs @@ -18,7 +18,9 @@ use axum::http::{HeaderMap, HeaderName, Method, StatusCode, Uri, header}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; -use dynamo_global_router::{InFlightTracker, PoolSelector, RoutingRequest, SelectionCoreSelector}; +use dynamo_global_router::{ + InFlightTracker, MAX_SESSION_ID_BYTES, PoolSelector, RoutingRequest, SelectionCoreSelector, +}; use dynamo_llm::protocols::agents::HEADER_DYNAMO_SESSION_ID; use dynamo_runtime::config::environment_names::llm::DYN_HTTP_BODY_LIMIT_MB; use dynamo_runtime::metrics::prometheus_names::frontend_service::error_type; @@ -224,6 +226,15 @@ async fn forward( let session = headers .get(HEADER_DYNAMO_SESSION_ID) .and_then(|session| session.to_str().ok()); + // A 400, not the 503 the core's refusal would become. + if let Selector::SelectionCore(_) = proxy.selector + && session.is_some_and(|session| session.len() > MAX_SESSION_ID_BYTES) + { + metrics.rejected(error_type::VALIDATION); + let message = + format!("{HEADER_DYNAMO_SESSION_ID} is longer than {MAX_SESSION_ID_BYTES} bytes"); + return error(StatusCode::BAD_REQUEST, message); + } let request = RoutingRequest::new(&model).with_session(session); let selecting = Instant::now(); let selected = match &proxy.selector { diff --git a/lib/global-frontend/tests/forwarding.rs b/lib/global-frontend/tests/forwarding.rs index e00bc1bc2942..e26aefdd87c8 100644 --- a/lib/global-frontend/tests/forwarding.rs +++ b/lib/global-frontend/tests/forwarding.rs @@ -13,8 +13,8 @@ use axum::http::HeaderMap; use dynamo_global_common::PoolKey; use dynamo_global_frontend::{Metrics, PoolSnapshot}; use dynamo_global_router::{ - InFlightTracker, Pool, PoolSelector, PoolTarget, RoundRobinSelector, RoutingRequest, - SelectionCorePolicy, SelectionCoreSelector, + AffinityConfig, InFlightTracker, Pool, PoolSelector, PoolSignals, PoolTarget, + RoundRobinSelector, RoutingRequest, SelectionCorePolicy, SelectionCoreSelector, }; use dynamo_llm::endpoint_type::EndpointType; use dynamo_llm::engines::{StreamingEngineAdapter, make_echo_engine}; @@ -175,18 +175,27 @@ async fn global_frontend_with( /// A Global Frontend that routes round robin through the router core. async fn frontend_with_selection_core(pools: Vec) -> GlobalFrontend { - let in_flight = Arc::new(InFlightTracker::new()); - let selection_core = - SelectionCoreSelector::new(SelectionCorePolicy::RoundRobin, in_flight.clone()) + frontend_with_selection_core_policy(SelectionCorePolicy::RoundRobin, pools, Arc::default()) + .await +} + +async fn frontend_with_selection_core_policy( + policy: SelectionCorePolicy, + pools: Vec, + in_flight: Arc, +) -> GlobalFrontend { + let selection_core = Arc::new( + SelectionCoreSelector::new(policy, in_flight.clone()) .await - .unwrap(); - let selection_core = Arc::new(selection_core); + .unwrap(), + ); let (sender, receiver) = watch::channel(Arc::new(PoolSnapshot::new(pools))); let (first_snapshot, reconcile) = selection_core.reconcile_catalog(receiver.clone(), CancellationToken::new()); tokio::spawn(reconcile); first_snapshot.await.unwrap(); let metrics = Metrics::new(Some(REGION.to_owned()), in_flight.clone()); + metrics.watch_selection_core(selection_core.clone()); let app = dynamo_global_frontend::app_with_selection_core( receiver, selection_core, @@ -767,3 +776,39 @@ async fn shutdown_lets_an_open_stream_finish_then_returns() { .unwrap(); drop(client); } + +/// One booking per new session; over-long session IDs are a 400. +#[tokio::test] +async fn affinity_on_the_core_counts_pin_bookings() { + let pool_frontend = echo_pool_frontend().await; + let in_flight = Arc::new(InFlightTracker::new()); + let policy = SelectionCorePolicy::Affinity(AffinityConfig { + signals: Arc::new(PoolSignals::new(in_flight.clone())), + region: REGION.to_owned(), + headroom_factor: 2.0, + idle: Duration::from_secs(60), + }); + let pools = vec![pool("blue", MODEL, &pool_frontend.url)]; + let global = frontend_with_selection_core_policy(policy, pools, in_flight).await; + + let send = |session: String| { + client() + .post(global.server.url.join("v1/chat/completions").unwrap()) + .header("x-dynamo-session-id", session) + .json(&chat(MODEL, "hi", false)) + .send() + }; + for session in ["a", "a", "b", "b"] { + assert_eq!( + send(session.to_owned()).await.unwrap().status(), + StatusCode::OK + ); + } + let response = send("s".repeat(257)).await.unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let bookings = "dynamo_global_frontend_pin_bookings_total"; + metrics_when(&global.server.url, |metrics| { + sample(metrics, bookings, &[]) == Some(2.0) + }) + .await; +} diff --git a/lib/global-router/src/lib.rs b/lib/global-router/src/lib.rs index 3ab9b54b5af5..09f8bd634695 100644 --- a/lib/global-router/src/lib.rs +++ b/lib/global-router/src/lib.rs @@ -23,7 +23,10 @@ mod signals; pub use affinity::AffinitySelector; pub use load::{InFlightGuard, InFlightTracker}; pub use pool::Pool; -pub use selection_core::{LeastInFlightConfig, SelectionCorePolicy, SelectionCoreSelector}; +pub use selection_core::{ + AffinityConfig, LeastInFlightConfig, MAX_SESSION_ID_BYTES, SelectionCorePolicy, + SelectionCoreSelector, +}; pub use selector::{ LeastInFlightSelector, PoolSelector, PoolTarget, RoundRobinSelector, RoutingRequest, }; diff --git a/lib/global-router/src/selection_core/mod.rs b/lib/global-router/src/selection_core/mod.rs index b71b8d952de0..95b2a5f55225 100644 --- a/lib/global-router/src/selection_core/mod.rs +++ b/lib/global-router/src/selection_core/mod.rs @@ -12,21 +12,27 @@ mod catalog; mod directory; mod policies; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; use anyhow::{Context, Result, anyhow}; use dynamo_global_common::PoolId; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::protocols::RoutingConstraints; use dynamo_kv_router::scheduling::RouterPolicyConfig; +use dynamo_kv_router::services::selection::affinity::{ + MAX_SESSION_AFFINITY_ENTRIES, SessionAffinityMode, +}; use dynamo_kv_router::services::selection::{ - CatalogReconciler, PromptRequest, SelectRequest, SelectResponse, SelectionError, - SelectionService, SelectionServiceBuilder, + CatalogReconciler, PromptRequest, SelectAndReserveRequest, SelectRequest, SelectResponse, + SelectionError, SelectionService, SelectionServiceBuilder, }; use dynamo_kv_router::{WorkerId, WorkerType}; +use parking_lot::Mutex; use tokio::sync::{oneshot, watch}; use tokio_util::sync::CancellationToken; @@ -38,6 +44,9 @@ use crate::{Pool, RoutingRequest}; use catalog::{PoolCatalogSource, ROUTING_GROUP, worker_id}; use directory::PoolDirectory; +/// The longest session id the core's session table accepts. +pub use dynamo_kv_router::services::selection::affinity::MAX_SESSION_AFFINITY_ID_BYTES as MAX_SESSION_ID_BYTES; + /// How the core picks among a model's pools. #[derive(Clone, Debug)] pub enum SelectionCorePolicy { @@ -45,6 +54,20 @@ pub enum SelectionCorePolicy { RoundRobin, /// As [`crate::LeastInFlightSelector`]. LeastInFlight(LeastInFlightConfig), + /// As [`crate::AffinitySelector`], with the core's session table holding pins. + Affinity(AffinityConfig), +} + +#[derive(Clone, Debug)] +pub struct AffinityConfig { + /// Pool occupancy and region round trips. + pub signals: Arc, + /// This process's region. + pub region: String, + /// Leave a pool only for one with this many times its spare room. + pub headroom_factor: f64, + /// How long a pin lasts without a request. + pub idle: Duration, } #[derive(Clone, Debug)] @@ -65,6 +88,58 @@ pub struct SelectionCoreSelector { frontends: Cursors, /// The frontend's in-flight counts and least-in-flight's pending picks. counts: Arc, + pins: Option, + pin_bookings: AtomicU64, +} + +/// When each session's pin was last booked. Only a booking sets or moves a +/// pin, and the core's binding expires `idle` after it, so a session books +/// when its pin is new, moved, or half `idle` old; the rest are advisory. +struct PinRefresh { + refresh: Duration, + capacity: usize, + booked: Mutex, +} + +#[derive(Default)] +struct Pins { + by_session: HashMap, + /// Last sweep of stale pins, at most one per refresh window. + swept: Option, +} + +impl PinRefresh { + fn new(idle: Duration) -> Self { + Self { + refresh: idle / 2, + capacity: MAX_SESSION_AFFINITY_ENTRIES, + booked: Mutex::default(), + } + } + + fn fresh(&self, session: &str, now: Instant) -> Option { + let booked = self.booked.lock(); + let &(worker, at) = booked.by_session.get(session)?; + (now.saturating_duration_since(at) < self.refresh).then_some(worker) + } + + /// A new session that finds the table full is not remembered, so it + /// books every request, as when the core's own table is full. + fn record(&self, session: &str, worker: WorkerId, now: Instant) { + let mut booked = self.booked.lock(); + let Pins { by_session, swept } = &mut *booked; + if by_session.len() >= self.capacity && !by_session.contains_key(session) { + let refresh = self.refresh; + if swept.is_none_or(|at| now.saturating_duration_since(at) >= refresh) { + by_session.retain(|_, (_, at)| now.saturating_duration_since(*at) < refresh); + *swept = Some(now); + } + if by_session.len() >= self.capacity { + return; + } + } + by_session.insert(session.to_owned(), (worker, now)); + } } impl SelectionCoreSelector { @@ -75,9 +150,17 @@ impl SelectionCoreSelector { let counts = Arc::new(FrontendSignals::new(in_flight)); let registry = policies::registry(&policy, directory.clone(), counts.clone()) .context("register the pool selection policies")?; - let builder = + let mut builder = SelectionServiceBuilder::new(config(&policy)?, WorkerType::Aggregated, registry) .indexer_threads(1); + let mut pins = None; + if let SelectionCorePolicy::Affinity(affinity) = &policy { + // Soft: the binding follows the pool picked. + builder = builder + .session_affinity(affinity.idle) + .session_affinity_mode(SessionAffinityMode::Soft); + pins = Some(PinRefresh::new(affinity.idle)); + } let service = Arc::new(builder.build().await.context("build the router core")?); Ok(Self { policy, @@ -85,9 +168,16 @@ impl SelectionCoreSelector { directory, frontends: Cursors::default(), counts, + pins, + pin_bookings: AtomicU64::new(0), }) } + /// Affinity pin bookings since start. + pub fn pin_bookings(&self) -> u64 { + self.pin_bookings.load(Ordering::Relaxed) + } + /// Keeps the core's workers in sync with `pools`, as the local router's /// `CatalogReconciler` does. The receiver resolves once the first /// snapshot is applied. @@ -114,15 +204,27 @@ impl SelectionCoreSelector { request: &RoutingRequest<'_>, candidates: &[&'a Pool], ) -> Option<(PoolTarget<'a>, InFlightGuard)> { - let allowed = candidates - .iter() - .filter(|pool| reachable(pool)) - .map(|pool| worker_id(&pool.id)) - .collect::>(); - if allowed.is_empty() { - return None; + let model = request.model; + let session = request.session.filter(|_| self.pins.is_some()); + let now = Instant::now(); + let pinned = session + .zip(self.pins.as_ref()) + .and_then(|(session, pins)| pins.fresh(session, now)); + // A session books when its pin is new or due for a refresh. + let mut book = session.is_some() && pinned.is_none(); + let mut selected = self + .pick(model, allowed(candidates)?, session, book) + .await?; + if !book && pinned.is_some_and(|pinned| pinned != selected.worker_id) { + // The plugin moved the session: only a booking moves its pin. + book = true; + selected = self + .pick(model, allowed(candidates)?, session, true) + .await?; + } + if let (true, Some(session), Some(pins)) = (book, session, &self.pins) { + pins.record(session, selected.worker_id, now); } - let selected = self.pick(request.model, allowed).await?; let Some(pool) = candidates .iter() .copied() @@ -137,22 +239,59 @@ impl SelectionCoreSelector { return None; }; let in_flight = match self.policy { - SelectionCorePolicy::RoundRobin => self.counts.start(&pool.id), + SelectionCorePolicy::RoundRobin | SelectionCorePolicy::Affinity(_) => { + self.counts.start(&pool.id) + } SelectionCorePolicy::LeastInFlight(_) => { self.counts.take_pending(&pool.id, selected.worker_id) } }; let target = target(&self.frontends, pool)?; - tracing::trace!(model = request.model, pool = %pool.name, "the router core picked a pool"); + tracing::trace!(model, session, pool = %pool.name, "the router core picked a pool"); Some((target, in_flight)) } - /// A least-in-flight pick stays pending until the caller takes it over - /// or gives it back. - async fn pick(&self, model: &str, allowed: HashSet) -> Option { - let request = advisory(model, allowed); - let selected = match self.policy { - SelectionCorePolicy::RoundRobin => self.service.select(request).await, + /// Books only if `book`, which commits the session's pin; a + /// least-in-flight pick stays pending until the caller takes it over. + async fn pick( + &self, + model: &str, + allowed: HashSet, + session: Option<&str>, + book: bool, + ) -> Option { + let booking = booking(model, allowed, session); + let selected = if book { + self.pin_bookings.fetch_add(1, Ordering::Relaxed); + self.book(booking).await + } else { + self.advise(advisory(booking)).await + }; + selected + .map_err(|error| tracing::debug!(model, %error, "the router core selected no pool")) + .ok() + } + + /// Frees the booking at once: the session's binding outlives it. + async fn book( + &self, + booking: SelectAndReserveRequest, + ) -> Result { + let selected = self.service.select_and_reserve(booking).await?; + if let Some(selection_id) = &selected.selection_id + && let Err(error) = self.service.free_reservation(selection_id).await + { + // The core expires an unfreed booking by age. + tracing::warn!(%selection_id, %error, "the router core did not free a pin booking"); + } + Ok(selected) + } + + async fn advise(&self, request: SelectRequest) -> Result { + match self.policy { + SelectionCorePolicy::RoundRobin | SelectionCorePolicy::Affinity(_) => { + self.service.select(request).await + } SelectionCorePolicy::LeastInFlight(_) => { let service = self.service.clone(); let select = Box::pin(async move { service.select(request).await }); @@ -160,10 +299,7 @@ impl SelectionCoreSelector { .finish() .await } - }; - selected - .map_err(|error| tracing::debug!(model, %error, "the router core selected no pool")) - .ok() + } } fn release(&self, worker: WorkerId) { @@ -217,9 +353,22 @@ impl Drop for PendingSelect { } } -/// An advisory select (books nothing) with no prompt: pools have no KV cache. -fn advisory(model: &str, allowed: HashSet) -> SelectRequest { - SelectRequest { +fn allowed(candidates: &[&Pool]) -> Option> { + let allowed = candidates + .iter() + .filter(|pool| reachable(pool)) + .map(|pool| worker_id(&pool.id)) + .collect::>(); + (!allowed.is_empty()).then_some(allowed) +} + +/// A booking with no prompt: pools have no KV cache. +fn booking( + model: &str, + allowed: HashSet, + session: Option<&str>, +) -> SelectAndReserveRequest { + SelectAndReserveRequest { model_name: model.to_owned(), routing_group: ROUTING_GROUP.to_owned(), selection_id: None, @@ -231,12 +380,32 @@ fn advisory(model: &str, allowed: HashSet) -> SelectRequest { expected_output_tokens: None, priority_jump: None, strict_priority: None, - session_id: None, + session_id: session.map(str::to_owned), session_context: None, affinity_target: None, pinned_worker: None, allowed_worker_ids: Some(allowed), routing_constraints: RoutingConstraints::default(), + } +} + +/// Books nothing; honours a session's pin without setting it. +fn advisory(booking: SelectAndReserveRequest) -> SelectRequest { + SelectRequest { + model_name: booking.model_name, + routing_group: booking.routing_group, + selection_id: None, + prompt: booking.prompt, + router_config_override: None, + expected_output_tokens: None, + priority_jump: None, + strict_priority: None, + session_id: booking.session_id, + session_context: None, + affinity_target: None, + pinned_worker: None, + allowed_worker_ids: booking.allowed_worker_ids, + routing_constraints: booking.routing_constraints, advisory: true, } } @@ -264,6 +433,7 @@ mod tests { use url::Url; + use crate::affinity::tests::{OHIO, fleet, set_occupancy}; use crate::selector::{self, tests::pool}; use crate::{PoolSelector, RoundRobinSelector}; @@ -480,4 +650,61 @@ mod tests { .expect("every abandoned pick was given back"); assert_eq!((in_flight.get(&blue.id), in_flight.get(&green.id)), (0, 0)); } + + fn affinity(signals: &Arc) -> SelectionCorePolicy { + SelectionCorePolicy::Affinity(AffinityConfig { + signals: signals.clone(), + region: OHIO.into(), + headroom_factor: 2.0, + idle: Duration::from_secs(60), + }) + } + + async fn session_s(selection_core: &SelectionCoreSelector, pools: &[Pool]) -> String { + let candidates = pools.iter().collect::>(); + let request = RoutingRequest::new("qwen").with_session(Some("s")); + let (target, _in_flight) = selection_core.select(&request, &candidates).await.unwrap(); + target.frontend.host_str().unwrap().to_owned() + } + + /// The pin holds after its booking is freed, and moves only for twice the room. + #[tokio::test] + async fn affinity_moves_for_twice_the_room_and_does_not_flip_back() { + let (pools, signals) = fleet(); + let (selection_core, _catalog) = + selection_core(affinity(&signals), &pools, &Arc::default()).await; + set_occupancy(&pools, &signals, [0.1, 0.5, 0.5, 0.5]); + assert_eq!(session_s(&selection_core, &pools).await, "ohio-a"); + // 50% against 30%: 0.5 spare against 0.7, under twice. + set_occupancy(&pools, &signals, [0.5, 0.3, 0.3, 0.3]); + assert_eq!(session_s(&selection_core, &pools).await, "ohio-a"); + // 80% against 10%: 0.2 against 0.9, so to the nearest such pool. + set_occupancy(&pools, &signals, [0.8, 0.1, 0.1, 0.1]); + assert_eq!(session_s(&selection_core, &pools).await, "ohio-b"); + // 0.4 spare against 0.6 is under twice, so no flip back. + set_occupancy(&pools, &signals, [0.4, 0.6, 0.6, 0.6]); + assert_eq!(session_s(&selection_core, &pools).await, "ohio-b"); + } + + #[tokio::test] + async fn affinity_books_only_to_set_or_move_a_pin() { + let (pools, signals) = fleet(); + let (selection_core, _catalog) = + selection_core(affinity(&signals), &pools, &Arc::default()).await; + set_occupancy(&pools, &signals, [0.1, 0.5, 0.5, 0.5]); + for _ in 0..3 { + assert_eq!(session_s(&selection_core, &pools).await, "ohio-a"); + } + let candidates = pools.iter().collect::>(); + selection_core + .select(&RoutingRequest::new("qwen"), &candidates) + .await + .unwrap(); + assert_eq!(selection_core.pin_bookings(), 1); + + set_occupancy(&pools, &signals, [0.8, 0.1, 0.1, 0.1]); + assert_eq!(session_s(&selection_core, &pools).await, "ohio-b"); + assert_eq!(session_s(&selection_core, &pools).await, "ohio-b"); + assert_eq!(selection_core.pin_bookings(), 2); + } } diff --git a/lib/global-router/src/selection_core/policies.rs b/lib/global-router/src/selection_core/policies.rs index 008b0f79ac03..ca62a0f30cb8 100644 --- a/lib/global-router/src/selection_core/policies.rs +++ b/lib/global-router/src/selection_core/policies.rs @@ -20,7 +20,7 @@ use dynamo_kv_router::{RoutingPartitionRef, WorkerType}; use dynamo_router_policies::router_modes; use super::directory::PoolDirectory; -use super::{LeastInFlightConfig, SelectionCorePolicy}; +use super::{AffinityConfig, LeastInFlightConfig, SelectionCorePolicy}; use crate::frontend_signals::FrontendSignals; use crate::placement::{self, Placed}; use crate::selector::in_flight_load; @@ -30,6 +30,7 @@ impl SelectionCorePolicy { match self { Self::RoundRobin => router_modes::ROUND_ROBIN, Self::LeastInFlight(_) => "global-least-in-flight", + Self::Affinity(_) => "global-affinity", } } } @@ -53,6 +54,13 @@ pub(super) fn registry( ties: AtomicUsize::new(0), }) } + SelectionCorePolicy::Affinity(config) => { + let config = config.clone(); + provider(move || HeadroomPicker { + config: config.clone(), + directory: directory.clone(), + }) + } }; registry.register_worker_selection(policy.name(), provider)?; Ok(registry) @@ -159,3 +167,56 @@ impl WorkerPicker for LeastInFlightPicker { Ok(row) } } + +/// [`crate::AffinitySelector`]'s rule; the default is the session's bound +/// pool (the core's affinity target), else the cheapest. +struct HeadroomPicker { + config: AffinityConfig, + directory: Arc, +} + +impl WorkerPicker for HeadroomPicker { + fn pick( + &self, + context: &WorkerSelectionContext<'_>, + input: WorkerInputView<'_>, + ) -> Result { + let pools = self.directory.snapshot(); + let signals = &self.config.signals; + let round_trips = signals.region_round_trips(); + let mut placed = input + .candidates() + .iter() + .enumerate() + .map(|(row, candidate)| { + let worker = candidate.worker().worker_id; + let pool = pools.get(&worker); + Placed { + item: (row, worker), + load: pool + .and_then(|pool| signals.occupancy(&pool.id)) + .map(|occupancy| occupancy.load), + distance: pool.and_then(|pool| { + placement::distance(&pool.region, &self.config.region, &round_trips) + }), + in_flight: 0, + } + }) + .collect::>(); + if placed.is_empty() { + return Err(WorkerSelectionPolicyError::failed("no candidates")); + } + placed.sort_by_key(|candidate| candidate.item.1); + placement::order(&mut placed); + let bound = context.affinity_target().map(|target| target.worker_id); + let default = bound + .and_then(|worker| { + placed + .iter() + .position(|candidate| candidate.item.1 == worker) + }) + .unwrap_or(0); + let choice = placement::headroom(&placed, default, self.config.headroom_factor); + Ok(placed[choice].item.0) + } +}