diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 898d3554230a..b603b9684e40 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -79,6 +79,10 @@ areas: - docs/fern/pages/reference/components/pool-relay-configuration.md - lib/bindings/python/rust/llm/pool_relay.rs - lib/global-discovery-client/ + - lib/global-discovery-service/ + - lib/global-ingress/ + - lib/global-router/ + - lib/global-router-shared/ - lib/llm/src/pool_relay.rs - lib/llm/src/pool_relay/ - lib/llm/tests/fixtures/relay-sources-canonical.json diff --git a/CODEOWNERS b/CODEOWNERS index e87af0496299..76b1ad2765e1 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -88,9 +88,13 @@ /components/src/dynamo/thunderagent_router/ @ai-dynamo/dynamo-router-codeowners # === global-router (@ai-dynamo/dynamo-global-router-codeowners) === -/examples/global-router/cdk/ @ai-dynamo/dynamo-global-router-codeowners +/lib/global-router/ @ai-dynamo/dynamo-global-router-codeowners +/lib/global-ingress/ @ai-dynamo/dynamo-global-router-codeowners +/lib/global-router-shared/ @ai-dynamo/dynamo-global-router-codeowners /lib/global-view-receiver/ @ai-dynamo/dynamo-global-router-codeowners +/examples/global-router/cdk/ @ai-dynamo/dynamo-global-router-codeowners /lib/global-discovery-client/ @ai-dynamo/dynamo-global-router-codeowners +/lib/global-discovery-service/ @ai-dynamo/dynamo-global-router-codeowners /components/src/dynamo/pool_relay/ @ai-dynamo/dynamo-global-router-codeowners # === backend-vllm (@ai-dynamo/dynamo-backend-vllm-codeowners) === diff --git a/Cargo.lock b/Cargo.lock index 84b57de13a12..202019c35d71 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2658,6 +2658,48 @@ dependencies = [ "url", ] +[[package]] +name = "dynamo-global-ingress" +version = "1.6.0" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "axum 0.8.4", + "clap", + "dynamo-global-discovery-client", + "dynamo-global-discovery-service", + "dynamo-global-router", + "dynamo-global-router-shared", + "dynamo-llm", + "dynamo-protocols", + "dynamo-runtime", + "futures", + "reqwest 0.12.28", + "serde", + "serde_json", + "tokio", + "tokio-util", + "tracing", + "url", +] + +[[package]] +name = "dynamo-global-router" +version = "1.6.0" +dependencies = [ + "axum 0.8.4", + "dashmap", + "dynamo-global-discovery-client", + "dynamo-global-router-shared", + "dynamo-runtime", + "serde_json", + "tokio", + "tokio-util", + "tracing", + "url", +] + [[package]] name = "dynamo-global-router-shared" version = "1.6.0" diff --git a/Cargo.toml b/Cargo.toml index 49b7805b5336..efb94862940b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,8 @@ members = [ "lib/kv-router", "lib/global-discovery-client", "lib/global-discovery-service", + "lib/global-ingress", + "lib/global-router", "lib/global-router-shared", "lib/router-plugins/catalog", "lib/router-plugins/builtin", diff --git a/lib/global-discovery-client/src/client.rs b/lib/global-discovery-client/src/client.rs index e7522b5abeb3..0f21e468e801 100644 --- a/lib/global-discovery-client/src/client.rs +++ b/lib/global-discovery-client/src/client.rs @@ -4,11 +4,13 @@ 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::watch::{Completeness, Snapshot}; +use crate::wire::{PoolListing, PoolRegistration, PoolsResponse, RoutersResponse}; #[derive(Debug, Error)] #[non_exhaustive] @@ -74,7 +76,28 @@ 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> { - let url = &self.routers_url; + let body: RoutersResponse = self.get_json(&self.routers_url).await?; + Ok(body.routers) + } + + /// Calls `GET /pools` and returns every pool the service lists. A Global + /// Router calls this to learn the pools and their frontends. + /// + /// Like [`Self::routers`], one entry that does not parse fails the whole + /// call. + pub async fn pools(&self) -> Result, DiscoveryClientError> { + let body: PoolsResponse = self.get_json(&self.pools_url).await?; + Ok(Snapshot { + items: body.pools, + completeness: if body.complete { + Completeness::Exact + } else { + Completeness::Estimated + }, + }) + } + + async fn get_json(&self, url: &Url) -> Result { let response = self.send(self.http.get(url.clone()), url).await?; let body = response .bytes() @@ -83,12 +106,10 @@ impl GlobalDiscoveryClient { url: url.clone(), source, })?; - let body: RoutersResponse = - serde_json::from_slice(&body).map_err(|source| DiscoveryClientError::Decode { - url: url.clone(), - source, - })?; - Ok(body.routers) + serde_json::from_slice(&body).map_err(|source| DiscoveryClientError::Decode { + url: url.clone(), + source, + }) } /// Calls `POST /pools` to register a pool or refresh its registration. diff --git a/lib/global-discovery-client/src/lib.rs b/lib/global-discovery-client/src/lib.rs index eaf066c270c1..b38a54adbda6 100644 --- a/lib/global-discovery-client/src/lib.rs +++ b/lib/global-discovery-client/src/lib.rs @@ -1,10 +1,10 @@ // 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`]). +//! Client for the Global Discovery Service. A PoolRelay lists every Global +//! Router replica ([`GlobalDiscoveryClient::routers`]) and registers its pool +//! ([`GlobalDiscoveryClient::register_pool`]). A Global Router lists every +//! pool ([`GlobalDiscoveryClient::pools`]). //! //! 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. @@ -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, PoolKey, PoolListing, PoolLocation, PoolRegistration}; diff --git a/lib/global-discovery-client/src/wire.rs b/lib/global-discovery-client/src/wire.rs index b58fd509f07b..f3afdfcc148d 100644 --- a/lib/global-discovery-client/src/wire.rs +++ b/lib/global-discovery-client/src/wire.rs @@ -14,6 +14,33 @@ pub(crate) struct RoutersResponse { pub routers: Vec, } +#[derive(Debug, Deserialize)] +pub(crate) struct PoolsResponse { + pub pools: Vec, + /// True when the service lists every pool. + pub complete: bool, +} + +/// One pool in the `GET /pools` response. +#[derive(Clone, Debug, PartialEq, Eq, Deserialize)] +pub struct PoolListing { + pub key: PoolKey, + pub location: PoolLocation, + /// `None` when the service knows only the pool's name, as with MCS. + pub model: Option, + #[serde(default)] + pub capabilities: Vec, + /// Local frontends the Global Router dispatches requests to. + #[serde(default)] + pub frontends: Vec, + #[serde(default)] + pub capacity: Vec, + /// The only Relay allowed to stream this pool's state, when the service + /// knows it. + #[serde(default)] + pub relay_identity: Option, +} + /// `POST /pools` request body. A PoolRelay sends it on every registration. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct PoolRegistration { diff --git a/lib/global-discovery-client/tests/router_calls.rs b/lib/global-discovery-client/tests/router_calls.rs new file mode 100644 index 000000000000..2567147780bd --- /dev/null +++ b/lib/global-discovery-client/tests/router_calls.rs @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::time::Duration; + +use axum::http::header; +use axum::routing::get; +use dynamo_global_discovery_client::{ + Capacity, Completeness, DiscoveryClientError, GlobalDiscoveryClient, PoolKey, PoolListing, + PoolLocation, +}; +use tokio::net::TcpListener; +use url::Url; + +/// A client for a service that answers `GET /discovery/pools` with `body`. +async fn client(body: &'static str) -> GlobalDiscoveryClient { + let app = axum::Router::new().route( + "/discovery/pools", + get(move || async move { ([(header::CONTENT_TYPE, "application/json")], body) }), + ); + 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() }); + let base_url = Url::parse(&format!("http://{address}/discovery")).unwrap(); + GlobalDiscoveryClient::new(base_url, Duration::from_secs(2)).unwrap() +} + +#[tokio::test] +async fn lists_every_pool() { + // The second pool is what the Kubernetes stores report: a name and a + // Relay, but no model or frontends. + let client = client( + r#"{ "complete": true, "pools": [ + { + "key": { "site_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1" }, + "location": { "region": "us-west-2", "availability_zone": null, "cluster": "aws-west", "datacenter": null }, + "model": "Qwen/Qwen3-0.6B", + "capabilities": ["aggregated"], + "frontends": ["http://10.1.22.20:8000/"], + "capacity": [{ "sku": "H100", "count": 8 }], + "relay_identity": null, + "added_later": "ignored" + }, + { + "key": { "site_id": "aws-ohio", "namespace": "dynamo", "dgd_name": "mocker-ohio-1" }, + "location": { "region": "us-east-2", "availability_zone": null, "cluster": null, "datacenter": null }, + "model": null, + "relay_identity": "10.0.3.4" + } + ] }"#, + ) + .await; + + let snapshot = client.pools().await.unwrap(); + assert_eq!(snapshot.completeness, Completeness::Exact); + assert_eq!( + snapshot.items, + [ + PoolListing { + 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: None, + cluster: Some("aws-west".into()), + datacenter: None, + }, + model: Some("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, + }], + relay_identity: None, + }, + PoolListing { + key: PoolKey { + site_id: "aws-ohio".into(), + namespace: "dynamo".into(), + dgd_name: "mocker-ohio-1".into(), + }, + location: PoolLocation { + region: "us-east-2".into(), + availability_zone: None, + cluster: None, + datacenter: None, + }, + model: None, + capabilities: Vec::new(), + frontends: Vec::new(), + capacity: Vec::new(), + relay_identity: Some("10.0.3.4".into()), + }, + ] + ); +} + +#[tokio::test] +async fn a_partial_list_is_estimated() { + let client = client(r#"{ "complete": false, "pools": [] }"#).await; + + let snapshot = client.pools().await.unwrap(); + assert_eq!(snapshot.completeness, Completeness::Estimated); + assert!(snapshot.items.is_empty()); +} + +#[tokio::test] +async fn one_bad_frontend_fails_the_whole_list() { + let client = client( + r#"{ "complete": true, "pools": [ { + "key": { "site_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1" }, + "location": { "region": "us-west-2", "availability_zone": null, "cluster": null, "datacenter": null }, + "model": "Qwen/Qwen3-0.6B", + "frontends": ["http://10.1.22.20:99999/"] + } ] }"#, + ) + .await; + + let error = client.pools().await.unwrap_err(); + assert!( + matches!(error, DiscoveryClientError::Decode { .. }), + "{error:?}" + ); +} diff --git a/lib/global-ingress/Cargo.toml b/lib/global-ingress/Cargo.toml new file mode 100644 index 000000000000..b6b2ddb85a1f --- /dev/null +++ b/lib/global-ingress/Cargo.toml @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "dynamo-global-ingress" +description = "Global Frontend: serves the OpenAI API and forwards each request to a pool frontend" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true + +[[bin]] +name = "global-ingress" +path = "src/main.rs" + +[dependencies] +dynamo-global-discovery-client = { path = "../global-discovery-client" } +dynamo-global-router = { path = "../global-router" } +dynamo-global-router-shared = { path = "../global-router-shared" } +# The block manager serves workers. A frontend that only forwards requests +# never uses it. +dynamo-llm = { path = "../llm", default-features = false } +dynamo-protocols = { workspace = true } +dynamo-runtime = { workspace = true } + +anyhow = { workspace = true } +async-stream = { workspace = true } +async-trait = { workspace = true } +clap = { version = "4.5.49", features = ["derive"] } +futures = { workspace = true } +reqwest = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tracing = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +dynamo-global-discovery-service = { path = "../global-discovery-service" } + +axum = { workspace = true } diff --git a/lib/global-ingress/src/forward.rs b/lib/global-ingress/src/forward.rs new file mode 100644 index 000000000000..44c78ece2af0 --- /dev/null +++ b/lib/global-ingress/src/forward.rs @@ -0,0 +1,348 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::io; +use std::sync::Arc; + +use async_trait::async_trait; +use dynamo_global_router::{PoolSelector, PoolSnapshot, PoolTarget, RoutingRequest}; +use dynamo_llm::http::service::error::HttpError; +use dynamo_llm::protocols::agents::HEADER_DYNAMO_SESSION_ID; +use dynamo_llm::protocols::codec::{Message, SseLineCodec}; +use dynamo_llm::protocols::common::extensions::{SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId}; +use dynamo_llm::protocols::openai::chat_completions::{ + NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse, +}; +use dynamo_llm::protocols::openai::completions::{ + NvCreateCompletionRequest, NvCreateCompletionResponse, +}; +use dynamo_protocols::types::ChatCompletionStreamOptions; +use dynamo_runtime::engine::{ + AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, ResponseStream, +}; +use dynamo_runtime::error::{DynamoError, ErrorType}; +use dynamo_runtime::pipeline::{Error, ManyOut, SingleIn}; +use dynamo_runtime::protocols::annotated::Annotated; +use futures::{Stream, StreamExt, TryStreamExt}; +use reqwest::header::ACCEPT; +use serde::Serialize; +use serde::de::DeserializeOwned; +use tokio::sync::watch; +use tokio_util::codec::FramedRead; +use tokio_util::io::StreamReader; +use url::Url; + +/// A request the engine can forward: the route it goes to on a pool +/// frontend, and the fields the engine reads or rewrites. +pub trait ForwardedRequest: Serialize + Send + Sync + 'static { + /// One event of the pool's streamed response. + type Response: DeserializeOwned + Send + Sync + 'static; + + /// Route below a pool frontend's root, as path segments. + const ROUTE: &'static [&'static str]; + + fn model(&self) -> &str; + + /// Ask the pool to stream its response. The HTTP handler consumes every + /// engine as a stream and folds it for clients that did not ask to + /// stream, so those clients' usage has to come back in the stream. + fn request_stream(&mut self); +} + +impl ForwardedRequest for NvCreateChatCompletionRequest { + type Response = NvCreateChatCompletionStreamResponse; + + const ROUTE: &'static [&'static str] = &["v1", "chat", "completions"]; + + fn model(&self) -> &str { + &self.inner.model + } + + fn request_stream(&mut self) { + if self.inner.stream != Some(true) { + self.inner.stream = Some(true); + self.inner.stream_options = Some(usage_in_stream()); + } + } +} + +impl ForwardedRequest for NvCreateCompletionRequest { + type Response = NvCreateCompletionResponse; + + const ROUTE: &'static [&'static str] = &["v1", "completions"]; + + fn model(&self) -> &str { + &self.inner.model + } + + fn request_stream(&mut self) { + if self.inner.stream != Some(true) { + self.inner.stream = Some(true); + self.inner.stream_options = Some(usage_in_stream()); + } + } +} + +fn usage_in_stream() -> ChatCompletionStreamOptions { + ChatCompletionStreamOptions { + include_usage: true, + continuous_usage_stats: false, + } +} + +/// Forwards each request to the pool frontend its selector picks and streams +/// the pool's response back. +/// +/// The engine reads the model from each request, so one instance serves +/// every model. +pub struct PoolForwardingEngine { + pools: watch::Receiver>, + selector: Arc, + http: reqwest::Client, +} + +impl PoolForwardingEngine { + /// `http` must not set a total request timeout, because a response + /// streams for as long as the pool generates. + pub fn new( + pools: watch::Receiver>, + selector: Arc, + http: reqwest::Client, + ) -> Self { + Self { + pools, + selector, + http, + } + } + + fn select(&self, model: &str) -> Result { + let snapshot = self.pools.borrow().clone(); + let candidates = snapshot.pools_for(model).collect::>(); + self.selector + .select(&RoutingRequest::new(model), &candidates) + .ok_or_else(|| unavailable(format!("no pool can serve model {model}"))) + } + + async fn forward( + &self, + request: SingleIn, + ) -> Result>, Error> { + // The HTTP handler moves the session header into the request context, + // so it has to be put back on the forwarded request. + let session = request + .get::(SESSION_AFFINITY_CONTEXT_KEY) + .ok(); + let (mut body, context) = request.into_parts(); + let ctx = context.context(); + let target = self.select(body.model())?; + let url = route_url(&target.frontend, R::ROUTE)?; + body.request_stream(); + + let mut pool_request = self + .http + .post(url.clone()) + .header(ACCEPT, "text/event-stream") + .json(&body); + if let Some(session) = session { + pool_request = pool_request.header(HEADER_DYNAMO_SESSION_ID, session.as_str()); + } + tracing::debug!(request_id = ctx.id(), pool = %target.pool, %url, "forwarding request"); + // TODO: retry the next pool when this one rejects the request before + // its first byte (HLD 2.2.2, priority retry). + let response = pool_request + .send() + .await + .map_err(|error| unavailable(format!("pool frontend {url} failed: {error}")))?; + let status = response.status(); + if !status.is_success() { + // The HTTP service maps this to the client response, keeping + // the pool's status. + let message = response.text().await.unwrap_or_default(); + return Err(HttpError { + code: status.as_u16(), + message, + } + .into()); + } + let events = pool_events::(response, ctx.clone()); + Ok(ResponseStream::new(Box::pin(events), ctx)) + } +} + +#[async_trait] +impl AsyncEngine, ManyOut>, Error> + for PoolForwardingEngine +{ + async fn generate( + &self, + request: SingleIn, + ) -> Result>, Error> { + self.forward(request).await + } +} + +fn unavailable(message: String) -> DynamoError { + DynamoError::builder() + .error_type(ErrorType::Unavailable) + .message(message) + .build() +} + +/// `frontend` with `route` appended, keeping any path it already has. +fn route_url(frontend: &Url, route: &[&str]) -> Result { + let mut url = frontend.clone(); + url.path_segments_mut() + .map_err(|()| { + DynamoError::builder() + .error_type(ErrorType::Internal) + .message(format!("pool frontend {frontend} cannot be a base URL")) + .build() + })? + .pop_if_empty() + .extend(route); + Ok(url) +} + +/// Decodes a pool frontend's SSE body into engine output. +/// +/// The stream ends with the body, at the first error, or when the client +/// goes away. Ending early drops the response body, which closes the +/// connection and so cancels the request at the pool. +/// +/// `SseLineCodec` consumes `data: [DONE]` without emitting it, so a pool that +/// closes its body cleanly before `[DONE]` looks like a finished stream. +fn pool_events( + response: reqwest::Response, + ctx: Arc, +) -> impl Stream> + Send { + let body = StreamReader::new(response.bytes_stream().map_err(io::Error::other)); + let mut messages = FramedRead::new(body, SseLineCodec::new()); + async_stream::stream! { + loop { + let message = tokio::select! { + biased; + _ = ctx.stopped() => break, + message = messages.next() => message, + }; + let message = match message { + Some(Ok(message)) => message, + Some(Err(error)) => { + yield Annotated::from_error(format!("reading the pool response failed: {error}")); + break; + } + None => break, + }; + // Keep-alive comments. The Global Frontend sends its own. + if message.data.is_none() && message.event.is_none() { + continue; + } + match decode(message) { + Ok(event) => yield event, + Err(error) => { + yield Annotated::from_error(error); + break; + } + } + } + } +} + +fn decode(message: Message) -> Result, String> { + if message.event.as_deref() == Some("error") { + return Err(message + .comments + .map(|comments| comments.join("\n")) + .unwrap_or_else(|| "the pool reported a stream error".to_owned())); + } + let data = match message.data.as_deref() { + None => None, + Some(data) => Some(serde_json::from_str(data).map_err(|error| { + error_frame_message(data).unwrap_or_else(|| { + format!("the pool sent an event this frontend cannot decode: {error}") + }) + })?), + }; + Ok(Annotated { + data, + id: message.id, + event: message.event, + comment: message.comments, + error: None, + }) +} + +/// The message of a `{"error": {"message": ...}}` frame, which a Dynamo +/// frontend sends when a stream fails after it started. +fn error_frame_message(data: &str) -> Option { + let frame: serde_json::Value = serde_json::from_str(data).ok()?; + let error = frame.get("error")?; + Some( + error + .get("message") + .and_then(serde_json::Value::as_str) + .map_or_else(|| error.to_string(), str::to_owned), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn route_url_keeps_the_frontend_path() { + for (frontend, expected) in [ + ("http://pool:8000", "http://pool:8000/v1/chat/completions"), + ("http://pool:8000/", "http://pool:8000/v1/chat/completions"), + ( + "http://gateway:8080/ohio/", + "http://gateway:8080/ohio/v1/chat/completions", + ), + ] { + let url = route_url( + &frontend.parse().unwrap(), + NvCreateChatCompletionRequest::ROUTE, + ); + assert_eq!(url.unwrap().as_str(), expected, "{frontend}"); + } + assert!(route_url(&"mailto:ops@example.com".parse().unwrap(), &["v1"]).is_err()); + } + + #[test] + fn non_streaming_requests_ask_the_pool_to_stream_usage() { + let mut request: NvCreateChatCompletionRequest = + serde_json::from_value(serde_json::json!({ + "model": "qwen", + "messages": [{"role": "user", "content": "hi"}], + })) + .unwrap(); + request.request_stream(); + assert_eq!(request.inner.stream, Some(true)); + assert!(request.inner.stream_options.unwrap().include_usage); + + let mut request: NvCreateChatCompletionRequest = + serde_json::from_value(serde_json::json!({ + "model": "qwen", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + })) + .unwrap(); + request.request_stream(); + assert_eq!( + request.inner.stream_options, None, + "a streaming client's own options must pass through" + ); + } + + #[test] + fn error_frames_surface_the_pool_message() { + let message = Message { + id: None, + event: None, + data: Some(r#"{"error":{"message":"worker died","type":"internal_error"}}"#.to_owned()), + comments: None, + }; + let error = decode::(message).unwrap_err(); + assert_eq!(error, "worker died"); + } +} diff --git a/lib/global-ingress/src/lib.rs b/lib/global-ingress/src/lib.rs new file mode 100644 index 000000000000..3f08164fb8c9 --- /dev/null +++ b/lib/global-ingress/src/lib.rs @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The Global Frontend: Dynamo's OpenAI-compatible HTTP service serving one +//! engine, [`PoolForwardingEngine`], that forwards each request to a pool +//! frontend chosen by the Global Router. +//! +//! The Global Frontend does not tokenize. The pool frontend preprocesses a +//! forwarded request exactly as it would one sent to it directly. +//! [`ModelSync`] keeps the HTTP service's models in step with the pool +//! catalog, so a model is served while at least one pool serves it. + +mod forward; +mod model_sync; + +pub use forward::{ForwardedRequest, PoolForwardingEngine}; +pub use model_sync::ModelSync; diff --git a/lib/global-ingress/src/main.rs b/lib/global-ingress/src/main.rs new file mode 100644 index 000000000000..c653145dbd8b --- /dev/null +++ b/lib/global-ingress/src/main.rs @@ -0,0 +1,157 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Global Frontend. See `dynamo_global_ingress`. +//! +//! ```bash +//! # Pools from the Global Discovery Service: +//! cargo run -p dynamo-global-ingress -- --port 8000 \ +//! --discovery-url http://global-discovery:8080/ +//! # Or pools given by hand: +//! cargo run -p dynamo-global-ingress -- --port 8000 \ +//! --pool ohio/dynamo/mocker-1=Qwen/Qwen3-0.6B@http://10.42.0.5:8000 +//! ``` + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{Context, Result, bail}; +use clap::{ArgGroup, Parser}; +use dynamo_global_discovery_client::GlobalDiscoveryClient; +use dynamo_global_ingress::{ModelSync, PoolForwardingEngine}; +use dynamo_global_router::{ + DiscoveryPoolCatalog, Pool, PoolCatalog, RoundRobinSelector, StaticPoolCatalog, +}; +use dynamo_global_router_shared::{PoolIdDeriver, PoolKey, PoolLocation, V1PoolIdDeriver}; +use dynamo_llm::http::service::service_v2::HttpService; +use tokio_util::sync::CancellationToken; +use url::Url; + +#[derive(Parser)] +#[command(about = "Serve the OpenAI API and forward each request to a pool frontend")] +#[command(group(ArgGroup::new("pool_source").required(true).args(["discovery_url", "pools"])))] +struct Args { + /// Address to serve HTTP on. + #[arg(long, default_value = "0.0.0.0")] + host: String, + #[arg(long, default_value_t = 8000)] + port: u16, + /// Root of the Global Discovery Service to read pools from, for example + /// `http://global-discovery:8080/`. + #[arg(long)] + discovery_url: Option, + /// How often to read pools from the Global Discovery Service. A read + /// that takes longer is abandoned. + #[arg(long, default_value_t = 5)] + discovery_interval_secs: u64, + /// A pool as `//=@`, in + /// place of `--discovery-url`. Separate several frontend URLs with commas, + /// and repeat the flag for each pool. The site is also the pool's region. + #[arg(long = "pool", value_parser = parse_pool)] + pools: Vec, + /// How long to wait for a pool frontend to accept a connection. + #[arg(long, default_value_t = 5)] + connect_timeout_secs: u64, +} + +fn parse_pool(value: &str) -> Result { + let shape = + || format!("expected //=@[,...], got {value}"); + let (key, served) = value.split_once('=').ok_or_else(shape)?; + let [site, namespace, dgd] = key.split('/').collect::>()[..] else { + return Err(shape()); + }; + let (model, frontends) = served.split_once('@').ok_or_else(shape)?; + let key = PoolKey::new(site, namespace, dgd).map_err(|error| error.to_string())?; + let frontends = frontends + .split(',') + .map(|url| { + url.parse::() + .map_err(|error| format!("{url}: {error}")) + }) + .collect::, _>>()?; + Ok(Pool { + id: V1PoolIdDeriver.derive(&key), + model: model.to_owned(), + location: PoolLocation { + region: site.to_owned(), + availability_zone: None, + cluster: None, + datacenter: None, + }, + frontends, + }) +} + +#[tokio::main] +async fn main() -> Result<()> { + dynamo_runtime::logging::init(); + let args = Args::parse(); + let mut ids = HashSet::new(); + if let Some(repeated) = args.pools.iter().find(|pool| !ids.insert(&pool.id)) { + bail!("--pool {} is given twice", repeated.id); + } + + let shutdown = CancellationToken::new(); + let signal = shutdown.clone(); + tokio::spawn(async move { + if let Err(error) = shutdown_signal().await { + tracing::warn!(%error, "shutdown signal handler failed"); + } + signal.cancel(); + }); + + let catalog: Box = match args.discovery_url { + Some(url) => { + let interval = Duration::from_secs(args.discovery_interval_secs); + let client = GlobalDiscoveryClient::new(url, interval) + .context("build the Global Discovery client")?; + Box::new(DiscoveryPoolCatalog::start(client, interval, shutdown.child_token()).await) + } + None => Box::new(StaticPoolCatalog::new(args.pools)), + }; + let service = HttpService::builder() + .host(args.host.clone()) + .port(args.port) + .cancel_token(Some(shutdown.clone())) + .build() + .context("build the HTTP service")?; + let http = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(args.connect_timeout_secs)) + .build() + .context("build the pool HTTP client")?; + let engine = Arc::new(PoolForwardingEngine::new( + catalog.subscribe(), + Arc::new(RoundRobinSelector::new()), + http, + )); + + // Register the models before serving, so the first request finds them. + let pools = catalog.subscribe(); + let mut models = ModelSync::new(service.clone(), engine); + models.apply(&pools.borrow()); + tokio::spawn(models.follow(pools, shutdown.child_token())); + + tracing::info!(host = %args.host, port = args.port, "Global Frontend listening"); + service + .run(shutdown) + .await + .context("serve the Global Frontend") +} + +async fn shutdown_signal() -> Result<()> { + #[cfg(unix)] + { + use tokio::signal::unix::{SignalKind, signal}; + + let mut terminate = signal(SignalKind::terminate()).context("register SIGTERM handler")?; + tokio::select! { + result = tokio::signal::ctrl_c() => result.context("wait for Ctrl-C")?, + _ = terminate.recv() => {} + } + } + #[cfg(not(unix))] + tokio::signal::ctrl_c().await.context("wait for Ctrl-C")?; + Ok(()) +} diff --git a/lib/global-ingress/src/model_sync.rs b/lib/global-ingress/src/model_sync.rs new file mode 100644 index 000000000000..3b92f20fc631 --- /dev/null +++ b/lib/global-ingress/src/model_sync.rs @@ -0,0 +1,122 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeSet; +use std::sync::Arc; + +use dynamo_global_router::PoolSnapshot; +use dynamo_llm::http::service::service_v2::HttpService; +use dynamo_llm::model_type::ModelType; +use dynamo_llm::types::openai::chat_completions::OpenAIChatCompletionsStreamingEngine; +use dynamo_llm::types::openai::completions::OpenAICompletionsStreamingEngine; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; + +use crate::forward::PoolForwardingEngine; + +/// A model registered by hand needs a card checksum. Every model here is the +/// same forwarding engine, so they share one. +const FORWARDING_CHECKSUM: &str = "global-ingress-forwarding"; + +/// Registers the forwarding engine under every model the pool catalog +/// serves, and enables the HTTP endpoints those models need. +/// +/// The Global Frontend is the only writer of its model manager. +pub struct ModelSync { + service: HttpService, + engine: Arc, + served: BTreeSet, + /// Whether the endpoints were last enabled or disabled. `None` until the + /// first snapshot, so that one always sets them. + endpoints_enabled: Option, +} + +impl ModelSync { + pub fn new(service: HttpService, engine: Arc) -> Self { + Self { + service, + engine, + served: BTreeSet::new(), + endpoints_enabled: None, + } + } + + /// Serve exactly the models in `snapshot`. + pub fn apply(&mut self, snapshot: &PoolSnapshot) { + let wanted = snapshot.models(); + let manager = self.service.model_manager(); + + let gone = self + .served + .iter() + .filter(|model| !wanted.contains(model.as_str())) + .cloned() + .collect::>(); + for model in gone { + manager.remove_model(&model); + self.served.remove(&model); + tracing::info!(model, "no pool serves the model any more"); + } + + for model in wanted { + if self.served.contains(model) { + continue; + } + match self.register(model) { + Ok(()) => { + self.served.insert(model.to_owned()); + tracing::info!(model, "serving model"); + } + Err(error) => { + // Drop a half-registered model so it is never served + // through only some of its APIs. + manager.remove_model(model); + tracing::error!(model, %error, "cannot serve model"); + } + } + } + + let enabled = !self.served.is_empty(); + if self.endpoints_enabled == Some(enabled) { + return; + } + self.endpoints_enabled = Some(enabled); + let endpoints = (ModelType::Chat | ModelType::Completions) + .as_endpoint_types_with_anthropic(self.service.anthropic_api_enabled()); + for endpoint in endpoints { + if let Err(error) = self.service.enable_model_endpoint(endpoint, enabled) { + tracing::error!(?endpoint, %error, "cannot update endpoint"); + } + } + } + + /// Apply every snapshot the catalog publishes until `cancel` fires or the + /// catalog is dropped. Models stay registered after either. + pub async fn follow( + mut self, + mut pools: watch::Receiver>, + cancel: CancellationToken, + ) { + loop { + let snapshot = pools.borrow_and_update().clone(); + self.apply(&snapshot); + tokio::select! { + _ = cancel.cancelled() => return, + changed = pools.changed() => { + if changed.is_err() { + return; + } + } + } + } + } + + fn register(&self, model: &str) -> anyhow::Result<()> { + let manager = self.service.model_manager(); + let chat: OpenAIChatCompletionsStreamingEngine = self.engine.clone(); + manager.add_chat_completions_model(model, FORWARDING_CHECKSUM, chat)?; + let completions: OpenAICompletionsStreamingEngine = self.engine.clone(); + manager.add_completions_model(model, FORWARDING_CHECKSUM, completions)?; + Ok(()) + } +} diff --git a/lib/global-ingress/tests/common/mod.rs b/lib/global-ingress/tests/common/mod.rs new file mode 100644 index 000000000000..6ac9b3f4c579 --- /dev/null +++ b/lib/global-ingress/tests/common/mod.rs @@ -0,0 +1,179 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Servers and requests shared by the Global Frontend tests. + +// Each test binary includes this module and uses only some of it. +#![allow(dead_code)] + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::Router; +use axum::http::HeaderMap; +use axum::routing; +use dynamo_global_ingress::{ModelSync, PoolForwardingEngine}; +use dynamo_global_router::{Pool, PoolSnapshot, RoundRobinSelector}; +use dynamo_global_router_shared::{PoolIdDeriver, PoolKey, PoolLocation, V1PoolIdDeriver}; +use dynamo_llm::http::service::service_v2::HttpService; +use reqwest::StatusCode; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; +use url::Url; + +pub const MODEL: &str = "echo-model"; + +/// A running HTTP server. Dropping it stops the server. +pub struct Server { + pub url: Url, + pub cancel: CancellationToken, +} + +impl Drop for Server { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +pub fn client() -> reqwest::Client { + reqwest::Client::builder().no_proxy().build().unwrap() +} + +pub async fn bind() -> (TcpListener, Url) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + ( + listener, + format!("http://127.0.0.1:{port}/").parse().unwrap(), + ) +} + +pub fn http_service(url: &Url) -> HttpService { + HttpService::builder() + .host("127.0.0.1") + .port(url.port().unwrap()) + .build() + .unwrap() +} + +/// Serves `service` and waits until it is healthy. +pub async fn serve(service: HttpService, listener: TcpListener, url: Url) -> Server { + let cancel = CancellationToken::new(); + service.spawn_with_listener(cancel.clone(), listener).await; + let health = url.join("health").unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while !client() + .get(health.clone()) + .send() + .await + .is_ok_and(|response| response.status().is_success()) + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("server did not become healthy"); + Server { url, cancel } +} + +/// A Global Frontend that routes round robin over the pools `pools` publishes. +pub async fn global_frontend(pools: watch::Receiver>) -> Server { + let (listener, url) = bind().await; + let service = http_service(&url); + let engine = Arc::new(PoolForwardingEngine::new( + pools.clone(), + Arc::new(RoundRobinSelector::new()), + client(), + )); + let mut models = ModelSync::new(service.clone(), engine); + models.apply(&pools.borrow()); + let server = serve(service, listener, url).await; + tokio::spawn(models.follow(pools, server.cancel.child_token())); + server +} + +pub fn pool(dgd: &str, model: &str, frontend: &Url) -> Pool { + Pool { + id: V1PoolIdDeriver.derive(&PoolKey::new("ohio", "dynamo", dgd).unwrap()), + model: model.to_owned(), + location: PoolLocation { + region: "us-east-2".to_owned(), + availability_zone: None, + cluster: None, + datacenter: None, + }, + frontends: vec![frontend.clone()], + } +} + +/// A pool frontend that records the headers of each chat request and answers +/// with a fixed SSE body. +pub async fn scripted_pool_frontend(sse: &'static str) -> (Server, Arc>>) { + let seen = Arc::new(Mutex::new(Vec::new())); + let recorder = seen.clone(); + let app = Router::new().route( + "/v1/chat/completions", + routing::post(move |headers: HeaderMap| { + recorder.lock().unwrap().push(headers); + async move { ([("content-type", "text/event-stream")], sse) } + }), + ); + let (listener, url) = bind().await; + let cancel = CancellationToken::new(); + let shutdown = cancel.clone(); + tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(shutdown.cancelled_owned()) + .await + }); + (Server { url, cancel }, seen) +} + +/// One chat completion chunk carrying `content`. +pub fn chat_chunk(content: &str) -> Value { + json!({ + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": MODEL, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}}], + }) +} + +pub fn chat(model: &str, prompt: &str, stream: bool) -> Value { + json!({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + }) +} + +pub async fn post(base: &Url, route: &str, body: &Value) -> reqwest::Response { + client() + .post(base.join(route).unwrap()) + .json(body) + .send() + .await + .unwrap() +} + +/// The JSON body of a successful response. On failure, the panic shows the +/// body, which carries the error. +pub async fn ok_json(response: reqwest::Response) -> Value { + let status = response.status(); + let body = response.text().await.unwrap(); + assert_eq!(status, StatusCode::OK, "{body}"); + serde_json::from_str(&body).unwrap() +} + +/// The `data:` payloads of an SSE body, up to `[DONE]`. +pub fn sse_data(body: &str) -> Vec { + body.lines() + .filter_map(|line| line.strip_prefix("data:")) + .map(str::trim) + .take_while(|data| *data != "[DONE]") + .map(|data| serde_json::from_str(data).unwrap()) + .collect() +} diff --git a/lib/global-ingress/tests/forwarding.rs b/lib/global-ingress/tests/forwarding.rs new file mode 100644 index 000000000000..4a05b4241515 --- /dev/null +++ b/lib/global-ingress/tests/forwarding.rs @@ -0,0 +1,273 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! A client calls the Global Frontend, which forwards to a pool frontend. +//! Most pools here are Dynamo's own HTTP service running the echo engine, +//! which streams the prompt back one character at a time. + +mod common; + +use std::sync::Arc; +use std::time::Duration; + +use common::{ + MODEL, Server, bind, chat, chat_chunk, client, http_service, ok_json, pool, post, + scripted_pool_frontend, serve, sse_data, +}; +use dynamo_global_router::{Pool, PoolSnapshot}; +use dynamo_llm::endpoint_type::EndpointType; +use dynamo_llm::engines::{StreamingEngineAdapter, make_echo_engine}; +use reqwest::StatusCode; +use serde_json::{Value, json}; +use tokio::sync::watch; + +struct GlobalFrontend { + server: Server, + pools: watch::Sender>, +} + +async fn echo_pool_frontend() -> Server { + let (listener, url) = bind().await; + let service = http_service(&url); + let engine = Arc::new(StreamingEngineAdapter::new(make_echo_engine())); + let manager = service.model_manager(); + manager + .add_chat_completions_model(MODEL, "echo", engine.clone()) + .unwrap(); + manager + .add_completions_model(MODEL, "echo", engine) + .unwrap(); + for endpoint in [EndpointType::Chat, EndpointType::Completion] { + service.enable_model_endpoint(endpoint, true).unwrap(); + } + serve(service, listener, url).await +} + +/// A Global Frontend over a catalog the test can change. +async fn global_frontend(pools: Vec) -> GlobalFrontend { + let (sender, receiver) = watch::channel(Arc::new(PoolSnapshot::new(pools))); + GlobalFrontend { + server: common::global_frontend(receiver).await, + pools: sender, + } +} + +#[tokio::test] +async fn streams_chat_completions_from_the_pool() { + let pool_frontend = echo_pool_frontend().await; + let global = global_frontend(vec![pool("blue", MODEL, &pool_frontend.url)]).await; + + let response = post( + &global.server.url, + "v1/chat/completions", + &chat(MODEL, "hello", true), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = response.text().await.unwrap(); + + let events = sse_data(&body); + assert!( + events.iter().all(|event| event.get("error").is_none()), + "{body}" + ); + let text = events + .iter() + .filter_map(|chunk| chunk["choices"][0]["delta"]["content"].as_str()) + .collect::(); + assert_eq!(text, "hello"); + assert!(body.trim_end().ends_with("data: [DONE]"), "{body}"); +} + +#[tokio::test] +async fn folds_the_pool_stream_for_non_streaming_clients() { + let pool_frontend = echo_pool_frontend().await; + let global = global_frontend(vec![pool("blue", MODEL, &pool_frontend.url)]).await; + + let response = post( + &global.server.url, + "v1/chat/completions", + &chat(MODEL, "hello", false), + ) + .await; + let body = ok_json(response).await; + assert_eq!(body["choices"][0]["message"]["content"], "hello", "{body}"); +} + +#[tokio::test] +async fn forwards_completions() { + let pool_frontend = echo_pool_frontend().await; + let global = global_frontend(vec![pool("blue", MODEL, &pool_frontend.url)]).await; + + let request = json!({"model": MODEL, "prompt": "hi there", "max_tokens": 16}); + let response = post(&global.server.url, "v1/completions", &request).await; + let body = ok_json(response).await; + assert_eq!(body["choices"][0]["text"], "hi there", "{body}"); +} + +/// The Responses API translates to chat completions at the Global Frontend, +/// so the pool sees a chat request. +#[tokio::test] +async fn serves_the_responses_api_through_chat() { + let pool_frontend = echo_pool_frontend().await; + let global = global_frontend(vec![pool("blue", MODEL, &pool_frontend.url)]).await; + + let request = json!({"model": MODEL, "input": "hello"}); + let response = post(&global.server.url, "v1/responses", &request).await; + let body = ok_json(response).await; + let text = body["output"] + .as_array() + .unwrap() + .iter() + .flat_map(|item| item["content"].as_array().cloned().unwrap_or_default()) + .filter_map(|part| part["text"].as_str().map(str::to_owned)) + .collect::(); + assert_eq!(text, "hello", "{body}"); +} + +#[tokio::test] +async fn models_follow_the_catalog() { + let pool_frontend = echo_pool_frontend().await; + let global = global_frontend(Vec::new()).await; + let request = chat(MODEL, "hi", false); + assert!(listed_models(&global).await.is_empty()); + assert_eq!( + post(&global.server.url, "v1/chat/completions", &request) + .await + .status(), + StatusCode::NOT_FOUND + ); + + global + .pools + .send_replace(Arc::new(PoolSnapshot::new(vec![pool( + "blue", + MODEL, + &pool_frontend.url, + )]))); + wait_for_models(&global, &[MODEL]).await; + assert_eq!( + post(&global.server.url, "v1/chat/completions", &request) + .await + .status(), + StatusCode::OK + ); + + global.pools.send_replace(Arc::new(PoolSnapshot::default())); + wait_for_models(&global, &[]).await; + assert_eq!( + post(&global.server.url, "v1/chat/completions", &request) + .await + .status(), + StatusCode::NOT_FOUND + ); +} + +async fn listed_models(global: &GlobalFrontend) -> Vec { + let response = client() + .get(global.server.url.join("v1/models").unwrap()) + .send() + .await + .unwrap(); + let body: Value = response.json().await.unwrap(); + body["data"] + .as_array() + .unwrap() + .iter() + .map(|model| model["id"].as_str().unwrap().to_owned()) + .collect() +} + +async fn wait_for_models(global: &GlobalFrontend, expected: &[&str]) { + tokio::time::timeout(Duration::from_secs(5), async { + while listed_models(global).await != expected { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the Global Frontend did not pick up the catalog change"); +} + +#[tokio::test] +async fn pool_errors_keep_their_status() { + // The catalog says the pool serves a model its frontend does not have. + let pool_frontend = echo_pool_frontend().await; + let global = global_frontend(vec![pool("blue", "ghost", &pool_frontend.url)]).await; + + let response = post( + &global.server.url, + "v1/chat/completions", + &chat("ghost", "hi", true), + ) + .await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn unreachable_pools_are_unavailable() { + let (listener, closed) = bind().await; + drop(listener); + let global = global_frontend(vec![pool("blue", MODEL, &closed)]).await; + + let response = post( + &global.server.url, + "v1/chat/completions", + &chat(MODEL, "hi", false), + ) + .await; + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); +} + +#[tokio::test] +async fn forwards_the_session_and_skips_pool_keep_alives() { + let sse = format!( + ": keep-alive\n\ndata: {}\n\n: keep-alive\n\ndata: {}\n\ndata: [DONE]\n\n", + chat_chunk("he"), + chat_chunk("llo") + ); + let (pool_frontend, seen) = scripted_pool_frontend(sse.leak()).await; + let global = global_frontend(vec![pool("blue", MODEL, &pool_frontend.url)]).await; + + let response = client() + .post(global.server.url.join("v1/chat/completions").unwrap()) + .header("x-dynamo-session-id", "session-123") + .json(&chat(MODEL, "hi", true)) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let text = sse_data(&response.text().await.unwrap()) + .iter() + .filter_map(|chunk| chunk["choices"][0]["delta"]["content"].as_str()) + .collect::(); + assert_eq!(text, "hello"); + + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0]["x-dynamo-session-id"], "session-123"); +} + +#[tokio::test] +async fn a_pool_error_mid_stream_ends_the_client_stream_with_an_error() { + let sse = format!( + "data: {}\n\ndata: {}\n\ndata: [DONE]\n\n", + chat_chunk("he"), + json!({"error": {"message": "worker died", "type": "internal_error", "code": 500}}) + ); + let (pool_frontend, _) = scripted_pool_frontend(sse.leak()).await; + let global = global_frontend(vec![pool("blue", MODEL, &pool_frontend.url)]).await; + + let response = post( + &global.server.url, + "v1/chat/completions", + &chat(MODEL, "hi", true), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let events = sse_data(&response.text().await.unwrap()); + assert_eq!(events[0]["choices"][0]["delta"]["content"], "he"); + assert!( + events.last().unwrap().get("error").is_some(), + "the client must see an error frame: {events:?}" + ); +} diff --git a/lib/global-ingress/tests/round_robin.rs b/lib/global-ingress/tests/round_robin.rs new file mode 100644 index 000000000000..a2da58024c25 --- /dev/null +++ b/lib/global-ingress/tests/round_robin.rs @@ -0,0 +1,133 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The Global Frontend reads pools from the real Global Discovery Service and +//! sends requests to them in turn. + +mod common; + +use std::collections::BTreeMap; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::http::HeaderMap; +use common::{MODEL, Server, chat, chat_chunk, post, scripted_pool_frontend}; +use dynamo_global_discovery_client::{ + GlobalDiscoveryClient, PoolKey, PoolLocation, PoolRegistration, +}; +use dynamo_global_discovery_service::backends::kubernetes::TrustedNetworkAuthenticator; +use dynamo_global_discovery_service::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; +use dynamo_global_discovery_service::pool_discovery::LeasePolicy; +use dynamo_global_discovery_service::service::{self, DiscoveryStore}; +use dynamo_global_router::{DiscoveryPoolCatalog, PoolCatalog}; +use reqwest::StatusCode; +use tokio_util::sync::CancellationToken; + +/// A Global Discovery Service with an in-memory store. +async fn discovery_service() -> (Server, GlobalDiscoveryClient) { + let lease = LeasePolicy { + heartbeat_interval: Duration::from_secs(5), + lease_duration: Duration::from_secs(15), + }; + let pools = Arc::new(MemoryPoolStore::new(lease).unwrap()); + let store = DiscoveryStore { + pools: pools.clone(), + routers: Arc::new(StaticReplicaDirectory::new(Vec::new())), + registrar: Some(pools), + authenticator: Arc::new(TrustedNetworkAuthenticator), + }; + let (listener, url) = common::bind().await; + let cancel = CancellationToken::new(); + tokio::spawn(service::serve(listener, store, cancel.clone())); + let client = GlobalDiscoveryClient::new(url.clone(), Duration::from_secs(2)).unwrap(); + (Server { url, cancel }, client) +} + +/// A pool frontend registered with the discovery service. +struct RegisteredPool { + _server: Server, + requests: Arc>>, +} + +async fn register_pool(discovery: &GlobalDiscoveryClient, dgd: &str) -> RegisteredPool { + let sse = format!("data: {}\n\ndata: [DONE]\n\n", chat_chunk(dgd)); + let (server, requests) = scripted_pool_frontend(sse.leak()).await; + discovery + .register_pool(&PoolRegistration { + key: PoolKey { + site_id: "aws-ohio".into(), + namespace: "dynamo".into(), + dgd_name: dgd.into(), + }, + location: PoolLocation { + region: "us-east-2".into(), + availability_zone: None, + cluster: Some("aws-ohio".into()), + datacenter: None, + }, + model: MODEL.into(), + capabilities: vec!["aggregated".into()], + frontends: vec![server.url.clone()], + capacity: Vec::new(), + }) + .await + .unwrap(); + RegisteredPool { + _server: server, + requests, + } +} + +/// Sends `count` chat requests through the Global Frontend. +async fn send_chats(global: &Server, count: usize) { + for _ in 0..count { + let response = post(&global.url, "v1/chat/completions", &chat(MODEL, "hi", true)).await; + assert_eq!(response.status(), StatusCode::OK); + response.text().await.unwrap(); + } +} + +/// How many requests each pool has received. +fn request_counts(pools: &[(&str, &RegisteredPool)]) -> BTreeMap { + pools + .iter() + .map(|(name, pool)| (name.to_string(), pool.requests.lock().unwrap().len())) + .collect() +} + +#[tokio::test] +async fn rotates_across_the_pools_discovery_lists() { + let (_discovery_server, discovery) = discovery_service().await; + let blue = register_pool(&discovery, "blue").await; + let green = register_pool(&discovery, "green").await; + + let cancel = CancellationToken::new(); + let catalog = + DiscoveryPoolCatalog::start(discovery.clone(), Duration::from_millis(50), cancel.clone()) + .await; + let global = common::global_frontend(catalog.subscribe()).await; + + send_chats(&global, 4).await; + assert_eq!( + request_counts(&[("blue", &blue), ("green", &green)]), + BTreeMap::from([("blue".into(), 2), ("green".into(), 2)]) + ); + + // A pool registered later joins the rotation without a restart. + let red = register_pool(&discovery, "red").await; + let mut pools = catalog.subscribe(); + tokio::time::timeout( + Duration::from_secs(5), + pools.wait_for(|snapshot| snapshot.pools().len() == 3), + ) + .await + .expect("the catalog did not pick up the new pool") + .unwrap(); + + send_chats(&global, 3).await; + assert_eq!( + request_counts(&[("blue", &blue), ("green", &green), ("red", &red)]), + BTreeMap::from([("blue".into(), 3), ("green".into(), 3), ("red".into(), 1)]) + ); + cancel.cancel(); +} diff --git a/lib/global-router/Cargo.toml b/lib/global-router/Cargo.toml new file mode 100644 index 000000000000..7e3fa728f351 --- /dev/null +++ b/lib/global-router/Cargo.toml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "dynamo-global-router" +description = "Pool catalog and pool selection for the Global Router" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true + +[dependencies] +dynamo-global-discovery-client = { path = "../global-discovery-client" } +dynamo-global-router-shared = { path = "../global-router-shared" } +dynamo-runtime = { workspace = true } + +dashmap = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tracing = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +axum = { workspace = true } +serde_json = { workspace = true } diff --git a/lib/global-router/src/catalog.rs b/lib/global-router/src/catalog.rs new file mode 100644 index 000000000000..e80d9c2fced9 --- /dev/null +++ b/lib/global-router/src/catalog.rs @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeSet; +use std::sync::Arc; + +use dynamo_global_router_shared::{PoolId, PoolLocation}; +use tokio::sync::watch; +use url::Url; + +/// One pool as the Global Router sees it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Pool { + pub id: PoolId, + /// The model the pool serves, as clients name it in requests. + pub model: String, + pub location: PoolLocation, + /// Frontends that accept forwarded requests. Each is the root of an + /// OpenAI-compatible HTTP API, for example `http://10.0.0.5:8000/`. + pub frontends: Vec, +} + +/// Every pool a catalog knows about at one point in time. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct PoolSnapshot { + pools: Vec, +} + +impl PoolSnapshot { + pub fn new(pools: Vec) -> Self { + Self { pools } + } + + pub fn pools(&self) -> &[Pool] { + &self.pools + } + + /// Pools serving `model`, in catalog order. + pub fn pools_for<'a>(&'a self, model: &'a str) -> impl Iterator { + self.pools.iter().filter(move |pool| pool.model == model) + } + + /// Models served by at least one pool. + pub fn models(&self) -> BTreeSet<&str> { + self.pools.iter().map(|pool| pool.model.as_str()).collect() + } +} + +/// A source of pools. An implementation publishes a new snapshot whenever a +/// pool joins, leaves, or changes its frontends. +pub trait PoolCatalog: Send + Sync { + /// The receiver holds the current snapshot, and its `changed()` waits for + /// the next one. It returns an error once the catalog is dropped. + fn subscribe(&self) -> watch::Receiver>; +} + +/// A catalog that never changes, for pools given on the command line and +/// for tests. +pub struct StaticPoolCatalog { + snapshot: watch::Sender>, +} + +impl StaticPoolCatalog { + pub fn new(pools: Vec) -> Self { + let (snapshot, _) = watch::channel(Arc::new(PoolSnapshot::new(pools))); + Self { snapshot } + } +} + +impl PoolCatalog for StaticPoolCatalog { + fn subscribe(&self) -> watch::Receiver> { + self.snapshot.subscribe() + } +} + +#[cfg(test)] +pub(crate) mod tests { + use dynamo_global_router_shared::{PoolIdDeriver, PoolKey, V1PoolIdDeriver}; + + use super::*; + + pub(crate) fn pool(dgd: &str, model: &str, frontends: &[&str]) -> Pool { + Pool { + id: V1PoolIdDeriver.derive(&PoolKey::new("ohio", "dynamo", dgd).unwrap()), + model: model.to_owned(), + location: PoolLocation { + region: "us-east-2".to_owned(), + availability_zone: None, + cluster: None, + datacenter: None, + }, + frontends: frontends.iter().map(|url| url.parse().unwrap()).collect(), + } + } + + #[test] + fn snapshot_groups_pools_by_model() { + let snapshot = PoolSnapshot::new(vec![ + pool("blue", "qwen", &["http://blue:8000"]), + pool("green", "qwen", &["http://green:8000"]), + pool("llama", "llama", &["http://llama:8000"]), + ]); + + assert_eq!( + snapshot.models().into_iter().collect::>(), + ["llama", "qwen"] + ); + let qwen = snapshot + .pools_for("qwen") + .map(|pool| pool.frontends[0].host_str()); + assert_eq!(qwen.collect::>(), [Some("blue"), Some("green")]); + assert_eq!(snapshot.pools_for("missing").count(), 0); + } + + #[test] + fn static_catalog_serves_its_pools_to_every_subscriber() { + let catalog = StaticPoolCatalog::new(vec![pool("blue", "qwen", &["http://blue:8000"])]); + let first = catalog.subscribe(); + let second = catalog.subscribe(); + + assert_eq!(first.borrow().pools().len(), 1); + assert_eq!(*first.borrow(), *second.borrow()); + } +} diff --git a/lib/global-router/src/discovery.rs b/lib/global-router/src/discovery.rs new file mode 100644 index 000000000000..096eaa5d8222 --- /dev/null +++ b/lib/global-router/src/discovery.rs @@ -0,0 +1,256 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::time::Duration; + +use dynamo_global_discovery_client::{Completeness, GlobalDiscoveryClient, PoolListing}; +use dynamo_global_router_shared::{PoolIdDeriver, PoolKey, PoolLocation, V1PoolIdDeriver}; +use tokio::sync::watch; +use tokio::time::MissedTickBehavior; +use tokio_util::sync::CancellationToken; + +use crate::catalog::{Pool, PoolCatalog, PoolSnapshot}; + +/// The pools the Global Discovery Service lists, polled on an interval. +/// +/// A failed poll keeps the last snapshot, so a discovery outage does not stop +/// routing to the pools already known. +pub struct DiscoveryPoolCatalog { + snapshot: watch::Sender>, +} + +impl DiscoveryPoolCatalog { + /// Polls once before returning, so the catalog starts with the service's + /// pools when it is reachable, then polls every `interval` until `cancel` + /// fires. + pub async fn start( + client: GlobalDiscoveryClient, + interval: Duration, + cancel: CancellationToken, + ) -> Self { + let (snapshot, _) = watch::channel(Arc::new(PoolSnapshot::default())); + refresh(&client, &snapshot).await; + let publisher = snapshot.clone(); + tokio::spawn(async move { + let mut ticks = tokio::time::interval(interval); + ticks.set_missed_tick_behavior(MissedTickBehavior::Delay); + // The first tick fires at once, and `start` has just polled. + ticks.tick().await; + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = ticks.tick() => refresh(&client, &publisher).await, + } + } + }); + Self { snapshot } + } +} + +impl PoolCatalog for DiscoveryPoolCatalog { + fn subscribe(&self) -> watch::Receiver> { + self.snapshot.subscribe() + } +} + +async fn refresh(client: &GlobalDiscoveryClient, snapshot: &watch::Sender>) { + let listing = match client.pools().await { + Ok(listing) => listing, + Err(error) => { + tracing::warn!(%error, "cannot list pools; keeping the last known pools"); + return; + } + }; + if listing.completeness == Completeness::Estimated { + tracing::debug!("the discovery service's pool list may be partial"); + } + let pools = PoolSnapshot::new(routable_pools(listing.items)); + snapshot.send_if_modified(|current| { + if **current == pools { + return false; + } + tracing::info!(pools = pools.pools().len(), "pool catalog changed"); + *current = Arc::new(pools); + true + }); +} + +/// The listed pools a router can route to. They are sorted by ID, so an +/// unchanged list compares equal and round robin keeps its place. +fn routable_pools(listings: Vec) -> Vec { + let mut pools = listings + .into_iter() + .filter_map(|listing| { + let key = &listing.key; + let Some(model) = listing.model else { + tracing::debug!(?key, "skipping a pool with no model"); + return None; + }; + let key = match PoolKey::new(&key.site_id, &key.namespace, &key.dgd_name) { + Ok(key) => key, + Err(error) => { + tracing::warn!(?key, %error, "skipping a pool with an invalid key"); + return None; + } + }; + let location = listing.location; + Some(Pool { + id: V1PoolIdDeriver.derive(&key), + model, + location: PoolLocation { + region: location.region, + availability_zone: location.availability_zone, + cluster: location.cluster, + datacenter: location.datacenter, + }, + frontends: listing.frontends, + }) + }) + .collect::>(); + pools.sort_by(|left, right| left.id.as_str().cmp(right.id.as_str())); + pools +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use axum::http::StatusCode; + use axum::routing::get; + use dynamo_global_discovery_client as wire; + use tokio::net::TcpListener; + use url::Url; + + use super::*; + + fn listing(dgd: &str, model: Option<&str>) -> PoolListing { + PoolListing { + key: wire::PoolKey { + site_id: "aws-west".into(), + namespace: "dynamo".into(), + dgd_name: dgd.into(), + }, + location: wire::PoolLocation { + region: "us-west-2".into(), + availability_zone: None, + cluster: Some("aws-west".into()), + datacenter: None, + }, + model: model.map(str::to_owned), + capabilities: Vec::new(), + frontends: vec![format!("http://{dgd}:8000/").parse().unwrap()], + capacity: Vec::new(), + relay_identity: None, + } + } + + #[test] + fn keeps_pools_a_router_can_route_to_in_id_order() { + let mut bad_key = listing("blank", Some("qwen")); + bad_key.key.namespace = String::new(); + let pools = routable_pools(vec![ + listing("green", Some("qwen")), + listing("unnamed", None), + bad_key, + listing("blue", Some("qwen")), + ]); + + let hosts = pools + .iter() + .map(|pool| pool.frontends[0].host_str().unwrap()) + .collect::>(); + assert_eq!(hosts.len(), 2); + assert!(hosts.contains(&"blue") && hosts.contains(&"green")); + assert!(pools[0].id.as_str() < pools[1].id.as_str()); + assert_eq!(pools[0].location.cluster.as_deref(), Some("aws-west")); + let reordered = routable_pools(vec![ + listing("blue", Some("qwen")), + listing("green", Some("qwen")), + ]); + assert_eq!(pools, reordered); + } + + /// A discovery service whose `GET /pools` answer the test can change. + struct FakeService { + response: Arc>, + client: GlobalDiscoveryClient, + } + + impl FakeService { + async fn start() -> Self { + let response = Arc::new(Mutex::new((StatusCode::OK, pools_body(&[])))); + let shared = response.clone(); + let app = axum::Router::new().route( + "/pools", + get(move || { + let (status, body) = shared.lock().unwrap().clone(); + async move { (status, body) } + }), + ); + 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() }); + let base_url = Url::parse(&format!("http://{address}/")).unwrap(); + let client = GlobalDiscoveryClient::new(base_url, Duration::from_secs(2)).unwrap(); + Self { response, client } + } + + fn answer(&self, status: StatusCode, body: String) { + *self.response.lock().unwrap() = (status, body); + } + } + + fn pools_body(dgds: &[&str]) -> String { + let pools = dgds + .iter() + .map(|dgd| { + serde_json::json!({ + "key": {"site_id": "aws-west", "namespace": "dynamo", "dgd_name": dgd}, + "location": {"region": "us-west-2", "availability_zone": null, "cluster": null, "datacenter": null}, + "model": "qwen", + "frontends": [format!("http://{dgd}:8000/")], + }) + }) + .collect::>(); + serde_json::json!({"complete": true, "pools": pools}).to_string() + } + + async fn next_snapshot(pools: &mut watch::Receiver>) -> Arc { + tokio::time::timeout(Duration::from_secs(5), pools.changed()) + .await + .expect("the catalog did not change") + .unwrap(); + pools.borrow_and_update().clone() + } + + #[tokio::test] + async fn follows_the_service_and_keeps_pools_through_an_outage() { + let service = FakeService::start().await; + service.answer(StatusCode::OK, pools_body(&["blue"])); + let cancel = CancellationToken::new(); + let catalog = DiscoveryPoolCatalog::start( + service.client.clone(), + Duration::from_millis(20), + cancel.clone(), + ) + .await; + let mut pools = catalog.subscribe(); + assert_eq!(pools.borrow_and_update().pools().len(), 1); + + service.answer(StatusCode::OK, pools_body(&["blue", "green"])); + assert_eq!(next_snapshot(&mut pools).await.pools().len(), 2); + + // An outage keeps the pools, and an unchanged list publishes nothing. + service.answer(StatusCode::SERVICE_UNAVAILABLE, String::new()); + tokio::time::sleep(Duration::from_millis(100)).await; + service.answer(StatusCode::OK, pools_body(&["green", "blue"])); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(!pools.has_changed().unwrap()); + assert_eq!(pools.borrow().pools().len(), 2); + + service.answer(StatusCode::OK, pools_body(&["green"])); + assert_eq!(next_snapshot(&mut pools).await.pools().len(), 1); + cancel.cancel(); + } +} diff --git a/lib/global-router/src/lib.rs b/lib/global-router/src/lib.rs new file mode 100644 index 000000000000..40027b6a9dbf --- /dev/null +++ b/lib/global-router/src/lib.rs @@ -0,0 +1,18 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Pool catalog and pool selection for the Global Router. +//! +//! The Global Frontend reads pools from a [`PoolCatalog`] and asks a +//! [`PoolSelector`] which pool serves each request. Neither knows about HTTP +//! or where pools are discovered, so a catalog backed by configuration, the +//! Global Discovery Service, or Pool Relay state plugs in behind the same +//! trait. + +mod catalog; +mod discovery; +mod selector; + +pub use catalog::{Pool, PoolCatalog, PoolSnapshot, StaticPoolCatalog}; +pub use discovery::DiscoveryPoolCatalog; +pub use selector::{PoolSelector, PoolTarget, RoundRobinSelector, RoutingRequest}; diff --git a/lib/global-router/src/selector.rs b/lib/global-router/src/selector.rs new file mode 100644 index 000000000000..fe48134ec2dd --- /dev/null +++ b/lib/global-router/src/selector.rs @@ -0,0 +1,166 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::borrow::Borrow; +use std::hash::Hash; + +use dashmap::DashMap; +use dynamo_global_router_shared::PoolId; +use dynamo_runtime::pipeline::BuiltinRoutePicker; +use url::Url; + +use crate::catalog::Pool; + +/// What a selector may know about a request. Routing inputs such as the +/// session ID or the client's region become fields here. +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub struct RoutingRequest<'a> { + pub model: &'a str, +} + +impl<'a> RoutingRequest<'a> { + pub fn new(model: &'a str) -> Self { + Self { model } + } +} + +/// Where to send one request. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PoolTarget { + pub pool: PoolId, + pub frontend: Url, +} + +/// Chooses a pool, and a frontend in it, for each request. +pub trait PoolSelector: Send + Sync { + /// `candidates` are the pools serving the request's model. `None` means + /// none of them can take the request. + fn select(&self, request: &RoutingRequest<'_>, candidates: &[&Pool]) -> Option; +} + +/// Rotates through the pools serving each model, then through the chosen +/// pool's frontends, so every pool gets an equal share of a model's requests +/// however many frontends it has. +/// +/// Each model and each pool has its own cursor, kept by the runtime's +/// round-robin picker. When the candidates change, the cursor carries on over +/// the new list. +#[derive(Debug, Default)] +pub struct RoundRobinSelector { + pools: DashMap, + frontends: DashMap, +} + +impl RoundRobinSelector { + pub fn new() -> Self { + Self::default() + } +} + +impl PoolSelector for RoundRobinSelector { + fn select(&self, request: &RoutingRequest<'_>, candidates: &[&Pool]) -> Option { + let reachable = candidates + .iter() + .copied() + .filter(|pool| !pool.frontends.is_empty()) + .collect::>(); + let pool = reachable[rotate(&self.pools, request.model, reachable.len())?]; + let frontend = &pool.frontends[rotate(&self.frontends, &pool.id, pool.frontends.len())?]; + Some(PoolTarget { + pool: pool.id.clone(), + frontend: frontend.clone(), + }) + } +} + +/// Advances `key`'s cursor over `count` candidates and returns the chosen +/// position, or `None` when there are no candidates. +fn rotate(pickers: &DashMap, key: &Q, count: usize) -> Option +where + K: Borrow + Eq + Hash, + Q: ToOwned + Eq + Hash + ?Sized, +{ + // The picker chooses among worker IDs. Candidate positions stand in for + // them, since round robin only counts. + let positions = (0..count as u64).collect::>(); + let pick = |picker: &BuiltinRoutePicker| picker.select_worker(&positions, |_| 0); + let position = match pickers.get(key) { + Some(picker) => pick(&picker), + None => pick( + &pickers + .entry(key.to_owned()) + .or_insert_with(BuiltinRoutePicker::round_robin), + ), + }; + position.map(|position| position as usize) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::catalog::tests::pool; + + fn picks(selector: &RoundRobinSelector, model: &str, candidates: &[&Pool]) -> String { + let target = selector + .select(&RoutingRequest::new(model), candidates) + .unwrap(); + target.frontend.host_str().unwrap().to_owned() + } + + #[test] + fn rotates_pools_then_their_frontends() { + let blue = pool( + "blue", + "qwen", + &["http://blue-0:8000", "http://blue-1:8000"], + ); + let green = pool("green", "qwen", &["http://green-0:8000"]); + let selector = RoundRobinSelector::new(); + + let order = (0..6) + .map(|_| picks(&selector, "qwen", &[&blue, &green])) + .collect::>(); + assert_eq!( + order, + [ + "blue-0", "green-0", "blue-1", "green-0", "blue-0", "green-0" + ] + ); + } + + #[test] + fn models_rotate_independently() { + let qwen = [ + pool("qwen-a", "qwen", &["http://qwen-a:8000"]), + pool("qwen-b", "qwen", &["http://qwen-b:8000"]), + ]; + let llama = [ + pool("llama-a", "llama", &["http://llama-a:8000"]), + pool("llama-b", "llama", &["http://llama-b:8000"]), + ]; + let qwen = qwen.iter().collect::>(); + let llama = llama.iter().collect::>(); + let selector = RoundRobinSelector::new(); + + assert_eq!(picks(&selector, "qwen", &qwen), "qwen-a"); + assert_eq!(picks(&selector, "llama", &llama), "llama-a"); + assert_eq!(picks(&selector, "qwen", &qwen), "qwen-b"); + assert_eq!(picks(&selector, "llama", &llama), "llama-b"); + } + + #[test] + fn skips_pools_without_frontends() { + let empty = pool("draining", "qwen", &[]); + let live = pool("live", "qwen", &["http://live:8000"]); + let selector = RoundRobinSelector::new(); + let request = RoutingRequest::new("qwen"); + + for _ in 0..3 { + let target = selector.select(&request, &[&empty, &live]).unwrap(); + assert_eq!(target.pool, live.id); + } + assert_eq!(selector.select(&request, &[&empty]), None); + assert_eq!(selector.select(&request, &[]), None); + } +}