diff --git a/Cargo.lock b/Cargo.lock index 6ea9c4f9795f..7f091414ba30 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2662,6 +2662,7 @@ dependencies = [ "tower 0.5.3", "tracing", "tracing-appender", + "url", "uuid", "xxhash-rust", "zeroize", diff --git a/lib/kv-router/Cargo.toml b/lib/kv-router/Cargo.toml index 053b84dec251..69bbdc51cddd 100644 --- a/lib/kv-router/Cargo.toml +++ b/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 } diff --git a/lib/kv-router/src/lib.rs b/lib/kv-router/src/lib.rs index c19b7274e810..72b42aee77bd 100644 --- a/lib/kv-router/src/lib.rs +++ b/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; diff --git a/lib/kv-router/src/pool_discovery/admission.rs b/lib/kv-router/src/pool_discovery/admission.rs new file mode 100644 index 000000000000..4509f55bb382 --- /dev/null +++ b/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 { + todo!() +} diff --git a/lib/kv-router/src/pool_discovery/connection.rs b/lib/kv-router/src/pool_discovery/connection.rs new file mode 100644 index 000000000000..d62e0c3c07d3 --- /dev/null +++ b/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), 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, 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, + dialer: Arc, + pool: PoolKey, +} + +impl StateBroadcaster { + pub fn new( + replicas: Arc, + dialer: Arc, + 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); +} diff --git a/lib/kv-router/src/pool_discovery/identity.rs b/lib/kv-router/src/pool_discovery/identity.rs new file mode 100644 index 000000000000..0d7b6fdb5289 --- /dev/null +++ b/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; +} + +pub async fn verify( + authenticator: &dyn RelayAuthenticator, + credentials: &PeerCredentials, +) -> Result { + authenticator + .authenticate(credentials) + .await + .map(VerifiedRelayIdentity) +} diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs new file mode 100644 index 000000000000..53b74e00be55 --- /dev/null +++ b/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: +//! . +//! +//! 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), +} diff --git a/lib/kv-router/src/pool_discovery/mode.rs b/lib/kv-router/src/pool_discovery/mode.rs new file mode 100644 index 000000000000..02e31c79989a --- /dev/null +++ b/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, + pub replicas: Arc, + /// `None` in Kubernetes and File modes. + pub registrar: Option>, + /// Incoming Relays go through `verify` with this, then + /// [`super::admit`] with `pools`. + pub authenticator: Arc, +} + +#[async_trait] +pub trait BuildPoolDiscovery { + async fn build(self) -> Result; +} diff --git a/lib/kv-router/src/pool_discovery/record.rs b/lib/kv-router/src/pool_discovery/record.rs new file mode 100644 index 000000000000..f29e14a2cb1f --- /dev/null +++ b/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, + /// Dynamo endpoint path of the pool's frontend, for example + /// `frontend.generate`. + #[builder(default)] + pub frontend_endpoint: Option, + #[builder(default)] + pub model: Option, + #[builder(default)] + pub capabilities: Vec, + /// URLs of the pool's local frontends, which a Global Router sends + /// requests to. + #[builder(default)] + pub frontends: Vec, + #[builder(default)] + pub capacity: Vec, + /// The only identity allowed to register this pool and stream its state. + #[builder(default)] + pub relay_identity: Option, + #[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; +pub type PoolEvent = WatchEvent; diff --git a/lib/kv-router/src/pool_discovery/registrar.rs b/lib/kv-router/src/pool_discovery/registrar.rs new file mode 100644 index 000000000000..d534a10d143d --- /dev/null +++ b/lib/kv-router/src/pool_discovery/registrar.rs @@ -0,0 +1,47 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::time::Duration; + +use async_trait::async_trait; + +use super::PoolDiscoveryError; +use super::identity::VerifiedRelayIdentity; +use super::record::PoolRecord; +use crate::global_view::PoolKey; + +/// A PoolRelay registers again every `heartbeat_interval`. A record expires +/// `lease_duration` after the last registration the store saw, measured on +/// the store's clock. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LeasePolicy { + pub heartbeat_interval: Duration, + pub lease_duration: Duration, +} + +/// Reply to `register`. It does not list router replicas: Relays get those +/// from a [`super::ReplicaDirectory`], served as `GET /routers`. +#[derive(Clone, Debug)] +pub struct Registration { + pub lease: LeasePolicy, +} + +/// Write side behind `RegisterPool` and `UnregisterPool`. A PoolRelay calls +/// `register` on a timer, and each call is also the pool's heartbeat: +/// registering an unchanged record again only renews it. gRPC mode keeps +/// records in memory. Hybrid writes them to the hub's Kubernetes API. +/// Kubernetes and File modes have no registrar. +#[async_trait] +pub trait PoolRegistrar: Send + Sync { + async fn register( + &self, + record: PoolRecord, + caller: &VerifiedRelayIdentity, + ) -> Result; + + async fn unregister( + &self, + key: &PoolKey, + caller: &VerifiedRelayIdentity, + ) -> Result<(), PoolDiscoveryError>; +} diff --git a/lib/kv-router/src/pool_discovery/replicas.rs b/lib/kv-router/src/pool_discovery/replicas.rs new file mode 100644 index 000000000000..3f5d87dfc4c4 --- /dev/null +++ b/lib/kv-router/src/pool_discovery/replicas.rs @@ -0,0 +1,35 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Router replicas a pool must connect to. Every replica needs every pool's +//! state, so a Relay opens one connection per replica. +//! +//! The operator deploys the Global Router, so each replica registers itself +//! with the hub's Dynamo runtime discovery, like a worker. Relays get the +//! list from the Global Discovery Service, as `GET /routers`. +//! +//! Runtime discovery lists only Ready pods, and Relays dial only listed +//! replicas. So a replica's readiness must not wait for its pools, or no +//! pool can ever reach it. Pool completeness is a separate signal +//! ([`super::Completeness`]). +//! +//! The Relay reads a `ReplicaDirectory` in every mode, filled from +//! `GET /routers`. In Kubernetes mode the service's own directory reads the +//! router's EndpointSlices that MCS imports. + +use url::Url; + +use super::watch::{Directory, WatchEvent}; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct ReplicaId(pub String); + +/// A replica's address as seen from workload clusters. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ReplicaEndpoint { + pub replica: ReplicaId, + pub url: Url, +} + +pub type ReplicaDirectory = dyn Directory; +pub type ReplicaEvent = WatchEvent; diff --git a/lib/kv-router/src/pool_discovery/watch.rs b/lib/kv-router/src/pool_discovery/watch.rs new file mode 100644 index 000000000000..c63591ea71bf --- /dev/null +++ b/lib/kv-router/src/pool_discovery/watch.rs @@ -0,0 +1,59 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The list-and-watch contract of `dynamo_runtime::discovery::Discovery`, +//! made generic so pools and router replicas share it. +//! +//! A watch starts with one `Added` per current item, then one `Resync` with the +//! same set, then changes. A `Resync` replaces all earlier state. + +use std::pin::Pin; + +use async_trait::async_trait; +use futures_util::Stream; +use tokio_util::sync::CancellationToken; + +use super::PoolDiscoveryError; + +#[derive(Clone, Debug)] +pub enum WatchEvent { + /// New or changed item. Replaces any earlier item with the same key. + Added(T), + Removed(K), + Resync(Snapshot), +} + +#[derive(Clone, Debug)] +pub struct Snapshot { + pub items: Vec, + pub completeness: Completeness, +} + +/// Lets a new router replica tell an exact "ready" from a best guess. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Completeness { + /// Read from a shared store: every existing item is listed. + Exact, + /// Built from what this replica has heard so far (gRPC mode). + Estimated, +} + +pub type WatchStream = + Pin, PoolDiscoveryError>> + Send>>; + +/// Read side. Method names follow `dynamo_runtime::discovery::Discovery`. +#[async_trait] +pub trait Directory: Send + Sync { + async fn check_connection(&self) -> Result<(), PoolDiscoveryError>; + + async fn list(&self) -> Result, PoolDiscoveryError>; + + /// One item, from the watch cache. Admission calls this per connection, + /// so it must not list. + async fn get(&self, key: &K) -> Result, PoolDiscoveryError>; + + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, PoolDiscoveryError>; +}