From 3dae6f8692d00cbeeb0a3f915699b689b8e04823 Mon Sep 17 00:00:00 2001 From: Sachal Malick Date: Sun, 25 Oct 2026 00:37:09 -0700 Subject: [PATCH 1/2] feat(global-router): add the Global Frontend forwarding engine Add dynamo-global-ingress, the Global Frontend. It runs the existing HttpService and registers one PoolForwardingEngine under every model in the pool catalog. The engine forwards each chat or completions request to the pool frontend a PoolSelector picks and streams the pool's SSE back. The Responses and Anthropic APIs translate to chat before the engine, so they forward too. The Global Frontend does not tokenize. Add dynamo-global-router with the PoolCatalog and PoolSelector traits, a static catalog for --pool flags and tests, and a first-pool selector that stands in until round robin. Claim the global-router crates, including the discovery service and the shared types, for the global-router codeowners. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Sachal Malick --- .github/codeowners/areas.yaml | 4 + CODEOWNERS | 6 +- Cargo.lock | 33 ++ Cargo.toml | 2 + lib/global-ingress/Cargo.toml | 41 +++ lib/global-ingress/src/forward.rs | 348 ++++++++++++++++++++ lib/global-ingress/src/lib.rs | 17 + lib/global-ingress/src/main.rs | 133 ++++++++ lib/global-ingress/src/model_sync.rs | 122 +++++++ lib/global-ingress/tests/forwarding.rs | 428 +++++++++++++++++++++++++ lib/global-router/Cargo.toml | 18 ++ lib/global-router/src/catalog.rs | 124 +++++++ lib/global-router/src/lib.rs | 16 + lib/global-router/src/selector.rs | 73 +++++ 14 files changed, 1364 insertions(+), 1 deletion(-) create mode 100644 lib/global-ingress/Cargo.toml create mode 100644 lib/global-ingress/src/forward.rs create mode 100644 lib/global-ingress/src/lib.rs create mode 100644 lib/global-ingress/src/main.rs create mode 100644 lib/global-ingress/src/model_sync.rs create mode 100644 lib/global-ingress/tests/forwarding.rs create mode 100644 lib/global-router/Cargo.toml create mode 100644 lib/global-router/src/catalog.rs create mode 100644 lib/global-router/src/lib.rs create mode 100644 lib/global-router/src/selector.rs 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..fc0013563bd9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2658,6 +2658,39 @@ dependencies = [ "url", ] +[[package]] +name = "dynamo-global-ingress" +version = "1.6.0" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "axum 0.8.4", + "clap", + "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 = [ + "dynamo-global-router-shared", + "tokio", + "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-ingress/Cargo.toml b/lib/global-ingress/Cargo.toml new file mode 100644 index 000000000000..7077936c798e --- /dev/null +++ b/lib/global-ingress/Cargo.toml @@ -0,0 +1,41 @@ +# 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-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] +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..b98ee80d0131 --- /dev/null +++ b/lib/global-ingress/src/main.rs @@ -0,0 +1,133 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Global Frontend. See `dynamo_global_ingress`. +//! +//! ```bash +//! 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::Parser; +use dynamo_global_ingress::{ModelSync, PoolForwardingEngine}; +use dynamo_global_router::{FirstPoolSelector, Pool, PoolCatalog, 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")] +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, + /// A pool as `//=@`. + /// 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, required = true)] + 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 = 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(FirstPoolSelector), + 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/forwarding.rs b/lib/global-ingress/tests/forwarding.rs new file mode 100644 index 000000000000..cc1791292675 --- /dev/null +++ b/lib/global-ingress/tests/forwarding.rs @@ -0,0 +1,428 @@ +// 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. + +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::{FirstPoolSelector, Pool, PoolSnapshot}; +use dynamo_global_router_shared::{PoolIdDeriver, PoolKey, PoolLocation, V1PoolIdDeriver}; +use dynamo_llm::endpoint_type::EndpointType; +use dynamo_llm::engines::{StreamingEngineAdapter, make_echo_engine}; +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; + +const MODEL: &str = "echo-model"; + +/// A running HTTP server. Dropping it stops the server. +struct Server { + url: Url, + cancel: CancellationToken, +} + +impl Drop for Server { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +struct GlobalFrontend { + server: Server, + pools: watch::Sender>, +} + +fn client() -> reqwest::Client { + reqwest::Client::builder().no_proxy().build().unwrap() +} + +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(), + ) +} + +fn http_service(url: &Url) -> HttpService { + HttpService::builder() + .host("127.0.0.1") + .port(url.port().unwrap()) + .build() + .unwrap() +} + +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 } +} + +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 +} + +async fn global_frontend(pools: Vec) -> GlobalFrontend { + let (sender, receiver) = watch::channel(Arc::new(PoolSnapshot::new(pools))); + let (listener, url) = bind().await; + let service = http_service(&url); + let engine = Arc::new(PoolForwardingEngine::new( + receiver.clone(), + Arc::new(FirstPoolSelector), + client(), + )); + let mut models = ModelSync::new(service.clone(), engine); + models.apply(&receiver.borrow()); + let server = serve(service, listener, url).await; + tokio::spawn(models.follow(receiver, server.cancel.child_token())); + GlobalFrontend { + server, + pools: sender, + } +} + +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()], + } +} + +fn chat(model: &str, prompt: &str, stream: bool) -> Value { + json!({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + }) +} + +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. +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]`. +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() +} + +#[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); +} + +/// A pool that records the headers it receives and answers with a fixed SSE +/// body. +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) +} + +#[tokio::test] +async fn forwards_the_session_and_skips_pool_keep_alives() { + let chunk = |content: &str| { + json!({ + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": MODEL, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}}], + }) + }; + let sse = format!( + ": keep-alive\n\ndata: {}\n\n: keep-alive\n\ndata: {}\n\ndata: [DONE]\n\n", + chunk("he"), + 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 chunk = json!({ + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 0, + "model": MODEL, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "he"}}], + }); + let sse = format!( + "data: {chunk}\n\ndata: {}\n\ndata: [DONE]\n\n", + 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-router/Cargo.toml b/lib/global-router/Cargo.toml new file mode 100644 index 000000000000..833120e711ba --- /dev/null +++ b/lib/global-router/Cargo.toml @@ -0,0 +1,18 @@ +# 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-router-shared = { path = "../global-router-shared" } + +tokio = { workspace = true } +url = { 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/lib.rs b/lib/global-router/src/lib.rs new file mode 100644 index 000000000000..da3f31a1eaec --- /dev/null +++ b/lib/global-router/src/lib.rs @@ -0,0 +1,16 @@ +// 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 selector; + +pub use catalog::{Pool, PoolCatalog, PoolSnapshot, StaticPoolCatalog}; +pub use selector::{FirstPoolSelector, PoolSelector, PoolTarget, RoutingRequest}; diff --git a/lib/global-router/src/selector.rs b/lib/global-router/src/selector.rs new file mode 100644 index 000000000000..25bb83ae14fb --- /dev/null +++ b/lib/global-router/src/selector.rs @@ -0,0 +1,73 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dynamo_global_router_shared::PoolId; +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; +} + +/// Sends every request to the first frontend of the first pool that has one. +/// A stand-in until the round-robin and load-aware policies land. +#[derive(Clone, Copy, Debug, Default)] +pub struct FirstPoolSelector; + +impl PoolSelector for FirstPoolSelector { + fn select(&self, _request: &RoutingRequest<'_>, candidates: &[&Pool]) -> Option { + candidates.iter().find_map(|pool| { + let frontend = pool.frontends.first()?; + Some(PoolTarget { + pool: pool.id.clone(), + frontend: frontend.clone(), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::catalog::tests::pool; + + #[test] + fn first_pool_selector_skips_pools_without_frontends() { + let empty = pool("draining", "qwen", &[]); + let live = pool("live", "qwen", &["http://live:8000", "http://live-2:8000"]); + let request = RoutingRequest::new("qwen"); + + let target = FirstPoolSelector + .select(&request, &[&empty, &live]) + .unwrap(); + assert_eq!(target.pool, live.id); + assert_eq!(target.frontend.as_str(), "http://live:8000/"); + assert_eq!(FirstPoolSelector.select(&request, &[&empty]), None); + assert_eq!(FirstPoolSelector.select(&request, &[]), None); + } +} From 74c2ec4471e8358ee8c21c9c0cdf493334fa5c9b Mon Sep 17 00:00:00 2001 From: Sachal Date: Sun, 25 Oct 2026 01:43:42 -0700 Subject: [PATCH 2/2] feat(global-router): route round robin over pools from discovery (#17273) Signed-off-by: Sachal Malick Co-authored-by: Claude Opus 5.5 --- Cargo.lock | 9 + lib/global-discovery-client/src/client.rs | 37 ++- lib/global-discovery-client/src/lib.rs | 10 +- lib/global-discovery-client/src/wire.rs | 27 ++ .../tests/router_calls.rs | 128 +++++++++ lib/global-ingress/Cargo.toml | 3 + lib/global-ingress/src/main.rs | 40 ++- lib/global-ingress/tests/common/mod.rs | 179 ++++++++++++ lib/global-ingress/tests/forwarding.rs | 183 +------------ lib/global-ingress/tests/round_robin.rs | 133 +++++++++ lib/global-router/Cargo.toml | 9 + lib/global-router/src/discovery.rs | 256 ++++++++++++++++++ lib/global-router/src/lib.rs | 4 +- lib/global-router/src/selector.rs | 137 ++++++++-- 14 files changed, 942 insertions(+), 213 deletions(-) create mode 100644 lib/global-discovery-client/tests/router_calls.rs create mode 100644 lib/global-ingress/tests/common/mod.rs create mode 100644 lib/global-ingress/tests/round_robin.rs create mode 100644 lib/global-router/src/discovery.rs diff --git a/Cargo.lock b/Cargo.lock index fc0013563bd9..202019c35d71 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2667,6 +2667,8 @@ dependencies = [ "async-trait", "axum 0.8.4", "clap", + "dynamo-global-discovery-client", + "dynamo-global-discovery-service", "dynamo-global-router", "dynamo-global-router-shared", "dynamo-llm", @@ -2686,8 +2688,15 @@ dependencies = [ 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", ] 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 index 7077936c798e..b6b2ddb85a1f 100644 --- a/lib/global-ingress/Cargo.toml +++ b/lib/global-ingress/Cargo.toml @@ -16,6 +16,7 @@ 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 @@ -38,4 +39,6 @@ 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/main.rs b/lib/global-ingress/src/main.rs index b98ee80d0131..c653145dbd8b 100644 --- a/lib/global-ingress/src/main.rs +++ b/lib/global-ingress/src/main.rs @@ -4,6 +4,10 @@ //! 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 //! ``` @@ -13,9 +17,12 @@ use std::sync::Arc; use std::time::Duration; use anyhow::{Context, Result, bail}; -use clap::Parser; +use clap::{ArgGroup, Parser}; +use dynamo_global_discovery_client::GlobalDiscoveryClient; use dynamo_global_ingress::{ModelSync, PoolForwardingEngine}; -use dynamo_global_router::{FirstPoolSelector, Pool, PoolCatalog, StaticPoolCatalog}; +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; @@ -23,16 +30,25 @@ 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, - /// A pool as `//=@`. - /// 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, required = true)] + /// 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)] @@ -86,7 +102,15 @@ async fn main() -> Result<()> { signal.cancel(); }); - let catalog = StaticPoolCatalog::new(args.pools); + 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) @@ -99,7 +123,7 @@ async fn main() -> Result<()> { .context("build the pool HTTP client")?; let engine = Arc::new(PoolForwardingEngine::new( catalog.subscribe(), - Arc::new(FirstPoolSelector), + Arc::new(RoundRobinSelector::new()), http, )); 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 index cc1791292675..4a05b4241515 100644 --- a/lib/global-ingress/tests/forwarding.rs +++ b/lib/global-ingress/tests/forwarding.rs @@ -5,84 +5,27 @@ //! Most pools here are Dynamo's own HTTP service running the echo engine, //! which streams the prompt back one character at a time. -use std::sync::{Arc, Mutex}; +mod common; + +use std::sync::Arc; use std::time::Duration; -use axum::Router; -use axum::http::HeaderMap; -use axum::routing; -use dynamo_global_ingress::{ModelSync, PoolForwardingEngine}; -use dynamo_global_router::{FirstPoolSelector, Pool, PoolSnapshot}; -use dynamo_global_router_shared::{PoolIdDeriver, PoolKey, PoolLocation, V1PoolIdDeriver}; +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 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; - -const MODEL: &str = "echo-model"; - -/// A running HTTP server. Dropping it stops the server. -struct Server { - url: Url, - cancel: CancellationToken, -} - -impl Drop for Server { - fn drop(&mut self) { - self.cancel.cancel(); - } -} struct GlobalFrontend { server: Server, pools: watch::Sender>, } -fn client() -> reqwest::Client { - reqwest::Client::builder().no_proxy().build().unwrap() -} - -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(), - ) -} - -fn http_service(url: &Url) -> HttpService { - HttpService::builder() - .host("127.0.0.1") - .port(url.port().unwrap()) - .build() - .unwrap() -} - -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 } -} - async fn echo_pool_frontend() -> Server { let (listener, url) = bind().await; let service = http_service(&url); @@ -100,75 +43,15 @@ async fn echo_pool_frontend() -> Server { 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))); - let (listener, url) = bind().await; - let service = http_service(&url); - let engine = Arc::new(PoolForwardingEngine::new( - receiver.clone(), - Arc::new(FirstPoolSelector), - client(), - )); - let mut models = ModelSync::new(service.clone(), engine); - models.apply(&receiver.borrow()); - let server = serve(service, listener, url).await; - tokio::spawn(models.follow(receiver, server.cancel.child_token())); GlobalFrontend { - server, + server: common::global_frontend(receiver).await, pools: sender, } } -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()], - } -} - -fn chat(model: &str, prompt: &str, stream: bool) -> Value { - json!({ - "model": model, - "messages": [{"role": "user", "content": prompt}], - "stream": stream, - }) -} - -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. -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]`. -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() -} - #[tokio::test] async fn streams_chat_completions_from_the_pool() { let pool_frontend = echo_pool_frontend().await; @@ -335,44 +218,12 @@ async fn unreachable_pools_are_unavailable() { assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); } -/// A pool that records the headers it receives and answers with a fixed SSE -/// body. -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) -} - #[tokio::test] async fn forwards_the_session_and_skips_pool_keep_alives() { - let chunk = |content: &str| { - json!({ - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 0, - "model": MODEL, - "choices": [{"index": 0, "delta": {"role": "assistant", "content": content}}], - }) - }; let sse = format!( ": keep-alive\n\ndata: {}\n\n: keep-alive\n\ndata: {}\n\ndata: [DONE]\n\n", - chunk("he"), - chunk("llo") + 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; @@ -398,15 +249,9 @@ async fn forwards_the_session_and_skips_pool_keep_alives() { #[tokio::test] async fn a_pool_error_mid_stream_ends_the_client_stream_with_an_error() { - let chunk = json!({ - "id": "chatcmpl-1", - "object": "chat.completion.chunk", - "created": 0, - "model": MODEL, - "choices": [{"index": 0, "delta": {"role": "assistant", "content": "he"}}], - }); let sse = format!( - "data: {chunk}\n\ndata: {}\n\ndata: [DONE]\n\n", + "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; 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 index 833120e711ba..7e3fa728f351 100644 --- a/lib/global-router/Cargo.toml +++ b/lib/global-router/Cargo.toml @@ -12,7 +12,16 @@ 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/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 index da3f31a1eaec..40027b6a9dbf 100644 --- a/lib/global-router/src/lib.rs +++ b/lib/global-router/src/lib.rs @@ -10,7 +10,9 @@ //! trait. mod catalog; +mod discovery; mod selector; pub use catalog::{Pool, PoolCatalog, PoolSnapshot, StaticPoolCatalog}; -pub use selector::{FirstPoolSelector, PoolSelector, PoolTarget, RoutingRequest}; +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 index 25bb83ae14fb..fe48134ec2dd 100644 --- a/lib/global-router/src/selector.rs +++ b/lib/global-router/src/selector.rs @@ -1,7 +1,12 @@ // 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; @@ -34,40 +39,128 @@ pub trait PoolSelector: Send + Sync { fn select(&self, request: &RoutingRequest<'_>, candidates: &[&Pool]) -> Option; } -/// Sends every request to the first frontend of the first pool that has one. -/// A stand-in until the round-robin and load-aware policies land. -#[derive(Clone, Copy, Debug, Default)] -pub struct FirstPoolSelector; - -impl PoolSelector for FirstPoolSelector { - fn select(&self, _request: &RoutingRequest<'_>, candidates: &[&Pool]) -> Option { - candidates.iter().find_map(|pool| { - let frontend = pool.frontends.first()?; - Some(PoolTarget { - pool: pool.id.clone(), - frontend: frontend.clone(), - }) +/// 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 first_pool_selector_skips_pools_without_frontends() { + 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", "http://live-2:8000"]); + let live = pool("live", "qwen", &["http://live:8000"]); + let selector = RoundRobinSelector::new(); let request = RoutingRequest::new("qwen"); - let target = FirstPoolSelector - .select(&request, &[&empty, &live]) - .unwrap(); - assert_eq!(target.pool, live.id); - assert_eq!(target.frontend.as_str(), "http://live:8000/"); - assert_eq!(FirstPoolSelector.select(&request, &[&empty]), None); - assert_eq!(FirstPoolSelector.select(&request, &[]), None); + 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); } }