feat: Abstractions for Global Router discovery fixes [DEP-1276]

lib/kv-router/Cargo.toml

@@ -54,6 +54,7 @@ tokio = { workspace = true }
 tokio-util = { workspace = true }
 tracing = { workspace = true }
 zmq = { version = "0.10", optional = true }
+url = { workspace = true }
 uuid = { workspace = true }
 xxhash-rust = { workspace = true }
 zeroize = { workspace = true }

lib/kv-router/src/lib.rs

@@ -16,6 +16,7 @@ pub mod identity;
 pub mod indexer;
 pub mod kv_hints;
 pub mod plugins;
+pub mod pool_discovery;
 pub mod protocols;
 pub mod recovery;
 pub mod scheduling;

lib/kv-router/src/pool_discovery/admission.rs

@@ -0,0 +1,23 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use super::PoolDiscoveryError;
+use super::identity::VerifiedRelayIdentity;
+use super::record::{PoolDirectory, PoolRecord};
+use crate::global_view::PoolKey;
+
+/// Decides whether an authenticated Relay may stream state for a pool. The
+/// rule is the same in every mode; only the source of `relay_identity`
+/// differs. A record without one admits any authenticated Relay.
+///
+/// Uses `PoolDirectory::get`, so it reads the watch cache and never lists.
+/// The runtime keys the stream by its `PlaneLease`, not by the claimed key.
+// Signature only; the body comes with the first mode's implementation.
+#[allow(unused_variables)]
+pub async fn admit(
+    pools: &PoolDirectory,
+    key: &PoolKey,
+    caller: &VerifiedRelayIdentity,
+) -> Result<PoolRecord, PoolDiscoveryError> {
+    todo!()
+}

lib/kv-router/src/pool_discovery/connection.rs

@@ -0,0 +1,78 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! The state connections a Relay opens to each router replica. Not part of
+//! discovery; shown so the flow is complete. One implementation serves every
+//! mode.
+//!
+//! Subject to change: the transport (a push RPC from the Relay, or a reverse
+//! tunnel) is not decided yet.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use tokio_util::sync::CancellationToken;
+
+use super::PoolDiscoveryError;
+use super::record::PoolRecord;
+use super::replicas::{ReplicaDirectory, ReplicaEndpoint};
+use crate::global_view::PoolKey;
+
+/// Router side. Checks each caller with `verify`, then `admit`.
+#[async_trait]
+pub trait StateReceiver: Send + Sync {
+    async fn accept(
+        &self,
+        cancel_token: CancellationToken,
+    ) -> Result<(PoolRecord, Box<dyn StateConnection>), PoolDiscoveryError>;
+}
+
+/// Relay side, one replica. The only part that depends on the transport.
+#[async_trait]
+pub trait StateDialer: Send + Sync {
+    async fn dial(
+        &self,
+        replica: &ReplicaEndpoint,
+        pool: &PoolKey,
+        cancel_token: CancellationToken,
+    ) -> Result<Box<dyn StateConnection>, PoolDiscoveryError>;
+}
+
+/// Relay side, every replica. Keeps one connection to each replica in the
+/// `ReplicaDirectory`: dials on `Added`, closes on `Removed`, and dials
+/// again when a connection closes. The same in every mode, because the
+/// directory already hides the mode.
+// Signature only; the body comes with the transport.
+#[allow(dead_code)]
+pub struct StateBroadcaster {
+    replicas: Arc<ReplicaDirectory>,
+    dialer: Arc<dyn StateDialer>,
+    pool: PoolKey,
+}
+
+impl StateBroadcaster {
+    pub fn new(
+        replicas: Arc<ReplicaDirectory>,
+        dialer: Arc<dyn StateDialer>,
+        pool: PoolKey,
+    ) -> Self {
+        Self {
+            replicas,
+            dialer,
+            pool,
+        }
+    }
+
+    /// Runs until `cancel_token` fires.
+    #[allow(unused_variables)]
+    pub async fn run(self, cancel_token: CancellationToken) -> Result<(), PoolDiscoveryError> {
+        todo!()
+    }
+}
+
+/// One open connection, on either side.
+#[async_trait]
+pub trait StateConnection: Send + Sync {
+    /// Resolves when the other side disconnects.
+    async fn closed(&self);
+}

lib/kv-router/src/pool_discovery/identity.rs

