feat(global-router): pools as router-core workers, with round-robin as a core plugin

docs/fern/pages/reference/observability/metrics-catalog.mdx

@@ -954,6 +954,10 @@ Emitted by the Global Frontend at `/metrics` on its serving port. It forwards ea
   **Experimental.** Requests this Global Frontend has sent to each pool and not yet finished. Other replicas count their own. Labeled by `pool` and `region`; read at scrape time, so a pool removed from the configuration loses its series.
 </ParamField>
 
+<ParamField path="dynamo_global_frontend_overhead_selection_ms" type="histogram">
+  **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>
+
 ## 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

@@ -128,8 +128,8 @@ new routers could keep Relays from ever finding them.
 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; and requests in flight per
-pool. The `pool` label is the pool's
+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
 `<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
@@ -150,6 +150,13 @@ for each family.
   Frontend's hosts, or each idle connection starts over with a small send
   window and a long prompt costs extra round trips.
 
+## Router core
+
+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. Only `round-robin` runs on the core for now, picking as the
+round-robin selector does.
+
 ## Limitations
 
 Pool registration remains the Relay deployment's responsibility. There is no

lib/global-frontend/src/catalog.rs

@@ -34,6 +34,12 @@ impl PoolSnapshot {
     }
 }
 
+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/lib.rs

@@ -24,7 +24,7 @@ mod round_trip;
 pub use catalog::PoolSnapshot;
 pub use discovery::discover_pools;
 pub use metrics::Metrics;
-pub use proxy::{app, serve};
+pub use proxy::{app, app_with_selection_core, serve};
 pub use relay::follow_relays;
 pub use relay_receiver::serve_relay_views;
 pub(crate) use relay_state::RelayPoolState;

lib/global-frontend/src/main.rs

@@ -11,17 +11,21 @@ use anyhow::{Context, Result, bail};
 use clap::{Parser, ValueEnum};
 use dynamo_global_common::{PoolIdDeriver, PoolKey, V1PoolIdDeriver};
 use dynamo_global_frontend::{
-    Metrics, PoolSnapshot, RelayGatedPoolCatalog, app, discover_pools, follow_relays,
-    measure_round_trips, serve, serve_relay_views,
+    Metrics, PoolSnapshot, RelayGatedPoolCatalog, app, app_with_selection_core, discover_pools,
+    follow_relays, measure_round_trips, serve, serve_relay_views,
 };
 use dynamo_global_router::{
     AffinitySelector, InFlightTracker, LeastInFlightSelector, Pool, PoolSelector, PoolSignals,
-    RoundRobinSelector,
+    RoundRobinSelector, SelectionCorePolicy, SelectionCoreSelector,
 };
+use dynamo_runtime::config::environment_names::logging as env_logging;
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
 use url::Url;
 
+/// The log filter without `DYN_LOG`: the router core logs each selection at
+/// INFO, which `--use-selection-core` would turn into a line per request.
+const DEFAULT_LOG: &str = "info,dynamo_kv_router::scheduling::selector=warn";
 /// How often the round trip to each pool is measured.
 const ROUND_TRIP_INTERVAL: Duration = Duration::from_secs(30);
 /// How often expired session pins are dropped.
@@ -67,6 +71,9 @@ 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 (round-robin only for now).
+    #[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"))]
@@ -154,8 +161,16 @@ fn parse_pool(value: &str) -> Result<Pool, String> {
     Ok(Pool::new(&key, model, key.cluster_id(), frontends))
 }
 
+fn main() -> Result<()> {
+    if std::env::var_os(env_logging::DYN_LOG).is_none() {
+        // SAFETY: no other thread has started yet.
+        unsafe { std::env::set_var(env_logging::DYN_LOG, DEFAULT_LOG) };
+    }
+    run()
+}
+
 #[tokio::main]
-async fn main() -> Result<()> {
+async fn run() -> Result<()> {
     dynamo_runtime::logging::init();
     let mut args = Args::parse();
     let mut ids = HashSet::new();
@@ -169,6 +184,9 @@ 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();
@@ -233,35 +251,49 @@ async fn main() -> Result<()> {
         .build()
         .context("build the pool HTTP client")?;
     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 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"
-                );
+    let app = if let Some(policy) = selection_core_policy {
+        let selection_core = Arc::new(SelectionCoreSelector::new(policy, in_flight.clone()).await?);
+        // Serve the pools the core has applied, from the first snapshot on,
+        // or early requests get a 503.
+        let (mut applied, reconcile) =
+            selection_core.reconcile_catalog(pools, shutdown.child_token());
+        tokio::spawn(reconcile);
+        applied
+            .changed()
+            .await
+            .context("the router core stopped before applying the first pool snapshot")?;
+        app_with_selection_core(applied, selection_core, in_flight, http, metrics)
+    } 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::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 idle = Duration::from_secs(args.session_idle_secs);
+                let selector = Arc::new(AffinitySelector::new(
+                    signals,
+                    region,
+                    args.headroom_factor,
+                    idle,
+                ));
+                tokio::spawn(sweep_pins(selector.clone(), shutdown.child_token()));
+                selector
             }
-            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 selector = Arc::new(AffinitySelector::new(
-                signals,
-                region,
-                args.headroom_factor,
-                idle,
-            ));
-            tokio::spawn(sweep_pins(selector.clone(), shutdown.child_token()));
-            selector
-        }
+        };
+        app(pools, selector, in_flight, http, metrics)
     };
-    let app = app(pools, selector, in_flight, http, metrics);
     let address = format!("{}:{}", args.host, args.port);
     let listener = tokio::net::TcpListener::bind(&address)
         .await
@@ -286,6 +318,15 @@ async fn main() -> Result<()> {
     }
 }
 
+fn selection_core_policy(policy: RouterPolicy) -> Result<SelectionCorePolicy> {
+    match policy {
+        RouterPolicy::RoundRobin => Ok(SelectionCorePolicy::RoundRobin),
+        RouterPolicy::LeastInFlight | RouterPolicy::Affinity => {
+            bail!("--use-selection-core supports only --router-policy round-robin for now")
+        }
+    }
+}
+
 /// 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) {

lib/global-frontend/src/metrics.rs

@@ -43,6 +43,7 @@ pub struct Metrics {
     time_to_first_byte: HistogramVec,
     duration: HistogramVec,
     in_flight: IntGaugeVec,
+    selection: Histogram,
     /// A scrape resets and refills the gauges, so scrapes must not overlap.
     scrape: Mutex<()>,
 }
@@ -121,6 +122,18 @@ impl Metrics {
                 &POOL_LABELS,
             ),
         );
+        // Mirrors the local router's `dynamo_router_overhead_*_ms`. Unlabelled:
+        // compare selection paths across runs.
+        let selection = register(
+            &registry,
+            Histogram::with_opts(
+                HistogramOpts::new(
+                    name("overhead_selection_ms"),
+                    "Time to select a pool for a request, in milliseconds",
+                )
+                .buckets(prometheus::exponential_buckets(0.01, 3.0, 17).expect("valid buckets")),
+            ),
+        );
         Self {
             registry,
             region,
@@ -131,10 +144,16 @@ impl Metrics {
             time_to_first_byte,
             duration,
             in_flight,
+            selection,
             scrape: Mutex::new(()),
         }
     }
 
+    pub(crate) fn selected(&self, started: Instant) {
+        self.selection
+            .observe(started.elapsed().as_secs_f64() * 1000.0);
+    }
+
     /// 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) {

lib/global-frontend/src/proxy.rs

@@ -18,7 +18,7 @@ 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};
+use dynamo_global_router::{InFlightTracker, 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;
@@ -45,12 +45,18 @@ const SHUTDOWN_GRACE: Duration = Duration::from_secs(30);
 
 struct Proxy {
     pools: watch::Receiver<Arc<PoolSnapshot>>,
-    selector: Arc<dyn PoolSelector>,
+    selector: Selector,
     in_flight: Arc<InFlightTracker>,
     http: reqwest::Client,
     metrics: Metrics,
 }
 
+/// A pool selector, or the router core.
+enum Selector {
+    Pool(Arc<dyn PoolSelector>),
+    SelectionCore(Arc<SelectionCoreSelector>),
+}
+
 /// The Global Frontend's routes over the pools `pools` publishes.
 ///
 /// Each request counts in `in_flight` from its pool choice until its
@@ -65,13 +71,35 @@ pub fn app(
     http: reqwest::Client,
     metrics: Metrics,
 ) -> Router {
-    let proxy = Arc::new(Proxy {
+    build_router(Proxy {
         pools,
-        selector,
+        selector: Selector::Pool(selector),
         in_flight,
         http,
         metrics,
-    });
+    })
+}
+
+/// [`app`], selecting through `selection_core`, which must count in `in_flight`.
+/// `pools` is the receiver its `reconcile_catalog` returns.
+pub fn app_with_selection_core(
+    pools: watch::Receiver<Arc<PoolSnapshot>>,
+    selection_core: Arc<SelectionCoreSelector>,
+    in_flight: Arc<InFlightTracker>,
+    http: reqwest::Client,
+    metrics: Metrics,
+) -> Router {
+    build_router(Proxy {
+        pools,
+        selector: Selector::SelectionCore(selection_core),
+        in_flight,
+        http,
+        metrics,
+    })
+}
+
+fn build_router(proxy: Proxy) -> Router {
+    let proxy = Arc::new(proxy);
     Router::new()
         .route("/v1/models", get(models))
         .route("/v1/{*route}", post(forward))
@@ -182,6 +210,11 @@ async fn forward(
         let message = "the request body must be a JSON object with a string `model`";
         return error(StatusCode::BAD_REQUEST, message.to_owned());
     };
+    // The router core's workers stay as `pools` shows them until the pick.
+    let catalog = match &proxy.selector {
+        Selector::Pool(_) => None,
+        Selector::SelectionCore(selection_core) => Some(selection_core.hold_catalog().await),
+    };
     let snapshot = proxy.pools.borrow().clone();
     let candidates = snapshot.pools_for(&model).collect::<Vec<_>>();
     if candidates.is_empty() {
@@ -198,13 +231,23 @@ async fn forward(
         .get(HEADER_DYNAMO_SESSION_ID)
         .and_then(|session| session.to_str().ok());
     let request = RoutingRequest::new(&model).with_session(session);
-    let Some(target) = proxy.selector.select(&request, &candidates) else {
+    let selecting = Instant::now();
+    let selected = match &proxy.selector {
+        Selector::Pool(selector) => selector
+            .select(&request, &candidates)
+            .map(|target| (target, proxy.in_flight.start(&target.pool.id))),
+        Selector::SelectionCore(selection_core) => {
+            selection_core.select(&request, &candidates).await
+        }
+    };
+    drop(catalog);
+    metrics.selected(selecting);
+    let Some((target, in_flight)) = selected else {
         metrics.rejected(error_type::UNAVAILABLE);
         let message = format!("no pool can serve model {model}");
         return error(StatusCode::SERVICE_UNAVAILABLE, message);
     };
     let pool = target.pool;
-    let in_flight = proxy.in_flight.start(&pool.id);
     let Some(url) = route_url(target.frontend, &uri) else {
         metrics.rejected(error_type::VALIDATION);
         return error(StatusCode::BAD_REQUEST, "invalid route".to_owned());