diff --git a/Cargo.lock b/Cargo.lock index 2931b3b97bdc..dbe6b17e27ab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2479,12 +2479,20 @@ dependencies = [ [[package]] name = "dynamo-global-common" version = "1.7.0" +dependencies = [ + "blake3", + "serde", + "serde_json", + "thiserror 2.0.18", + "url", +] [[package]] name = "dynamo-global-discovery-client" version = "1.7.0" dependencies = [ "axum 0.8.4", + "dynamo-global-common", "reqwest 0.12.28", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 1e713cc0e4ab..a3b300ec02a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -72,6 +72,7 @@ dynamo-tokens = { path = "lib/tokens", version = "1.7.0" } dynamo-truthy = { path = "lib/truthy", version = "1.7.0" } dynamo-data-gen = { path = "lib/data-gen", version = "1.7.0" } dynamo-kv-hashing = { path = "lib/kv-hashing", version = "1.7.0" } +dynamo-global-common = { path = "lib/global-common", version = "1.7.0" } dynamo-memory = { path = "lib/memory", version = "1.7.0" } dynamo-mocker = { path = "lib/mocker", version = "1.7.0" } dynamo-kv-router = { path = "lib/kv-router", version = "1.7.0" } diff --git a/lib/global-common/Cargo.toml b/lib/global-common/Cargo.toml index b6d4a810bddb..8338e78e2194 100644 --- a/lib/global-common/Cargo.toml +++ b/lib/global-common/Cargo.toml @@ -9,3 +9,12 @@ edition.workspace = true authors.workspace = true license.workspace = true repository.workspace = true + +[dependencies] +blake3 = { workspace = true } +serde = { workspace = true } +thiserror = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +serde_json = { workspace = true } diff --git a/lib/global-common/src/discovery.rs b/lib/global-common/src/discovery.rs new file mode 100644 index 000000000000..b838cb1c236d --- /dev/null +++ b/lib/global-common/src/discovery.rs @@ -0,0 +1,263 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The JSON of the Global Discovery Service's HTTP API. A reader ignores +//! fields it does not know, so the service can add a field before its clients +//! use it. + +use std::net::IpAddr; +use std::time::Duration; + +use serde::{Deserialize, Deserializer, Serialize}; +use thiserror::Error; +use url::Url; + +use crate::pool::{self, PoolKey, PoolLocation}; +use crate::replicas::RouterReplica; + +/// `POST /pools` body. A PoolRelay sends it on a timer, and each call is also +/// the pool's heartbeat. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PoolRegistration { + pub key: PoolKey, + pub location: PoolLocation, + /// The model or checkpoint the pool serves. Reading it checks it as + /// [`PoolKey::new`] checks the key's fields. + #[serde(deserialize_with = "checked_model")] + pub model: String, + #[serde(default)] + pub capabilities: Vec, + /// Local frontends the Global Router sends requests to. + #[serde(default)] + pub frontends: Vec, + #[serde(default)] + pub capacity: Vec, +} + +fn checked_model<'de, D: Deserializer<'de>>(deserializer: D) -> Result { + let model = String::deserialize(deserializer)?; + pool::validate(&model, "model").map_err(serde::de::Error::custom)?; + Ok(model) +} + +/// The workers of one hardware SKU. +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct Capacity { + pub sku: String, + pub workers: u32, +} + +/// `POST /pools` reply: register again after the heartbeat interval. If the +/// lease passes with no successful registration, the pool drops out of +/// `GET /pools`. +/// +/// Both are whole seconds, and the interval is above zero and below the +/// lease. A reply that breaks this does not parse, so it cannot make a +/// PoolRelay register in a tight loop, and every `Lease` can be written as +/// JSON and read back. There is no upper bound: a service bounds the lease it +/// grants, and a client follows any lease it is given. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(try_from = "LeaseFields")] +#[must_use = "register again after `heartbeat_interval`"] +pub struct Lease { + heartbeat_interval_s: u64, + lease_duration_s: u64, +} + +impl Lease { + pub fn new(heartbeat_interval_s: u64, lease_duration_s: u64) -> Result { + if heartbeat_interval_s == 0 || heartbeat_interval_s >= lease_duration_s { + return Err(LeaseError { + heartbeat_interval_s, + lease_duration_s, + }); + } + Ok(Self { + heartbeat_interval_s, + lease_duration_s, + }) + } + + pub fn heartbeat_interval(&self) -> Duration { + Duration::from_secs(self.heartbeat_interval_s) + } + + pub fn duration(&self) -> Duration { + Duration::from_secs(self.lease_duration_s) + } +} + +#[derive(Debug, Error, PartialEq, Eq)] +#[error( + "heartbeat interval {heartbeat_interval_s}s must be above zero and below the \ + {lease_duration_s}s lease" +)] +pub struct LeaseError { + pub heartbeat_interval_s: u64, + pub lease_duration_s: u64, +} + +#[derive(Deserialize)] +struct LeaseFields { + heartbeat_interval_s: u64, + lease_duration_s: u64, +} + +impl TryFrom for Lease { + type Error = LeaseError; + + fn try_from(fields: LeaseFields) -> Result { + Self::new(fields.heartbeat_interval_s, fields.lease_duration_s) + } +} + +/// `GET /pools` reply. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PoolsResponse { + pub pools: Vec, + /// True when the service lists every pool, so a new router can tell a + /// full list from a partial one. + pub complete: bool, +} + +/// One pool in `GET /pools`. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PoolListing { + pub key: PoolKey, + pub location: PoolLocation, + /// `None` when the service does not know the pool's model. + pub model: Option, + #[serde(default)] + pub capabilities: Vec, + /// Local frontends the Global Router sends requests to. + #[serde(default)] + pub frontends: Vec, + #[serde(default)] + pub capacity: Vec, + /// The only PoolRelay allowed to stream this pool's state, by each + /// identity it can connect with. On a trusted network these are its pod's + /// addresses, as [`address_identity`] writes them: two on a dual-stack + /// cluster. Empty when the service does not know the PoolRelay. + #[serde(default)] + pub relay_identities: Vec, +} + +/// Names a PoolRelay by one of its addresses on a trusted network, as +/// `ip:
`. An IPv4-mapped IPv6 address is written as IPv4, as +/// EndpointSlices list it. A router matches a connecting PoolRelay against +/// [`PoolListing::relay_identities`] with this. +pub fn address_identity(address: IpAddr) -> String { + format!("ip:{}", address.to_canonical()) +} + +/// `GET /routers` reply: every Global Router replica. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct RoutersResponse { + pub routers: Vec, +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + + use super::*; + + fn round_trip Deserialize<'de>>(json: Value) -> T { + let parsed: T = serde_json::from_value(json.clone()).unwrap(); + assert_eq!(serde_json::to_value(&parsed).unwrap(), json); + parsed + } + + #[test] + fn json_matches_the_service_and_its_clients() { + let key = + json!({"cluster_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1"}); + let location = + json!({"region": "us-west-2", "availability_zone": null, "datacenter": null}); + let mut registration = json!({ + "key": key, + "location": location, + "model": "Qwen/Qwen3-0.6B", + "capabilities": ["aggregated"], + "frontends": ["http://10.43.1.2:8000/"], + "capacity": [{"sku": "H100", "workers": 8}] + }); + round_trip::(registration.clone()); + registration["model"] = json!(" "); + assert!(serde_json::from_value::(registration).is_err()); + round_trip::(json!({ + "pools": [{ + "key": key, + "location": location, + "model": null, + "capabilities": [], + "frontends": [], + "capacity": [], + "relay_identities": ["ip:10.43.66.175", "ip:fd00::5"] + }], + "complete": true + })); + round_trip::(json!({ + "routers": [{"replica": "aws-ohio/global-router-0", "url": "http://10.42.0.5:50051/"}] + })); + let mapped: IpAddr = "::ffff:10.43.66.175".parse().unwrap(); + assert_eq!(address_identity(mapped), "ip:10.43.66.175"); + assert_eq!(address_identity("fd00::2".parse().unwrap()), "ip:fd00::2"); + } + + #[test] + fn a_lease_parses_unless_it_would_make_a_relay_spin() { + let lease = round_trip::(json!({"heartbeat_interval_s": 5, "lease_duration_s": 15})); + assert_eq!(lease.heartbeat_interval(), Duration::from_secs(5)); + assert_eq!(lease.duration(), Duration::from_secs(15)); + // The longest lease a client could be given. + let longest = + round_trip::(json!({"heartbeat_interval_s": 1, "lease_duration_s": u64::MAX})); + assert_eq!(longest.duration(), Duration::from_secs(u64::MAX)); + for (interval, lease) in [(0, 15), (15, 15), (20, 15)] { + let body = json!({"heartbeat_interval_s": interval, "lease_duration_s": lease}); + assert!( + serde_json::from_value::(body).is_err(), + "{interval}s / {lease}s" + ); + } + } + + #[test] + fn a_reply_without_its_list_does_not_parse() { + // `{}` must not read as "no routers", or a PoolRelay would drop every + // router connection. + assert!(serde_json::from_value::(json!({})).is_err()); + assert!(serde_json::from_value::(json!({"complete": true})).is_err()); + } + + #[test] + fn a_reader_skips_unknown_fields_and_fills_missing_lists() { + // Either side can add a field, at any depth, before the other reads it. + let key = json!({"cluster_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1", "new": 1}); + let location = + json!({"region": "us-west-2", "availability_zone": null, "datacenter": null, "new": 1}); + serde_json::from_value::(json!({ + "key": key, "location": location, "model": "Qwen/Qwen3-0.6B", "new": 1 + })) + .unwrap(); + serde_json::from_value::(json!({ + "pools": [{ + "key": key, "location": location, "model": null, + "capacity": [{"sku": "H100", "workers": 8, "new": 1}], "new": 1 + }], + "complete": true, + "new": 1 + })) + .unwrap(); + let _: Lease = serde_json::from_value( + json!({"heartbeat_interval_s": 5, "lease_duration_s": 15, "new": 1}), + ) + .unwrap(); + serde_json::from_value::(json!({ + "routers": [{"replica": "aws-ohio/global-router-0", "url": "http://10.42.0.5:50051/", "new": 1}], + "new": 1 + })) + .unwrap(); + } +} diff --git a/lib/global-common/src/lib.rs b/lib/global-common/src/lib.rs index 3dcb805a6d82..7adb7a080391 100644 --- a/lib/global-common/src/lib.rs +++ b/lib/global-common/src/lib.rs @@ -2,3 +2,33 @@ // SPDX-License-Identifier: Apache-2.0 //! Types shared by the Global Router components. +//! +//! - A pool's identity and location: [`PoolKey`], [`PoolId`] and its deriver +//! [`V1PoolIdDeriver`], and [`PoolLocation`]. +//! - The JSON of the Global Discovery Service's HTTP API, which the service +//! and its clients exchange: +//! +//! | Call | Body | Reply | +//! |---|---|---| +//! | `POST /pools` | [`PoolRegistration`] | [`Lease`] | +//! | `GET /pools` | | [`PoolsResponse`] | +//! | `GET /routers` | | [`RoutersResponse`] | +//! +//! A router matches a connecting PoolRelay against a listing's +//! `relay_identities` with [`address_identity`]. +//! - Global Router replicas ([`RouterReplica`]), and list-and-watch events +//! and snapshots ([`WatchEvent`], [`Snapshot`]), which the service's watch +//! contract and its clients share. + +mod discovery; +mod pool; +mod replicas; +mod watch; + +pub use discovery::{ + Capacity, Lease, LeaseError, PoolListing, PoolRegistration, PoolsResponse, RoutersResponse, + address_identity, +}; +pub use pool::{PoolId, PoolIdDeriver, PoolKey, PoolKeyError, PoolLocation, V1PoolIdDeriver}; +pub use replicas::{RouterReplica, RouterReplicaEvent, RouterReplicaId}; +pub use watch::{Completeness, Snapshot, WatchEvent}; diff --git a/lib/global-common/src/pool.rs b/lib/global-common/src/pool.rs new file mode 100644 index 000000000000..9de5616acf29 --- /dev/null +++ b/lib/global-common/src/pool.rs @@ -0,0 +1,225 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Runtime-independent identities for pools in a Global View. +//! +//! A routing pool is one DGD in the initial deployment. Its identity is stable +//! across process and relay restarts. + +use std::fmt; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +/// An opaque routing identity. Callers must not parse its representation. +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct PoolId(String); + +impl PoolId { + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl fmt::Display for PoolId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +/// Stable deployment coordinates that a [`PoolIdDeriver`] turns into a [`PoolId`]. +/// +/// `cluster_id` is the cluster ID the pool's PoolRelay runs with +/// (`--cluster-id`, `DYN_CLUSTER_ID`). Where MCS also reports the pool, that +/// ID must be the cluster's name in the ClusterSet, so a registration and the +/// pool MCS reports have one key. `namespace` is the Kubernetes namespace +/// containing the DGD, not the Dynamo runtime namespace derived by the +/// operator. In JSON it is `{"cluster_id", "namespace", "dgd_name"}`, and +/// reading it checks each field as [`PoolKey::new`] does. +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(try_from = "PoolKeyFields", into = "PoolKeyFields")] +pub struct PoolKey { + cluster_id: String, + namespace: String, + dgd_name: String, +} + +#[derive(Serialize, Deserialize)] +struct PoolKeyFields { + cluster_id: String, + namespace: String, + dgd_name: String, +} + +impl TryFrom for PoolKey { + type Error = PoolKeyError; + + fn try_from(fields: PoolKeyFields) -> Result { + Self::new(fields.cluster_id, fields.namespace, fields.dgd_name) + } +} + +impl From for PoolKeyFields { + fn from(key: PoolKey) -> Self { + Self { + cluster_id: key.cluster_id, + namespace: key.namespace, + dgd_name: key.dgd_name, + } + } +} + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum PoolKeyError { + #[error("{field} must be non-empty without surrounding whitespace or control characters")] + InvalidField { field: &'static str }, +} + +impl PoolKey { + pub fn new( + cluster_id: impl Into, + namespace: impl Into, + dgd_name: impl Into, + ) -> Result { + let cluster_id = cluster_id.into(); + let namespace = namespace.into(); + let dgd_name = dgd_name.into(); + validate(&cluster_id, "cluster_id")?; + validate(&namespace, "namespace")?; + validate(&dgd_name, "dgd_name")?; + Ok(Self { + cluster_id, + namespace, + dgd_name, + }) + } + + pub fn cluster_id(&self) -> &str { + &self.cluster_id + } + + pub fn namespace(&self) -> &str { + &self.namespace + } + + pub fn dgd_name(&self) -> &str { + &self.dgd_name + } +} + +pub(crate) fn validate(value: &str, field: &'static str) -> Result<(), PoolKeyError> { + if value.is_empty() || value.trim() != value || value.chars().any(char::is_control) { + return Err(PoolKeyError::InvalidField { field }); + } + Ok(()) +} + +/// Global View owns assignment of routing pool IDs. Implementations can +/// change without leaking derivation details into catalog or routing code. +pub trait PoolIdDeriver: Send + Sync { + fn derive(&self, key: &PoolKey) -> PoolId; +} + +/// Initial deterministic derivation, shared by every Global Router replica. +/// +/// Changing this algorithm changes pool IDs and requires a deployment +/// migration. The API continues treating IDs as opaque strings. +#[derive(Clone, Copy, Debug, Default)] +pub struct V1PoolIdDeriver; + +impl PoolIdDeriver for V1PoolIdDeriver { + fn derive(&self, key: &PoolKey) -> PoolId { + let mut hasher = blake3::Hasher::new_derive_key("dynamo.global_view.pool_id.v1"); + for part in [&key.cluster_id, &key.namespace, &key.dgd_name] { + let bytes = part.as_bytes(); + hasher.update(&(bytes.len() as u64).to_be_bytes()); + hasher.update(bytes); + } + PoolId(format!("pool-v1-{}", hasher.finalize().to_hex())) + } +} + +/// The cluster a pool runs in is [`PoolKey::cluster_id`], not a field here. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PoolLocation { + pub region: String, + pub availability_zone: Option, + pub datacenter: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(cluster: &str, namespace: &str, dgd: &str) -> PoolKey { + PoolKey::new(cluster, namespace, dgd).unwrap() + } + + #[test] + fn identity_survives_deriver_restart() { + let pool = key("ohio", "dynamo", "mocker-1"); + // Pinned: changing the derivation changes every deployed pool's ID. + assert_eq!( + V1PoolIdDeriver.derive(&pool).as_str(), + "pool-v1-2d83e849329ee9362567e384ac43e4d13db27a2b6a00cb79a8be9024504c56ae" + ); + } + + #[test] + fn moving_or_renaming_a_dgd_changes_identity() { + let deriver = V1PoolIdDeriver; + let original = deriver.derive(&key("ohio", "dynamo", "mocker-1")); + for changed in [ + key("west", "dynamo", "mocker-1"), + key("ohio", "other", "mocker-1"), + key("ohio", "dynamo", "mocker-2"), + ] { + assert_ne!(original, deriver.derive(&changed)); + } + } + + #[test] + fn component_boundaries_cannot_alias() { + let deriver = V1PoolIdDeriver; + assert_ne!( + deriver.derive(&key("ab", "c", "d")), + deriver.derive(&key("a", "bc", "d")), + ); + } + + #[test] + fn rejects_invalid_coordinates() { + assert_eq!( + PoolKey::new("", "dynamo", "mocker-1"), + Err(PoolKeyError::InvalidField { + field: "cluster_id" + }) + ); + assert_eq!( + PoolKey::new("ohio", " dynamo", "mocker-1"), + Err(PoolKeyError::InvalidField { field: "namespace" }) + ); + assert_eq!( + PoolKey::new("ohio", "dynamo", "mocker\n1"), + Err(PoolKeyError::InvalidField { field: "dgd_name" }) + ); + } + + #[test] + fn api_representation_is_an_opaque_string() { + let id = V1PoolIdDeriver.derive(&key("ohio", "dynamo", "mocker-1")); + assert_eq!(serde_json::to_string(&id).unwrap(), format!("\"{id}\"")); + } + + #[test] + fn key_json_is_checked_when_read() { + let pool = key("ohio", "dynamo", "mocker-1"); + let json = serde_json::json!({"cluster_id": "ohio", "namespace": "dynamo", "dgd_name": "mocker-1"}); + assert_eq!(serde_json::to_value(&pool).unwrap(), json); + assert_eq!(serde_json::from_value::(json).unwrap(), pool); + let invalid = + serde_json::json!({"cluster_id": " ", "namespace": "dynamo", "dgd_name": "mocker-1"}); + assert!(serde_json::from_value::(invalid).is_err()); + } +} diff --git a/lib/global-common/src/replicas.rs b/lib/global-common/src/replicas.rs new file mode 100644 index 000000000000..4bea24c2fd14 --- /dev/null +++ b/lib/global-common/src/replicas.rs @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Global Router replicas. Every replica needs every pool's state, so a +//! PoolRelay opens one connection per replica. + +use serde::{Deserialize, Serialize}; +use url::Url; + +use crate::watch::WatchEvent; + +/// In JSON, a plain string. +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct RouterReplicaId(pub String); + +/// A Global Router replica and the address workload clusters reach it at. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct RouterReplica { + pub replica: RouterReplicaId, + pub url: Url, +} + +pub type RouterReplicaEvent = WatchEvent; diff --git a/lib/global-discovery-client/src/watch.rs b/lib/global-common/src/watch.rs similarity index 54% rename from lib/global-discovery-client/src/watch.rs rename to lib/global-common/src/watch.rs index 362ebd5fa292..12ad773fd94b 100644 --- a/lib/global-discovery-client/src/watch.rs +++ b/lib/global-common/src/watch.rs @@ -1,12 +1,17 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! List-and-watch events. A `Resync` replaces all earlier state. +//! List-and-watch, as `dynamo_runtime::discovery::Discovery` does it: a watch +//! starts with one `Upserted` per current item (the runtime's `Added`), then +//! one `Resync` with the same set, then changes. A `Resync` replaces all +//! earlier state. #[derive(Clone, Debug)] pub enum WatchEvent { /// A new or changed item. It replaces any earlier item with the same key. - Added(T), + Upserted(T), + /// An item went away. It can name a key the watcher never got, for an + /// item that went away just before the watch started; ignore those. Removed(K), Resync(Snapshot), } @@ -21,8 +26,8 @@ pub struct Snapshot { /// complete list from a partial one. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum Completeness { - /// Read from a shared store, so it lists every existing item. + /// Read from a shared store: every existing item is listed. Exact, - /// Built from the updates this replica has received so far. + /// Built from what this process has heard so far. Estimated, } diff --git a/lib/global-discovery-client/Cargo.toml b/lib/global-discovery-client/Cargo.toml index 31d5920254e9..0445b966a691 100644 --- a/lib/global-discovery-client/Cargo.toml +++ b/lib/global-discovery-client/Cargo.toml @@ -12,6 +12,7 @@ homepage.workspace = true repository.workspace = true [dependencies] +dynamo-global-common = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/lib/global-discovery-client/src/client.rs b/lib/global-discovery-client/src/client.rs index 0c6b77b76a69..672616103e24 100644 --- a/lib/global-discovery-client/src/client.rs +++ b/lib/global-discovery-client/src/client.rs @@ -3,14 +3,12 @@ use std::time::Duration; +use dynamo_global_common::{Lease, PoolRegistration, RouterReplica, RoutersResponse}; use reqwest::StatusCode; use serde::de::DeserializeOwned; use thiserror::Error; use url::Url; -use crate::replicas::ReplicaEndpoint; -use crate::wire::{Lease, PoolRegistration, RoutersResponse}; - #[derive(Debug, Error)] #[non_exhaustive] pub enum DiscoveryClientError { @@ -75,7 +73,7 @@ impl GlobalDiscoveryClient { /// /// Each call returns the full current set. If any entry does not parse, /// the whole call fails, so a caller never acts on a partial list. - pub async fn routers(&self) -> Result, DiscoveryClientError> { + pub async fn routers(&self) -> Result, DiscoveryClientError> { let url = &self.routers_url; let response = self.send(self.http.get(url.clone()), url).await?; let body: RoutersResponse = read_json(response, url).await?; diff --git a/lib/global-discovery-client/src/lib.rs b/lib/global-discovery-client/src/lib.rs index de0081896992..404703d2d050 100644 --- a/lib/global-discovery-client/src/lib.rs +++ b/lib/global-discovery-client/src/lib.rs @@ -6,15 +6,13 @@ //! ([`GlobalDiscoveryClient::routers`]) and registering its pool //! ([`GlobalDiscoveryClient::register_pool`]). //! -//! The crate also defines the router replica types and the list-and-watch -//! events. +//! The types it sends and receives are `dynamo-global-common`'s, re-exported +//! here, so the client reads and writes the JSON as the service does. mod client; -mod replicas; -mod watch; -mod wire; pub use client::{DiscoveryClientError, GlobalDiscoveryClient}; -pub use replicas::{ReplicaEndpoint, ReplicaEvent, ReplicaId}; -pub use watch::{Completeness, Snapshot, WatchEvent}; -pub use wire::{Capacity, Lease, PoolKey, PoolLocation, PoolRegistration}; +pub use dynamo_global_common::{ + Capacity, Lease, LeaseError, PoolKey, PoolKeyError, PoolLocation, PoolRegistration, + RouterReplica, RouterReplicaId, +}; diff --git a/lib/global-discovery-client/src/replicas.rs b/lib/global-discovery-client/src/replicas.rs deleted file mode 100644 index b34b5e526a1d..000000000000 --- a/lib/global-discovery-client/src/replicas.rs +++ /dev/null @@ -1,22 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Global Router replicas a PoolRelay connects to. Every replica needs every -//! pool's state, so a PoolRelay opens one connection per replica. - -use serde::{Deserialize, Serialize}; -use url::Url; - -use crate::watch::WatchEvent; - -#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct ReplicaId(pub String); - -/// A replica's address as seen from workload clusters. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct ReplicaEndpoint { - pub replica: ReplicaId, - pub url: Url, -} - -pub type ReplicaEvent = WatchEvent; diff --git a/lib/global-discovery-client/src/wire.rs b/lib/global-discovery-client/src/wire.rs deleted file mode 100644 index 2f1a550a6c6e..000000000000 --- a/lib/global-discovery-client/src/wire.rs +++ /dev/null @@ -1,117 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! JSON exchanged with the Global Discovery Service. The client ignores -//! response fields it does not know, so the service can add fields first. - -use std::time::Duration; - -use serde::{Deserialize, Serialize}; -use url::Url; - -use crate::replicas::ReplicaEndpoint; - -#[derive(Debug, Deserialize)] -pub(crate) struct RoutersResponse { - pub routers: Vec, -} - -/// `POST /pools` reply: the service's timing for registrations. Register again -/// after `heartbeat_interval`. If `lease_duration` passes with no successful -/// registration, the pool drops out of `GET /pools`. A missing body, a zero -/// interval, or an interval not below the lease does not parse, so a bad reply -/// cannot make a PoolRelay register in a tight loop. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)] -#[serde(try_from = "LeaseBody")] -#[must_use = "wait `heartbeat_interval` before the next `register_pool` call"] -pub struct Lease { - pub heartbeat_interval: Duration, - pub lease_duration: Duration, -} - -#[derive(Deserialize)] -struct LeaseBody { - heartbeat_interval_s: u64, - lease_duration_s: u64, -} - -impl TryFrom for Lease { - type Error = String; - - fn try_from(body: LeaseBody) -> Result { - if body.heartbeat_interval_s == 0 || body.heartbeat_interval_s >= body.lease_duration_s { - return Err(format!( - "heartbeat interval {}s must be above zero and below the {}s lease", - body.heartbeat_interval_s, body.lease_duration_s - )); - } - Ok(Self { - heartbeat_interval: Duration::from_secs(body.heartbeat_interval_s), - lease_duration: Duration::from_secs(body.lease_duration_s), - }) - } -} - -/// `POST /pools` request body. A PoolRelay sends it on every registration. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct PoolRegistration { - pub key: PoolKey, - pub location: PoolLocation, - /// The model or checkpoint the pool serves. - pub model: String, - pub capabilities: Vec, - /// Local frontends the Global Router dispatches requests to. - pub frontends: Vec, - pub capacity: Vec, -} - -/// `namespace` is the Kubernetes namespace containing the DGD, not the -/// Dynamo runtime namespace derived by the operator. -#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] -pub struct PoolKey { - pub site_id: String, - pub namespace: String, - pub dgd_name: String, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct PoolLocation { - pub region: String, - pub availability_zone: Option, - pub cluster: Option, - pub datacenter: Option, -} - -/// The workers of one hardware SKU. -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -pub struct Capacity { - pub sku: String, - pub count: u32, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn a_lease_that_would_spin_does_not_parse() { - let lease: Lease = - serde_json::from_str(r#"{ "heartbeat_interval_s": 5, "lease_duration_s": 15 }"#) - .unwrap(); - assert_eq!( - lease, - Lease { - heartbeat_interval: Duration::from_secs(5), - lease_duration: Duration::from_secs(15), - } - ); - for body in [ - "", - r#"{ "heartbeat_interval_s": 0, "lease_duration_s": 15 }"#, - r#"{ "heartbeat_interval_s": 15, "lease_duration_s": 15 }"#, - r#"{ "heartbeat_interval_s": 20, "lease_duration_s": 15 }"#, - ] { - assert!(serde_json::from_str::(body).is_err(), "{body:?}"); - } - } -} diff --git a/lib/global-discovery-client/tests/relay_calls.rs b/lib/global-discovery-client/tests/relay_calls.rs index b59a004af1ec..740c518f44f0 100644 --- a/lib/global-discovery-client/tests/relay_calls.rs +++ b/lib/global-discovery-client/tests/relay_calls.rs @@ -11,7 +11,7 @@ use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use dynamo_global_discovery_client::{ Capacity, DiscoveryClientError, GlobalDiscoveryClient, Lease, PoolKey, PoolLocation, - PoolRegistration, ReplicaEndpoint, ReplicaId, + PoolRegistration, RouterReplica, RouterReplicaId, }; use serde_json::{Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -77,24 +77,19 @@ async fn serve_register( Json(json!({ "heartbeat_interval_s": 5, "lease_duration_s": 15 })) } -fn replica(id: &str, url: &str) -> ReplicaEndpoint { - ReplicaEndpoint { - replica: ReplicaId(id.into()), +fn replica(id: &str, url: &str) -> RouterReplica { + RouterReplica { + replica: RouterReplicaId(id.into()), url: url.parse().unwrap(), } } fn registration() -> PoolRegistration { PoolRegistration { - key: PoolKey { - site_id: "aws-west".into(), - namespace: "dynamo".into(), - dgd_name: "mocker-west-1".into(), - }, + key: PoolKey::new("aws-west", "dynamo", "mocker-west-1").unwrap(), location: PoolLocation { region: "us-west-2".into(), availability_zone: Some("usw2-az1".into()), - cluster: Some("aws-west".into()), datacenter: None, }, model: "Qwen/Qwen3-0.6B".into(), @@ -102,7 +97,7 @@ fn registration() -> PoolRegistration { frontends: vec!["http://10.1.22.20:8000".parse().unwrap()], capacity: vec![Capacity { sku: "H100".into(), - count: 8, + workers: 8, }], } } @@ -188,29 +183,22 @@ async fn registers_a_pool() { let lease = client.register_pool(®istration()).await.unwrap(); - assert_eq!( - lease, - Lease { - heartbeat_interval: Duration::from_secs(5), - lease_duration: Duration::from_secs(15), - } - ); + assert_eq!(lease, Lease::new(5, 15).unwrap()); // The field names are the contract with the service, so compare the JSON // it received rather than decoding it back into the same struct. assert_eq!( fake.0.lock().unwrap().registrations, [json!({ - "key": { "site_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1" }, + "key": { "cluster_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1" }, "location": { "region": "us-west-2", "availability_zone": "usw2-az1", - "cluster": "aws-west", "datacenter": null }, "model": "Qwen/Qwen3-0.6B", "capabilities": ["aggregated"], "frontends": ["http://10.1.22.20:8000/"], - "capacity": [{ "sku": "H100", "count": 8 }] + "capacity": [{ "sku": "H100", "workers": 8 }] })] ); }