diff --git a/Cargo.lock b/Cargo.lock index 71d4342e2b5f..da69c28705ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2482,6 +2482,15 @@ version = "1.7.0" [[package]] name = "dynamo-global-discovery-client" version = "1.7.0" +dependencies = [ + "axum 0.8.4", + "reqwest 0.12.28", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "url", +] [[package]] name = "dynamo-global-discovery-service" diff --git a/lib/global-discovery-client/Cargo.toml b/lib/global-discovery-client/Cargo.toml index 2bce64562194..31d5920254e9 100644 --- a/lib/global-discovery-client/Cargo.toml +++ b/lib/global-discovery-client/Cargo.toml @@ -8,4 +8,16 @@ version.workspace = true edition.workspace = true authors.workspace = true license.workspace = true +homepage.workspace = true repository.workspace = true + +[dependencies] +reqwest = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +axum = { workspace = true } +tokio = { workspace = true } diff --git a/lib/global-discovery-client/src/client.rs b/lib/global-discovery-client/src/client.rs new file mode 100644 index 000000000000..0c6b77b76a69 --- /dev/null +++ b/lib/global-discovery-client/src/client.rs @@ -0,0 +1,176 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::time::Duration; + +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 { + /// The base URL cannot have a path appended, for example `mailto:`. + #[error("{0} cannot be a base URL")] + InvalidBaseUrl(Url), + #[error("building the HTTP client failed")] + Build(#[source] reqwest::Error), + #[error("request to {url} failed")] + Request { + url: Url, + #[source] + source: reqwest::Error, + }, + #[error("{url} returned {status}: {body}")] + Status { + url: Url, + status: StatusCode, + body: String, + }, + /// The body is not in the shape this client expects, for example a + /// renamed field, a router URL that does not parse, or an invalid + /// [`Lease`]. + #[error("{url} returned a body this client cannot parse")] + Decode { + url: Url, + #[source] + source: serde_json::Error, + }, +} + +/// Cloning is cheap: clones share one connection pool, so give each task its +/// own clone instead of wrapping the client in an `Arc`. +#[derive(Clone, Debug)] +pub struct GlobalDiscoveryClient { + http: reqwest::Client, + routers_url: Url, + pools_url: Url, +} + +impl GlobalDiscoveryClient { + /// `base_url` is the root of the service, for example + /// `http://global-discovery:8080/`. The client keeps any path in it as a + /// prefix. `request_timeout` bounds each call. + pub fn new(base_url: Url, request_timeout: Duration) -> Result { + let http = reqwest::Client::builder() + .timeout(request_timeout) + // A followed redirect can turn the registration POST into a GET, + // so a 3xx must fail instead of looking like a success. + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(DiscoveryClientError::Build)?; + Ok(Self { + http, + routers_url: endpoint(&base_url, "routers")?, + pools_url: endpoint(&base_url, "pools")?, + }) + } + + /// Calls `GET /routers` and returns every Global Router replica the + /// service lists. + /// + /// 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> { + let url = &self.routers_url; + let response = self.send(self.http.get(url.clone()), url).await?; + let body: RoutersResponse = read_json(response, url).await?; + Ok(body.routers) + } + + /// Calls `POST /pools` to register a pool or refresh its registration, + /// and returns the service's [`Lease`]. Each call is also the pool's + /// heartbeat, so the PoolRelay calls this again after + /// `Lease::heartbeat_interval`. + pub async fn register_pool( + &self, + registration: &PoolRegistration, + ) -> Result { + let url = &self.pools_url; + let response = self + .send(self.http.post(url.clone()).json(registration), url) + .await?; + read_json(response, url).await + } + + async fn send( + &self, + request: reqwest::RequestBuilder, + url: &Url, + ) -> Result { + let response = request + .send() + .await + .map_err(|source| DiscoveryClientError::Request { + url: url.clone(), + source, + })?; + let status = response.status(); + if status.is_success() { + return Ok(response); + } + Err(DiscoveryClientError::Status { + url: url.clone(), + status, + body: response.text().await.unwrap_or_default(), + }) + } +} + +/// A body cut off mid-transfer is a `Request` error. Only JSON that does not +/// match `T` is a `Decode` error. +async fn read_json( + response: reqwest::Response, + url: &Url, +) -> Result { + let body = response + .bytes() + .await + .map_err(|source| DiscoveryClientError::Request { + url: url.clone(), + source, + })?; + serde_json::from_slice(&body).map_err(|source| DiscoveryClientError::Decode { + url: url.clone(), + source, + }) +} + +/// `base_url` with `segment` appended, keeping any path it already has. +fn endpoint(base_url: &Url, segment: &str) -> Result { + let mut url = base_url.clone(); + url.path_segments_mut() + .map_err(|()| DiscoveryClientError::InvalidBaseUrl(base_url.clone()))? + .pop_if_empty() + .push(segment); + Ok(url) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn endpoint_keeps_the_base_path() { + for (base, expected) in [ + ("http://hub:8080", "http://hub:8080/pools"), + ( + "http://hub:8080/discovery", + "http://hub:8080/discovery/pools", + ), + ( + "http://hub:8080/discovery/", + "http://hub:8080/discovery/pools", + ), + ] { + let url = endpoint(&base.parse().unwrap(), "pools").unwrap(); + assert_eq!(url.as_str(), expected, "{base}"); + } + let error = endpoint(&"mailto:ops@example.com".parse().unwrap(), "pools").unwrap_err(); + assert!(matches!(error, DiscoveryClientError::InvalidBaseUrl(_))); + } +} diff --git a/lib/global-discovery-client/src/lib.rs b/lib/global-discovery-client/src/lib.rs index 423fc3a9ea72..de0081896992 100644 --- a/lib/global-discovery-client/src/lib.rs +++ b/lib/global-discovery-client/src/lib.rs @@ -1,4 +1,20 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! HTTP client for the Global Router's Global Discovery Service. +//! Client for the two calls a PoolRelay makes to the Global Discovery +//! Service: listing every Global Router replica +//! ([`GlobalDiscoveryClient::routers`]) and registering its pool +//! ([`GlobalDiscoveryClient::register_pool`]). +//! +//! The crate also defines the router replica types and the list-and-watch +//! events. + +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}; diff --git a/lib/global-discovery-client/src/replicas.rs b/lib/global-discovery-client/src/replicas.rs new file mode 100644 index 000000000000..b34b5e526a1d --- /dev/null +++ b/lib/global-discovery-client/src/replicas.rs @@ -0,0 +1,22 @@ +// 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/watch.rs b/lib/global-discovery-client/src/watch.rs new file mode 100644 index 000000000000..362ebd5fa292 --- /dev/null +++ b/lib/global-discovery-client/src/watch.rs @@ -0,0 +1,28 @@ +// 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. + +#[derive(Clone, Debug)] +pub enum WatchEvent { + /// A new or changed item. It 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, +} + +/// Whether a snapshot lists every item, so a new router replica can tell a +/// 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. + Exact, + /// Built from the updates this replica has received so far. + Estimated, +} diff --git a/lib/global-discovery-client/src/wire.rs b/lib/global-discovery-client/src/wire.rs new file mode 100644 index 000000000000..2f1a550a6c6e --- /dev/null +++ b/lib/global-discovery-client/src/wire.rs @@ -0,0 +1,117 @@ +// 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 new file mode 100644 index 000000000000..b59a004af1ec --- /dev/null +++ b/lib/global-discovery-client/tests/relay_calls.rs @@ -0,0 +1,237 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::Json; +use axum::extract::State; +use axum::http::{StatusCode, header}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use dynamo_global_discovery_client::{ + Capacity, DiscoveryClientError, GlobalDiscoveryClient, Lease, PoolKey, PoolLocation, + PoolRegistration, ReplicaEndpoint, ReplicaId, +}; +use serde_json::{Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; +use url::Url; + +#[derive(Clone)] +struct FakeService(Arc>); + +struct FakeState { + routers_status: StatusCode, + routers_body: String, + /// Registrations as the JSON the service received. + registrations: Vec, +} + +impl FakeService { + fn new(routers_body: &str) -> Self { + Self(Arc::new(Mutex::new(FakeState { + routers_status: StatusCode::OK, + routers_body: routers_body.to_owned(), + registrations: Vec::new(), + }))) + } + + /// Serves under `/discovery`, so every test also checks that the client + /// keeps the base URL's path. + async fn client(&self) -> GlobalDiscoveryClient { + let app = axum::Router::new() + .route("/discovery/routers", get(serve_routers)) + .route("/discovery/pools", post(serve_register)) + .with_state(self.clone()); + let address = serve(app).await; + let base_url = Url::parse(&format!("http://{address}/discovery")).unwrap(); + GlobalDiscoveryClient::new(base_url, Duration::from_secs(2)).unwrap() + } +} + +async fn serve(app: axum::Router) -> std::net::SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + address +} + +async fn serve_routers(State(fake): State) -> Response { + let state = fake.0.lock().unwrap(); + if state.routers_status != StatusCode::OK { + return (state.routers_status, "not synced").into_response(); + } + ( + [(header::CONTENT_TYPE, "application/json")], + state.routers_body.clone(), + ) + .into_response() +} + +async fn serve_register( + State(fake): State, + Json(registration): Json, +) -> Json { + fake.0.lock().unwrap().registrations.push(registration); + Json(json!({ "heartbeat_interval_s": 5, "lease_duration_s": 15 })) +} + +fn replica(id: &str, url: &str) -> ReplicaEndpoint { + ReplicaEndpoint { + replica: ReplicaId(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(), + }, + 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(), + capabilities: vec!["aggregated".into()], + frontends: vec!["http://10.1.22.20:8000".parse().unwrap()], + capacity: vec![Capacity { + sku: "H100".into(), + count: 8, + }], + } +} + +#[tokio::test] +async fn lists_every_router() { + let fake = FakeService::new( + r#"{ "routers": [ + { "replica": "aws-ohio/global-router-0", "url": "http://10.0.12.7:50061" }, + { "replica": "aws-west/global-router-0", "url": "http://[fd00::9]:50061", "zone": "usw2-az1" } + ] }"#, + ); + let client = fake.client().await; + + assert_eq!( + client.routers().await.unwrap(), + [ + replica("aws-ohio/global-router-0", "http://10.0.12.7:50061"), + replica("aws-west/global-router-0", "http://[fd00::9]:50061"), + ] + ); +} + +#[tokio::test] +async fn one_bad_url_fails_the_whole_list() { + let fake = FakeService::new( + r#"{ "routers": [ + { "replica": "aws-ohio/global-router-0", "url": "http://10.0.12.7:50061" }, + { "replica": "aws-west/global-router-0", "url": "http://10.1.40.3:99999" } + ] }"#, + ); + let client = fake.client().await; + + let error = client.routers().await.unwrap_err(); + assert!( + matches!(error, DiscoveryClientError::Decode { .. }), + "{error:?}" + ); +} + +#[tokio::test] +async fn reports_a_non_2xx_status() { + let fake = FakeService::new(r#"{ "routers": [] }"#); + fake.0.lock().unwrap().routers_status = StatusCode::SERVICE_UNAVAILABLE; + let client = fake.client().await; + + let error = client.routers().await.unwrap_err(); + assert!( + matches!(&error, DiscoveryClientError::Status { status, body, .. } if status.as_u16() == 503 && body == "not synced"), + "{error:?}" + ); +} + +#[tokio::test] +async fn a_cut_off_body_is_a_request_error() { + // The body ends before its Content-Length. That is a transport failure, + // not a change in the JSON shape, so it must not be reported as `Decode`. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + let _ = socket.read(&mut request).await.unwrap(); + socket + .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 100\r\n\r\n{\"routers\": [") + .await + .unwrap(); + }); + let base_url = Url::parse(&format!("http://{address}/")).unwrap(); + let client = GlobalDiscoveryClient::new(base_url, Duration::from_secs(2)).unwrap(); + + let error = client.routers().await.unwrap_err(); + assert!( + matches!(error, DiscoveryClientError::Request { .. }), + "{error:?}" + ); +} + +#[tokio::test] +async fn registers_a_pool() { + let fake = FakeService::new(r#"{ "routers": [] }"#); + let client = fake.client().await; + + let lease = client.register_pool(®istration()).await.unwrap(); + + assert_eq!( + lease, + Lease { + heartbeat_interval: Duration::from_secs(5), + lease_duration: Duration::from_secs(15), + } + ); + // 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" }, + "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 }] + })] + ); +} + +#[tokio::test] +async fn a_redirect_is_an_error() { + // Following a 301 would resend the registration as a GET without its + // body, so the client must report the 3xx instead. + let app = axum::Router::new().fallback(|| async { + ( + StatusCode::MOVED_PERMANENTLY, + [(header::LOCATION, "/elsewhere")], + ) + }); + let address = serve(app).await; + let base_url = Url::parse(&format!("http://{address}/")).unwrap(); + let client = GlobalDiscoveryClient::new(base_url, Duration::from_secs(2)).unwrap(); + + let error = client.register_pool(®istration()).await.unwrap_err(); + assert!( + matches!(&error, DiscoveryClientError::Status { status, .. } if *status == StatusCode::MOVED_PERMANENTLY), + "{error:?}" + ); +}