@@ -0,0 +1,59 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::net::IpAddr;
+
+use async_trait::async_trait;
+
+use super::PoolDiscoveryError;
+
+/// Expected identity of a pool's Relay, for example a SPIFFE ID.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct RelayIdentity(pub String);
+
+/// An identity proven by a [`RelayAuthenticator`]. Only [`verify`] creates
+/// one, so writes cannot use a claimed identity.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct VerifiedRelayIdentity(pub(super) RelayIdentity);
+
+impl VerifiedRelayIdentity {
+    /// Read access for other crates, for example to compare with
+    /// `PoolRecord::relay_identity` or to log a denied caller.
+    pub fn identity(&self) -> &RelayIdentity {
+        &self.0
+    }
+}
+
+/// Raw credentials the transport took from a Relay's connection.
+/// No `Debug`: it can hold a secret.
+#[derive(Clone)]
+#[non_exhaustive]
+pub enum PeerCredentials {
+    /// From a client certificate that TLS already checked.
+    Spiffe(String),
+    BearerToken(String),
+    /// The connection's source address, on a trusted network with no mesh.
+    /// Names the caller only if traffic between clusters is not
+    /// source-NATed.
+    PeerAddress(IpAddr),
+}
+
+/// Implementations live in other crates, so they return a plain
+/// `RelayIdentity` and callers go through [`verify`].
+#[async_trait]
+pub trait RelayAuthenticator: Send + Sync {
+    async fn authenticate(
+        &self,
+        credentials: &PeerCredentials,
+    ) -> Result<RelayIdentity, PoolDiscoveryError>;
+}
+
+pub async fn verify(
+    authenticator: &dyn RelayAuthenticator,
+    credentials: &PeerCredentials,
+) -> Result<VerifiedRelayIdentity, PoolDiscoveryError> {
+    authenticator
+        .authenticate(credentials)
+        .await
+        .map(VerifiedRelayIdentity)
+}

lib/kv-router/src/pool_discovery/mod.rs

