From 49dcb5060421749e7558b97f27009137a87ba8a2 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Thu, 8 Oct 2026 15:12:13 -0700 Subject: [PATCH 01/12] Initial Interfaces for the Discovery Abstractions Signed-off-by: Anna Tchernych --- Cargo.lock | 1 + lib/kv-router/Cargo.toml | 1 + lib/kv-router/src/lib.rs | 1 + lib/kv-router/src/pool_discovery/admission.rs | 21 ++++++ lib/kv-router/src/pool_discovery/announce.rs | 28 ++++++++ lib/kv-router/src/pool_discovery/identity.rs | 33 ++++++++++ lib/kv-router/src/pool_discovery/mod.rs | 53 +++++++++++++++ lib/kv-router/src/pool_discovery/mode.rs | 64 +++++++++++++++++++ lib/kv-router/src/pool_discovery/record.rs | 34 ++++++++++ lib/kv-router/src/pool_discovery/registrar.rs | 52 +++++++++++++++ lib/kv-router/src/pool_discovery/replicas.rs | 26 ++++++++ lib/kv-router/src/pool_discovery/watch.rs | 55 ++++++++++++++++ .../src/kv_dc_relay/global_view_consumer.rs | 1 + .../global_view_consumer/inbound.rs | 37 +++++++++++ lib/llm/src/kv_dc_relay/wan.rs | 1 + lib/llm/src/kv_dc_relay/wan/outbound.rs | 40 ++++++++++++ 16 files changed, 448 insertions(+) create mode 100644 lib/kv-router/src/pool_discovery/admission.rs create mode 100644 lib/kv-router/src/pool_discovery/announce.rs create mode 100644 lib/kv-router/src/pool_discovery/identity.rs create mode 100644 lib/kv-router/src/pool_discovery/mod.rs create mode 100644 lib/kv-router/src/pool_discovery/mode.rs create mode 100644 lib/kv-router/src/pool_discovery/record.rs create mode 100644 lib/kv-router/src/pool_discovery/registrar.rs create mode 100644 lib/kv-router/src/pool_discovery/replicas.rs create mode 100644 lib/kv-router/src/pool_discovery/watch.rs create mode 100644 lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs create mode 100644 lib/llm/src/kv_dc_relay/wan/outbound.rs 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..7ac62ece1855 --- /dev/null +++ b/lib/kv-router/src/pool_discovery/admission.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use async_trait::async_trait; + +use super::PoolDiscoveryError; +use super::identity::VerifiedRelayIdentity; +use super::record::PoolRecord; +use crate::global_view::PoolKey; + +/// Decides whether an authenticated Relay may stream state for a pool. +/// Usually a `PoolDirectory` lookup plus a `relay_identity` check. +/// The runtime keys the stream by its `PlaneLease`, not by the claimed key. +#[async_trait] +pub trait PoolAdmission: Send + Sync { + async fn admit( + &self, + key: &PoolKey, + caller: &VerifiedRelayIdentity, + ) -> Result; +} diff --git a/lib/kv-router/src/pool_discovery/announce.rs b/lib/kv-router/src/pool_discovery/announce.rs new file mode 100644 index 000000000000..46e49a964608 --- /dev/null +++ b/lib/kv-router/src/pool_discovery/announce.rs @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Pool side, running next to the Relay in a workload cluster. + +use async_trait::async_trait; +use tokio::sync::watch; + +use super::PoolDiscoveryError; +use super::record::PoolRecord; +use super::replicas::ReplicaEndpoint; + +/// gRPC and Hybrid modes call `RegisterPool` on the regional router address. +/// Kubernetes mode writes the record to the hub and renews its Lease. +#[async_trait] +pub trait PoolAnnouncer: Send + Sync { + async fn announce( + &self, + record: PoolRecord, + ) -> Result, PoolDiscoveryError>; +} + +/// RAII guard. Keeps the pool registered while alive; dropping it stops +/// heartbeats and unregisters the pool. +pub trait Announcement: Send + Sync { + /// Router replicas this pool must open state connections to. + fn replicas(&self) -> watch::Receiver>; +} 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..ee3041d55161 --- /dev/null +++ b/lib/kv-router/src/pool_discovery/identity.rs @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +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 authenticators in +/// this module tree can create one, so writes cannot use a claimed identity. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct VerifiedRelayIdentity(pub(super) RelayIdentity); + +/// 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), +} + +#[async_trait] +pub trait RelayAuthenticator: Send + Sync { + async fn authenticate( + &self, + credentials: &PeerCredentials, + ) -> Result; +} 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..ad9f1f721d1c --- /dev/null +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -0,0 +1,53 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Pluggable pool discovery for the Global Router (DEP-1276). +//! +//! 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`]. +//! +//! | Mode | Pool list stored in | Pools call | +//! |---|---|---| +//! | gRPC | Router memory | `RegisterPool` on the router | +//! | Hybrid | Hub Kubernetes API, written by the router | `RegisterPool` on the router | +//! | Kubernetes | Hub Kubernetes API, written by the pool | Hub Kubernetes API | +//! | File | Static file (tests, first deployments) | Nothing | + +pub mod admission; +pub mod announce; +pub mod identity; +pub mod mode; +pub mod record; +pub mod registrar; +pub mod replicas; +pub mod watch; + +pub use admission::PoolAdmission; +pub use announce::{Announcement, PoolAnnouncer}; +pub use identity::{PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity}; +pub use mode::{BuildPoolDiscovery, PoolDiscovery, PoolDiscoveryMode, SettlePolicy}; +pub use record::{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..fbc532f3d76d --- /dev/null +++ b/lib/kv-router/src/pool_discovery/mode.rs @@ -0,0 +1,64 @@ +// 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 std::time::Duration; + +use async_trait::async_trait; + +use super::PoolDiscoveryError; +use super::admission::PoolAdmission; +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`. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub enum PoolDiscoveryMode { + Grpc { + lease: LeasePolicy, + settle: SettlePolicy, + }, + Hybrid { + lease: LeasePolicy, + hub_namespace: String, + }, + Kubernetes { + lease_duration: Duration, + pool_namespace_selector: String, + }, + File { + path: PathBuf, + }, +} + +/// When a gRPC-mode replica treats its pool list as settled (soft readiness). +#[derive(Clone, Copy, Debug)] +pub enum SettlePolicy { + FixedDelay(Duration), + /// No new pool for this long. + QuietPeriod(Duration), + /// Copy a peer replica's list; fall back to a delay if none answers. + PeerSnapshot { + fallback: Duration, + }, +} + +/// 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>, + pub admission: Arc, + 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..8c15db28f9e2 --- /dev/null +++ b/lib/kv-router/src/pool_discovery/record.rs @@ -0,0 +1,34 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use derive_builder::Builder; + +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, + /// With `frontend_endpoint`, becomes `RelayPoolScope` in dynamo-llm. + pub runtime_namespace: String, + pub frontend_endpoint: String, + pub model: String, + /// The only identity allowed to register this pool and stream its state. + pub relay_identity: RelayIdentity, + #[builder(default)] + pub revision: Revision, +} + +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..327ab9902475 --- /dev/null +++ b/lib/kv-router/src/pool_discovery/registrar.rs @@ -0,0 +1,52 @@ +// 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 super::replicas::ReplicaEndpoint; +use crate::global_view::PoolKey; + +/// A record expires `lease_duration` after the last heartbeat this router saw, +/// measured on the router's clock. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LeasePolicy { + pub heartbeat_interval: Duration, + pub lease_duration: Duration, +} + +/// Reply to `register` and `heartbeat`. +#[derive(Clone, Debug)] +pub struct Registration { + pub lease: LeasePolicy, + /// Every router replica. The Relay opens a state connection to each one. + pub replicas: Vec, +} + +/// Write side behind `RegisterPool`, `Heartbeat` and `UnregisterPool`. +/// 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 heartbeat( + &self, + key: &PoolKey, + 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..226431cf3357 --- /dev/null +++ b/lib/kv-router/src/pool_discovery/replicas.rs @@ -0,0 +1,26 @@ +// 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. +//! +//! Replicas register themselves, so a `ReplicaDirectory` can be backed by +//! `dynamo_runtime::discovery::Discovery`. A replica must be listed before it +//! is ready, or pools can never reach it. + +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..8d27341e2dad --- /dev/null +++ b/lib/kv-router/src/pool_discovery/watch.rs @@ -0,0 +1,55 @@ +// 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, settled by a `SettlePolicy`. + 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>; + + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, PoolDiscoveryError>; +} diff --git a/lib/llm/src/kv_dc_relay/global_view_consumer.rs b/lib/llm/src/kv_dc_relay/global_view_consumer.rs index 9f149c525835..dbe705bf55e2 100644 --- a/lib/llm/src/kv_dc_relay/global_view_consumer.rs +++ b/lib/llm/src/kv_dc_relay/global_view_consumer.rs @@ -5,6 +5,7 @@ //! Catalog, CKF, stats, and reconnect lifecycle are owned by this module. mod coordinator; +pub mod inbound; mod runtime; mod scorer; pub mod stats; diff --git a/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs b/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs new file mode 100644 index 000000000000..2773372b2e0f --- /dev/null +++ b/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Dial-in replacement for [`super::RelayDgdSource`] (DEP-1276). +//! +//! Pool details come from discovery. The Relay-opened connection replaces the +//! Relay and stats channels the router used to dial. + +use async_trait::async_trait; +use dynamo_kv_router::pool_discovery::{PoolDiscoveryError, PoolRecord, VerifiedRelayIdentity}; +use tokio_util::sync::CancellationToken; + +pub struct InboundPoolSource { + pub record: PoolRecord, + pub connection: Box, +} + +/// A connection a Relay opened to this router replica, already authenticated +/// and admitted. The state protocol on it is outside DEP-1276. +#[async_trait] +pub trait RelayConnection: Send + Sync { + fn caller(&self) -> &VerifiedRelayIdentity; + + /// Resolves when the Relay disconnects, so the runtime can release the + /// pool's `PlaneLease`s. + async fn closed(&self); +} + +/// The router's gRPC server side. Authenticates with `RelayAuthenticator` and +/// checks with `PoolAdmission` before yielding a source. +#[async_trait] +pub trait RelayListener: Send + Sync { + async fn accept( + &self, + cancel_token: CancellationToken, + ) -> Result; +} diff --git a/lib/llm/src/kv_dc_relay/wan.rs b/lib/llm/src/kv_dc_relay/wan.rs index 4d0d5c67ca13..6a6e47bf5ed1 100644 --- a/lib/llm/src/kv_dc_relay/wan.rs +++ b/lib/llm/src/kv_dc_relay/wan.rs @@ -4,3 +4,4 @@ //! WAN adapters over the transport-neutral Relay publication source. pub mod grpc; +pub mod outbound; diff --git a/lib/llm/src/kv_dc_relay/wan/outbound.rs b/lib/llm/src/kv_dc_relay/wan/outbound.rs new file mode 100644 index 000000000000..b06493d82d57 --- /dev/null +++ b/lib/llm/src/kv_dc_relay/wan/outbound.rs @@ -0,0 +1,40 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Dial-in replacement for the Relay's gRPC server (DEP-1276). +//! +//! The Relay opens one connection per pool to each router replica listed by +//! its `Announcement`, and sends the same publications the server sends +//! today. Push RPC vs reverse tunnel is outside DEP-1276. + +use std::sync::Arc; + +use async_trait::async_trait; +use dynamo_kv_router::global_view::PoolKey; +use dynamo_kv_router::pool_discovery::{PoolDiscoveryError, ReplicaEndpoint, ReplicaId}; +use tokio_util::sync::CancellationToken; + +use crate::kv_dc_relay::RelayPublicationSource; + +/// Counterpart of `global_view_consumer::inbound::RelayListener`. +#[async_trait] +pub trait RelayDialer: Send + Sync { + async fn dial( + &self, + replica: &ReplicaEndpoint, + pool: &PoolKey, + publication: Arc, + cancel_token: CancellationToken, + ) -> Result, PoolDiscoveryError>; +} + +/// The router issues `PlaneLease`s when it accepts. A new Relay incarnation +/// must dial again. +#[async_trait] +pub trait OutboundConnection: Send + Sync { + fn replica(&self) -> &ReplicaId; + + /// Resolves when the router replica disconnects. The Relay redials while + /// the replica is still listed. + async fn closed(&self); +} From 84f0470d648971ce5593e4e31d51992e4ed7c3df Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Thu, 8 Oct 2026 17:09:34 -0700 Subject: [PATCH 02/12] fix(router): adjust pool discovery traits found by the Kubernetes sketch - RelayAuthenticator returns a plain RelayIdentity; verify() wraps it in VerifiedRelayIdentity, so authenticators in other crates can implement the trait. - Directory gains get(&K), so admission does not list per connection. - Announcement gains async withdraw(); Drop cannot run async cleanup. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/admission.rs | 2 +- lib/kv-router/src/pool_discovery/announce.rs | 8 ++++++-- lib/kv-router/src/pool_discovery/identity.rs | 18 +++++++++++++++--- lib/kv-router/src/pool_discovery/mod.rs | 4 +++- lib/kv-router/src/pool_discovery/watch.rs | 4 ++++ .../global_view_consumer/inbound.rs | 2 +- 6 files changed, 30 insertions(+), 8 deletions(-) diff --git a/lib/kv-router/src/pool_discovery/admission.rs b/lib/kv-router/src/pool_discovery/admission.rs index 7ac62ece1855..0fffe0715d1f 100644 --- a/lib/kv-router/src/pool_discovery/admission.rs +++ b/lib/kv-router/src/pool_discovery/admission.rs @@ -9,7 +9,7 @@ use super::record::PoolRecord; use crate::global_view::PoolKey; /// Decides whether an authenticated Relay may stream state for a pool. -/// Usually a `PoolDirectory` lookup plus a `relay_identity` check. +/// Usually `PoolDirectory::get` plus a `relay_identity` check. /// The runtime keys the stream by its `PlaneLease`, not by the claimed key. #[async_trait] pub trait PoolAdmission: Send + Sync { diff --git a/lib/kv-router/src/pool_discovery/announce.rs b/lib/kv-router/src/pool_discovery/announce.rs index 46e49a964608..67d023d03e4b 100644 --- a/lib/kv-router/src/pool_discovery/announce.rs +++ b/lib/kv-router/src/pool_discovery/announce.rs @@ -20,9 +20,13 @@ pub trait PoolAnnouncer: Send + Sync { ) -> Result, PoolDiscoveryError>; } -/// RAII guard. Keeps the pool registered while alive; dropping it stops -/// heartbeats and unregisters the pool. +/// RAII guard. Keeps the pool registered while alive. Call `withdraw` on a +/// clean shutdown; dropping only stops heartbeats, and the lease expires. +#[async_trait] pub trait Announcement: Send + Sync { /// Router replicas this pool must open state connections to. fn replicas(&self) -> watch::Receiver>; + + /// Unregisters the pool and waits for the backend to confirm. + async fn withdraw(self: Box) -> Result<(), PoolDiscoveryError>; } diff --git a/lib/kv-router/src/pool_discovery/identity.rs b/lib/kv-router/src/pool_discovery/identity.rs index ee3041d55161..2a6ae6866924 100644 --- a/lib/kv-router/src/pool_discovery/identity.rs +++ b/lib/kv-router/src/pool_discovery/identity.rs @@ -9,8 +9,8 @@ use super::PoolDiscoveryError; #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct RelayIdentity(pub String); -/// An identity proven by a [`RelayAuthenticator`]. Only authenticators in -/// this module tree can create one, so writes cannot use a claimed identity. +/// 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); @@ -24,10 +24,22 @@ pub enum PeerCredentials { BearerToken(String), } +/// 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; + ) -> 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 index ad9f1f721d1c..6f019428e9ca 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -26,7 +26,9 @@ pub mod watch; pub use admission::PoolAdmission; pub use announce::{Announcement, PoolAnnouncer}; -pub use identity::{PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity}; +pub use identity::{ + PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, +}; pub use mode::{BuildPoolDiscovery, PoolDiscovery, PoolDiscoveryMode, SettlePolicy}; pub use record::{PoolDirectory, PoolEvent, PoolRecord, PoolRecordBuilder, Revision}; pub use registrar::{LeasePolicy, PoolRegistrar, Registration}; diff --git a/lib/kv-router/src/pool_discovery/watch.rs b/lib/kv-router/src/pool_discovery/watch.rs index 8d27341e2dad..469407eccab3 100644 --- a/lib/kv-router/src/pool_discovery/watch.rs +++ b/lib/kv-router/src/pool_discovery/watch.rs @@ -48,6 +48,10 @@ pub trait Directory: Send + Sync { 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, diff --git a/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs b/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs index 2773372b2e0f..c60281ceda39 100644 --- a/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs +++ b/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs @@ -26,7 +26,7 @@ pub trait RelayConnection: Send + Sync { async fn closed(&self); } -/// The router's gRPC server side. Authenticates with `RelayAuthenticator` and +/// The router's gRPC server side. Authenticates through `pool_discovery::verify` and /// checks with `PoolAdmission` before yielding a source. #[async_trait] pub trait RelayListener: Send + Sync { From ca4c9ee0e7f1a2048b949150bb470039c71e58ac Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Fri, 9 Oct 2026 13:43:04 -0700 Subject: [PATCH 03/12] fix(router): fit pool discovery traits to the SIG Multicluster mode Kubernetes mode now follows SIG Multicluster (MCS) on a trusted network with reachable pod addresses. Discovery there only knows a pool's name and whether its PoolRelay endpoint is Ready, so: - PoolRecord: runtime_namespace, frontend_endpoint, model and relay_identity are Option, filled later from the Relay's catalog. - PoolDiscoveryMode::Kubernetes drops the lease and takes the MCS implementation (Standard or Karmada), a PoolRelay Service name suffix, and the router Service. A name rule replaces exportedLabels, which not every implementation supports. - PeerCredentials::PeerAddress names a caller by its source address, for trusted networks without a mesh. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/announce.rs | 2 +- lib/kv-router/src/pool_discovery/identity.rs | 6 +++++ lib/kv-router/src/pool_discovery/mod.rs | 6 +++-- lib/kv-router/src/pool_discovery/mode.rs | 23 ++++++++++++++++++-- lib/kv-router/src/pool_discovery/record.rs | 14 ++++++++---- 5 files changed, 42 insertions(+), 9 deletions(-) diff --git a/lib/kv-router/src/pool_discovery/announce.rs b/lib/kv-router/src/pool_discovery/announce.rs index 67d023d03e4b..a0ea30d86fd1 100644 --- a/lib/kv-router/src/pool_discovery/announce.rs +++ b/lib/kv-router/src/pool_discovery/announce.rs @@ -11,7 +11,7 @@ use super::record::PoolRecord; use super::replicas::ReplicaEndpoint; /// gRPC and Hybrid modes call `RegisterPool` on the regional router address. -/// Kubernetes mode writes the record to the hub and renews its Lease. +/// Kubernetes mode has no announcer: the operator creates a `ServiceExport`. #[async_trait] pub trait PoolAnnouncer: Send + Sync { async fn announce( diff --git a/lib/kv-router/src/pool_discovery/identity.rs b/lib/kv-router/src/pool_discovery/identity.rs index 2a6ae6866924..d4bfd846c79d 100644 --- a/lib/kv-router/src/pool_discovery/identity.rs +++ b/lib/kv-router/src/pool_discovery/identity.rs @@ -1,6 +1,8 @@ // 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; @@ -22,6 +24,10 @@ 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 diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index 6f019428e9ca..fb162341ce67 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -12,7 +12,7 @@ //! |---|---|---| //! | gRPC | Router memory | `RegisterPool` on the router | //! | Hybrid | Hub Kubernetes API, written by the router | `RegisterPool` on the router | -//! | Kubernetes | Hub Kubernetes API, written by the pool | Hub Kubernetes API | +//! | 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; @@ -29,7 +29,9 @@ pub use announce::{Announcement, PoolAnnouncer}; pub use identity::{ PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, }; -pub use mode::{BuildPoolDiscovery, PoolDiscovery, PoolDiscoveryMode, SettlePolicy}; +pub use mode::{ + BuildPoolDiscovery, McsImplementation, PoolDiscovery, PoolDiscoveryMode, SettlePolicy, +}; pub use record::{PoolDirectory, PoolEvent, PoolRecord, PoolRecordBuilder, Revision}; pub use registrar::{LeasePolicy, PoolRegistrar, Registration}; pub use replicas::{ReplicaDirectory, ReplicaEndpoint, ReplicaEvent, ReplicaId}; diff --git a/lib/kv-router/src/pool_discovery/mode.rs b/lib/kv-router/src/pool_discovery/mode.rs index fbc532f3d76d..81bb8a347707 100644 --- a/lib/kv-router/src/pool_discovery/mode.rs +++ b/lib/kv-router/src/pool_discovery/mode.rs @@ -26,15 +26,34 @@ pub enum PoolDiscoveryMode { lease: LeasePolicy, hub_namespace: String, }, + /// SIG Multicluster: pools are MCS `ServiceExport`s. Liveness is endpoint + /// readiness, so there is no lease. Assumes a trusted network with pod + /// addresses that are reachable between clusters. Kubernetes { - lease_duration: Duration, - pool_namespace_selector: String, + mcs: McsImplementation, + /// The operator names each PoolRelay Service ``. + /// A name rule works with every MCS implementation; `exportedLabels` + /// does not. + pool_service_suffix: String, + /// The Global Router's Service, exported to every workload cluster. + router_service: String, }, File { path: PathBuf, }, } +/// How the MCS implementation names imported EndpointSlices. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum McsImplementation { + /// The `multicluster.kubernetes.io` labels (Submariner, Cilium, GKE). + Standard, + /// Karmada's own names: a `derived-` Service, and the source + /// cluster in the `work.karmada.io/namespace` annotation. + Karmada, +} + /// When a gRPC-mode replica treats its pool list as settled (soft readiness). #[derive(Clone, Copy, Debug)] pub enum SettlePolicy { diff --git a/lib/kv-router/src/pool_discovery/record.rs b/lib/kv-router/src/pool_discovery/record.rs index 8c15db28f9e2..4fdad332e8aa 100644 --- a/lib/kv-router/src/pool_discovery/record.rs +++ b/lib/kv-router/src/pool_discovery/record.rs @@ -20,12 +20,18 @@ pub struct Revision(pub u64); pub struct PoolRecord { pub key: PoolKey, pub location: PoolLocation, + // The fields below are `None` until the Relay's catalog reports them, in + // modes where discovery only knows the pool's name (Kubernetes / MCS). /// With `frontend_endpoint`, becomes `RelayPoolScope` in dynamo-llm. - pub runtime_namespace: String, - pub frontend_endpoint: String, - pub model: String, + #[builder(default)] + pub runtime_namespace: Option, + #[builder(default)] + pub frontend_endpoint: Option, + #[builder(default)] + pub model: Option, /// The only identity allowed to register this pool and stream its state. - pub relay_identity: RelayIdentity, + #[builder(default)] + pub relay_identity: Option, #[builder(default)] pub revision: Revision, } From 79b29e283a5e920316309cbb5b82ef76b02ed720 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 11 Oct 2026 15:41:19 -0700 Subject: [PATCH 04/12] refactor(router): keep only the pool discovery traits that differ by mode Admission, the Relay's RegisterPool client, and the Relay connections work the same in every discovery mode, so they need no traits: - PoolAdmission becomes the function `admit`, with unit tests. - PoolAnnouncer and Announcement are removed. The Relay reads a ReplicaDirectory in every mode; gRPC and Hybrid fill it from the RegisterPool and Heartbeat replies. - RelayListener and RelayDialer move out to the transport work, which still has to choose between a push RPC and a reverse tunnel. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/admission.rs | 135 ++++++++++++++++-- lib/kv-router/src/pool_discovery/announce.rs | 32 ----- lib/kv-router/src/pool_discovery/mod.rs | 9 +- lib/kv-router/src/pool_discovery/mode.rs | 4 +- lib/kv-router/src/pool_discovery/replicas.rs | 5 + .../src/kv_dc_relay/global_view_consumer.rs | 1 - .../global_view_consumer/inbound.rs | 37 ----- lib/llm/src/kv_dc_relay/wan.rs | 1 - lib/llm/src/kv_dc_relay/wan/outbound.rs | 40 ------ 9 files changed, 136 insertions(+), 128 deletions(-) delete mode 100644 lib/kv-router/src/pool_discovery/announce.rs delete mode 100644 lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs delete mode 100644 lib/llm/src/kv_dc_relay/wan/outbound.rs diff --git a/lib/kv-router/src/pool_discovery/admission.rs b/lib/kv-router/src/pool_discovery/admission.rs index 0fffe0715d1f..b651404ddbca 100644 --- a/lib/kv-router/src/pool_discovery/admission.rs +++ b/lib/kv-router/src/pool_discovery/admission.rs @@ -1,21 +1,132 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use async_trait::async_trait; - use super::PoolDiscoveryError; use super::identity::VerifiedRelayIdentity; -use super::record::PoolRecord; +use super::record::{PoolDirectory, PoolRecord}; use crate::global_view::PoolKey; -/// Decides whether an authenticated Relay may stream state for a pool. -/// Usually `PoolDirectory::get` plus a `relay_identity` check. +/// 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. -#[async_trait] -pub trait PoolAdmission: Send + Sync { - async fn admit( - &self, - key: &PoolKey, - caller: &VerifiedRelayIdentity, - ) -> Result; +pub async fn admit( + pools: &PoolDirectory, + key: &PoolKey, + caller: &VerifiedRelayIdentity, +) -> Result { + let record = pools + .get(key) + .await? + .ok_or_else(|| PoolDiscoveryError::UnknownPool(key.clone()))?; + match &record.relay_identity { + Some(expected) if *expected != caller.0 => { + Err(PoolDiscoveryError::PermissionDenied(key.clone())) + } + _ => Ok(record), + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use async_trait::async_trait; + use tokio_util::sync::CancellationToken; + + use super::*; + use crate::global_view::state::PoolLocation; + use crate::pool_discovery::identity::RelayIdentity; + use crate::pool_discovery::record::PoolRecordBuilder; + use crate::pool_discovery::watch::{Directory, Snapshot, WatchStream}; + + struct Fixed(HashMap); + + #[async_trait] + impl Directory for Fixed { + async fn check_connection(&self) -> Result<(), PoolDiscoveryError> { + unimplemented!() + } + + async fn list(&self) -> Result, PoolDiscoveryError> { + unimplemented!() + } + + async fn get(&self, key: &PoolKey) -> Result, PoolDiscoveryError> { + Ok(self.0.get(key).cloned()) + } + + async fn list_and_watch( + &self, + _cancel_token: Option, + ) -> Result, PoolDiscoveryError> { + unimplemented!() + } + } + + fn key(dgd: &str) -> PoolKey { + PoolKey::new("site-a", "ns", dgd).unwrap() + } + + fn record(dgd: &str, relay: Option<&str>) -> PoolRecord { + PoolRecordBuilder::default() + .key(key(dgd)) + .location(PoolLocation { + region: "r1".into(), + availability_zone: None, + cluster: None, + datacenter: None, + }) + .relay_identity(relay.map(|id| RelayIdentity(id.into()))) + .build() + .unwrap() + } + + fn caller(id: &str) -> VerifiedRelayIdentity { + VerifiedRelayIdentity(RelayIdentity(id.into())) + } + + fn pools() -> Fixed { + Fixed( + [record("bound", Some("relay-a")), record("open", None)] + .into_iter() + .map(|r| (r.key.clone(), r)) + .collect(), + ) + } + + #[tokio::test] + async fn admits_the_bound_relay() { + let record = admit(&pools(), &key("bound"), &caller("relay-a")) + .await + .unwrap(); + assert_eq!(record.key, key("bound")); + } + + #[tokio::test] + async fn rejects_another_relay() { + let err = admit(&pools(), &key("bound"), &caller("relay-b")) + .await + .unwrap_err(); + assert!(matches!(err, PoolDiscoveryError::PermissionDenied(k) if k == key("bound"))); + } + + #[tokio::test] + async fn admits_any_relay_without_a_bound_identity() { + assert!( + admit(&pools(), &key("open"), &caller("relay-b")) + .await + .is_ok() + ); + } + + #[tokio::test] + async fn rejects_an_unknown_pool() { + let err = admit(&pools(), &key("missing"), &caller("relay-a")) + .await + .unwrap_err(); + assert!(matches!(err, PoolDiscoveryError::UnknownPool(k) if k == key("missing"))); + } } diff --git a/lib/kv-router/src/pool_discovery/announce.rs b/lib/kv-router/src/pool_discovery/announce.rs deleted file mode 100644 index a0ea30d86fd1..000000000000 --- a/lib/kv-router/src/pool_discovery/announce.rs +++ /dev/null @@ -1,32 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Pool side, running next to the Relay in a workload cluster. - -use async_trait::async_trait; -use tokio::sync::watch; - -use super::PoolDiscoveryError; -use super::record::PoolRecord; -use super::replicas::ReplicaEndpoint; - -/// gRPC and Hybrid modes call `RegisterPool` on the regional router address. -/// Kubernetes mode has no announcer: the operator creates a `ServiceExport`. -#[async_trait] -pub trait PoolAnnouncer: Send + Sync { - async fn announce( - &self, - record: PoolRecord, - ) -> Result, PoolDiscoveryError>; -} - -/// RAII guard. Keeps the pool registered while alive. Call `withdraw` on a -/// clean shutdown; dropping only stops heartbeats, and the lease expires. -#[async_trait] -pub trait Announcement: Send + Sync { - /// Router replicas this pool must open state connections to. - fn replicas(&self) -> watch::Receiver>; - - /// Unregisters the pool and waits for the backend to confirm. - async fn withdraw(self: Box) -> Result<(), PoolDiscoveryError>; -} diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index fb162341ce67..45d38b7e4efd 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -8,6 +8,11 @@ //! 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, the Relay's `RegisterPool` client, and the connections that +//! carry pool state work the same in every mode, so they are concrete code. +//! //! | Mode | Pool list stored in | Pools call | //! |---|---|---| //! | gRPC | Router memory | `RegisterPool` on the router | @@ -16,7 +21,6 @@ //! | File | Static file (tests, first deployments) | Nothing | pub mod admission; -pub mod announce; pub mod identity; pub mod mode; pub mod record; @@ -24,8 +28,7 @@ pub mod registrar; pub mod replicas; pub mod watch; -pub use admission::PoolAdmission; -pub use announce::{Announcement, PoolAnnouncer}; +pub use admission::admit; pub use identity::{ PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, }; diff --git a/lib/kv-router/src/pool_discovery/mode.rs b/lib/kv-router/src/pool_discovery/mode.rs index 81bb8a347707..2362d8a5bee5 100644 --- a/lib/kv-router/src/pool_discovery/mode.rs +++ b/lib/kv-router/src/pool_discovery/mode.rs @@ -8,7 +8,6 @@ use std::time::Duration; use async_trait::async_trait; use super::PoolDiscoveryError; -use super::admission::PoolAdmission; use super::identity::RelayAuthenticator; use super::record::PoolDirectory; use super::registrar::{LeasePolicy, PoolRegistrar}; @@ -73,7 +72,8 @@ pub struct PoolDiscovery { pub replicas: Arc, /// `None` in Kubernetes and File modes. pub registrar: Option>, - pub admission: Arc, + /// Incoming Relays go through `verify` with this, then + /// [`super::admit`] with `pools`. pub authenticator: Arc, } diff --git a/lib/kv-router/src/pool_discovery/replicas.rs b/lib/kv-router/src/pool_discovery/replicas.rs index 226431cf3357..c28dd5073286 100644 --- a/lib/kv-router/src/pool_discovery/replicas.rs +++ b/lib/kv-router/src/pool_discovery/replicas.rs @@ -7,6 +7,11 @@ //! Replicas register themselves, so a `ReplicaDirectory` can be backed by //! `dynamo_runtime::discovery::Discovery`. A replica must be listed before it //! is ready, or pools can never reach it. +//! +//! The Relay reads a `ReplicaDirectory` in every mode. In gRPC and Hybrid +//! modes it is filled from the `RegisterPool` and `Heartbeat` replies +//! ([`super::Registration::replicas`]). In Kubernetes mode it reads the +//! router's EndpointSlices that MCS imports into the Relay's cluster. use url::Url; diff --git a/lib/llm/src/kv_dc_relay/global_view_consumer.rs b/lib/llm/src/kv_dc_relay/global_view_consumer.rs index dbe705bf55e2..9f149c525835 100644 --- a/lib/llm/src/kv_dc_relay/global_view_consumer.rs +++ b/lib/llm/src/kv_dc_relay/global_view_consumer.rs @@ -5,7 +5,6 @@ //! Catalog, CKF, stats, and reconnect lifecycle are owned by this module. mod coordinator; -pub mod inbound; mod runtime; mod scorer; pub mod stats; diff --git a/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs b/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs deleted file mode 100644 index c60281ceda39..000000000000 --- a/lib/llm/src/kv_dc_relay/global_view_consumer/inbound.rs +++ /dev/null @@ -1,37 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Dial-in replacement for [`super::RelayDgdSource`] (DEP-1276). -//! -//! Pool details come from discovery. The Relay-opened connection replaces the -//! Relay and stats channels the router used to dial. - -use async_trait::async_trait; -use dynamo_kv_router::pool_discovery::{PoolDiscoveryError, PoolRecord, VerifiedRelayIdentity}; -use tokio_util::sync::CancellationToken; - -pub struct InboundPoolSource { - pub record: PoolRecord, - pub connection: Box, -} - -/// A connection a Relay opened to this router replica, already authenticated -/// and admitted. The state protocol on it is outside DEP-1276. -#[async_trait] -pub trait RelayConnection: Send + Sync { - fn caller(&self) -> &VerifiedRelayIdentity; - - /// Resolves when the Relay disconnects, so the runtime can release the - /// pool's `PlaneLease`s. - async fn closed(&self); -} - -/// The router's gRPC server side. Authenticates through `pool_discovery::verify` and -/// checks with `PoolAdmission` before yielding a source. -#[async_trait] -pub trait RelayListener: Send + Sync { - async fn accept( - &self, - cancel_token: CancellationToken, - ) -> Result; -} diff --git a/lib/llm/src/kv_dc_relay/wan.rs b/lib/llm/src/kv_dc_relay/wan.rs index 6a6e47bf5ed1..4d0d5c67ca13 100644 --- a/lib/llm/src/kv_dc_relay/wan.rs +++ b/lib/llm/src/kv_dc_relay/wan.rs @@ -4,4 +4,3 @@ //! WAN adapters over the transport-neutral Relay publication source. pub mod grpc; -pub mod outbound; diff --git a/lib/llm/src/kv_dc_relay/wan/outbound.rs b/lib/llm/src/kv_dc_relay/wan/outbound.rs deleted file mode 100644 index b06493d82d57..000000000000 --- a/lib/llm/src/kv_dc_relay/wan/outbound.rs +++ /dev/null @@ -1,40 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Dial-in replacement for the Relay's gRPC server (DEP-1276). -//! -//! The Relay opens one connection per pool to each router replica listed by -//! its `Announcement`, and sends the same publications the server sends -//! today. Push RPC vs reverse tunnel is outside DEP-1276. - -use std::sync::Arc; - -use async_trait::async_trait; -use dynamo_kv_router::global_view::PoolKey; -use dynamo_kv_router::pool_discovery::{PoolDiscoveryError, ReplicaEndpoint, ReplicaId}; -use tokio_util::sync::CancellationToken; - -use crate::kv_dc_relay::RelayPublicationSource; - -/// Counterpart of `global_view_consumer::inbound::RelayListener`. -#[async_trait] -pub trait RelayDialer: Send + Sync { - async fn dial( - &self, - replica: &ReplicaEndpoint, - pool: &PoolKey, - publication: Arc, - cancel_token: CancellationToken, - ) -> Result, PoolDiscoveryError>; -} - -/// The router issues `PlaneLease`s when it accepts. A new Relay incarnation -/// must dial again. -#[async_trait] -pub trait OutboundConnection: Send + Sync { - fn replica(&self) -> &ReplicaId; - - /// Resolves when the router replica disconnects. The Relay redials while - /// the replica is still listed. - async fn closed(&self); -} From cb1b722bf28745c1d4691c79569bd21326da42d4 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 11 Oct 2026 15:58:58 -0700 Subject: [PATCH 05/12] refactor(router): keep the pool discovery PR to abstractions only - `admit` keeps its signature; the body and tests come with the first mode's implementation (they are in 79b29e283a). - PoolDiscoveryMode keeps one variant per mode. Mode-only settings (McsImplementation, the Kubernetes Service names, SettlePolicy, Hybrid's hub namespace) move to each mode's PR. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/admission.rs | 115 +----------------- lib/kv-router/src/pool_discovery/mod.rs | 4 +- lib/kv-router/src/pool_discovery/mode.rs | 43 +------ lib/kv-router/src/pool_discovery/watch.rs | 2 +- 4 files changed, 11 insertions(+), 153 deletions(-) diff --git a/lib/kv-router/src/pool_discovery/admission.rs b/lib/kv-router/src/pool_discovery/admission.rs index b651404ddbca..4509f55bb382 100644 --- a/lib/kv-router/src/pool_discovery/admission.rs +++ b/lib/kv-router/src/pool_discovery/admission.rs @@ -12,121 +12,12 @@ use crate::global_view::PoolKey; /// /// 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 { - let record = pools - .get(key) - .await? - .ok_or_else(|| PoolDiscoveryError::UnknownPool(key.clone()))?; - match &record.relay_identity { - Some(expected) if *expected != caller.0 => { - Err(PoolDiscoveryError::PermissionDenied(key.clone())) - } - _ => Ok(record), - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use async_trait::async_trait; - use tokio_util::sync::CancellationToken; - - use super::*; - use crate::global_view::state::PoolLocation; - use crate::pool_discovery::identity::RelayIdentity; - use crate::pool_discovery::record::PoolRecordBuilder; - use crate::pool_discovery::watch::{Directory, Snapshot, WatchStream}; - - struct Fixed(HashMap); - - #[async_trait] - impl Directory for Fixed { - async fn check_connection(&self) -> Result<(), PoolDiscoveryError> { - unimplemented!() - } - - async fn list(&self) -> Result, PoolDiscoveryError> { - unimplemented!() - } - - async fn get(&self, key: &PoolKey) -> Result, PoolDiscoveryError> { - Ok(self.0.get(key).cloned()) - } - - async fn list_and_watch( - &self, - _cancel_token: Option, - ) -> Result, PoolDiscoveryError> { - unimplemented!() - } - } - - fn key(dgd: &str) -> PoolKey { - PoolKey::new("site-a", "ns", dgd).unwrap() - } - - fn record(dgd: &str, relay: Option<&str>) -> PoolRecord { - PoolRecordBuilder::default() - .key(key(dgd)) - .location(PoolLocation { - region: "r1".into(), - availability_zone: None, - cluster: None, - datacenter: None, - }) - .relay_identity(relay.map(|id| RelayIdentity(id.into()))) - .build() - .unwrap() - } - - fn caller(id: &str) -> VerifiedRelayIdentity { - VerifiedRelayIdentity(RelayIdentity(id.into())) - } - - fn pools() -> Fixed { - Fixed( - [record("bound", Some("relay-a")), record("open", None)] - .into_iter() - .map(|r| (r.key.clone(), r)) - .collect(), - ) - } - - #[tokio::test] - async fn admits_the_bound_relay() { - let record = admit(&pools(), &key("bound"), &caller("relay-a")) - .await - .unwrap(); - assert_eq!(record.key, key("bound")); - } - - #[tokio::test] - async fn rejects_another_relay() { - let err = admit(&pools(), &key("bound"), &caller("relay-b")) - .await - .unwrap_err(); - assert!(matches!(err, PoolDiscoveryError::PermissionDenied(k) if k == key("bound"))); - } - - #[tokio::test] - async fn admits_any_relay_without_a_bound_identity() { - assert!( - admit(&pools(), &key("open"), &caller("relay-b")) - .await - .is_ok() - ); - } - - #[tokio::test] - async fn rejects_an_unknown_pool() { - let err = admit(&pools(), &key("missing"), &caller("relay-a")) - .await - .unwrap_err(); - assert!(matches!(err, PoolDiscoveryError::UnknownPool(k) if k == key("missing"))); - } + todo!() } diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index 45d38b7e4efd..17bd41abe818 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -32,9 +32,7 @@ pub use admission::admit; pub use identity::{ PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, }; -pub use mode::{ - BuildPoolDiscovery, McsImplementation, PoolDiscovery, PoolDiscoveryMode, SettlePolicy, -}; +pub use mode::{BuildPoolDiscovery, PoolDiscovery, PoolDiscoveryMode}; pub use record::{PoolDirectory, PoolEvent, PoolRecord, PoolRecordBuilder, Revision}; pub use registrar::{LeasePolicy, PoolRegistrar, Registration}; pub use replicas::{ReplicaDirectory, ReplicaEndpoint, ReplicaEvent, ReplicaId}; diff --git a/lib/kv-router/src/pool_discovery/mode.rs b/lib/kv-router/src/pool_discovery/mode.rs index 2362d8a5bee5..f353fb728103 100644 --- a/lib/kv-router/src/pool_discovery/mode.rs +++ b/lib/kv-router/src/pool_discovery/mode.rs @@ -3,7 +3,6 @@ use std::path::PathBuf; use std::sync::Arc; -use std::time::Duration; use async_trait::async_trait; @@ -14,57 +13,27 @@ 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, - settle: SettlePolicy, }, + /// Pools call `RegisterPool`. The router writes records to the hub's + /// Kubernetes API. Hybrid { lease: LeasePolicy, - hub_namespace: String, }, /// SIG Multicluster: pools are MCS `ServiceExport`s. Liveness is endpoint - /// readiness, so there is no lease. Assumes a trusted network with pod - /// addresses that are reachable between clusters. - Kubernetes { - mcs: McsImplementation, - /// The operator names each PoolRelay Service ``. - /// A name rule works with every MCS implementation; `exportedLabels` - /// does not. - pool_service_suffix: String, - /// The Global Router's Service, exported to every workload cluster. - router_service: String, - }, + /// readiness, so there is no lease. + Kubernetes, File { path: PathBuf, }, } -/// How the MCS implementation names imported EndpointSlices. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -#[non_exhaustive] -pub enum McsImplementation { - /// The `multicluster.kubernetes.io` labels (Submariner, Cilium, GKE). - Standard, - /// Karmada's own names: a `derived-` Service, and the source - /// cluster in the `work.karmada.io/namespace` annotation. - Karmada, -} - -/// When a gRPC-mode replica treats its pool list as settled (soft readiness). -#[derive(Clone, Copy, Debug)] -pub enum SettlePolicy { - FixedDelay(Duration), - /// No new pool for this long. - QuietPeriod(Duration), - /// Copy a peer replica's list; fall back to a delay if none answers. - PeerSnapshot { - fallback: Duration, - }, -} - /// Everything the router needs from the selected mode. The router code is /// the same for every mode. pub struct PoolDiscovery { diff --git a/lib/kv-router/src/pool_discovery/watch.rs b/lib/kv-router/src/pool_discovery/watch.rs index 469407eccab3..c63591ea71bf 100644 --- a/lib/kv-router/src/pool_discovery/watch.rs +++ b/lib/kv-router/src/pool_discovery/watch.rs @@ -34,7 +34,7 @@ pub struct Snapshot { pub enum Completeness { /// Read from a shared store: every existing item is listed. Exact, - /// Built from what this replica has heard so far, settled by a `SettlePolicy`. + /// Built from what this replica has heard so far (gRPC mode). Estimated, } From 333fdbe49aa4005d0a9135309e7b7540d6aaa66d Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 11 Oct 2026 16:09:59 -0700 Subject: [PATCH 06/12] feat(router): sketch the Relay state connection traits RelayListener (router side), RelayDialer (Relay side) and StateConnection, as signatures only, so reviewers see the whole flow. They are subject to change until the transport (push RPC or reverse tunnel) is decided. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- .../src/pool_discovery/connection.rs | 44 +++++++++++++++++++ lib/kv-router/src/pool_discovery/mod.rs | 7 ++- 2 files changed, 49 insertions(+), 2 deletions(-) create mode 100644 lib/kv-router/src/pool_discovery/connection.rs 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..5d4a315b8eef --- /dev/null +++ b/lib/kv-router/src/pool_discovery/connection.rs @@ -0,0 +1,44 @@ +// 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 async_trait::async_trait; +use tokio_util::sync::CancellationToken; + +use super::PoolDiscoveryError; +use super::record::PoolRecord; +use super::replicas::ReplicaEndpoint; +use crate::global_view::PoolKey; + +/// Router side. Checks each caller with `verify`, then `admit`. +#[async_trait] +pub trait RelayListener: Send + Sync { + async fn accept( + &self, + cancel_token: CancellationToken, + ) -> Result<(PoolRecord, Box), PoolDiscoveryError>; +} + +/// Relay side. Dials every replica in its `ReplicaDirectory`. +#[async_trait] +pub trait RelayDialer: Send + Sync { + async fn dial( + &self, + replica: &ReplicaEndpoint, + pool: &PoolKey, + cancel_token: CancellationToken, + ) -> Result, PoolDiscoveryError>; +} + +/// 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/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index 17bd41abe818..157f508dd939 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -10,8 +10,9 @@ //! //! Only the parts that differ between modes are traits: [`PoolDirectory`], //! [`PoolRegistrar`], [`ReplicaDirectory`] and [`RelayAuthenticator`]. -//! Admission, the Relay's `RegisterPool` client, and the connections that -//! carry pool state work the same in every mode, so they are concrete code. +//! 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. //! //! | Mode | Pool list stored in | Pools call | //! |---|---|---| @@ -21,6 +22,7 @@ //! | File | Static file (tests, first deployments) | Nothing | pub mod admission; +pub mod connection; pub mod identity; pub mod mode; pub mod record; @@ -29,6 +31,7 @@ pub mod replicas; pub mod watch; pub use admission::admit; +pub use connection::{RelayDialer, RelayListener, StateConnection}; pub use identity::{ PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, }; From 2ae63912b8d993b055fabc8c8dc220ad6088a778 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Sun, 11 Oct 2026 16:15:50 -0700 Subject: [PATCH 07/12] docs(router): link the public pool discovery design, not a Linear ID Addresses the CodeRabbit review: internal Linear IDs do not resolve for external readers. No GitHub issue tracks this work, so link the public design proposal (ai-dynamo/enhancements#103) instead. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index 157f508dd939..8da34daa37e1 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -1,7 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Pluggable pool discovery for the Global Router (DEP-1276). +//! 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 From 53a71d919721467423953bf3e166613dd4433c78 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Mon, 12 Oct 2026 16:43:41 -0700 Subject: [PATCH 08/12] docs(router): router replicas register with the Dynamo runtime Assume the operator deploys the Global Router, so each replica registers with the hub's runtime discovery like a worker. Runtime discovery lists only Ready pods, so a replica's readiness must not wait for its pools. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/replicas.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/lib/kv-router/src/pool_discovery/replicas.rs b/lib/kv-router/src/pool_discovery/replicas.rs index c28dd5073286..1c49ef49a3c9 100644 --- a/lib/kv-router/src/pool_discovery/replicas.rs +++ b/lib/kv-router/src/pool_discovery/replicas.rs @@ -4,9 +4,14 @@ //! Router replicas a pool must connect to. Every replica needs every pool's //! state, so a Relay opens one connection per replica. //! -//! Replicas register themselves, so a `ReplicaDirectory` can be backed by -//! `dynamo_runtime::discovery::Discovery`. A replica must be listed before it -//! is ready, or pools can never reach it. +//! The operator deploys the Global Router, so each replica registers itself +//! with the hub's Dynamo runtime discovery, like a worker, and the router +//! reads that list back. It needs the list only for the `RegisterPool` reply. +//! +//! 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. In gRPC and Hybrid //! modes it is filled from the `RegisterPool` and `Heartbeat` replies From e32d4ff81242ec58b9025abfb3863bddc9ec33eb Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Thu, 15 Oct 2026 17:39:56 -0700 Subject: [PATCH 09/12] refactor(router): rename RelayListener and RelayDialer to StateReceiver and StateBroadcaster Review feedback: the router side receives pool state and the Relay side broadcasts it, like a radio broadcaster and its receivers. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/connection.rs | 4 ++-- lib/kv-router/src/pool_discovery/mod.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/kv-router/src/pool_discovery/connection.rs b/lib/kv-router/src/pool_discovery/connection.rs index 5d4a315b8eef..8bb33450bde9 100644 --- a/lib/kv-router/src/pool_discovery/connection.rs +++ b/lib/kv-router/src/pool_discovery/connection.rs @@ -18,7 +18,7 @@ use crate::global_view::PoolKey; /// Router side. Checks each caller with `verify`, then `admit`. #[async_trait] -pub trait RelayListener: Send + Sync { +pub trait StateReceiver: Send + Sync { async fn accept( &self, cancel_token: CancellationToken, @@ -27,7 +27,7 @@ pub trait RelayListener: Send + Sync { /// Relay side. Dials every replica in its `ReplicaDirectory`. #[async_trait] -pub trait RelayDialer: Send + Sync { +pub trait StateBroadcaster: Send + Sync { async fn dial( &self, replica: &ReplicaEndpoint, diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index 8da34daa37e1..6185f4e15cf6 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -32,7 +32,7 @@ pub mod replicas; pub mod watch; pub use admission::admit; -pub use connection::{RelayDialer, RelayListener, StateConnection}; +pub use connection::{StateBroadcaster, StateConnection, StateReceiver}; pub use identity::{ PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, }; From bd16741df7c78698338dceb3737c22d349289943 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Thu, 15 Oct 2026 17:44:12 -0700 Subject: [PATCH 10/12] feat(router): add a read accessor to VerifiedRelayIdentity Code in other crates can now read the verified identity, to compare it with a pool record or log it, but still cannot create one. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/identity.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/lib/kv-router/src/pool_discovery/identity.rs b/lib/kv-router/src/pool_discovery/identity.rs index d4bfd846c79d..0d7b6fdb5289 100644 --- a/lib/kv-router/src/pool_discovery/identity.rs +++ b/lib/kv-router/src/pool_discovery/identity.rs @@ -16,6 +16,14 @@ pub struct RelayIdentity(pub String); #[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)] From 241a7467fc12f16aff5435f83adfb3ee40fa861b Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Thu, 15 Oct 2026 18:10:29 -0700 Subject: [PATCH 11/12] refactor(router): split StateDialer from StateBroadcaster Review feedback: dialing is scoped to one replica and broadcasting to the whole directory. StateDialer opens one connection and is the only part that depends on the transport. StateBroadcaster keeps one connection per replica in the ReplicaDirectory and is the same in every mode. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- .../src/pool_discovery/connection.rs | 40 +++++++++++++++++-- lib/kv-router/src/pool_discovery/mod.rs | 2 +- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/lib/kv-router/src/pool_discovery/connection.rs b/lib/kv-router/src/pool_discovery/connection.rs index 8bb33450bde9..d62e0c3c07d3 100644 --- a/lib/kv-router/src/pool_discovery/connection.rs +++ b/lib/kv-router/src/pool_discovery/connection.rs @@ -8,12 +8,14 @@ //! 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::ReplicaEndpoint; +use super::replicas::{ReplicaDirectory, ReplicaEndpoint}; use crate::global_view::PoolKey; /// Router side. Checks each caller with `verify`, then `admit`. @@ -25,9 +27,9 @@ pub trait StateReceiver: Send + Sync { ) -> Result<(PoolRecord, Box), PoolDiscoveryError>; } -/// Relay side. Dials every replica in its `ReplicaDirectory`. +/// Relay side, one replica. The only part that depends on the transport. #[async_trait] -pub trait StateBroadcaster: Send + Sync { +pub trait StateDialer: Send + Sync { async fn dial( &self, replica: &ReplicaEndpoint, @@ -36,6 +38,38 @@ pub trait StateBroadcaster: Send + Sync { ) -> 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 { diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index 6185f4e15cf6..070ac89c7861 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -32,7 +32,7 @@ pub mod replicas; pub mod watch; pub use admission::admit; -pub use connection::{StateBroadcaster, StateConnection, StateReceiver}; +pub use connection::{StateBroadcaster, StateConnection, StateDialer, StateReceiver}; pub use identity::{ PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, }; From 764137296cb8bed81beb6c01cd71b1d192cbf2b1 Mon Sep 17 00:00:00 2001 From: Anna Tchernych Date: Fri, 23 Oct 2026 14:43:16 -0700 Subject: [PATCH 12/12] feat(router): make each pool registration its heartbeat Registration now goes through the Global Discovery Service: a PoolRelay calls RegisterPool on a timer, and each call is also the pool's heartbeat. Relays list router replicas separately, as GET /routers. - PoolRegistrar drops `heartbeat`. Registering an unchanged record again only renews it. - Registration drops `replicas` and keeps the lease policy. - PoolRecord gains the fields RegisterPool carries: capabilities, the local frontend URLs, and capacity per hardware SKU. - The docs say who hosts these calls in the Hybrid and Kubernetes modes. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Anna Tchernych --- lib/kv-router/src/pool_discovery/mod.rs | 8 ++++-- lib/kv-router/src/pool_discovery/mode.rs | 4 +-- lib/kv-router/src/pool_discovery/record.rs | 22 ++++++++++++++-- lib/kv-router/src/pool_discovery/registrar.rs | 25 ++++++++----------- lib/kv-router/src/pool_discovery/replicas.rs | 11 ++++---- 5 files changed, 43 insertions(+), 27 deletions(-) diff --git a/lib/kv-router/src/pool_discovery/mod.rs b/lib/kv-router/src/pool_discovery/mod.rs index 070ac89c7861..53b74e00be55 100644 --- a/lib/kv-router/src/pool_discovery/mod.rs +++ b/lib/kv-router/src/pool_discovery/mod.rs @@ -15,10 +15,14 @@ //! 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 router | `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 | @@ -37,7 +41,7 @@ pub use identity::{ PeerCredentials, RelayAuthenticator, RelayIdentity, VerifiedRelayIdentity, verify, }; pub use mode::{BuildPoolDiscovery, PoolDiscovery, PoolDiscoveryMode}; -pub use record::{PoolDirectory, PoolEvent, PoolRecord, PoolRecordBuilder, Revision}; +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}; diff --git a/lib/kv-router/src/pool_discovery/mode.rs b/lib/kv-router/src/pool_discovery/mode.rs index f353fb728103..02e31c79989a 100644 --- a/lib/kv-router/src/pool_discovery/mode.rs +++ b/lib/kv-router/src/pool_discovery/mode.rs @@ -21,8 +21,8 @@ pub enum PoolDiscoveryMode { Grpc { lease: LeasePolicy, }, - /// Pools call `RegisterPool`. The router writes records to the hub's - /// Kubernetes API. + /// Pools call `RegisterPool`. The Global Discovery Service writes records + /// to the hub's Kubernetes API. Hybrid { lease: LeasePolicy, }, diff --git a/lib/kv-router/src/pool_discovery/record.rs b/lib/kv-router/src/pool_discovery/record.rs index 4fdad332e8aa..f29e14a2cb1f 100644 --- a/lib/kv-router/src/pool_discovery/record.rs +++ b/lib/kv-router/src/pool_discovery/record.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use derive_builder::Builder; +use url::Url; use super::identity::RelayIdentity; use super::watch::{Directory, WatchEvent}; @@ -20,15 +21,25 @@ pub struct Revision(pub u64); pub struct PoolRecord { pub key: PoolKey, pub location: PoolLocation, - // The fields below are `None` until the Relay's catalog reports them, in - // modes where discovery only knows the pool's name (Kubernetes / MCS). + // 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, @@ -36,5 +47,12 @@ pub struct PoolRecord { 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 index 327ab9902475..d534a10d143d 100644 --- a/lib/kv-router/src/pool_discovery/registrar.rs +++ b/lib/kv-router/src/pool_discovery/registrar.rs @@ -8,28 +8,29 @@ use async_trait::async_trait; use super::PoolDiscoveryError; use super::identity::VerifiedRelayIdentity; use super::record::PoolRecord; -use super::replicas::ReplicaEndpoint; use crate::global_view::PoolKey; -/// A record expires `lease_duration` after the last heartbeat this router saw, -/// measured on the router's clock. +/// 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` and `heartbeat`. +/// 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, - /// Every router replica. The Relay opens a state connection to each one. - pub replicas: Vec, } -/// Write side behind `RegisterPool`, `Heartbeat` and `UnregisterPool`. -/// gRPC mode keeps records in memory. Hybrid writes them to the hub's -/// Kubernetes API. Kubernetes and File modes have no registrar. +/// 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( @@ -38,12 +39,6 @@ pub trait PoolRegistrar: Send + Sync { caller: &VerifiedRelayIdentity, ) -> Result; - async fn heartbeat( - &self, - key: &PoolKey, - caller: &VerifiedRelayIdentity, - ) -> Result; - async fn unregister( &self, key: &PoolKey, diff --git a/lib/kv-router/src/pool_discovery/replicas.rs b/lib/kv-router/src/pool_discovery/replicas.rs index 1c49ef49a3c9..3f5d87dfc4c4 100644 --- a/lib/kv-router/src/pool_discovery/replicas.rs +++ b/lib/kv-router/src/pool_discovery/replicas.rs @@ -5,18 +5,17 @@ //! 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, and the router -//! reads that list back. It needs the list only for the `RegisterPool` reply. +//! 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. In gRPC and Hybrid -//! modes it is filled from the `RegisterPool` and `Heartbeat` replies -//! ([`super::Registration::replicas`]). In Kubernetes mode it reads the -//! router's EndpointSlices that MCS imports into the Relay's cluster. +//! 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;