From afde31a2f506c550c9824ab14bea00ed0ad22afb Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Tue, 15 Sep 2026 22:30:31 -0700 Subject: [PATCH] [router] Speak cleartext h2c on both edges: serve it inbound, forward it outbound (#39006) Co-authored-by: Kangyan Zhou Co-authored-by: Claude Opus 5 (1M context) --- experimental/sgl-router/Cargo.lock | 3 + experimental/sgl-router/Cargo.toml | 23 +- experimental/sgl-router/README.md | 45 ++- experimental/sgl-router/src/main.rs | 4 + experimental/sgl-router/src/proxy/mod.rs | 83 +++++- .../sgl-router/src/server/routes/cache.rs | 8 +- .../sgl-router/src/server/routes/chat.rs | 6 + .../sgl-router/src/workers/manager.rs | 21 +- experimental/sgl-router/src/workers/worker.rs | 18 +- .../k8s_integration/Dockerfile.fake_worker | 2 +- .../tests/e2e/k8s_integration/fake_worker.py | 57 +++- .../k8s_integration/test_h2c_forwarding.py | 198 +++++++++++++ .../sgl-router/tests/proxy/chat_routing.rs | 8 +- .../sgl-router/tests/proxy/h2c_forward.rs | 195 +++++++++++++ .../sgl-router/tests/proxy/inbound_h2c.rs | 116 ++++++++ experimental/sgl-router/tests/proxy/main.rs | 3 + .../tests/proxy/pd_protocol_binding.rs | 260 ++++++++++++++++++ 17 files changed, 1009 insertions(+), 41 deletions(-) create mode 100644 experimental/sgl-router/tests/e2e/k8s_integration/test_h2c_forwarding.py create mode 100644 experimental/sgl-router/tests/proxy/h2c_forward.rs create mode 100644 experimental/sgl-router/tests/proxy/inbound_h2c.rs create mode 100644 experimental/sgl-router/tests/proxy/pd_protocol_binding.rs diff --git a/experimental/sgl-router/Cargo.lock b/experimental/sgl-router/Cargo.lock index 15f474db1..81b322705 100644 --- a/experimental/sgl-router/Cargo.lock +++ b/experimental/sgl-router/Cargo.lock @@ -2887,6 +2887,7 @@ dependencies = [ "bytes", "futures-core", "futures-util", + "h2", "http", "http-body", "http-body-util", @@ -3294,6 +3295,8 @@ dependencies = [ "futures", "hf-hub", "http-body-util", + "hyper", + "hyper-util", "k8s-openapi", "kube", "minijinja", diff --git a/experimental/sgl-router/Cargo.toml b/experimental/sgl-router/Cargo.toml index 7a11cb670..b24e9d24b 100644 --- a/experimental/sgl-router/Cargo.toml +++ b/experimental/sgl-router/Cargo.toml @@ -38,10 +38,21 @@ chrono = { version = "0.4", default-features = false, features = ["clock"] } # Async runtime + http tokio = { version = "1.42", features = ["full"] } -axum = { version = "0.8", features = ["macros", "tracing"] } +# http2 on both sides: axum serves cleartext h2c on the inbound listener, reqwest +# dials it outbound. +# +# axum's `http2` looks redundant and is not — please don't drop it. `axum::serve` +# runs on hyper-util's auto Builder, whose HTTP/2 branch is compiled by +# `hyper-util/http2`, and that arrives today only transitively, via +# sgl-kv-indexer -> tonic "transport" -> tonic/server -> hyper-util/server-auto +# (which is server + http1 + http2). So the listener speaks h2c whether or not +# this line exists. Declaring it forwards `hyper-util/http2` from axum directly, +# so inbound h2c no longer rides on a gRPC dependency it has nothing to do with, +# and it enables HTTP/2 extended CONNECT for h2 websockets. +axum = { version = "0.8", features = ["macros", "tracing", "http2"] } tower = { version = "0.5", features = ["full"] } tower-http = { version = "0.6", features = ["trace", "compression-gzip", "cors", "timeout", "request-id"] } -reqwest = { version = "0.12", features = ["stream", "json", "rustls-tls"], default-features = false } +reqwest = { version = "0.12", features = ["stream", "json", "rustls-tls", "http2"], default-features = false } sgl-kv-indexer = { path = "sgl-kv-indexer" } # Serialization @@ -81,6 +92,14 @@ zeromq = { version = "0.6", default-features = false, features = ["tokio-runtime [dev-dependencies] dirs = "5" http-body-util = "0.1" +# HTTP/2-only mock server for the proxy h2c forwarding tests (tests/proxy/h2c_forward.rs). +# `http1` as well as `http2`: tests/proxy/pd_protocol_binding.rs drives both +# `hyper::server::conn::http1` and `::http2` to build single-protocol mocks. +# `http1` arrives transitively via axum's defaults today; declaring it here +# means an axum `default-features = false` later fails at the dependency +# rather than as an unresolved import in a test. +hyper = { version = "1", features = ["server", "http1", "http2"] } +hyper-util = { version = "0.1", features = ["tokio"] } serde = { version = "1", features = ["derive"] } serde_json = "1" tempfile = "3" diff --git a/experimental/sgl-router/README.md b/experimental/sgl-router/README.md index 15801cf11..bac246372 100644 --- a/experimental/sgl-router/README.md +++ b/experimental/sgl-router/README.md @@ -6,7 +6,8 @@ Serves a single model and routes across its workers. Exposes `/v1/tokenize`, `/v1/detokenize`, `/v1/models`, `/v1/chat/completions` (buffered and SSE), plus `/healthz` / `/readyz` and `/metrics`. Worker pools come from either a static URL list or Kubernetes EndpointSlice -discovery. +discovery. Both edges speak cleartext HTTP/2 where the peer does — see +[HTTP/2](#http2). ## Building @@ -68,6 +69,48 @@ The Indexer replaces the Router-local radix tree as the native Cache-Aware signal. Query timeouts and local concurrency are bounded by the two Indexer options, which default to 100 ms and 32 respectively. +## HTTP/2 + +There is nothing to configure. The router negotiates per connection inbound and +resolves the protocol per worker outbound; every combination below is reached +automatically, and HTTP/1.1 remains a supported peer on both edges. + +**Inbound.** The listener accepts cleartext HTTP/2 (h2c, prior knowledge) and +HTTP/1.1 on the same `--port`, chosen per connection. A mesh sidecar or +load balancer that prefers h2c multiplexes over one connection instead of +opening one per request; an HTTP/1.1 client is unaffected. + +**Outbound.** At registration the router reads each worker's `/server_info` and +forwards over cleartext h2c only when that worker reports `--enable-http2` (the +engine's Granian server, which serves h2c and HTTP/1.1 together) **and** the +worker URL is cleartext. Everything else uses the default client, which speaks +HTTP/1.1 in cleartext and negotiates ALPN `h2, http/1.1` over TLS: + +| `/server_info` | worker URL | router forwards over | +|---|---|---| +| `enable_http2: true` | `http://` | cleartext h2c | +| `enable_http2: true` | `https://` | HTTP/2 over TLS, by ALPN | +| `enable_http2: false`, or absent | any | HTTP/1.1 (cleartext) or ALPN (TLS) | + +The choice is per worker, not fleet-wide, so a mixed fleet works and one +worker's state never changes another's. The flag comes from the `/server_info` +launch record, which has reported it all along, so no engine change is needed. +A worker whose `/server_info` does not answer has no readable protocol and +forwards over HTTP/1.1 for as long as it stays registered — a throughput cost, +never a correctness one. Admin fan-out (`/flush_cache`) always uses the default +client, because it addresses every worker at once rather than a selected one. + +Two things worth knowing when debugging. h2c is prior-knowledge only — there is +no negotiation and no fallback — which is why the router requires the engine's +own `enable_http2` report before using it. And a worker's protocol is fixed +for as long as that worker stays registered: it is read once, at registration, +and never re-read. Under the K8s backend a restarting engine flips its +EndpointSlice to `ready=false`, which is a `Removed` → `Added` cycle and so a +fresh reading; under `--worker-urls` the fan-out happens once at startup and +nothing re-registers. So a worker registered over h2c that later stops serving +it (a proxy interposed on its port, say) is not detected until it +is re-registered; its circuit breaker will open in the meantime. + ## Upgrading from `cache_aware_zmq` The `cache_aware_zmq` policy has been removed. Configurations using it should diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index dcaf35fea..759b2e14f 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -167,6 +167,10 @@ async fn main() -> Result<()> { sgl_router::policies::active_load::spawn_janitor(Arc::clone(&active_load), sweep_interval); // Spawn discovery + manager tasks. + // The manager resolves each worker's wire protocol from its `/server_info` + // and stamps it onto the registered worker. The proxy holds one client per + // protocol and selects by the worker's protocol per request, so the manager + // needs no proxy handle. let (event_rx, discovery_handle) = sgl_router::discovery::spawn_discovery(&cfg) .await .context("spawn discovery")?; diff --git a/experimental/sgl-router/src/proxy/mod.rs b/experimental/sgl-router/src/proxy/mod.rs index 7477f6f3d..5efda106d 100644 --- a/experimental/sgl-router/src/proxy/mod.rs +++ b/experimental/sgl-router/src/proxy/mod.rs @@ -8,6 +8,7 @@ pub mod sse; use crate::health::circuit_breaker::CircuitBreaker; use crate::server::error::ApiError; use crate::server::header_utils::should_forward_request_header; +use crate::workers::WireProtocol; use anyhow::Context; use axum::body::Body; use axum::http::{HeaderMap, HeaderName, HeaderValue, Response}; @@ -32,28 +33,68 @@ fn parse_worker_url(worker_url: &str, breaker: &CircuitBreaker) -> Result Result { + let builder = Client::builder() + .pool_max_idle_per_host(64) + .connect_timeout(Duration::from_secs(5)); + match protocol { + WireProtocol::Http1 => builder, + WireProtocol::H2c => builder.http2_prior_knowledge(), + } + .build() + .context("build reqwest client") +} + impl Proxy { /// Build a proxy. `request_timeout` is the per-request wall-clock budget for /// non-streaming forwards. Connect timeout is hard-coded to 5 s — even a /// streaming request fails fast at TCP setup if the worker is unreachable. + /// + /// WHY both clients up front: protocol is a per-worker property resolved + /// from each engine's `/server_info`, so the request path must be able to + /// pick either one per request. Building them here reduces that to a + /// selection — no per-request client construction, and no single shared + /// client whose first writer decides the protocol for the whole fleet. pub fn new(request_timeout: Duration) -> Result { - let client = Client::builder() - .pool_max_idle_per_host(64) - .connect_timeout(Duration::from_secs(5)) - .build() - .context("build reqwest client")?; Ok(Self { - client, + default_client: build_client(WireProtocol::Http1)?, + h2c_client: build_client(WireProtocol::H2c)?, request_timeout, }) } + /// The forwarding client for `protocol`, taken from the selected worker's + /// [`crate::workers::Worker::protocol`]. + fn client_for(&self, protocol: WireProtocol) -> &Client { + match protocol { + WireProtocol::Http1 => &self.default_client, + WireProtocol::H2c => &self.h2c_client, + } + } + + /// The client for side-channel admin traffic (e.g. `/flush_cache`), which + /// fans out across workers and so cannot use any one worker's protocol. + pub fn admin_client(&self) -> &Client { + &self.default_client + } + /// Classify a reqwest error into the right `ApiError` variant, given an /// explicit worker URL. Called from the breaker-gated `forward_*_to` /// methods, which carry per-request worker URLs (not a single proxy-level @@ -91,6 +132,7 @@ impl Proxy { pub async fn forward_json_to( &self, worker_url: &str, + protocol: WireProtocol, breaker: &CircuitBreaker, path: &str, headers: &HeaderMap, @@ -105,7 +147,7 @@ impl Proxy { let url = worker_url.join(path).map_err(|e| { ApiError::Internal(anyhow::Error::new(e).context(format!("join worker path {path}"))) })?; - let mut req = self.client.post(url.clone()).body(body); + let mut req = self.client_for(protocol).post(url.clone()).body(body); for (k, v) in headers { if should_forward_request_header(k) { req = req.header(k, v); @@ -164,13 +206,14 @@ impl Proxy { /// for the full streaming lifetime — without which a long-running SSE /// response would under-report load. // Each parameter is a distinct, required input to a single upstream - // forward (target, breaker, path, headers, body, plus the + // forward (target, protocol, breaker, path, headers, body, plus the // streaming-lifetime callbacks). Bundling them into a struct purely to // satisfy the arg-count heuristic would add indirection without clarity. #[allow(clippy::too_many_arguments)] pub async fn forward_streaming_to( &self, worker_url: &str, + protocol: WireProtocol, breaker: &Arc, path: &str, headers: &HeaderMap, @@ -188,7 +231,7 @@ impl Proxy { let url = worker_url.join(path).map_err(|e| { ApiError::Internal(anyhow::Error::new(e).context(format!("join worker path {path}"))) })?; - let mut req = self.client.post(url.clone()).body(body); + let mut req = self.client_for(protocol).post(url.clone()).body(body); for (k, v) in headers { if should_forward_request_header(k) { req = req.header(k, v); @@ -274,4 +317,24 @@ mod tests { let p = Proxy::new(Duration::from_secs(5)).unwrap(); assert_eq!(p.request_timeout, Duration::from_secs(5)); } + + /// `client_for` routes each protocol to its own field, and admin traffic + /// shares the default client. Asserting the two clients differ by address + /// would be vacuous — they are distinct struct fields, so that holds even + /// if `build_client` ignored its argument. What the selector must get right + /// is the mapping, so pin that instead; the on-the-wire difference between + /// the two clients is covered by tests/proxy/h2c_forward.rs. + #[tokio::test] + async fn client_for_maps_each_protocol_to_its_own_client() { + let p = Proxy::new(Duration::from_secs(5)).unwrap(); + assert!(std::ptr::eq( + p.client_for(WireProtocol::Http1), + &p.default_client + )); + assert!(std::ptr::eq(p.client_for(WireProtocol::H2c), &p.h2c_client)); + assert!(std::ptr::eq( + p.client_for(WireProtocol::Http1), + p.admin_client() + )); + } } diff --git a/experimental/sgl-router/src/server/routes/cache.rs b/experimental/sgl-router/src/server/routes/cache.rs index ef136533f..6fcc494cb 100644 --- a/experimental/sgl-router/src/server/routes/cache.rs +++ b/experimental/sgl-router/src/server/routes/cache.rs @@ -100,8 +100,12 @@ pub async fn flush_cache(State(ctx): State>) -> Response { .into_response(); } - let (successful, failed) = - fan_out_flush(&workers, &ctx.proxy.client, ctx.proxy.request_timeout).await; + let (successful, failed) = fan_out_flush( + &workers, + ctx.proxy.admin_client(), + ctx.proxy.request_timeout, + ) + .await; // Partial failure is an operational event an operator needs to see at the // common production log level — match the rest of the router, which warns diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index fcfe06d63..f49564810 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -539,6 +539,7 @@ pub async fn chat_completions( let bootstrap_room = bootstrap_room.expect("PD dispatch implies a resolved bootstrap room"); let prefill_url = worker.url.clone(); + let prefill_protocol = worker.protocol(); let prefill_breaker = Arc::clone(&worker.breaker); let prefill_headers = headers.clone(); let prefill_body = outgoing_body.clone(); @@ -555,6 +556,7 @@ pub async fn chat_completions( match prefill_proxy .forward_json_to( &prefill_url, + prefill_protocol, &prefill_breaker, "/v1/chat/completions", &prefill_headers, @@ -588,6 +590,7 @@ pub async fn chat_completions( Box::new((decode_guard, decode_active_guard, make_duration_guard())); let fetch = ctx.proxy.forward_streaming_to( &decode_worker.url, + decode_worker.protocol(), &decode_worker.breaker, "/v1/chat/completions", &headers, @@ -605,6 +608,7 @@ pub async fn chat_completions( let _decode_hold = (decode_guard, decode_active_guard); let fetch = ctx.proxy.forward_json_to( &decode_worker.url, + decode_worker.protocol(), &decode_worker.breaker, "/v1/chat/completions", &headers, @@ -624,6 +628,7 @@ pub async fn chat_completions( Box::new((guard, active_guard, make_duration_guard())); let fetch = ctx.proxy.forward_streaming_to( &worker.url, + worker.protocol(), &worker.breaker, "/v1/chat/completions", &headers, @@ -653,6 +658,7 @@ pub async fn chat_completions( let _holds: (LoadGuard, _) = (guard, active_guard); let fetch = ctx.proxy.forward_json_to( &worker.url, + worker.protocol(), &worker.breaker, "/v1/chat/completions", &headers, diff --git a/experimental/sgl-router/src/workers/manager.rs b/experimental/sgl-router/src/workers/manager.rs index 037936262..3de03a638 100644 --- a/experimental/sgl-router/src/workers/manager.rs +++ b/experimental/sgl-router/src/workers/manager.rs @@ -82,11 +82,11 @@ fn dials_cleartext(worker_url: &str) -> Option { /// that risk. /// /// [`WireProtocol::Http1`] is the fallback for everything else, and it does -/// mean HTTP/1.1 on the wire: reqwest is built here without its `http2` -/// feature, so the forwarding client advertises only `http/1.1` even on TLS. -/// An `https://` worker running `--enable-http2` therefore stays on HTTP/1.1 -/// — correct, because h2c cannot be sent to a TLS endpoint either way, but -/// not the ALPN upgrade the flag might suggest. See [`WireProtocol`]. +/// not mean "HTTP/1.1 on the wire". It selects the negotiating client, which +/// advertises ALPN `h2, http/1.1`, so an `https://` worker running +/// `--enable-http2` reaches HTTP/2 over TLS on its own — the correct outcome, +/// arrived at by negotiation rather than by assumption. Only cleartext workers +/// need prior-knowledge h2c. See [`WireProtocol`]. fn resolve_protocol(enable_http2: Option, cleartext: Option) -> WireProtocol { match (enable_http2, cleartext) { (Some(true), Some(true)) => WireProtocol::H2c, @@ -449,12 +449,12 @@ fn log_protocol_resolution(worker_url: &str, enable_http2: Option, clearte worker_url = %worker_url, "/server_info reports --enable-http2 on a cleartext worker; forwarding over h2c", ), - // A TLS worker. h2c is unsendable there and the client does not - // negotiate, so this worker stays on HTTP/1.1. + // A TLS worker. h2c is unsendable there, but the negotiating client + // advertises ALPN h2, so a TLS engine still reaches HTTP/2 on its own. (Some(true), Some(false)) => tracing::info!( worker_url = %worker_url, - "/server_info reports --enable-http2 on a TLS worker; \ - forwarding over HTTP/1.1", + "/server_info reports --enable-http2 on a TLS worker; using the \ + negotiating client, which reaches HTTP/2 over TLS via ALPN", ), // `dials_cleartext` could not parse the URL. Reaching this at all means // `/server_info` answered over a URL the scheme check then rejected, so @@ -473,7 +473,8 @@ fn log_protocol_resolution(worker_url: &str, enable_http2: Option, clearte // revisiting it, so this reading is the only one it will ever get. (None, _) => tracing::info!( worker_url = %worker_url, - "no --enable-http2 reading from /server_info; forwarding over HTTP/1.1", + "no --enable-http2 reading from /server_info; using the negotiating \ + client (HTTP/1.1 in cleartext)", ), // The engine explicitly disabled it. Nothing to explain. (Some(false), _) => {} diff --git a/experimental/sgl-router/src/workers/worker.rs b/experimental/sgl-router/src/workers/worker.rs index e045f40db..27583d0c8 100644 --- a/experimental/sgl-router/src/workers/worker.rs +++ b/experimental/sgl-router/src/workers/worker.rs @@ -13,19 +13,19 @@ use std::time::Instant; /// Fixed for the worker's lifetime: it is derived by `manager::resolve_protocol` /// from the engine's `--enable-http2` launch flag and the dialed URL scheme, /// neither of which changes while the process runs. The asymmetry that drives -/// the default: HTTP/1.1 is accepted by every engine, while h2c is -/// prior-knowledge only and fails outright against an engine that does not +/// the default: the negotiating client is accepted by every engine, while h2c +/// is prior-knowledge only and fails outright against an engine that does not /// serve it. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum WireProtocol { - /// HTTP/1.1, in cleartext and over TLS alike. Safe for every engine, so it - /// is also the fallback. + /// The negotiating client. Safe for every engine, so it is also the + /// fallback. /// - /// Not ALPN-negotiated: `Cargo.toml` builds reqwest with - /// `default-features = false` and does not enable `http2`, so the - /// forwarding client advertises only `http/1.1` and a TLS engine running - /// `--enable-http2` still gets HTTP/1.1. Enabling that feature is what - /// would make this variant negotiate. + /// HTTP/1.1 in cleartext, ALPN-negotiated over TLS: this crate enables + /// reqwest's `http2` feature (see `Cargo.toml`), so the client advertises + /// `h2, http/1.1` and a TLS engine running `--enable-http2` reaches HTTP/2 + /// on its own. Dropping that feature silently reduces this variant to + /// HTTP/1.1 everywhere. #[default] Http1, /// Cleartext HTTP/2 with prior knowledge (h2c). Used only when a worker diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.fake_worker b/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.fake_worker index 8cff23e91..22edadf20 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.fake_worker +++ b/experimental/sgl-router/tests/e2e/k8s_integration/Dockerfile.fake_worker @@ -1,6 +1,6 @@ FROM python:3.12-slim WORKDIR /app -RUN pip install --no-cache-dir fastapi uvicorn +RUN pip install --no-cache-dir fastapi uvicorn granian COPY fake_worker.py . EXPOSE 30000 CMD ["python", "fake_worker.py"] diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/fake_worker.py b/experimental/sgl-router/tests/e2e/k8s_integration/fake_worker.py index 658556bac..4cdc1e7c3 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/fake_worker.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/fake_worker.py @@ -2,9 +2,21 @@ Responds to: GET /health -> {"status": "ok"} - GET /server_info -> {"served_model_name": MODEL_ID} + GET /server_info -> {"served_model_name": MODEL_ID, ...} GET /v1/models -> list with a single MODEL_ID model entry - POST /v1/chat/completions -> echoes the last user message back + POST /v1/chat/completions -> echoes the last user message back, plus + the HTTP version the request arrived on + +Set `FAKE_WORKER_HTTP2=1` to imitate an engine launched with +`--enable-http2`: `/server_info` advertises the flag and the app is served by +Granian in `HTTPModes.auto`, which is what the real engine runs, so one port +serves cleartext h2c alongside HTTP/1.1. The default is uvicorn, which speaks +HTTP/1.1 only. + +`x_http_version` on the chat response is the load-bearing part for +`test_h2c_forwarding.py`: a chat completion returns 200 over either protocol, +so without the worker reporting what it actually received, an h2c test passes +whether or not h2c was used. """ from __future__ import annotations @@ -17,6 +29,16 @@ from fastapi import FastAPI, Request app = FastAPI() MODEL_ID = os.environ.get("MODEL_ID", "tiny") +# Normalised before comparing: a manifest that writes the Python-idiomatic +# "False" must not silently turn this worker into a Granian/h2c one, which +# would fail test_h2c_forwarding with a message about a router bug. +ENABLE_HTTP2 = os.environ.get("FAKE_WORKER_HTTP2", "").strip().lower() not in ( + "", + "0", + "false", + "no", + "off", +) @app.get("/health") @@ -27,8 +49,12 @@ async def health(): @app.get("/server_info") async def server_info(): # The sgl-router worker manager fetches this on every Added event and - # uses `served_model_name` to populate the registry's model index. - return {"served_model_name": MODEL_ID} + # uses `served_model_name` to populate the registry's model index, and + # `enable_http2` to resolve the worker's forwarding protocol. + info = {"served_model_name": MODEL_ID} + if ENABLE_HTTP2: + info["enable_http2"] = True + return info @app.get("/v1/models") @@ -55,6 +81,11 @@ async def chat_completions(request: Request): "id": "chatcmpl-mock", "object": "chat.completion", "model": payload.get("model", MODEL_ID), + # Non-standard, and deliberately so: the router returns the upstream + # body verbatim (`proxy::forward_*` hands back `resp.bytes()`), so this + # is how a test on the other side of the router learns which protocol + # the forward leg actually used. "1.1" or "2". + "x_http_version": request.scope.get("http_version"), "choices": [ { "index": 0, @@ -70,4 +101,20 @@ async def chat_completions(request: Request): if __name__ == "__main__": - uvicorn.run(app, host="0.0.0.0", port=30000) + if ENABLE_HTTP2: + # Mirrors the engine's own server (`_run_granian_server` in + # sglang/srt/entrypoints/http_server.py): HTTPModes.auto dispatches per + # connection on the first bytes, so h2c prior-knowledge and HTTP/1.1 + # share one cleartext port. + from granian import Granian + from granian.constants import HTTPModes, Interfaces + + Granian( + target="fake_worker:app", + address="0.0.0.0", + port=30000, + interface=Interfaces.ASGI, + http=HTTPModes.auto, + ).serve() + else: + uvicorn.run(app, host="0.0.0.0", port=30000) diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_h2c_forwarding.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_h2c_forwarding.py new file mode 100644 index 000000000..c9d818e3c --- /dev/null +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_h2c_forwarding.py @@ -0,0 +1,198 @@ +"""E2E: the router forwards to an h2c-capable worker over cleartext HTTP/2. + +The in-process tests (`tests/proxy/h2c_forward.rs`, `inbound_h2c.rs`) already +drive real HTTP/2 sockets, and dropping reqwest's `http2` feature fails the +build outright, so neither the client nor the framing needs covering again +here. What no in-process test can assemble is the *chain*: a worker discovered +through a real EndpointSlice, introspected over the network, resolved to +`WireProtocol::H2c` from its own `/server_info`, and then actually forwarded to +over h2c. + +The fleet is deliberately mixed. `setup.sh` leaves three uvicorn workers +(HTTP/1.1 only) behind the `app=sglang` Service; this module adds one Granian +worker reporting `enable_http2: true` to the same Service, so both protocols +must be in use simultaneously. That is the e2e form of the per-worker-protocol +property: a router that resolved one protocol fleet-wide would either fail +against the h2c worker or break the three HTTP/1.1 ones, and either way this +test fails. +""" + +from __future__ import annotations + +import httpx +import pytest +from conftest import ( + NAMESPACE, + _apply_from_stdin, + _kubectl, + _poll_until, + _wait_for_deployment_ready, + logger, +) + +H2C_DEPLOYMENT = "fake-worker-h2c" + +# Round-robin over a 4-worker pool: 12 requests give every worker ~3 turns, so +# a miss means a routing or resolution failure rather than an unlucky draw. +_PROBE_REQUESTS = 12 + +# Kept low enough that a whole probe round (12 x 5 s worst case) fits inside the +# 90 s poll budget below. A fake worker answers instantly; a request that needs +# more than 5 s is already a failure, and letting a round outlast its own poll +# would make that timeout non-binding. +_CHAT_TIMEOUT = 5.0 +_CONVERGE_TIMEOUT = 90 + +_H2C_WORKER_MANIFEST = f""" +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {H2C_DEPLOYMENT} + namespace: {NAMESPACE} +spec: + replicas: 1 + selector: + matchLabels: + app: sglang-h2c + template: + metadata: + labels: + app: sglang-h2c + spec: + containers: + - name: worker + image: sgl-router-fake-worker:e2e + imagePullPolicy: Never + env: + - name: FAKE_WORKER_HTTP2 + value: "1" + - name: MODEL_ID + value: "tiny" + ports: + - containerPort: 30000 + readinessProbe: + httpGet: + path: /health + port: 30000 + initialDelaySeconds: 2 + periodSeconds: 3 +--- +apiVersion: v1 +kind: Service +metadata: + name: {H2C_DEPLOYMENT} + namespace: {NAMESPACE} + # The router watches ENDPOINTSLICES whose labels match `--selector + # app=sglang`, and Kubernetes mirrors a Service's labels onto the slices it + # manages -- so this label, not the pods', is what puts these workers in the + # router's view. + labels: + app: sglang +spec: + # Pods are labelled `app: sglang-h2c`, deliberately NOT `app: sglang`: the + # fake-worker Deployment's selector is a bare `app=sglang`, so sharing that + # label would put these pods inside another controller's selector and into + # the HTTP/1.1 Service as well. + selector: + app: sglang-h2c + ports: + - port: 30000 + targetPort: 30000 +""" + + +@pytest.fixture(scope="module") +def h2c_worker(k8s_cluster): + """Add one Granian/h2c worker, behind its own Service, to the router's view. + + Its own Service rather than the existing one: the router selects + EndpointSlices, so a second Service labelled `app: sglang` is watched just + the same, while its pods stay out of the `fake-worker` Deployment's bare + `app=sglang` selector. Torn down afterwards so the suite's other modules + see the three-worker fleet they expect. + """ + _apply_from_stdin(_H2C_WORKER_MANIFEST) + try: + _wait_for_deployment_ready(H2C_DEPLOYMENT) + yield + finally: + for kind in ("deployment", "service"): + _kubectl( + "delete", + kind, + H2C_DEPLOYMENT, + "-n", + NAMESPACE, + "--ignore-not-found", + "--wait=true", + check=False, + ) + + +def _chat(router_url: str, content: str) -> httpx.Response: + return httpx.post( + f"{router_url}/v1/chat/completions", + json={ + "model": "tiny", + "messages": [{"role": "user", "content": content}], + "stream": False, + }, + timeout=_CHAT_TIMEOUT, + ) + + +def _observed_protocols(router_url: str, *, strict: bool) -> set[str]: + """Fan out round-robin and collect the HTTP version each worker saw. + + `x_http_version` is reported by the worker itself, not inferred from the + client side: the test's own connection to the router is a separate hop, so + only the worker can say what the forward leg used. Distinct content per + request keeps any content-derived routing from collapsing onto one worker. + + `strict=False` while converging. The router runs `--cb-threshold 1` + (manifests/router.yaml), so one refused connection to the still-starting h2c + pod opens its breaker and round-robin hands back a 502 for that turn. That + is precisely what the poll is meant to wait out — and `conftest._poll_until` + retries only transport-level errors, so an `AssertionError` raised here + would escape the retry budget and fail the test on the first blip. Skip + non-200s while converging; assert on them once converged. + """ + seen: set[str] = set() + for i in range(_PROBE_REQUESTS): + r = _chat(router_url, f"h2c-probe-{i}") + if r.status_code != 200: + if strict: + raise AssertionError(f"request {i} failed {r.status_code}: {r.text}") + continue + version = r.json().get("x_http_version") + # Fatal either way: this is a stale fake-worker image or a router that + # stopped returning the upstream body verbatim, neither of which a retry + # fixes, and without it the test cannot tell h2c from HTTP/1.1 at all. + assert version is not None, ( + "worker did not report `x_http_version` — the fake-worker image is " + "stale, or the router stopped returning the upstream body verbatim; " + "either way this test cannot tell h2c from HTTP/1.1" + ) + seen.add(version) + logger.info("protocols observed across %d requests: %s", _PROBE_REQUESTS, seen) + return seen + + +def test_router_forwards_over_h2c_to_an_http2_worker(router_url, h2c_worker): + """A worker advertising `enable_http2` is reached over HTTP/2, and the + HTTP/1.1 workers alongside it keep their own protocol.""" + # The router must first see the new pod's EndpointSlice entry and + # introspect it; until then every response comes back "1.1". + _poll_until( + lambda: "2" in _observed_protocols(router_url, strict=False), + "router forwards to the h2c worker over HTTP/2", + timeout=_CONVERGE_TIMEOUT, + interval=5, + ) + + seen = _observed_protocols(router_url, strict=True) + assert "2" in seen, f"expected an HTTP/2 forward, saw {seen}" + assert "1.1" in seen, ( + f"expected the three uvicorn workers to stay on HTTP/1.1, saw {seen}; " + "a fleet-wide protocol would have taken them with it" + ) diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 2ac07d6ed..1c87765a7 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -12,7 +12,7 @@ use sgl_router::server::app::build_router; use sgl_router::server::app_context::AppContext; use sgl_router::server::routes::chat::MAX_CHAT_BODY_BYTES; use sgl_router::tokenizer::TokenizerRegistry; -use sgl_router::workers::{Worker, WorkerRegistry}; +use sgl_router::workers::{WireProtocol, Worker, WorkerRegistry}; use axum::body::Body; use axum::http::{Request, StatusCode}; @@ -839,6 +839,7 @@ async fn forward_json_to_records_failure_on_body_drop() { let res: Result<_, ApiError> = proxy .forward_json_to( &worker.url, + WireProtocol::Http1, &breaker, "/v1/chat/completions", &headers, @@ -895,6 +896,7 @@ async fn forward_json_to_records_success_only_after_body_completes() { let res: Result<_, ApiError> = proxy .forward_json_to( &ok_worker.url, + WireProtocol::Http1, &breaker, "/v1/chat/completions", &headers, @@ -945,6 +947,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() { let res: Result<_, ApiError> = proxy .forward_streaming_to( &worker.url, + WireProtocol::Http1, &breaker, "/v1/chat/completions", &headers, @@ -995,6 +998,7 @@ async fn forward_json_to_records_failure_on_5xx() { let _: Result<_, ApiError> = proxy .forward_json_to( &worker.url, + WireProtocol::Http1, &breaker, "/v1/chat/completions", &headers, @@ -1028,6 +1032,7 @@ async fn forward_json_to_rejects_when_breaker_open() { let res = proxy .forward_json_to( &worker.url, + WireProtocol::Http1, &breaker, "/v1/chat/completions", &headers, @@ -1065,6 +1070,7 @@ async fn forward_json_to_malformed_url_returns_worker_misconfigured_and_trips_br let res = proxy .forward_json_to( "not-a-url", + WireProtocol::Http1, &breaker, "/v1/chat/completions", &headers, diff --git a/experimental/sgl-router/tests/proxy/h2c_forward.rs b/experimental/sgl-router/tests/proxy/h2c_forward.rs new file mode 100644 index 000000000..0cbca58ab --- /dev/null +++ b/experimental/sgl-router/tests/proxy/h2c_forward.rs @@ -0,0 +1,195 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Cleartext-HTTP/2 (h2c) forwarding tests for [`Proxy`]. +//! +//! These spin up an **HTTP/2-only** server (hyper's `http2::Builder`, no +//! HTTP/1.1 path) and drive `forward_json_to` against it with an explicit +//! per-request [`WireProtocol`]. Together the two tests prove what a value-only +//! unit test cannot: that the proxy's h2c client genuinely speaks HTTP/2 on the +//! wire when a request selects [`WireProtocol::H2c`] (not just that the value +//! was threaded through), and that the HTTP/1.1 client cannot reach an h2c +//! worker — so a regression that dropped the `http2` Cargo feature, removed +//! `http2_prior_knowledge()`, or mismatched `build_client`'s arms would fail +//! here rather than slip through. They also pin the per-worker design: the +//! protocol is chosen per `forward_json_to` call, not committed fleet-wide. + +use std::convert::Infallible; +use std::sync::Arc; +use std::time::Duration; + +use bytes::Bytes; +use http_body_util::{BodyExt, Full, StreamBody}; +use hyper::body::Frame; +use hyper::server::conn::http2; +use hyper::service::service_fn; +use hyper::{Request, Response}; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use sgl_router::health::circuit_breaker::CircuitBreaker; +use sgl_router::proxy::Proxy; +use sgl_router::workers::WireProtocol; +use tokio::net::TcpListener; + +/// Serve HTTP/2 only. `http2::Builder::serve_connection` speaks the HTTP/2 +/// framing protocol with no HTTP/1.1 fallback, so a client that does not +/// send the HTTP/2 connection preface cannot complete a request. Returns the +/// base URL; the accept loop is aborted when the test runtime shuts down. +async fn spawn_h2c_only_server() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + tokio::spawn(async move { + let _ = http2::Builder::new(TokioExecutor::new()) + .serve_connection( + TokioIo::new(stream), + service_fn(|_req: Request| async { + Ok::<_, Infallible>(Response::new(Full::new(Bytes::from_static( + b"{\"ok\":true}", + )))) + }), + ) + .await; + }); + } + }); + format!("http://{addr}") +} + +#[tokio::test] +async fn h2c_client_reaches_http2_only_worker() { + let url = spawn_h2c_only_server().await; + let proxy = Proxy::new(Duration::from_secs(5)).unwrap(); + let breaker = CircuitBreaker::new(); + + // Select h2c per request, as the chat handler does for a worker whose + // `/server_info` reported --enable-http2 on a cleartext URL. + let resp = proxy + .forward_json_to( + &url, + WireProtocol::H2c, + &breaker, + "/v1/chat/completions", + &axum::http::HeaderMap::new(), + Bytes::from_static(b"{}"), + ) + .await + .expect("h2c client must reach an HTTP/2-only worker"); + assert_eq!(resp.status(), 200); +} + +#[tokio::test] +async fn http1_client_cannot_reach_http2_only_worker() { + let url = spawn_h2c_only_server().await; + let proxy = Proxy::new(Duration::from_secs(5)).unwrap(); + let breaker = CircuitBreaker::new(); + + // The HTTP/1.1 client never sends the HTTP/2 preface, so the h2c-only + // server cannot serve it. This is what makes selecting the protocol + // meaningful: Http1 and H2c are different protocols on the wire, not the + // same client pointed at the same endpoint. + let res = proxy + .forward_json_to( + &url, + WireProtocol::Http1, + &breaker, + "/v1/chat/completions", + &axum::http::HeaderMap::new(), + Bytes::from_static(b"{}"), + ) + .await; + assert!( + res.is_err(), + "HTTP/1.1 client must not complete a request against an HTTP/2-only worker, got {res:?}", + ); +} + +/// Serve HTTP/2 only, answering with a multi-frame SSE body. Each `data:` +/// chunk is its own HTTP/2 DATA frame, which is the framing the real engine +/// produces for a streaming generation. +async fn spawn_h2c_only_sse_server() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + tokio::spawn(async move { + let _ = http2::Builder::new(TokioExecutor::new()) + .serve_connection( + TokioIo::new(stream), + service_fn(|_req: Request| async { + let chunks: Vec, Infallible>> = vec![ + Ok(Frame::data(Bytes::from_static(b"data: {\"i\":0}\n\n"))), + Ok(Frame::data(Bytes::from_static(b"data: {\"i\":1}\n\n"))), + Ok(Frame::data(Bytes::from_static(b"data: [DONE]\n\n"))), + ]; + let body = StreamBody::new(futures::stream::iter(chunks)); + Ok::<_, Infallible>( + Response::builder() + .header("content-type", "text/event-stream") + .body(body) + .unwrap(), + ) + }), + ) + .await; + }); + } + }); + format!("http://{addr}") +} + +/// The streaming forward — the path every chat generation actually takes — must +/// work over h2c, and must deliver the whole multi-frame body. +/// +/// `forward_streaming_to` differs from `forward_json_to` in ways HTTP/2 reaches +/// differently: it omits the request timeout, and it hands `bytes_stream()` to +/// the SSE pump rather than buffering. A break confined to SSE-over-h2c — a +/// truncated body, a stall, a mid-stream reset surfacing as success — leaves +/// the buffered tests green, so cover it directly. +#[tokio::test] +async fn h2c_client_streams_sse_from_http2_only_worker() { + let url = spawn_h2c_only_sse_server().await; + let proxy = Proxy::new(Duration::from_secs(5)).unwrap(); + let breaker = Arc::new(CircuitBreaker::new()); + + let first_byte_seen = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let flag = Arc::clone(&first_byte_seen); + + let resp = proxy + .forward_streaming_to( + &url, + WireProtocol::H2c, + &breaker, + "/v1/chat/completions", + &axum::http::HeaderMap::new(), + Bytes::from_static(b"{}"), + None, + Some(Box::new(move || { + flag.store(true, std::sync::atomic::Ordering::SeqCst); + })), + None, + ) + .await + .expect("h2c client must stream from an HTTP/2-only worker"); + assert_eq!(resp.status(), 200); + + let body = resp + .into_body() + .collect() + .await + .expect("streaming body must complete over h2c") + .to_bytes(); + let text = String::from_utf8(body.to_vec()).unwrap(); + + // Every frame arrives, in order — not just the first. + assert!(text.contains("\"i\":0"), "missing first chunk: {text}"); + assert!(text.contains("\"i\":1"), "missing second chunk: {text}"); + assert!( + text.contains("[DONE]"), + "stream truncated before DONE: {text}" + ); + assert!( + first_byte_seen.load(std::sync::atomic::Ordering::SeqCst), + "on_first_byte must fire for an h2c stream (TTFT accounting depends on it)", + ); +} diff --git a/experimental/sgl-router/tests/proxy/inbound_h2c.rs b/experimental/sgl-router/tests/proxy/inbound_h2c.rs new file mode 100644 index 000000000..fbf6c9feb --- /dev/null +++ b/experimental/sgl-router/tests/proxy/inbound_h2c.rs @@ -0,0 +1,116 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Inbound cleartext-HTTP/2 (h2c) listener tests for the router's own server. +//! +//! Where `h2c_forward.rs` proves the router's *outbound* client speaks h2c to +//! an HTTP/2-only worker, these prove the *inbound* side: the router's +//! `axum::serve` listener accepts an h2c prior-knowledge client on the same +//! cleartext port it serves HTTP/1.1 on, auto-negotiating per connection. +//! +//! Protocol negotiation happens at the connection layer, below the tower +//! service — so unlike the `oneshot` tests in `chat_routing.rs`, these must +//! drive a real socket via `axum::serve` (mirroring `main.rs`). +//! +//! What actually compiles the listener's h2 path is `hyper-util/http2`, reached +//! through `server-auto`, and several dependencies enable it. So these tests do +//! NOT fail if axum's own `http2` feature is dropped — verified by removing it. +//! Their job is the property an operator depends on, whatever the feature graph +//! happens to look like: one cleartext port serving both h2c prior-knowledge +//! and HTTP/1.1. That is what would break if a dependency change silently took +//! `hyper-util/http2` away. + +use std::sync::Arc; +use std::time::Duration; + +use sgl_router::config::PolicyKind; +use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry; +use sgl_router::proxy::Proxy; +use sgl_router::server::app::build_router; +use sgl_router::server::app_context::AppContext; +use sgl_router::tokenizer::TokenizerRegistry; +use sgl_router::workers::WorkerRegistry; + +use axum::http::Version; + +const TEST_TIMEOUT: Duration = Duration::from_secs(5); + +fn build_ctx() -> Arc { + // Reuse the shared fixture rather than restating the whole `Config`: this + // test is about the listener, not the routing policy, and a local literal + // would be one more site to edit every time `ModelConfig` gains a field. + let mut cfg = crate::common::cache_aware_fixture::config(); + cfg.model.policy = PolicyKind::RoundRobin; + cfg.model.cache_aware = None; + + let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); + let registry = Arc::new(WorkerRegistry::default()); + let policies = Arc::new(build_policy_registry(&cfg).unwrap()); + let proxy = Arc::new(Proxy::new(TEST_TIMEOUT).unwrap()); + Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) +} + +/// Spawn the real router (`build_router`) behind `axum::serve` on an ephemeral +/// port — the exact serve path used in `main.rs`. `/healthz` returns 200 +/// unconditionally, so no worker is needed. Returns the base URL; the accept +/// loop is dropped when the test runtime shuts down. +async fn spawn_router() -> String { + let ctx = build_ctx(); + ctx.mark_ready(); + let app = build_router(ctx); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + format!("http://{addr}") +} + +#[tokio::test] +async fn inbound_accepts_h2c_prior_knowledge() { + let base = spawn_router().await; + + // `http2_prior_knowledge()` sends the HTTP/2 connection preface directly + // over cleartext (no ALPN, no h1 upgrade) — the same way a service-mesh + // sidecar dials h2c. The listener must recognize the preface and serve h2. + let client = reqwest::Client::builder() + .http2_prior_knowledge() + .build() + .unwrap(); + + let resp = client + .get(format!("{base}/healthz")) + .send() + .await + .expect("h2c prior-knowledge client must reach the router listener"); + + assert_eq!(resp.status(), 200); + assert_eq!( + resp.version(), + Version::HTTP_2, + "listener must negotiate HTTP/2 for an h2c prior-knowledge client", + ); +} + +#[tokio::test] +async fn inbound_still_accepts_http1() { + let base = spawn_router().await; + + // The default reqwest client speaks HTTP/1.1 over cleartext. Enabling h2c + // must not break existing HTTP/1.1 callers (load balancers, probes, curl) + // — the auto builder serves both on the same port. + let client = reqwest::Client::new(); + + let resp = client + .get(format!("{base}/healthz")) + .send() + .await + .expect("HTTP/1.1 client must still reach the router listener"); + + assert_eq!(resp.status(), 200); + assert_eq!( + resp.version(), + Version::HTTP_11, + "an HTTP/1.1 client must keep negotiating HTTP/1.1 on the same port", + ); +} diff --git a/experimental/sgl-router/tests/proxy/main.rs b/experimental/sgl-router/tests/proxy/main.rs index 179cfa2f5..cd360f173 100644 --- a/experimental/sgl-router/tests/proxy/main.rs +++ b/experimental/sgl-router/tests/proxy/main.rs @@ -16,9 +16,12 @@ mod chat_routing; mod external_indexer_routing; mod failover; mod graceful_shutdown; +mod h2c_forward; mod header_forwarding; +mod inbound_h2c; mod pd_bootstrap_injection; mod pd_pool_isolation; +mod pd_protocol_binding; mod radix_tree_routing; mod roundrobin_input_ids; mod shared_prefill_admission; diff --git a/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs b/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs new file mode 100644 index 000000000..d7065e448 --- /dev/null +++ b/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs @@ -0,0 +1,260 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Binds protocol *resolution* to protocol *use*, across the PD split. +//! +//! Everywhere else the two halves are tested apart: the manager tests assert +//! what lands on the registry, and `h2c_forward.rs` passes a `WireProtocol` to +//! the proxy by hand. Nothing asserts that the protocol a worker resolved to is +//! the one its own forward actually uses — and in the PD arm of +//! `chat_completions` that join is three separate expressions, a +//! `prefill_protocol` captured before a `tokio::spawn` plus two live +//! `decode_worker.protocol()` reads. Passing the wrong one of those compiles, +//! and every other test in the suite stays green. +//! +//! So make the two workers disagree and let the wire enforce it: the prefill +//! mock speaks **HTTP/2 only**, the decode mock speaks **HTTP/1.1 only**, and +//! each is registered with the matching protocol. Any mix-up sends a client at +//! a server that cannot answer it. This is also why neither mock can be the +//! shared `MockWorker` — that one runs on `axum::serve`, which answers both +//! protocols and would pass no matter which was selected. + +use std::convert::Infallible; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use bytes::Bytes; +use http_body_util::{BodyExt, Full}; +use hyper::server::conn::{http1, http2}; +use hyper::service::service_fn; +use hyper::Response as HyperResponse; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use sgl_router::config::{ + ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind, + ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig, +}; +use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; +use sgl_router::policies::factory::build_registry_with_defaults; +use sgl_router::proxy::Proxy; +use sgl_router::server::app::build_router; +use sgl_router::server::app_context::AppContext; +use sgl_router::tokenizer::TokenizerRegistry; +use sgl_router::workers::{WireProtocol, WorkerRegistry}; +use tokio::net::TcpListener; +use tower::ServiceExt; + +/// A chat-completions response shaped enough for the handler to return 200. +fn chat_completion_json() -> Bytes { + Bytes::from( + serde_json::json!({ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "tiny", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }) + .to_string(), + ) +} + +/// Serve HTTP/2 only, recording that a request body arrived. An HTTP/1.1 +/// client cannot complete a request here — it never sends the HTTP/2 preface. +async fn spawn_h2_only_worker(hits: Arc>) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + let hits = Arc::clone(&hits); + tokio::spawn(async move { + let _ = http2::Builder::new(TokioExecutor::new()) + .serve_connection( + TokioIo::new(stream), + service_fn(move |_req| { + let hits = Arc::clone(&hits); + async move { + *hits.lock().unwrap() += 1; + Ok::<_, Infallible>(HyperResponse::new(Full::new( + chat_completion_json(), + ))) + } + }), + ) + .await; + }); + } + }); + format!("http://{addr}") +} + +/// Serve HTTP/1.1 only. `http1::Builder` speaks no HTTP/2, so an h2c client +/// sending the preface cannot be served here. +async fn spawn_http1_only_worker(hits: Arc>) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + let hits = Arc::clone(&hits); + tokio::spawn(async move { + let _ = http1::Builder::new() + .serve_connection( + TokioIo::new(stream), + service_fn(move |_req| { + let hits = Arc::clone(&hits); + async move { + *hits.lock().unwrap() += 1; + Ok::<_, Infallible>(HyperResponse::new(Full::new( + chat_completion_json(), + ))) + } + }), + ) + .await; + }); + } + }); + format!("http://{addr}") +} + +fn config() -> Config { + Config { + server: ServerConfig { + host: "0".into(), + port: 0, + }, + observability: ObservabilityConfig::default(), + model: ModelConfig { + id: "tiny".into(), + tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(), + policy: PolicyKind::RoundRobin, + decode_policy: Default::default(), + bucket_config: None, + circuit_breaker: None, + cache_aware: None, + sticky: None, + affinity: None, + fused: None, + eligibility: None, + }, + discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { + urls: vec!["http://placeholder:0".into()], + }), + proxy: ProxyConfig::default(), + active_load: ActiveLoadConfig::default(), + } +} + +/// Register both workers with their own resolved protocol, the way +/// `manager::register_one` does after introspecting `/server_info`. +fn build_ctx(prefill_url: String, decode_url: String) -> Arc { + let cfg = config(); + let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); + let registry = Arc::new(WorkerRegistry::default()); + + let prefill_id = WorkerId("p1".into()); + let decode_id = WorkerId("d1".into()); + // The two disagree on purpose. This is the state a mixed fleet reaches + // when only some engines run with --enable-http2. + registry + .add_with_cb( + WorkerSpec { + id: prefill_id.clone(), + url: prefill_url, + mode: WorkerMode::Prefill, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: Some(8997), + }, + None, + WireProtocol::H2c, + ) + .unwrap(); + registry + .add_with_cb( + WorkerSpec { + id: decode_id.clone(), + url: decode_url, + mode: WorkerMode::Decode, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + }, + None, + WireProtocol::Http1, + ) + .unwrap(); + + let policies = Arc::new(build_registry_with_defaults(&cfg).unwrap()); + let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()); + Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies)) +} + +fn chat_request() -> Request { + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + serde_json::to_vec(&serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + "stream": false + })) + .unwrap(), + )) + .unwrap() +} + +/// Each PD leg must forward over the protocol *its own* worker resolved. +/// +/// The decode leg is awaited, so a protocol mix-up there fails the response +/// outright. The prefill leg is detached, so its mix-up shows up as a body that +/// never arrives — poll for it rather than reading once. +#[tokio::test] +async fn pd_legs_each_use_their_own_workers_protocol() { + let prefill_hits = Arc::new(Mutex::new(0usize)); + let decode_hits = Arc::new(Mutex::new(0usize)); + let prefill_url = spawn_h2_only_worker(Arc::clone(&prefill_hits)).await; + let decode_url = spawn_http1_only_worker(Arc::clone(&decode_hits)).await; + + let app = build_router(build_ctx(prefill_url, decode_url)); + let res = app.oneshot(chat_request()).await.unwrap(); + + // Decode answered, so the decode leg used HTTP/1.1 — had it been handed the + // prefill worker's H2c, this HTTP/1.1-only server could not have replied. + assert_eq!( + res.status(), + StatusCode::OK, + "decode leg must forward over the decode worker's own protocol (HTTP/1.1)", + ); + let body = res.into_body().collect().await.unwrap().to_bytes(); + assert!( + !body.is_empty(), + "decode response body must reach the client", + ); + + // Prefill is spawn-and-forget; it races the response back to the client. + let reached = tokio::time::timeout(Duration::from_secs(3), async { + loop { + if *prefill_hits.lock().unwrap() > 0 { + return true; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await; + assert!( + reached.is_ok(), + "prefill leg must forward over the prefill worker's own protocol (h2c); \ + an HTTP/1.1 client cannot reach this HTTP/2-only worker", + ); + assert_eq!( + *decode_hits.lock().unwrap(), + 1, + "decode worker should be hit exactly once", + ); +}