feat(global-router): session affinity on the router core

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.
 </ParamField>
 
+The following is emitted with `--use-selection-core`. It is read from the embedded router core at scrape time.
+
+<ParamField path="dynamo_global_frontend_pin_bookings_total" type="counter">
+  **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.
+</ParamField>
+
 ## 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.

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
 `<cluster>/<namespace>/<DGD name>` 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
 

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),
+        })),
     }
 }
 

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<WatchedSelectionCore>,
+    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(
+            &registry,
+            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<SelectionCoreSelector>) {
+        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<String> {
         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<SelectionCoreSelector>);
+
+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}")
 }

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 {

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<Pool>) -> 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<Pool>,
+    in_flight: Arc<InFlightTracker>,
+) -> 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;
+}

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