From 27200dfa0fff195880bcf171d98b0f8d47bcd5f6 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Thu, 22 Oct 2026 18:54:48 -0400 Subject: [PATCH 1/9] feat(router): add a Global Discovery Service client crate Add dynamo-global-discovery-client with the two calls a PoolRelay makes to the Global Discovery Service in the Global Router design: - GlobalDiscoveryClient::routers lists every Global Router replica (GET /routers) as ReplicaEndpoint values. A URL that does not parse fails the whole call. - GlobalDiscoveryClient::register_pool registers or refreshes a pool (POST /pools). The body carries the pool's key and location, with the field names of PoolKey and PoolLocation, plus its model, capabilities, frontend URLs, and capacity. The PoolRelay calls it on a schedule, and each call is also the pool's heartbeat. The crate also defines the read-side types from the pool_discovery draft (#15933): ReplicaId, ReplicaEndpoint, WatchEvent, Snapshot, and Completeness. The client and the Global Discovery Service share these definitions. The paths and JSON are a strawman. The service does not serve these calls yet, and their shape is not agreed with its owners. Tests run the client against a local HTTP server. Refs: LLM-212 Signed-off-by: Jacob Ioffe --- .github/codeowners/areas.yaml | 1 + CODEOWNERS | 1 + Cargo.lock | 12 ++ Cargo.toml | 1 + lib/global-discovery-client/Cargo.toml | 22 +++ lib/global-discovery-client/src/client.rs | 154 +++++++++++++++++ lib/global-discovery-client/src/lib.rs | 20 +++ lib/global-discovery-client/src/replicas.rs | 22 +++ lib/global-discovery-client/src/watch.rs | 32 ++++ lib/global-discovery-client/src/wire.rs | 53 ++++++ .../tests/relay_calls.rs | 158 ++++++++++++++++++ 11 files changed, 476 insertions(+) create mode 100644 lib/global-discovery-client/Cargo.toml create mode 100644 lib/global-discovery-client/src/client.rs create mode 100644 lib/global-discovery-client/src/lib.rs create mode 100644 lib/global-discovery-client/src/replicas.rs create mode 100644 lib/global-discovery-client/src/watch.rs create mode 100644 lib/global-discovery-client/src/wire.rs create mode 100644 lib/global-discovery-client/tests/relay_calls.rs diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 57f8b743f723..068f09e7eaab 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -76,6 +76,7 @@ areas: - docs/fern/pages/kubernetes/kv-aware-routing/kv-dc-relay.md - docs/fern/pages/reference/components/kv-dc-relay-configuration.md - lib/bindings/python/rust/llm/kv_dc_relay.rs + - lib/global-discovery-client/ - lib/llm/src/kv_dc_relay.rs - lib/llm/src/kv_dc_relay/ - lib/llm/tests/fixtures/relay-sources-canonical.json diff --git a/CODEOWNERS b/CODEOWNERS index 17aec522c0ad..e2f15c1a8b4c 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -88,6 +88,7 @@ /components/src/dynamo/thunderagent_router/ @ai-dynamo/dynamo-router-codeowners # === global-router (@ai-dynamo/dynamo-global-router-codeowners) === +/lib/global-discovery-client/ @ai-dynamo/dynamo-global-router-codeowners /components/src/dynamo/kv_dc_relay/ @ai-dynamo/dynamo-global-router-codeowners # === backend-vllm (@ai-dynamo/dynamo-backend-vllm-codeowners) === diff --git a/Cargo.lock b/Cargo.lock index dcd7355de6fd..b86a145e7da9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2475,6 +2475,18 @@ dependencies = [ "validator", ] +[[package]] +name = "dynamo-global-discovery-client" +version = "1.6.0" +dependencies = [ + "axum 0.8.4", + "reqwest 0.12.28", + "serde", + "thiserror 2.0.18", + "tokio", + "url", +] + [[package]] name = "dynamo-kv-hashing" version = "1.6.0" diff --git a/Cargo.toml b/Cargo.toml index b7c3ac9f6e81..995d8f890c5c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "lib/kv-hashing", "lib/mocker", "lib/kv-router", + "lib/global-discovery-client", "lib/router-plugins/catalog", "lib/router-plugins/builtin", "examples/router/custom-policy-example/soft-pin-repin", diff --git a/lib/global-discovery-client/Cargo.toml b/lib/global-discovery-client/Cargo.toml new file mode 100644 index 000000000000..52bf0b948f5c --- /dev/null +++ b/lib/global-discovery-client/Cargo.toml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "dynamo-global-discovery-client" +description = "HTTP client for the Global Router's Global Discovery Service" +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 } +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..d139ea6dd901 --- /dev/null +++ b/lib/global-discovery-client/src/client.rs @@ -0,0 +1,154 @@ +// 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 thiserror::Error; +use url::Url; + +use crate::replicas::ReplicaEndpoint; +use crate::wire::{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 or a router URL that does not parse. + #[error("{url} returned a body this client cannot parse")] + Decode { + url: Url, + #[source] + source: reqwest::Error, + }, +} + +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) + .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 = + response + .json() + .await + .map_err(|source| DiscoveryClientError::Decode { + url: url.clone(), + source, + })?; + Ok(body.routers) + } + + /// Calls `POST /pools` to register a pool or refresh its registration. + /// The PoolRelay calls this on a schedule, and each call is also the + /// pool's heartbeat. + pub async fn register_pool( + &self, + registration: &PoolRegistration, + ) -> Result<(), DiscoveryClientError> { + let url = &self.pools_url; + self.send(self.http.post(url.clone()).json(registration), url) + .await?; + Ok(()) + } + + 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(), + }) + } +} + +/// `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 new file mode 100644 index 000000000000..eaf066c270c1 --- /dev/null +++ b/lib/global-discovery-client/src/lib.rs @@ -0,0 +1,20 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! 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 read-side types that the client and the service +//! share: 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, 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..c3a7d1127f1b --- /dev/null +++ b/lib/global-discovery-client/src/watch.rs @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! List-and-watch events, shared by the Global Discovery Service and its +//! clients. +//! +//! A watch starts with one `Added` per current item, 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), + 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..a2cc9d0f18fc --- /dev/null +++ b/lib/global-discovery-client/src/wire.rs @@ -0,0 +1,53 @@ +// 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 serde::{Deserialize, Serialize}; +use url::Url; + +use crate::replicas::ReplicaEndpoint; + +/// `GET /routers` response body. +#[derive(Debug, Deserialize)] +pub(crate) struct RoutersResponse { + pub routers: Vec, +} + +/// `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, +} 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..8d3b7721f30a --- /dev/null +++ b/lib/global-discovery-client/tests/relay_calls.rs @@ -0,0 +1,158 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Tests the client against a fake Global Discovery Service that runs in the +//! test process. + +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, PoolKey, PoolLocation, PoolRegistration, + ReplicaEndpoint, ReplicaId, +}; +use url::Url; + +#[derive(Clone)] +struct FakeService(Arc>); + +struct FakeState { + routers_status: StatusCode, + routers_body: String, + 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 listener = tokio::net::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() }); + let base_url = Url::parse(&format!("http://{address}/discovery")).unwrap(); + GlobalDiscoveryClient::new(base_url, Duration::from_secs(2)).unwrap() + } +} + +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, +) -> StatusCode { + fake.0.lock().unwrap().registrations.push(registration); + StatusCode::NO_CONTENT +} + +fn replica(id: &str, url: &str) -> ReplicaEndpoint { + ReplicaEndpoint { + replica: ReplicaId(id.into()), + url: url.parse().unwrap(), + } +} + +#[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, .. } if status.as_u16() == 503), + "{error:?}" + ); +} + +#[tokio::test] +async fn registers_a_pool() { + let fake = FakeService::new(r#"{ "routers": [] }"#); + let client = fake.client().await; + let registration = 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, + }], + }; + + client.register_pool(®istration).await.unwrap(); + + assert_eq!(fake.0.lock().unwrap().registrations, [registration]); +} From 49731a75382a76cac28c05420e816f4f65a7063a Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Fri, 23 Oct 2026 15:33:55 -0400 Subject: [PATCH 2/9] test(router): drop a redundant module comment in relay_calls Signed-off-by: Jacob Ioffe --- lib/global-discovery-client/tests/relay_calls.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/lib/global-discovery-client/tests/relay_calls.rs b/lib/global-discovery-client/tests/relay_calls.rs index 8d3b7721f30a..839ee01bc375 100644 --- a/lib/global-discovery-client/tests/relay_calls.rs +++ b/lib/global-discovery-client/tests/relay_calls.rs @@ -1,9 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Tests the client against a fake Global Discovery Service that runs in the -//! test process. - use std::sync::{Arc, Mutex}; use std::time::Duration; From 9e39b5447624e61bfd8372d55a9901a5e230f868 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Fri, 23 Oct 2026 15:55:28 -0400 Subject: [PATCH 3/9] fix(router): reject redirects and separate body read errors in the discovery client Disable automatic redirects. A followed 301, 302, or 303 turns the registration POST into a GET, so register_pool returned Ok even though nothing was registered. A 3xx now returns a Status error. Read the /routers body before parsing it. A body cut off mid-transfer now returns a Request error, and Decode only covers JSON that does not match the expected shape. reqwest reports both cases as decode errors, so the client parses with serde_json itself. Remove the redundant RoutersResponse doc comment. Signed-off-by: Jacob Ioffe --- Cargo.lock | 1 + lib/global-discovery-client/Cargo.toml | 1 + lib/global-discovery-client/src/client.rs | 23 +++++++++++++++-------- lib/global-discovery-client/src/wire.rs | 1 - 4 files changed, 17 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b86a145e7da9..b5a7cccabb92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2482,6 +2482,7 @@ dependencies = [ "axum 0.8.4", "reqwest 0.12.28", "serde", + "serde_json", "thiserror 2.0.18", "tokio", "url", diff --git a/lib/global-discovery-client/Cargo.toml b/lib/global-discovery-client/Cargo.toml index 52bf0b948f5c..31d5920254e9 100644 --- a/lib/global-discovery-client/Cargo.toml +++ b/lib/global-discovery-client/Cargo.toml @@ -14,6 +14,7 @@ repository.workspace = true [dependencies] reqwest = { workspace = true } serde = { workspace = true } +serde_json = { workspace = true } thiserror = { workspace = true } url = { workspace = true } diff --git a/lib/global-discovery-client/src/client.rs b/lib/global-discovery-client/src/client.rs index d139ea6dd901..d2bdaef57c8f 100644 --- a/lib/global-discovery-client/src/client.rs +++ b/lib/global-discovery-client/src/client.rs @@ -36,7 +36,7 @@ pub enum DiscoveryClientError { Decode { url: Url, #[source] - source: reqwest::Error, + source: serde_json::Error, }, } @@ -53,6 +53,9 @@ impl GlobalDiscoveryClient { 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 { @@ -70,14 +73,18 @@ impl GlobalDiscoveryClient { 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 = response + .bytes() + .await + .map_err(|source| DiscoveryClientError::Request { + url: url.clone(), + source, + })?; let body: RoutersResponse = - response - .json() - .await - .map_err(|source| DiscoveryClientError::Decode { - url: url.clone(), - source, - })?; + serde_json::from_slice(&body).map_err(|source| DiscoveryClientError::Decode { + url: url.clone(), + source, + })?; Ok(body.routers) } diff --git a/lib/global-discovery-client/src/wire.rs b/lib/global-discovery-client/src/wire.rs index a2cc9d0f18fc..b58fd509f07b 100644 --- a/lib/global-discovery-client/src/wire.rs +++ b/lib/global-discovery-client/src/wire.rs @@ -9,7 +9,6 @@ use url::Url; use crate::replicas::ReplicaEndpoint; -/// `GET /routers` response body. #[derive(Debug, Deserialize)] pub(crate) struct RoutersResponse { pub routers: Vec, From 9c6c1c231ea0bc91c940aae7bc732ac9efe2b0eb Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Fri, 23 Oct 2026 16:16:50 -0400 Subject: [PATCH 4/9] feat(router): derive Clone and Debug for GlobalDiscoveryClient Every other public type in the crate implements Debug. Clone is cheap because reqwest::Client is reference-counted, so a PoolRelay can give its registration loop and its routers poll their own clones instead of wrapping the client in an Arc. Signed-off-by: Jacob Ioffe --- lib/global-discovery-client/src/client.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/lib/global-discovery-client/src/client.rs b/lib/global-discovery-client/src/client.rs index d2bdaef57c8f..e7522b5abeb3 100644 --- a/lib/global-discovery-client/src/client.rs +++ b/lib/global-discovery-client/src/client.rs @@ -40,6 +40,9 @@ pub enum DiscoveryClientError { }, } +/// 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, From 0738961174ac0dfa99eecc358235c2559aec8a32 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Fri, 23 Oct 2026 16:16:55 -0400 Subject: [PATCH 5/9] test(router): pin the registration JSON and cover the redirect fix - registers_a_pool now compares the JSON the service received with a literal, so renaming a field fails the test. It used to decode the body into the same struct the client encoded. - a_redirect_is_an_error checks that a 301 comes back as a Status error instead of being followed. Signed-off-by: Jacob Ioffe --- .../tests/relay_calls.rs | 103 +++++++++++++----- 1 file changed, 76 insertions(+), 27 deletions(-) diff --git a/lib/global-discovery-client/tests/relay_calls.rs b/lib/global-discovery-client/tests/relay_calls.rs index 839ee01bc375..05620df5aabd 100644 --- a/lib/global-discovery-client/tests/relay_calls.rs +++ b/lib/global-discovery-client/tests/relay_calls.rs @@ -13,6 +13,8 @@ use dynamo_global_discovery_client::{ Capacity, DiscoveryClientError, GlobalDiscoveryClient, PoolKey, PoolLocation, PoolRegistration, ReplicaEndpoint, ReplicaId, }; +use serde_json::{Value, json}; +use tokio::net::TcpListener; use url::Url; #[derive(Clone)] @@ -21,7 +23,8 @@ struct FakeService(Arc>); struct FakeState { routers_status: StatusCode, routers_body: String, - registrations: Vec, + /// Registrations as the JSON the service received. + registrations: Vec, } impl FakeService { @@ -40,14 +43,19 @@ impl FakeService { .route("/discovery/routers", get(serve_routers)) .route("/discovery/pools", post(serve_register)) .with_state(self.clone()); - let listener = tokio::net::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() }); + 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 { @@ -62,7 +70,7 @@ async fn serve_routers(State(fake): State) -> Response { async fn serve_register( State(fake): State, - Json(registration): Json, + Json(registration): Json, ) -> StatusCode { fake.0.lock().unwrap().registrations.push(registration); StatusCode::NO_CONTENT @@ -75,6 +83,29 @@ fn replica(id: &str, url: &str) -> ReplicaEndpoint { } } +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( @@ -128,28 +159,46 @@ async fn reports_a_non_2xx_status() { async fn registers_a_pool() { let fake = FakeService::new(r#"{ "routers": [] }"#); let client = fake.client().await; - let registration = 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, - }], - }; - client.register_pool(®istration).await.unwrap(); + client.register_pool(®istration()).await.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" }, + "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 }] + })] + ); +} - assert_eq!(fake.0.lock().unwrap().registrations, [registration]); +#[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:?}" + ); } From 4fa71a6ea09df2e179acaa5b11f93d88e65ef886 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Mon, 26 Oct 2026 08:10:33 -0400 Subject: [PATCH 6/9] feat(router): return the service's lease from register_pool The Global Discovery Service answers POST /pools with {"heartbeat_interval_s": 5, "lease_duration_s": 15}. register_pool now parses that reply and returns it as a Lease, so the PoolRelay can take its heartbeat interval from the service instead of configuring it separately. If the interval ever reached the lease, pools would drop out of GET /pools between heartbeats. A missing body, a zero interval, or an interval not below the lease fails to parse and returns Decode, so a bad reply cannot make the relay register in a tight loop. routers() and register_pool() share one helper that reads the body and then parses it. Signed-off-by: Jacob Ioffe --- lib/global-discovery-client/src/client.rs | 49 ++++++++------ lib/global-discovery-client/src/lib.rs | 2 +- lib/global-discovery-client/src/wire.rs | 64 +++++++++++++++++++ .../tests/relay_calls.rs | 17 +++-- 4 files changed, 107 insertions(+), 25 deletions(-) diff --git a/lib/global-discovery-client/src/client.rs b/lib/global-discovery-client/src/client.rs index e7522b5abeb3..49d8e09d9635 100644 --- a/lib/global-discovery-client/src/client.rs +++ b/lib/global-discovery-client/src/client.rs @@ -4,11 +4,12 @@ use std::time::Duration; use reqwest::StatusCode; +use serde::de::DeserializeOwned; use thiserror::Error; use url::Url; use crate::replicas::ReplicaEndpoint; -use crate::wire::{PoolRegistration, RoutersResponse}; +use crate::wire::{Lease, PoolRegistration, RoutersResponse}; #[derive(Debug, Error)] #[non_exhaustive] @@ -76,32 +77,23 @@ impl GlobalDiscoveryClient { 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 = response - .bytes() - .await - .map_err(|source| DiscoveryClientError::Request { - url: url.clone(), - source, - })?; - let body: RoutersResponse = - serde_json::from_slice(&body).map_err(|source| DiscoveryClientError::Decode { - url: url.clone(), - source, - })?; + let body: RoutersResponse = read_json(response, url).await?; Ok(body.routers) } - /// Calls `POST /pools` to register a pool or refresh its registration. - /// The PoolRelay calls this on a schedule, and each call is also the - /// pool's heartbeat. + /// 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<(), DiscoveryClientError> { + ) -> Result { let url = &self.pools_url; - self.send(self.http.post(url.clone()).json(registration), url) + let response = self + .send(self.http.post(url.clone()).json(registration), url) .await?; - Ok(()) + read_json(response, url).await } async fn send( @@ -128,6 +120,25 @@ impl GlobalDiscoveryClient { } } +/// 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(); diff --git a/lib/global-discovery-client/src/lib.rs b/lib/global-discovery-client/src/lib.rs index eaf066c270c1..61e3cc216a10 100644 --- a/lib/global-discovery-client/src/lib.rs +++ b/lib/global-discovery-client/src/lib.rs @@ -17,4 +17,4 @@ mod wire; pub use client::{DiscoveryClientError, GlobalDiscoveryClient}; pub use replicas::{ReplicaEndpoint, ReplicaEvent, ReplicaId}; pub use watch::{Completeness, Snapshot, WatchEvent}; -pub use wire::{Capacity, PoolKey, PoolLocation, PoolRegistration}; +pub use wire::{Capacity, Lease, PoolKey, PoolLocation, PoolRegistration}; diff --git a/lib/global-discovery-client/src/wire.rs b/lib/global-discovery-client/src/wire.rs index b58fd509f07b..4a0383f57c76 100644 --- a/lib/global-discovery-client/src/wire.rs +++ b/lib/global-discovery-client/src/wire.rs @@ -4,6 +4,8 @@ //! 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; @@ -14,6 +16,41 @@ pub(crate) struct RoutersResponse { pub routers: Vec, } +/// `POST /pools` reply: the service's timing for registrations. Register again +/// after `heartbeat_interval`; without one, the registration lapses after +/// `lease_duration`. 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")] +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 { @@ -50,3 +87,30 @@ 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 05620df5aabd..881d4fb04566 100644 --- a/lib/global-discovery-client/tests/relay_calls.rs +++ b/lib/global-discovery-client/tests/relay_calls.rs @@ -10,8 +10,8 @@ use axum::http::{StatusCode, header}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use dynamo_global_discovery_client::{ - Capacity, DiscoveryClientError, GlobalDiscoveryClient, PoolKey, PoolLocation, PoolRegistration, - ReplicaEndpoint, ReplicaId, + Capacity, DiscoveryClientError, GlobalDiscoveryClient, Lease, PoolKey, PoolLocation, + PoolRegistration, ReplicaEndpoint, ReplicaId, }; use serde_json::{Value, json}; use tokio::net::TcpListener; @@ -71,9 +71,9 @@ async fn serve_routers(State(fake): State) -> Response { async fn serve_register( State(fake): State, Json(registration): Json, -) -> StatusCode { +) -> Json { fake.0.lock().unwrap().registrations.push(registration); - StatusCode::NO_CONTENT + Json(json!({ "heartbeat_interval_s": 5, "lease_duration_s": 15 })) } fn replica(id: &str, url: &str) -> ReplicaEndpoint { @@ -160,8 +160,15 @@ async fn registers_a_pool() { let fake = FakeService::new(r#"{ "routers": [] }"#); let client = fake.client().await; - client.register_pool(®istration()).await.unwrap(); + 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!( From 9f0f234635214f618585a9fcf2621f4d3262e346 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Mon, 26 Oct 2026 10:43:41 -0400 Subject: [PATCH 7/9] fix(router): warn when a caller drops the registration Lease `register_pool(®).await?;` compiled without a warning and threw away the service's timing, which is the failure the review on #17208 asked to prevent. #[must_use] on Lease makes that line warn. The Lease doc now says what happens when lease_duration passes with no successful registration: the pool drops out of GET /pools. Signed-off-by: Jacob Ioffe --- lib/global-discovery-client/src/wire.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/lib/global-discovery-client/src/wire.rs b/lib/global-discovery-client/src/wire.rs index 4a0383f57c76..2f1a550a6c6e 100644 --- a/lib/global-discovery-client/src/wire.rs +++ b/lib/global-discovery-client/src/wire.rs @@ -17,12 +17,13 @@ pub(crate) struct RoutersResponse { } /// `POST /pools` reply: the service's timing for registrations. Register again -/// after `heartbeat_interval`; without one, the registration lapses after -/// `lease_duration`. 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. +/// 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, From 36a9bbce1a7e9b9e3e766036908c06ce2ce06db8 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Mon, 26 Oct 2026 10:46:13 -0400 Subject: [PATCH 8/9] docs(router): remove discovery client doc claims nothing backs up The crate and watch module docs said the service shares these types, but the service does not depend on this crate. The watch doc also fixed a start-up event order that nothing here emits and that is not agreed yet. The Decode doc now lists an invalid Lease as a cause. Signed-off-by: Jacob Ioffe --- lib/global-discovery-client/src/client.rs | 3 ++- lib/global-discovery-client/src/lib.rs | 4 ++-- lib/global-discovery-client/src/watch.rs | 6 +----- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/lib/global-discovery-client/src/client.rs b/lib/global-discovery-client/src/client.rs index 49d8e09d9635..0c6b77b76a69 100644 --- a/lib/global-discovery-client/src/client.rs +++ b/lib/global-discovery-client/src/client.rs @@ -32,7 +32,8 @@ pub enum DiscoveryClientError { body: String, }, /// The body is not in the shape this client expects, for example a - /// renamed field or a router URL that does not parse. + /// 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, diff --git a/lib/global-discovery-client/src/lib.rs b/lib/global-discovery-client/src/lib.rs index 61e3cc216a10..de0081896992 100644 --- a/lib/global-discovery-client/src/lib.rs +++ b/lib/global-discovery-client/src/lib.rs @@ -6,8 +6,8 @@ //! ([`GlobalDiscoveryClient::routers`]) and registering its pool //! ([`GlobalDiscoveryClient::register_pool`]). //! -//! The crate also defines the read-side types that the client and the service -//! share: the router replica types and the list-and-watch events. +//! The crate also defines the router replica types and the list-and-watch +//! events. mod client; mod replicas; diff --git a/lib/global-discovery-client/src/watch.rs b/lib/global-discovery-client/src/watch.rs index c3a7d1127f1b..362ebd5fa292 100644 --- a/lib/global-discovery-client/src/watch.rs +++ b/lib/global-discovery-client/src/watch.rs @@ -1,11 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! List-and-watch events, shared by the Global Discovery Service and its -//! clients. -//! -//! A watch starts with one `Added` per current item, then one `Resync` with the -//! same set, then changes. A `Resync` replaces all earlier state. +//! List-and-watch events. A `Resync` replaces all earlier state. #[derive(Clone, Debug)] pub enum WatchEvent { From e4f6dcde29aef68cf28648dabb3f7f99001d3b3b Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Wed, 28 Oct 2026 11:47:22 -0400 Subject: [PATCH 9/9] test(global-discovery-client): cover a cut-off body and the Status error body Signed-off-by: Jacob Ioffe --- .../tests/relay_calls.rs | 28 ++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/lib/global-discovery-client/tests/relay_calls.rs b/lib/global-discovery-client/tests/relay_calls.rs index 881d4fb04566..b59a004af1ec 100644 --- a/lib/global-discovery-client/tests/relay_calls.rs +++ b/lib/global-discovery-client/tests/relay_calls.rs @@ -14,6 +14,7 @@ use dynamo_global_discovery_client::{ PoolRegistration, ReplicaEndpoint, ReplicaId, }; use serde_json::{Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use url::Url; @@ -150,7 +151,32 @@ async fn reports_a_non_2xx_status() { let error = client.routers().await.unwrap_err(); assert!( - matches!(&error, DiscoveryClientError::Status { status, .. } if status.as_u16() == 503), + 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:?}" ); }