[router] Speak cleartext h2c on both edges: serve it inbound, forward it outbound (#39006)
Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Kangyan Zhou
Claude Opus 5
parent
c1f5b4736a
commit
afde31a2f5
@@ -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,
|
||||
|
||||
@@ -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<hyper::body::Incoming>| 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<hyper::body::Incoming>| async {
|
||||
let chunks: Vec<Result<Frame<Bytes>, 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)",
|
||||
);
|
||||
}
|
||||
@@ -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<AppContext> {
|
||||
// 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",
|
||||
);
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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<Mutex<usize>>) -> 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<Mutex<usize>>) -> 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<AppContext> {
|
||||
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<Body> {
|
||||
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",
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user