@@ -0,0 +1,66 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Pluggable pool discovery for the Global Router. Design:
+//! <https://github.com/ai-dynamo/enhancements/pull/103>.
+//!
+//! Pools dial into the hub. Discovery answers three questions: which pools
+//! exist, who may speak for each pool, and which router replicas a pool must
+//! connect to. Pool state still arrives on the Relay's own connections and is
+//! assembled by [`crate::global_view::source::PoolObservationAssembler`].
+//!
+//! Only the parts that differ between modes are traits: [`PoolDirectory`],
+//! [`PoolRegistrar`], [`ReplicaDirectory`] and [`RelayAuthenticator`].
+//! Admission and the Relay's `RegisterPool` client work the same in every
+//! mode, so they are concrete code. The state connections in [`connection`]
+//! are shown for the full picture and are subject to change.
+//!
+//! In the Kubernetes modes, the Global Discovery Service serves these traits
+//! over HTTP: `POST /pools` registers a pool, and `GET /pools` and
+//! `GET /routers` list pools and router replicas.
+//!
+//! | Mode | Pool list stored in | Pools call |
+//! |---|---|---|
+//! | gRPC | Router memory | `RegisterPool` on the router |
+//! | Hybrid | Hub Kubernetes API, written by the Global Discovery Service | `RegisterPool` on the service |
+//! | Kubernetes (SIG Multicluster) | Imported EndpointSlices, written by the MCS implementation | Nothing: the operator creates a `ServiceExport` |
+//! | File | Static file (tests, first deployments) | Nothing |
+
+pub mod admission;
+pub mod connection;
+pub mod identity;
+pub mod mode;
+pub mod record;
+pub mod registrar;
+pub mod replicas;
+pub mod watch;
+
+pub use admission::admit;
+pub use connection::{StateBroadcaster, StateConnection, StateDialer, StateReceiver};
+pub use identity::{
+    PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify,
+};
+pub use mode::{BuildPoolDiscovery, PoolDiscovery, PoolDiscoveryMode};
+pub use record::{Capacity, PoolDirectory, PoolEvent, PoolRecord, PoolRecordBuilder, Revision};
+pub use registrar::{LeasePolicy, PoolRegistrar, Registration};
+pub use replicas::{ReplicaDirectory, ReplicaEndpoint, ReplicaEvent, ReplicaId};
+pub use watch::{Completeness, Directory, Snapshot, WatchEvent, WatchStream};
+
+use thiserror::Error;
+
+use crate::global_view::PoolKey;
+
+/// Shared by every discovery interface so the gRPC server can map each case
+/// to one status code.
+#[derive(Debug, Error)]
+#[non_exhaustive]
+pub enum PoolDiscoveryError {
+    #[error("caller is not authenticated")]
+    Unauthenticated,
+    #[error("caller may not act for pool {0:?}")]
+    PermissionDenied(PoolKey),
+    #[error("pool {0:?} is not registered")]
+    UnknownPool(PoolKey),
+    #[error("discovery backend is unavailable")]
+    Unavailable(#[source] anyhow::Error),
+}

lib/kv-router/src/pool_discovery/mode.rs

@@ -0,0 +1,52 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+
+use super::PoolDiscoveryError;
+use super::identity::RelayAuthenticator;
+use super::record::PoolDirectory;
+use super::registrar::{LeasePolicy, PoolRegistrar};
+use super::replicas::ReplicaDirectory;
+
+/// Chosen once at startup, like `dynamo_runtime::distributed::DiscoveryBackend`.
+/// Each mode's own settings are added with its implementation.
+#[derive(Clone, Debug)]
+#[non_exhaustive]
+pub enum PoolDiscoveryMode {
+    /// Pools call `RegisterPool`. Records live in router memory.
+    Grpc {
+        lease: LeasePolicy,
+    },
+    /// Pools call `RegisterPool`. The Global Discovery Service writes records
+    /// to the hub's Kubernetes API.
+    Hybrid {
+        lease: LeasePolicy,
+    },
+    /// SIG Multicluster: pools are MCS `ServiceExport`s. Liveness is endpoint
+    /// readiness, so there is no lease.
+    Kubernetes,
+    File {
+        path: PathBuf,
+    },
+}
+
+/// Everything the router needs from the selected mode. The router code is
+/// the same for every mode.
+pub struct PoolDiscovery {
+    pub pools: Arc<PoolDirectory>,
+    pub replicas: Arc<ReplicaDirectory>,
+    /// `None` in Kubernetes and File modes.
+    pub registrar: Option<Arc<dyn PoolRegistrar>>,
+    /// Incoming Relays go through `verify` with this, then
+    /// [`super::admit`] with `pools`.
+    pub authenticator: Arc<dyn RelayAuthenticator>,
+}
+
+#[async_trait]
+pub trait BuildPoolDiscovery {
+    async fn build(self) -> Result<PoolDiscovery, PoolDiscoveryError>;
+}

lib/kv-router/src/pool_discovery/record.rs

@@ -0,0 +1,58 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use derive_builder::Builder;
+use url::Url;
+
+use super::identity::RelayIdentity;
+use super::watch::{Directory, WatchEvent};
+use crate::global_view::PoolKey;
+use crate::global_view::state::PoolLocation;
+
+/// Orders updates to one pool's record.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct Revision(pub u64);
+
+/// One pool as discovery knows it. Pool state (load, KV, readiness) is not
+/// part of the record. Router replicas derive the routing `PoolId` from `key`
+/// with the shared `PoolIdDeriver`.
+#[derive(Clone, Debug, PartialEq, Eq, Builder)]
+#[builder(setter(into))]
+pub struct PoolRecord {
+    pub key: PoolKey,
+    pub location: PoolLocation,
+    // The fields below are empty until the pool registers them or the Relay's
+    // catalog reports them. Kubernetes (MCS) mode knows only the pool's name.
+    /// With `frontend_endpoint`, becomes `RelayPoolScope` in dynamo-llm.
+    #[builder(default)]
+    pub runtime_namespace: Option<String>,
+    /// Dynamo endpoint path of the pool's frontend, for example
+    /// `frontend.generate`.
+    #[builder(default)]
+    pub frontend_endpoint: Option<String>,
+    #[builder(default)]
+    pub model: Option<String>,
+    #[builder(default)]
+    pub capabilities: Vec<String>,
+    /// URLs of the pool's local frontends, which a Global Router sends
+    /// requests to.
+    #[builder(default)]
+    pub frontends: Vec<Url>,
+    #[builder(default)]
+    pub capacity: Vec<Capacity>,
+    /// The only identity allowed to register this pool and stream its state.
+    #[builder(default)]
+    pub relay_identity: Option<RelayIdentity>,
+    #[builder(default)]
+    pub revision: Revision,
+}
+
+/// The workers of one hardware SKU.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Capacity {
+    pub sku: String,
+    pub count: u32,
+}
+
+pub type PoolDirectory = dyn Directory<PoolRecord, PoolKey>;
+pub type PoolEvent = WatchEvent<PoolRecord, PoolKey>;