feat(global-router): replicate bookings and session pins between frontends

lib/global-frontend/src/catalog.rs

@@ -34,6 +34,13 @@ impl PoolSnapshot {
     }
 }
 
+/// The router core follows the catalog as a list of pools.
+impl AsRef<[Pool]> for PoolSnapshot {
+    fn as_ref(&self) -> &[Pool] {
+        &self.pools
+    }
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use dynamo_global_common::PoolKey;

lib/global-frontend/src/main.rs

@@ -2,6 +2,21 @@
 // SPDX-License-Identifier: Apache-2.0
 
 //! Global Frontend. See `dynamo_global_frontend` and the crate's README.
+//!
+//! `--selection-core` selects through the embedded router core instead of
+//! the Global Router's own selectors, with each pool as one of the core's
+//! workers and every `--router-policy` as one of its plugins. With
+//! `--replica-sync`, the core shares its requests in flight and session
+//! pins with the other Global Frontends (`--peer`) over its ZMQ mesh,
+//! publishing on `--replica-sync-port`:
+//!
+//! ```bash
+//! cargo run -p dynamo-global-frontend -- --port 8000 --selection-core \
+//!     --router-policy affinity --region ohio \
+//!     --pool ohio/dynamo/mocker-1=Qwen/Qwen3-0.6B@http://10.42.0.5:8000 \
+//!     --pool-relay ohio/dynamo/mocker-1=http://10.42.0.6:50051 \
+//!     --replica-sync --replica-sync-port 5600 --peer tcp://10.42.1.5:5600
+//! ```
 
 use std::collections::HashSet;
 use std::sync::Arc;
@@ -14,8 +29,8 @@ use dynamo_global_frontend::{
     Metrics, PoolSnapshot, app, follow_relays, measure_round_trips, serve,
 };
 use dynamo_global_router::{
-    AffinitySelector, InFlightTracker, LeastInFlightSelector, Pool, PoolSelector, PoolSignals,
-    RoundRobinSelector,
+    AffinityConfig, AffinitySelector, CorePolicy, CoreSelector, InFlightTracker,
+    LeastInFlightSelector, Pool, PoolSignals, ReplicaSync, RequestSelector, RoundRobinSelector,
 };
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
@@ -58,6 +73,10 @@ struct Args {
     connect_timeout_secs: u64,
     #[arg(long, value_enum, default_value_t = RouterPolicy::RoundRobin)]
     router_policy: RouterPolicy,
+    /// Select through the embedded router core, with each pool as one of
+    /// its workers, instead of the Global Router's own selectors.
+    #[arg(long)]
+    selection_core: bool,
     /// The region this process runs in, named as `--pool` names clusters.
     /// `affinity` tries pools in it first, and metrics tell requests kept in
     /// it from the rest.
@@ -74,8 +93,22 @@ struct Args {
     #[arg(long, default_value_t = 2.0)]
     headroom_factor: f64,
     /// `affinity`: how long a session keeps its pool without a request.
+    /// The core keeps at most 65,536 sessions per model; while the table is
+    /// full, a new session is placed without a pin until idle ones expire.
     #[arg(long, default_value_t = 1800)]
     session_idle_secs: u64,
+    /// Share this replica's requests in flight and session pins with the
+    /// other Global Frontends over the router core's ZMQ mesh. Needs
+    /// `--selection-core` and `--replica-sync-port`.
+    #[arg(long, requires_all = ["selection_core", "replica_sync_port"])]
+    replica_sync: bool,
+    /// `--replica-sync`: the port this replica publishes on.
+    #[arg(long, requires = "replica_sync")]
+    replica_sync_port: Option<u16>,
+    /// `--replica-sync`: another Global Frontend's `--replica-sync-port`
+    /// endpoint, as `tcp://host:port`. Repeat the flag for each peer.
+    #[arg(long = "peer", value_parser = parse_peer, requires = "replica_sync")]
+    peers: Vec<String>,
 }
 
 fn parse_key(key: &str) -> Result<PoolKey, String> {
@@ -109,6 +142,20 @@ fn parse_pool_relay(value: &str) -> Result<(PoolKey, Url), String> {
     Ok((parse_key(key)?, relay))
 }
 
+fn parse_peer(value: &str) -> Result<String, String> {
+    let Some(address) = value.strip_prefix("tcp://") else {
+        return Err(format!("expected tcp://host:port, got {value}"));
+    };
+    match address.rsplit_once(':') {
+        Some((host, port))
+            if !host.is_empty() && port.parse::<u16>().is_ok_and(|port| port > 0) =>
+        {
+            Ok(value.to_owned())
+        }
+        _ => Err(format!("expected tcp://host:port, got {value}")),
+    }
+}
+
 /// Gives each pool its Relay. Every Relay must name a configured pool, and
 /// no two pools may share one.
 fn attach_relays(pools: &mut [Pool], relays: &[(PoolKey, Url)]) -> Result<()> {
@@ -146,6 +193,45 @@ fn parse_pool(value: &str) -> Result<Pool, String> {
     Ok(Pool::new(&key, model, key.cluster_id(), frontends))
 }
 
+fn replica_sync(args: &Args) -> Option<ReplicaSync> {
+    // The parser makes `--replica-sync` bring its port.
+    let port = args.replica_sync_port.filter(|_| args.replica_sync)?;
+    Some(ReplicaSync {
+        port,
+        peers: args.peers.clone(),
+    })
+}
+
+/// Starts following each pool's Relay and measuring each region's round
+/// trip, and returns the signals they fill, for the affinity policy on
+/// either path.
+fn load_signals(
+    args: &Args,
+    pools: &watch::Receiver<Arc<PoolSnapshot>>,
+    in_flight: Arc<InFlightTracker>,
+    shutdown: &CancellationToken,
+) -> Arc<PoolSignals> {
+    if args.pool_relays.is_empty() {
+        tracing::warn!(
+            "no --pool-relay is given, so affinity knows no pool's load \
+             and keeps each session on its nearest pool"
+        );
+    }
+    let signals = Arc::new(PoolSignals::new(in_flight));
+    tokio::spawn(follow_relays(
+        pools.clone(),
+        signals.clone(),
+        shutdown.child_token(),
+    ));
+    tokio::spawn(measure_round_trips(
+        pools.clone(),
+        signals.clone(),
+        ROUND_TRIP_INTERVAL,
+        shutdown.child_token(),
+    ));
+    signals
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     dynamo_runtime::logging::init();
@@ -173,37 +259,44 @@ async fn main() -> Result<()> {
 
     // The pools never change; the sender must live as long as the process,
     // since the Relay followers stop once it drops.
-    let (_catalog, pools) = watch::channel(Arc::new(PoolSnapshot::new(args.pools)));
+    let (_catalog, pools) =
+        watch::channel(Arc::new(PoolSnapshot::new(std::mem::take(&mut args.pools))));
     let http = reqwest::Client::builder()
         .connect_timeout(Duration::from_secs(args.connect_timeout_secs))
         .build()
         .context("build the pool HTTP client")?;
     let in_flight = Arc::new(InFlightTracker::new());
     let metrics = Metrics::new(args.region.clone(), in_flight.clone());
-    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::Affinity => {
+    let idle = Duration::from_secs(args.session_idle_secs);
+    let selector: Arc<dyn RequestSelector> = match (args.selection_core, args.router_policy) {
+        (true, policy) => {
+            let policy = match policy {
+                RouterPolicy::RoundRobin => CorePolicy::RoundRobin,
+                RouterPolicy::LeastInFlight => CorePolicy::LeastRequests,
+                RouterPolicy::Affinity => CorePolicy::Affinity(AffinityConfig {
+                    signals: load_signals(&args, &pools, in_flight.clone(), &shutdown),
+                    region: args.region.clone().context("--region is required")?,
+                    headroom_factor: args.headroom_factor,
+                    idle,
+                }),
+            };
+            let core = Arc::new(CoreSelector::new(policy, replica_sync(&args)).await?);
+            // Serve only once the configured pools are the core's workers;
+            // a request before that would get a spurious 503.
+            let (ready, follower) = core.follow(pools.clone(), shutdown.child_token());
+            tokio::spawn(follower);
+            ready
+                .await
+                .context("the router core stopped before taking the first pool snapshot")?;
+            core
+        }
+        (false, RouterPolicy::RoundRobin) => Arc::new(RoundRobinSelector::new()),
+        (false, RouterPolicy::LeastInFlight) => {
+            Arc::new(LeastInFlightSelector::new(in_flight.clone()))
+        }
+        (false, RouterPolicy::Affinity) => {
             let region = args.region.clone().context("--region is required")?;
-            if args.pool_relays.is_empty() {
-                tracing::warn!(
-                    "no --pool-relay is given, so affinity knows no pool's load \
-                     and keeps each session on its nearest pool"
-                );
-            }
-            let signals = Arc::new(PoolSignals::new(in_flight.clone()));
-            tokio::spawn(follow_relays(
-                pools.clone(),
-                signals.clone(),
-                shutdown.child_token(),
-            ));
-            tokio::spawn(measure_round_trips(
-                pools.clone(),
-                signals.clone(),
-                ROUND_TRIP_INTERVAL,
-                shutdown.child_token(),
-            ));
-            let idle = Duration::from_secs(args.session_idle_secs);
+            let signals = load_signals(&args, &pools, in_flight.clone(), &shutdown);
             let selector = Arc::new(AffinitySelector::new(
                 signals,
                 region,
@@ -301,6 +394,62 @@ mod tests {
         assert_eq!(args.router_policy, RouterPolicy::RoundRobin);
     }
 
+    #[test]
+    fn the_core_and_replica_sync_are_off_unless_asked_for() {
+        let pool = "--pool=ohio/dynamo/blue=qwen@http://blue:8000";
+        let args = Args::try_parse_from(["global-frontend", pool]).unwrap();
+        assert!(!args.selection_core);
+        assert!(!args.replica_sync);
+        assert_eq!(replica_sync(&args), None);
+        let args = Args::try_parse_from(["global-frontend", pool, "--selection-core"]).unwrap();
+        assert_eq!(replica_sync(&args), None);
+
+        let args = Args::try_parse_from([
+            "global-frontend",
+            pool,
+            "--selection-core",
+            "--replica-sync",
+            "--replica-sync-port",
+            "5600",
+            "--peer",
+            "tcp://10.0.0.2:5600",
+            "--peer",
+            "tcp://10.0.0.3:5600",
+        ])
+        .unwrap();
+        assert!(args.selection_core);
+        assert_eq!(
+            replica_sync(&args),
+            Some(ReplicaSync {
+                port: 5600,
+                peers: vec![
+                    "tcp://10.0.0.2:5600".to_owned(),
+                    "tcp://10.0.0.3:5600".to_owned()
+                ],
+            })
+        );
+
+        // Replica sync needs the core and the port this replica publishes
+        // on; its port and peers need replica sync.
+        for rejected in [
+            &["--replica-sync", "--replica-sync-port", "5600"][..],
+            &["--selection-core", "--replica-sync"],
+            &["--selection-core", "--replica-sync-port", "5600"],
+            &["--selection-core", "--peer", "tcp://10.0.0.2:5600"],
+        ] {
+            let args = [&["global-frontend", pool][..], rejected].concat();
+            assert!(Args::try_parse_from(&args).is_err(), "{rejected:?}");
+        }
+        // The mesh dials tcp:// endpoints only.
+        assert!(parse_peer("http://10.0.0.2:5600").is_err());
+        assert!(parse_peer("tcp://10.0.0.2").is_err());
+        assert!(parse_peer("tcp://:5600").is_err());
+        assert_eq!(
+            parse_peer("tcp://peer:5600").as_deref(),
+            Ok("tcp://peer:5600")
+        );
+    }
+
     #[test]
     fn requires_a_pool() {
         assert!(Args::try_parse_from(["global-frontend"]).is_err());

lib/global-frontend/src/proxy.rs

@@ -19,7 +19,9 @@ use axum::response::{IntoResponse, Response};
 use axum::routing::{get, post};
 use axum::serve::ListenerExt;
 use axum::{Json, Router};
-use dynamo_global_router::{InFlightTracker, PoolSelector, RoutingRequest};
+use dynamo_global_router::{
+    InFlightTracker, MAX_SESSION_ID_BYTES, RequestSelector, RoutingRequest, Selection,
+};
 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;
@@ -42,7 +44,7 @@ const SHUTDOWN_GRACE: Duration = Duration::from_secs(30);
 
 struct Proxy {
     pools: watch::Receiver<Arc<PoolSnapshot>>,
-    selector: Arc<dyn PoolSelector>,
+    selector: Arc<dyn RequestSelector>,
     in_flight: Arc<InFlightTracker>,
     http: reqwest::Client,
     metrics: Metrics,
@@ -51,13 +53,14 @@ struct Proxy {
 /// The Global Frontend's routes over the pools `pools` publishes.
 ///
 /// Each request counts in `in_flight` from its pool choice until its
-/// response ends; selectors and signals that weigh this process's own
-/// requests must read the same tracker. `http` must not set a total request
-/// timeout, because a response streams for as long as the pool generates.
-/// `/metrics` serves `metrics`.
+/// response ends, and whatever `selector` holds for it is held as long;
+/// selectors and signals that weigh this process's own requests must read
+/// the same tracker. `http` must not set a total request timeout, because a
+/// response streams for as long as the pool generates. `/metrics` serves
+/// `metrics`.
 pub fn app(
     pools: watch::Receiver<Arc<PoolSnapshot>>,
-    selector: Arc<dyn PoolSelector>,
+    selector: Arc<dyn RequestSelector>,
     in_flight: Arc<InFlightTracker>,
     http: reqwest::Client,
     metrics: Metrics,
@@ -150,8 +153,21 @@ async fn forward(
     let session = headers
         .get(HEADER_DYNAMO_SESSION_ID)
         .and_then(|session| session.to_str().ok());
+    // The core's session table refuses longer ids; refusing them here keeps
+    // that a client error rather than an unavailable pool.
+    if session.is_some_and(|session| session.len() > MAX_SESSION_ID_BYTES) {
+        metrics.rejected(error_type::VALIDATION);
+        return error(
+            StatusCode::BAD_REQUEST,
+            format!("{HEADER_DYNAMO_SESSION_ID} is longer than {MAX_SESSION_ID_BYTES} bytes"),
+        );
+    }
     let request = RoutingRequest::new(&model).with_session(session);
-    let Some(target) = proxy.selector.select(&request, &candidates) else {
+    // `lifecycle` is what the selector holds for the request. It is dropped
+    // wherever this handler or the response stream is, which ends the
+    // request for the selector too.
+    let Some(Selection { target, lifecycle }) = proxy.selector.route(&request, &candidates).await
+    else {
         metrics.rejected(error_type::UNAVAILABLE);
         let message = format!("no pool can serve model {model}");
         return error(StatusCode::SERVICE_UNAVAILABLE, message);
@@ -212,6 +228,7 @@ async fn forward(
     // ends or the client drops it; dropping it closes the pool connection.
     let body = upstream.bytes_stream().map(move |chunk| {
         let _in_flight = &in_flight;
+        let _lifecycle = &lifecycle;
         if let Some(timing) = &mut timing {
             timing.chunk(&chunk);
         }