[router] Resolve a wire protocol per worker at registration (#39004)

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:
Kangyan-Zhou
2026-09-15 16:54:37 -07:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 5
parent 0dabef3d30
commit 63845a1bb2
6 changed files with 768 additions and 104 deletions
@@ -6,10 +6,11 @@
//! Two concurrent requests, because the worker answers two different //! Two concurrent requests, because the worker answers two different
//! questions on two different endpoints: `/model_info` reports the identity //! questions on two different endpoints: `/model_info` reports the identity
//! the worker currently serves under (a weight update moves it), while //! the worker currently serves under (a weight update moves it), while
//! `/server_info` reports its launch configuration — kv-event publisher and //! `/server_info` reports its launch configuration — kv-event publisher,
//! disaggregation role. The result is dispatched by the manager: registry //! disaggregation role, and whether the engine serves cleartext h2c. The
//! consumes `served_model_name`, the optional `KvEventIndex` consumes the //! result is dispatched by the manager: registry consumes
//! resolved `EventConfig`. //! `served_model_name` and `enable_http2`, the optional `KvEventIndex`
//! consumes the resolved `EventConfig`.
//! //!
//! `served_model_name` is taken from `/model_info`, falling back to //! `served_model_name` is taken from `/model_info`, falling back to
//! `/server_info` for workers that predate the field there. //! `/server_info` for workers that predate the field there.
@@ -56,6 +57,14 @@ pub struct ServerInfo {
pub served_model_name: Option<String>, pub served_model_name: Option<String>,
pub event_config: Option<EventConfig>, pub event_config: Option<EventConfig>,
pub disaggregation_role: Option<DisaggregationRole>, pub disaggregation_role: Option<DisaggregationRole>,
/// Whether the engine was launched with `--enable-http2` (Granian,
/// serving cleartext h2c + HTTP/1.1), from the `/server_info` launch
/// record. `Some(true)` ⇒ the router may forward over h2c;
/// `Some(false)` / `None` ⇒ stay on HTTP/1.1 — which is also what a
/// worker whose `/server_info` did not answer gets, so an unread
/// protocol costs throughput and never correctness. Consumed by
/// `manager::register_one` to set [`crate::workers::WireProtocol`].
pub enable_http2: Option<bool>,
} }
/// PD classification derived from a worker's `/server_info` response. /// PD classification derived from a worker's `/server_info` response.
@@ -157,6 +166,7 @@ impl WorkerIntrospector {
served_model_name, served_model_name,
event_config, event_config,
disaggregation_role, disaggregation_role,
enable_http2: parsed.enable_http2,
} }
} }
@@ -327,7 +337,8 @@ pub(crate) fn resolve_event_config(
} }
/// Projection of `/model_info` used by the introspector: the identity the /// Projection of `/model_info` used by the introspector: the identity the
/// worker currently serves under, which a weight update moves. /// worker currently serves under. `#[serde(default)]` so an engine that
/// predates the field still deserialises.
#[derive(Debug, Default, Deserialize)] #[derive(Debug, Default, Deserialize)]
struct ModelInfoBody { struct ModelInfoBody {
#[serde(default)] #[serde(default)]
@@ -360,6 +371,12 @@ struct ServerInfoBody {
/// bootstrap server binds to exactly this port (no internal offset). /// bootstrap server binds to exactly this port (no internal offset).
#[serde(default)] #[serde(default)]
disaggregation_bootstrap_port: Option<u16>, disaggregation_bootstrap_port: Option<u16>,
/// `ServerArgs.enable_http2`. `true` ⇒ the engine runs Granian and
/// serves cleartext h2c alongside HTTP/1.1, so the router may forward
/// to it with prior knowledge. Absent on older SGLang versions that
/// predate the flag.
#[serde(default)]
enable_http2: Option<bool>,
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
@@ -706,6 +723,41 @@ mod tests {
assert_eq!(got.disaggregation_role, Some(DisaggregationRole::Plain)); assert_eq!(got.disaggregation_role, Some(DisaggregationRole::Plain));
} }
/// `enable_http2: true` is surfaced so the manager forwards over h2c.
#[tokio::test]
async fn fetch_surfaces_enable_http2_true() {
let (url, _shutdown) = spawn_fake_worker_with_model_info(
json!({"served_model_name": "m", "enable_http2": true}),
Some(json!({"served_model_name": "m"})),
)
.await;
let got = fast_introspector().fetch(&url).await;
assert_eq!(got.enable_http2, Some(true));
}
/// An explicit `enable_http2: false` (HTTP/1.1-only engine) is surfaced
/// as `Some(false)`, distinct from the older-SGLang absent case.
#[tokio::test]
async fn fetch_surfaces_enable_http2_false() {
let (url, _shutdown) = spawn_fake_worker_with_model_info(
json!({"served_model_name": "m", "enable_http2": false}),
Some(json!({"served_model_name": "m"})),
)
.await;
let got = fast_introspector().fetch(&url).await;
assert_eq!(got.enable_http2, Some(false));
}
/// Older SGLang predates `enable_http2`; its absence must read as
/// `None` (the manager then keeps the safe HTTP/1.1 default), not as a
/// parse failure.
#[tokio::test]
async fn fetch_enable_http2_absent_is_none() {
let (url, _shutdown) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
let got = fast_introspector().fetch(&url).await;
assert_eq!(got.enable_http2, None);
}
/// Partial data (`prefill` mode with no bootstrap port) returns /// Partial data (`prefill` mode with no bootstrap port) returns
/// `None` so the manager keeps the discovery backend's /// `None` so the manager keeps the discovery backend's
/// classification. The alternative — forcing Plain — would silently /// classification. The alternative — forcing Plain — would silently
+409 -79
View File
@@ -7,7 +7,7 @@ use crate::health::circuit_breaker::CircuitBreakerConfig;
use crate::policies::active_load::ActiveLoadRegistry; use crate::policies::active_load::ActiveLoadRegistry;
use crate::policies::kv_events::KvEventIndex; use crate::policies::kv_events::KvEventIndex;
use crate::workers::introspect::{DisaggregationRole, WorkerIntrospector}; use crate::workers::introspect::{DisaggregationRole, WorkerIntrospector};
use crate::workers::WorkerRegistry; use crate::workers::{WireProtocol, WorkerRegistry};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
@@ -15,14 +15,13 @@ use tokio::sync::mpsc;
use tokio::task::JoinHandle; use tokio::task::JoinHandle;
/// Production reconcile cadence. Workers that register without resolving /// Production reconcile cadence. Workers that register without resolving
/// their model IDs (a `/server_info` introspection that failed at `Added` /// their model IDs (an introspection that failed at `Added` time — e.g. the
/// time — e.g. the EndpointSlice flipped `ready=true` before the engine's /// EndpointSlice flipped `ready=true` before the engine's HTTP server could
/// scheduler-backed `/server_info` could answer) are re-introspected on /// answer) are re-introspected on this interval until they join their model
/// this interval until they join their model pool. The worst-case /// pool. The worst-case "registered but invisible" window is about one
/// "registered but invisible" window is about one interval plus the /// interval plus the introspection round-trip; steady state costs one cheap
/// introspection round-trip; steady state costs one cheap registry scan /// registry scan per interval. See `reconcile_unresolved_workers` for the
/// per interval. See `reconcile_unresolved_workers` for the (benign) /// (benign) case of a worker that answers but never advertises a model name.
/// case of a worker that answers but never advertises a model name.
const RECONCILE_INTERVAL: Duration = Duration::from_secs(30); const RECONCILE_INTERVAL: Duration = Duration::from_secs(30);
/// Resolve the circuit-breaker config for all model IDs carried by a spec. /// Resolve the circuit-breaker config for all model IDs carried by a spec.
@@ -42,6 +41,59 @@ fn cb_config_for_spec(spec: &WorkerSpec, cfg: &Config) -> Option<CircuitBreakerC
None None
} }
/// Whether `worker_url` is one the router will dial in cleartext.
///
/// The scheme test cannot be `Url::parse(worker_url)?.scheme() == "http"`: a
/// schemeless `worker-0.engines.svc:8000` parses with `worker-0.engines.svc`
/// as the *scheme*, so that test answers "not cleartext" for a shape that is
/// nothing of the sort. Normalize the way `parse_bootstrap_host` does — retry
/// under an explicit `http://` — and read the scheme off that.
///
/// A schemeless URL therefore reads as cleartext here. That answer never
/// reaches the wire: `proxy::parse_worker_url` is a bare `Url::parse`, which
/// *accepts* `worker-0.engines.svc:8000` (scheme `worker-0.engines.svc`,
/// opaque path `8000`), and the request then dies at `worker_url.join(path)`
/// with `RelativeUrlWithCannotBeABaseBase` before any client is used. So
/// `Some(true)` costs nothing today and stays correct if that path ever grows
/// the same normalization; answering `None` would only mislabel the reason in
/// the log.
fn dials_cleartext(worker_url: &str) -> Option<bool> {
if worker_url.contains("://") {
return url::Url::parse(worker_url)
.ok()
.map(|u| u.scheme() == "http");
}
url::Url::parse(&format!("http://{worker_url}"))
.ok()
.map(|_| true)
}
/// Resolve the wire protocol for a worker from its `/server_info`
/// (`enable_http2`) and whether the router dials it in cleartext
/// (`dials_cleartext`, `None` when the URL did not parse).
///
/// Both inputs are fixed for the worker's lifetime, so this runs once per
/// worker, before it is registered.
///
/// Upgrades to [`WireProtocol::H2c`] only when the engine self-reports
/// `--enable-http2` **and** the worker URL is cleartext. h2c is HTTP/2 with
/// prior knowledge — no negotiation — so sending it anywhere that is not
/// known to serve it fails every request; the cleartext gate is what bounds
/// 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`].
fn resolve_protocol(enable_http2: Option<bool>, cleartext: Option<bool>) -> WireProtocol {
match (enable_http2, cleartext) {
(Some(true), Some(true)) => WireProtocol::H2c,
_ => WireProtocol::Http1,
}
}
pub async fn run(rx: mpsc::Receiver<DiscoveryEvent>, registry: Arc<WorkerRegistry>) { pub async fn run(rx: mpsc::Receiver<DiscoveryEvent>, registry: Arc<WorkerRegistry>) {
run_with_config(rx, registry, None, None, None).await; run_with_config(rx, registry, None, None, None).await;
} }
@@ -51,6 +103,11 @@ pub async fn run(rx: mpsc::Receiver<DiscoveryEvent>, registry: Arc<WorkerRegistr
/// on every worker add / remove, and an optional active-load registry /// on every worker add / remove, and an optional active-load registry
/// that is asked to forget per-worker counters on `Removed`. /// that is asked to forget per-worker counters on `Removed`.
/// ///
/// The forwarding wire protocol (HTTP/1.1 vs cleartext h2c) is resolved per
/// worker from its `/server_info` and carried into the registered
/// [`crate::workers::Worker`] at construction (see `register_one`), so the
/// manager does not need a handle to the proxy.
///
/// When `kv_index` is `None`, KV-event and load-subscriber state is disabled; when /// When `kv_index` is `None`, KV-event and load-subscriber state is disabled; when
/// `active_load` is `None` the active-load bookkeeping is not pruned /// `active_load` is `None` the active-load bookkeeping is not pruned
/// on worker removal (leaks one `WorkerCounters` slot per departed /// on worker removal (leaks one `WorkerCounters` slot per departed
@@ -58,7 +115,7 @@ pub async fn run(rx: mpsc::Receiver<DiscoveryEvent>, registry: Arc<WorkerRegistr
/// `cfg` is `None` the default CB config is used for every worker /// `cfg` is `None` the default CB config is used for every worker
/// (threshold = 3). /// (threshold = 3).
/// ///
/// Uses the default HTTP client (2-second timeout) for `/server_info` /// Uses the default HTTP client (2-second timeout) for worker
/// introspection. Tests that want a tighter timeout call /// introspection. Tests that want a tighter timeout call
/// [`run_with_introspector`] directly. /// [`run_with_introspector`] directly.
pub async fn run_with_config( pub async fn run_with_config(
@@ -303,17 +360,15 @@ async fn handle_discovery_event(
} }
} }
/// Re-introspect workers that registered without resolving their model /// Re-introspect workers that registered without resolving their model IDs.
/// IDs.
/// ///
/// A worker lands in the registry with empty `model_ids` when its /// A worker lands in the registry with empty `model_ids` when its
/// `/server_info` introspection failed at `Added` time (e.g. the /// introspection failed at `Added` time (e.g. the EndpointSlice flips
/// EndpointSlice flips `ready=true` before the engine can answer the /// `ready=true` before the engine's HTTP server can answer, and the
/// scheduler-backed `/server_info` round-trip, and the introspector's /// introspector's bounded retry budget is exhausted). Such a worker is
/// bounded retry budget is exhausted). Such a worker is present in /// present in `by_id` but absent from every `by_model` pool, so it gets zero
/// `by_id` but absent from every `by_model` pool, so it gets zero traffic /// traffic — and it never recovers on its own: discovery re-lists carry the
/// — and it never recovers on its own: discovery re-lists carry the same /// same empty `model_ids` the backend always emits, so they produce no new
/// empty `model_ids` the backend always emits, so they produce no new
/// `Added` event. /// `Added` event.
/// ///
/// This pass re-runs `register_one` (an idempotent registry upsert + /// This pass re-runs `register_one` (an idempotent registry upsert +
@@ -325,12 +380,11 @@ async fn handle_discovery_event(
/// mid-reconcile awaits the in-flight handle before clearing the /// mid-reconcile awaits the in-flight handle before clearing the
/// registry, so a worker that genuinely left is not resurrected. /// registry, so a worker that genuinely left is not resurrected.
/// ///
/// A worker that answers `/server_info` but never advertises a /// A worker that answers but never advertises a `served_model_name` also
/// `served_model_name` also stays in this set and is re-introspected /// stays in this set and is re-introspected every interval — a benign,
/// every interval. That is a benign, bounded poll (one cheap round-trip /// bounded poll, not a leak. It is never escalated, so the per-attempt
/// per worker per interval), not a leak — but it is also never escalated, /// logging stays at `debug!`; the introspector emits the `warn!` that
/// so the per-attempt failure logging stays at `debug!`; the introspector /// surfaces a persistently failing worker.
/// itself emits the `warn!` that surfaces a persistently failing worker.
fn reconcile_unresolved_workers( fn reconcile_unresolved_workers(
registry: &Arc<WorkerRegistry>, registry: &Arc<WorkerRegistry>,
cfg: &Option<Arc<Config>>, cfg: &Option<Arc<Config>>,
@@ -348,31 +402,33 @@ fn reconcile_unresolved_workers(
// finish rather than racing a second introspection. // finish rather than racing a second introspection.
continue; continue;
} }
// Rebuild a discovery-shaped spec: empty `model_ids` so let registry_t = registry.clone();
// `register_one` re-resolves them from `/server_info`; current let introspector_t = introspector.clone();
// mode + bootstrap_port as the seed (`register_one` re-applies let worker_url = worker.url.clone();
// any `/server_info` override). // Rebuild a discovery-shaped spec: empty `model_ids` so `register_one`
// re-resolves them; current mode + bootstrap_port as the seed
// (`register_one` re-applies any `/server_info` override).
let spec = WorkerSpec { let spec = WorkerSpec {
id: id.clone(), id: id.clone(),
url: worker.url.clone(), url: worker_url.clone(),
mode: worker.mode(), mode: worker.mode(),
model_ids: Vec::new(), model_ids: Vec::new(),
bootstrap_port: worker.bootstrap_port(), bootstrap_port: worker.bootstrap_port(),
}; };
// `debug!` not `info!`: this fires every interval for each // `debug!` not `info!`: this fires every interval for each
// still-unresolved worker, so info-level would spam for a worker // still-unresolved worker, so info-level would spam for one that is
// that is permanently model-less. The introspector logs the // permanently model-less. The introspector logs the underlying failure
// underlying `/server_info` failure at `warn!` on each attempt, // at `warn!` on each attempt, which is the operator-facing signal.
// which is the operator-facing signal.
tracing::debug!( tracing::debug!(
worker_id = %id, worker_id = %id,
worker_url = %worker.url, worker_url = %worker_url,
"reconcile: re-introspecting worker that registered without model_ids", "reconcile: re-introspecting worker that registered without model_ids",
); );
let registry_t = registry.clone();
let cfg_t = cfg.clone(); let cfg_t = cfg.clone();
let kv_index_t = kv_index.clone(); let kv_index_t = kv_index.clone();
let introspector_t = introspector.clone(); // Safe to go back through the registry upsert only because this worker
// is in no model pool: the fresh `Worker` it builds discards a breaker
// and load counters that a model-less worker has never accumulated.
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
register_one(spec, registry_t, cfg_t, kv_index_t, introspector_t).await; register_one(spec, registry_t, cfg_t, kv_index_t, introspector_t).await;
}); });
@@ -380,6 +436,50 @@ fn reconcile_unresolved_workers(
} }
} }
/// Explain a resolved protocol at the level an operator needs: the h2c upgrade
/// is a behaviour change worth an `info!`, an engine that asked for HTTP/2 and
/// did not get it should say why, and a worker whose flag was never read should
/// not look the same as one that reported `false`.
///
/// Takes the same `cleartext` that `resolve_protocol` was given, so the log and
/// the decision cannot disagree about what the router would dial.
fn log_protocol_resolution(worker_url: &str, enable_http2: Option<bool>, cleartext: Option<bool>) {
match (enable_http2, cleartext) {
(Some(true), Some(true)) => tracing::info!(
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.
(Some(true), Some(false)) => tracing::info!(
worker_url = %worker_url,
"/server_info reports --enable-http2 on a TLS worker; \
forwarding over HTTP/1.1",
),
// `dials_cleartext` could not parse the URL. Reaching this at all means
// `/server_info` answered over a URL the scheme check then rejected, so
// the worker is misconfigured rather than merely un-upgradable — every
// forward to it will fail in `proxy::parse_worker_url` or at
// `worker_url.join(path)`. Warn rather than inform.
(Some(true), None) => tracing::warn!(
worker_url = %worker_url,
"/server_info reports --enable-http2 but the worker URL did not \
parse; cannot classify the endpoint, and forwards to it are \
expected to fail",
),
// Never read: `/server_info` did not answer, or the engine predates the
// flag. Distinct from an explicit `false`, and worth saying out loud —
// once a worker has a model id `reconcile_unresolved_workers` stops
// 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",
),
// The engine explicitly disabled it. Nothing to explain.
(Some(false), _) => {}
}
}
/// Onboard a single worker: introspect once, then dispatch the result /// Onboard a single worker: introspect once, then dispatch the result
/// to the registry and (if enabled) the KV-event index. Failure of any /// to the registry and (if enabled) the KV-event index. Failure of any
/// step is logged inside the call chain; we still register the worker /// step is logged inside the call chain; we still register the worker
@@ -427,8 +527,19 @@ async fn register_one(
spec.bootstrap_port = new_port; spec.bootstrap_port = new_port;
} }
} }
let cleartext = dials_cleartext(&worker_url);
let protocol = resolve_protocol(info.enable_http2, cleartext);
// Captured before the insert: `reconcile_unresolved_workers` re-runs this
// function every interval for a worker that never advertises a model name,
// so logging unconditionally would repeat the same line for the life of the
// process. Logging only a new or changed resolution keeps the reconcile
// path quiet, matching why its own progress message stays at `debug!`.
let previous_protocol = registry.get(&spec.id).map(|w| w.protocol());
let cb = cfg.as_ref().and_then(|c| cb_config_for_spec(&spec, c)); let cb = cfg.as_ref().and_then(|c| cb_config_for_spec(&spec, c));
if let Err(e) = registry.add_with_cb(spec, cb) { // `protocol` rides beside the spec rather than on it: `WorkerSpec` is the
// serde wire type for `DiscoveryEvent`, and no discovery backend can know
// a worker's protocol.
if let Err(e) = registry.add_with_cb(spec, cb, protocol) {
// Mixed PD + plain on the same model is rejected at registration // Mixed PD + plain on the same model is rejected at registration
// time. Log loudly so the operator notices the conflicting // time. Log loudly so the operator notices the conflicting
// worker — the alternative (silently dropping into either pool) // worker — the alternative (silently dropping into either pool)
@@ -442,6 +553,11 @@ async fn register_one(
); );
return; return;
} }
// After the insert, so the log describes a worker that is actually taking
// traffic — a spec refused above never reaches the wire at all.
if previous_protocol != Some(protocol) {
log_protocol_resolution(&worker_url, info.enable_http2, cleartext);
}
if let Some(idx) = kv_index { if let Some(idx) = kv_index {
// Pass the pre-resolved EventConfig so the KvEventIndex does // Pass the pre-resolved EventConfig so the KvEventIndex does
// not issue a second `/server_info` round-trip. // not issue a second `/server_info` round-trip.
@@ -509,19 +625,54 @@ mod tests {
assert_eq!(cb.cool_down, Duration::from_secs(60)); assert_eq!(cb.cool_down, Duration::from_secs(60));
} }
/// Helper: spawn a tiny fake worker that returns the supplied JSON body #[test]
/// on `GET /server_info`. Returns the worker URL + a shutdown channel. fn resolve_protocol_upgrades_to_h2c_only_for_cleartext_http2_engine() {
async fn spawn_fake_server_info_worker(body: Value) -> (String, oneshot::Sender<()>) { // The one case that gets h2c: engine self-reports --enable-http2 and
let body = Arc::new(body); // we dial cleartext http://.
assert_eq!(
resolve_protocol(Some(true), dials_cleartext("http://10.0.0.1:30000")),
WireProtocol::H2c,
);
}
#[test]
fn resolve_protocol_stays_http1_for_https_even_with_http2() {
// A TLS engine with --enable-http2 serves h2-over-TLS, not cleartext
// h2c; the router dials cleartext, so it must stay on HTTP/1.1
// (which negotiates fine over TLS) rather than break every request.
assert_eq!(
resolve_protocol(Some(true), dials_cleartext("https://10.0.0.1:30000")),
WireProtocol::Http1,
);
}
#[test]
fn resolve_protocol_stays_http1_when_http2_disabled_or_unknown() {
// Explicit false (HTTP/1.1-only Uvicorn) and absent field (older
// SGLang) both keep the safe default.
assert_eq!(
resolve_protocol(Some(false), dials_cleartext("http://10.0.0.1:30000")),
WireProtocol::Http1,
);
assert_eq!(
resolve_protocol(None, dials_cleartext("http://10.0.0.1:30000")),
WireProtocol::Http1,
);
}
#[test]
fn resolve_protocol_stays_http1_for_unparsable_url() {
assert_eq!(
resolve_protocol(Some(true), dials_cleartext("not a url")),
WireProtocol::Http1,
);
}
/// Serve `app` on an ephemeral port. Returns its base URL + a shutdown
/// channel; every fake worker in this module is a `Router` plus this.
async fn serve(app: Router) -> (String, oneshot::Sender<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port(); let port = listener.local_addr().unwrap().port();
let app = Router::new().route(
"/server_info",
get(move || {
let body = body.clone();
async move { Json((*body).clone()) }
}),
);
let (tx, rx) = oneshot::channel::<()>(); let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move { tokio::spawn(async move {
let _ = axum::serve(listener, app) let _ = axum::serve(listener, app)
@@ -533,6 +684,19 @@ mod tests {
(format!("http://127.0.0.1:{port}"), tx) (format!("http://127.0.0.1:{port}"), tx)
} }
/// Answers the supplied JSON body on `GET /server_info` only.
async fn spawn_fake_server_info_worker(body: Value) -> (String, oneshot::Sender<()>) {
let body = Arc::new(body);
serve(Router::new().route(
"/server_info",
get(move || {
let body = body.clone();
async move { Json((*body).clone()) }
}),
))
.await
}
/// Reserve a TCP port and immediately drop the listener so subsequent /// Reserve a TCP port and immediately drop the listener so subsequent
/// connection attempts during the test fail fast with /// connection attempts during the test fail fast with
/// ConnectionRefused. /// ConnectionRefused.
@@ -941,43 +1105,122 @@ mod tests {
let _ = manager_handle.await; let _ = manager_handle.await;
} }
/// Fake `/server_info` worker whose readiness is switchable at /// Fake worker whose readiness is switchable at runtime. While `ready` is
/// runtime. While `ready` is false it answers `503` (mimicking an /// false both introspection endpoints answer `503` (mimicking an engine
/// engine whose EndpointSlice flipped `ready=true` before its /// whose EndpointSlice flipped `ready=true` before its HTTP server could
/// scheduler-backed `/server_info` could answer); flip `ready` to /// answer anything); flip `ready` to true and both serve `body`.
/// true and it serves `body`. async fn spawn_switchable_worker(
async fn spawn_switchable_server_info_worker(
body: Value, body: Value,
ready: Arc<std::sync::atomic::AtomicBool>, ready: Arc<std::sync::atomic::AtomicBool>,
) -> (String, oneshot::Sender<()>) { ) -> (String, oneshot::Sender<()>) {
use axum::http::StatusCode; use axum::http::StatusCode;
use axum::response::IntoResponse; use axum::response::IntoResponse;
let body = Arc::new(body); let body = Arc::new(body);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let when_ready = move || {
let port = listener.local_addr().unwrap().port(); let body = body.clone();
let app = Router::new().route( let ready = ready.clone();
"/server_info", async move {
get(move || { if ready.load(std::sync::atomic::Ordering::SeqCst) {
let body = body.clone(); Json((*body).clone()).into_response()
let ready = ready.clone(); } else {
async move { StatusCode::SERVICE_UNAVAILABLE.into_response()
if ready.load(std::sync::atomic::Ordering::SeqCst) { }
Json((*body).clone()).into_response() }
} else { };
StatusCode::SERVICE_UNAVAILABLE.into_response() serve(
Router::new()
.route("/model_info", get(when_ready.clone()))
.route("/server_info", get(when_ready)),
)
.await
}
/// Fake worker that resolves a model name but 503s `/server_info`: the
/// warming-engine shape, where the scheduler round-trip behind
/// `/server_info` is not answerable yet.
async fn spawn_worker_without_server_info(
model_info_body: Value,
) -> (String, oneshot::Sender<()>) {
use axum::http::StatusCode;
let body = Arc::new(model_info_body);
serve(
Router::new()
.route(
"/model_info",
get(move || {
let body = body.clone();
async move { Json((*body).clone()) }
}),
)
.route(
"/server_info",
get(|| async { (StatusCode::SERVICE_UNAVAILABLE, "warming up") }),
),
)
.await
}
/// A worker whose `/server_info` never answers joins its model pool on
/// the HTTP/1.1 default.
///
/// `enable_http2` comes from `/server_info`, so a worker that cannot
/// serve it has no readable protocol. It still resolves a model name and
/// becomes routable, and `reconcile_unresolved_workers` keys on empty
/// `model_ids`, so it is never revisited — the worker forwards over
/// HTTP/1.1 for its lifetime. A throughput cost, never a correctness one.
///
/// The reconcile interval here is far longer than the timeout, so the
/// result cannot be the repair loop arriving late either way.
#[tokio::test]
async fn warming_worker_without_server_info_registers_on_http1() {
let (worker_url, _shutdown) =
spawn_worker_without_server_info(json!({"served_model_name": "m"})).await;
let registry = Arc::new(WorkerRegistry::default());
let (tx, rx) = mpsc::channel::<DiscoveryEvent>(8);
let manager_handle = tokio::spawn(run_with_introspector_and_reconcile(
rx,
registry.clone(),
None,
None,
None,
fast_introspector(),
Duration::from_secs(600),
));
let id = WorkerId("w-warming".into());
tx.send(DiscoveryEvent::Added(WorkerSpec {
id: id.clone(),
url: worker_url,
mode: WorkerMode::Plain,
model_ids: Vec::new(),
bootstrap_port: None,
}))
.await
.unwrap();
let resolved = tokio::time::timeout(Duration::from_secs(3), async {
loop {
if let Some(w) = registry.get(&id) {
if !w.model_ids.is_empty() {
return w.protocol();
} }
} }
}), tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await;
assert_eq!(
resolved.expect("worker must join its model pool"),
WireProtocol::Http1,
"a worker whose /server_info never answered has no readable \
protocol, so it must take traffic on the HTTP/1.1 default rather \
than on a guess",
); );
let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move { drop(tx);
let _ = axum::serve(listener, app) let _ = manager_handle.await;
.with_graceful_shutdown(async move {
let _ = rx.await;
})
.await;
});
(format!("http://127.0.0.1:{port}"), tx)
} }
/// A worker that registers with empty `model_ids` because /// A worker that registers with empty `model_ids` because
@@ -993,8 +1236,7 @@ mod tests {
let ready = Arc::new(AtomicBool::new(false)); let ready = Arc::new(AtomicBool::new(false));
let (worker_url, _shutdown) = let (worker_url, _shutdown) =
spawn_switchable_server_info_worker(json!({"served_model_name": "m"}), ready.clone()) spawn_switchable_worker(json!({"served_model_name": "m"}), ready.clone()).await;
.await;
let registry = Arc::new(WorkerRegistry::default()); let registry = Arc::new(WorkerRegistry::default());
let (tx, rx) = mpsc::channel::<DiscoveryEvent>(8); let (tx, rx) = mpsc::channel::<DiscoveryEvent>(8);
@@ -1066,6 +1308,94 @@ mod tests {
let _ = manager_handle.await; let _ = manager_handle.await;
} }
/// A worker whose introspection failed entirely (so it registered with no
/// model IDs and the HTTP/1.1 default) must come back as h2c once the
/// engine is ready and reports `enable_http2: true`. The repair is the
/// re-registration `reconcile_unresolved_workers` already performs for a
/// model-less worker: the fresh `Worker` it builds carries the protocol
/// resolved by that same fetch, so a transient startup failure does not
/// strand an h2c-capable engine on HTTP/1.1 forever.
#[tokio::test]
async fn reconcile_upgrades_worker_to_h2c_after_transient_failure() {
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::time::timeout;
// While not-ready the worker answers 503 (introspection fails → Http1,
// empty model_ids); once ready it reports a model AND enable_http2.
let ready = Arc::new(AtomicBool::new(false));
let (worker_url, _shutdown) = spawn_switchable_worker(
json!({"served_model_name": "m", "enable_http2": true}),
ready.clone(),
)
.await;
let registry = Arc::new(WorkerRegistry::default());
let (tx, rx) = mpsc::channel::<DiscoveryEvent>(8);
let manager_handle = tokio::spawn(run_with_introspector_and_reconcile(
rx,
registry.clone(),
None,
None,
None,
fast_introspector(),
Duration::from_millis(150),
));
let id = WorkerId("w-warming".into());
tx.send(DiscoveryEvent::Added(WorkerSpec {
id: id.clone(),
url: worker_url,
mode: WorkerMode::Plain,
model_ids: Vec::new(),
bootstrap_port: None,
}))
.await
.unwrap();
// Phase 1: the warming worker is registered on the safe HTTP/1.1
// default (introspection failed → empty model_ids, no h2c).
let stuck = timeout(Duration::from_secs(2), async {
loop {
if let Some(w) = registry.get(&id) {
if w.model_ids.is_empty() && w.protocol() == WireProtocol::Http1 {
return true;
}
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await;
assert!(
stuck.is_ok(),
"a worker that failed initial introspection must register on the HTTP/1.1 default",
);
// The engine finishes coming up.
ready.store(true, Ordering::SeqCst);
// Phase 2: reconcile re-introspects and upgrades the worker to h2c.
let upgraded = timeout(Duration::from_secs(3), async {
loop {
if registry
.get(&id)
.is_some_and(|w| w.protocol() == WireProtocol::H2c)
{
return true;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await;
assert!(
upgraded.is_ok(),
"reconcile must upgrade the worker to h2c once /server_info reports enable_http2; got {:?}",
registry.get(&id).map(|w| w.protocol()),
);
drop(tx);
let _ = manager_handle.await;
}
/// Resurrection safety: a `Removed` that arrives while a reconcile /// Resurrection safety: a `Removed` that arrives while a reconcile
/// re-introspection for the same id is in-flight must NOT resurrect /// re-introspection for the same id is in-flight must NOT resurrect
/// the worker. The `Removed` handler awaits the in-flight handle (which /// the worker. The `Removed` handler awaits the in-flight handle (which
@@ -9,4 +9,5 @@ pub mod worker;
pub use introspect::{ServerInfo, WorkerIntrospector}; pub use introspect::{ServerInfo, WorkerIntrospector};
pub use registry::WorkerRegistry; pub use registry::WorkerRegistry;
pub use worker::LoadGuard; pub use worker::LoadGuard;
pub use worker::WireProtocol;
pub use worker::Worker; pub use worker::Worker;
@@ -3,7 +3,7 @@
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use crate::health::circuit_breaker::CircuitBreakerConfig; use crate::health::circuit_breaker::CircuitBreakerConfig;
use crate::workers::worker::Worker; use crate::workers::worker::{WireProtocol, Worker};
use dashmap::DashMap; use dashmap::DashMap;
use std::collections::HashSet; use std::collections::HashSet;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
@@ -51,11 +51,12 @@ pub struct WorkerRegistry {
impl WorkerRegistry { impl WorkerRegistry {
pub fn add(&self, spec: WorkerSpec) -> Result<(), AddWorkerError> { pub fn add(&self, spec: WorkerSpec) -> Result<(), AddWorkerError> {
self.add_with_cb(spec, None) self.add_with_cb(spec, None, WireProtocol::default())
} }
/// Add a worker, optionally supplying a circuit-breaker config. /// Add a worker, optionally supplying a circuit-breaker config, and with
/// Pass `None` to use the circuit-breaker default (threshold = 3). /// the forwarding protocol resolved for it. Pass `None` to use the
/// circuit-breaker default (threshold = 3).
/// ///
/// Re-adding an existing `WorkerId` is an upsert: the prior entry's /// Re-adding an existing `WorkerId` is an upsert: the prior entry's
/// `by_model` memberships are cleared first so a model that the new /// `by_model` memberships are cleared first so a model that the new
@@ -83,6 +84,7 @@ impl WorkerRegistry {
&self, &self,
spec: WorkerSpec, spec: WorkerSpec,
cb: Option<CircuitBreakerConfig>, cb: Option<CircuitBreakerConfig>,
protocol: WireProtocol,
) -> Result<(), AddWorkerError> { ) -> Result<(), AddWorkerError> {
let incoming_mode = spec.mode; let incoming_mode = spec.mode;
// Hold the write lock for the entire validate→insert sequence. // Hold the write lock for the entire validate→insert sequence.
@@ -121,7 +123,7 @@ impl WorkerRegistry {
} }
} }
} }
let w = Arc::new(Worker::with_cb_config(spec, cb)); let w = Arc::new(Worker::with_cb_config(spec, cb, protocol));
let id = w.id.clone(); let id = w.id.clone();
self.remove_locked(&id); self.remove_locked(&id);
for m in &w.model_ids { for m in &w.model_ids {
@@ -310,7 +312,11 @@ mod tests {
use std::time::Duration; use std::time::Duration;
let r = WorkerRegistry::default(); let r = WorkerRegistry::default();
let _ = r.add_with_cb(spec("ok", WorkerMode::Plain, &["m"]), None); let _ = r.add_with_cb(
spec("ok", WorkerMode::Plain, &["m"]),
None,
WireProtocol::default(),
);
// Give "bad" a threshold=1 breaker so a single record_failure // Give "bad" a threshold=1 breaker so a single record_failure
// flips it to Open. // flips it to Open.
let _ = r.add_with_cb( let _ = r.add_with_cb(
@@ -319,6 +325,7 @@ mod tests {
threshold: NonZeroU32::new(1).unwrap(), threshold: NonZeroU32::new(1).unwrap(),
cool_down: Duration::from_secs(30), cool_down: Duration::from_secs(30),
}), }),
WireProtocol::default(),
); );
let bad = r.get(&WorkerId("bad".into())).expect("bad worker present"); let bad = r.get(&WorkerId("bad".into())).expect("bad worker present");
bad.breaker.record_failure(); bad.breaker.record_failure();
+58 -3
View File
@@ -8,6 +8,31 @@ use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::Instant; use std::time::Instant;
/// Which forwarding client the proxy uses for a worker.
///
/// 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
/// 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.
///
/// 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.
#[default]
Http1,
/// Cleartext HTTP/2 with prior knowledge (h2c). Used only when a worker
/// reports `--enable-http2` on a cleartext URL.
H2c,
}
/// Parse a host from a worker URL. Matches SMG's `worker_builder.rs` /// Parse a host from a worker URL. Matches SMG's `worker_builder.rs`
/// fallback chain: parse as-is, retry with `http://` prefix if missing, /// fallback chain: parse as-is, retry with `http://` prefix if missing,
/// fall back to `"localhost"` if both fail. The fallback is defensive — /// fall back to `"localhost"` if both fail. The fallback is defensive —
@@ -134,6 +159,9 @@ pub struct Worker {
/// Interior-mutable mode so `ModeChanged` can update in place without /// Interior-mutable mode so `ModeChanged` can update in place without
/// dropping the Worker (which would reset `active_requests` + breaker). /// dropping the Worker (which would reset `active_requests` + breaker).
mode: AtomicU8, mode: AtomicU8,
/// Forwarding wire protocol, resolved from `/server_info` before this
/// worker was constructed. Immutable: see [`WireProtocol`].
protocol: WireProtocol,
pub model_ids: Vec<ModelId>, pub model_ids: Vec<ModelId>,
pub breaker: Arc<CircuitBreaker>, pub breaker: Arc<CircuitBreaker>,
pub active_requests: Arc<AtomicUsize>, pub active_requests: Arc<AtomicUsize>,
@@ -155,14 +183,16 @@ pub struct Worker {
impl Worker { impl Worker {
pub fn new(spec: crate::discovery::WorkerSpec) -> Self { pub fn new(spec: crate::discovery::WorkerSpec) -> Self {
Self::with_cb_config(spec, None) Self::with_cb_config(spec, None, WireProtocol::default())
} }
/// Construct a worker with an explicit circuit-breaker configuration. /// Construct a worker with an explicit circuit-breaker configuration and
/// Pass `None` to use the default config (threshold = 3, cool_down = 30 s). /// forwarding protocol. Pass `None` for the default breaker config
/// (threshold = 3, cool_down = 30 s).
pub fn with_cb_config( pub fn with_cb_config(
spec: crate::discovery::WorkerSpec, spec: crate::discovery::WorkerSpec,
cb: Option<CircuitBreakerConfig>, cb: Option<CircuitBreakerConfig>,
protocol: WireProtocol,
) -> Self { ) -> Self {
let breaker = match cb { let breaker = match cb {
Some(cfg) => Arc::new(CircuitBreaker::with_config(cfg)), Some(cfg) => Arc::new(CircuitBreaker::with_config(cfg)),
@@ -175,6 +205,7 @@ impl Worker {
id: spec.id, id: spec.id,
url: spec.url, url: spec.url,
mode: AtomicU8::new(spec.mode.as_u8()), mode: AtomicU8::new(spec.mode.as_u8()),
protocol,
model_ids: spec.model_ids, model_ids: spec.model_ids,
breaker, breaker,
active_requests, active_requests,
@@ -210,6 +241,11 @@ impl Worker {
self.mode.store(m.as_u8(), Ordering::Relaxed); self.mode.store(m.as_u8(), Ordering::Relaxed);
} }
/// The wire protocol the proxy uses when forwarding to this worker.
pub fn protocol(&self) -> WireProtocol {
self.protocol
}
pub fn active_load(&self) -> usize { pub fn active_load(&self) -> usize {
self.active_requests.load(Ordering::Relaxed) self.active_requests.load(Ordering::Relaxed)
} }
@@ -247,6 +283,7 @@ impl std::fmt::Debug for Worker {
.field("id", &self.id) .field("id", &self.id)
.field("url", &self.url) .field("url", &self.url)
.field("mode", &self.mode()) .field("mode", &self.mode())
.field("protocol", &self.protocol)
.field("active_load", &self.active_load()) .field("active_load", &self.active_load())
.finish() .finish()
} }
@@ -335,6 +372,24 @@ mod tests {
assert_eq!(w.mode(), WorkerMode::Plain); assert_eq!(w.mode(), WorkerMode::Plain);
} }
#[test]
fn protocol_is_carried_from_construction() {
let spec = || WorkerSpec {
id: WorkerId("w".into()),
url: "http://x".into(),
mode: WorkerMode::Plain,
model_ids: vec![],
bootstrap_port: None,
};
// `new` takes the always-safe default; the resolved protocol reaches a
// worker only through the constructor the registry uses.
assert_eq!(Worker::new(spec()).protocol(), WireProtocol::Http1);
assert_eq!(
Worker::with_cb_config(spec(), None, WireProtocol::H2c).protocol(),
WireProtocol::H2c,
);
}
#[test] #[test]
fn bootstrap_port_returns_spec_value_for_prefill() { fn bootstrap_port_returns_spec_value_for_prefill() {
let w = Worker::new(WorkerSpec { let w = Worker::new(WorkerSpec {
@@ -4,26 +4,43 @@
use axum::{routing::get, Json, Router}; use axum::{routing::get, Json, Router};
use serde_json::{json, Value}; use serde_json::{json, Value};
use sgl_router::discovery::{DiscoveryEvent, ModelId, WorkerId, WorkerMode, WorkerSpec}; use sgl_router::discovery::{DiscoveryEvent, ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::workers::{manager, WorkerRegistry}; use sgl_router::workers::{manager, WireProtocol, WorkerRegistry};
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc; use std::sync::Arc;
use std::time::Duration; use std::time::{Duration, Instant};
use tokio::net::TcpListener; use tokio::net::TcpListener;
use tokio::sync::{mpsc, oneshot, Barrier}; use tokio::sync::{mpsc, oneshot, Barrier};
/// Spin up a tiny fake worker that returns `body` on `GET /server_info`. /// Spin up a tiny fake worker that returns `body` on both introspection
/// Returns the worker base URL and a shutdown channel. /// endpoints. Returns the worker base URL and a shutdown channel.
///
/// One body for both because a real engine reports `served_model_name` on
/// each.
async fn spawn_fake_worker(body: Value) -> (String, oneshot::Sender<()>) { async fn spawn_fake_worker(body: Value) -> (String, oneshot::Sender<()>) {
spawn_worker_serving(body, true).await
}
/// A worker that answers `/server_info` only — an SGLang predating
/// `served_model_name` on `/model_info`, and the shape the kind e2e fleet
/// still has (`tests/e2e/k8s_integration/fake_worker.py` defines no
/// `/model_info`). Registration must resolve the model name from the
/// `/server_info` fallback for this worker.
async fn spawn_server_info_only_worker(body: Value) -> (String, oneshot::Sender<()>) {
spawn_worker_serving(body, false).await
}
async fn spawn_worker_serving(body: Value, with_model_info: bool) -> (String, oneshot::Sender<()>) {
let body = Arc::new(body); let body = Arc::new(body);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port(); let port = listener.local_addr().unwrap().port();
let app = Router::new().route( let serve_body = move || {
"/server_info", let body = body.clone();
get(move || { async move { Json((*body).clone()) }
let body = body.clone(); };
async move { Json((*body).clone()) } let mut app = Router::new().route("/server_info", get(serve_body.clone()));
}), if with_model_info {
); app = app.route("/model_info", get(serve_body));
}
let (tx, rx) = oneshot::channel::<()>(); let (tx, rx) = oneshot::channel::<()>();
tokio::spawn(async move { tokio::spawn(async move {
let _ = axum::serve(listener, app) let _ = axum::serve(listener, app)
@@ -176,6 +193,208 @@ async fn manager_handles_mode_changed() {
h.await.unwrap(); h.await.unwrap();
} }
/// A worker that reports `enable_http2: true` is registered with
/// [`WireProtocol::H2c`], so the proxy forwards to it over cleartext h2c.
#[tokio::test]
async fn manager_resolves_h2c_protocol_from_introspection() {
let (url, _s) =
spawn_fake_worker(json!({"served_model_name": "m", "enable_http2": true})).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url,
WorkerMode::Plain,
)))
.await
.unwrap();
let resolved = wait_for_protocol(&registry, "w1", WireProtocol::H2c).await;
assert!(
resolved,
"worker reporting enable_http2 must resolve to h2c, got {:?}",
registry.get(&WorkerId("w1".into())).map(|w| w.protocol()),
);
drop(tx);
h.await.unwrap();
}
/// A worker that omits `enable_http2` (older SGLang) keeps the safe HTTP/1.1
/// default on its registry entry.
#[tokio::test]
async fn manager_defaults_http1_when_enable_http2_absent() {
let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w1",
&url,
WorkerMode::Plain,
)))
.await
.unwrap();
// Non-empty `model_ids` is what proves introspection ran: `spawn_fake_worker`
// answers from one body, so a resolved model id means `/server_info`
// answered too and its silence on `enable_http2` is the engine's, not the
// fixture's. The positive direction — that `enable_http2` is actually read
// — is pinned by `manager_resolves_protocol_per_worker_no_fleet_lock`
// below.
let registered = wait_for(Duration::from_secs(2), || {
registry
.get(&WorkerId("w1".into()))
.is_some_and(|w| !w.model_ids.is_empty())
})
.await;
assert!(registered, "worker registered");
let w = registry.get(&WorkerId("w1".into())).unwrap();
assert_eq!(w.protocol(), WireProtocol::Http1);
drop(tx);
h.await.unwrap();
}
/// A worker with no `/model_info` at all still registers, resolving its model
/// name from the `/server_info` fallback and its protocol from the same
/// response.
///
/// Every other test here uses a fixture that answers both endpoints, so without
/// this one the fallback chain in `introspect::fetch` has no component-level
/// coverage — while the kind e2e fleet runs exactly this shape.
#[tokio::test]
async fn manager_registers_worker_that_serves_server_info_only() {
let (url, _s) =
spawn_server_info_only_worker(json!({"served_model_name": "m", "enable_http2": true}))
.await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w-no-model-info",
&url,
WorkerMode::Plain,
)))
.await
.unwrap();
assert!(
wait_for_protocol(&registry, "w-no-model-info", WireProtocol::H2c).await,
"a worker without /model_info must still resolve its model name and \
protocol from /server_info; got {:?}",
registry
.get(&WorkerId("w-no-model-info".into()))
.map(|w| (w.model_ids.clone(), w.protocol())),
);
drop(tx);
h.await.unwrap();
}
/// The load-bearing regression for the production h2c bug: per-worker protocol
/// means one worker that resolved HTTP/1.1 first does NOT lock the rest of the
/// fleet off h2c. A mixed fleet (one h2c-capable worker registered AFTER a
/// plain HTTP/1.1 worker) ends with each worker on its own protocol — the old
/// single-client first-write-wins design forced both to HTTP/1.1.
#[tokio::test]
async fn manager_resolves_protocol_per_worker_no_fleet_lock() {
// w-http1 registers first and reports no enable_http2 (resolves Http1);
// w-h2c registers second and reports enable_http2: true (must still get
// h2c — it is not dragged down by w-http1's earlier Http1 resolution).
let (url_http1, _s1) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
let (url_h2c, _s2) =
spawn_fake_worker(json!({"served_model_name": "m", "enable_http2": true})).await;
let (tx, rx) = mpsc::channel(16);
let registry = Arc::new(WorkerRegistry::default());
let h = tokio::spawn(manager::run(rx, registry.clone()));
tx.send(DiscoveryEvent::Added(spec_for(
"w-http1",
&url_http1,
WorkerMode::Plain,
)))
.await
.unwrap();
// Let the HTTP/1.1 worker resolve first so it is the one that would have
// "won" the old fleet-wide client.
assert!(
wait_for_protocol(&registry, "w-http1", WireProtocol::Http1).await,
"first worker should resolve Http1",
);
tx.send(DiscoveryEvent::Added(spec_for(
"w-h2c",
&url_h2c,
WorkerMode::Plain,
)))
.await
.unwrap();
assert!(
wait_for_protocol(&registry, "w-h2c", WireProtocol::H2c).await,
"an h2c-capable worker registered after an Http1 worker must still resolve h2c \
(per-worker protocol, no fleet-wide lock); got {:?}",
registry
.get(&WorkerId("w-h2c".into()))
.map(|w| w.protocol()),
);
// The earlier worker is untouched.
let first = registry.get(&WorkerId("w-http1".into())).unwrap();
assert_eq!(
first.protocol(),
WireProtocol::Http1,
"the Http1 worker must stay Http1",
);
drop(tx);
h.await.unwrap();
}
/// Block until `worker_id`'s registry entry reports `expected`, or 2 s elapse.
/// Returns whether it converged.
///
/// The `model_ids` check is half the predicate on purpose. `H2c` proves itself
/// — it is reachable only through a successful `/server_info` read — but
/// `Http1` is also the default for a worker built without a resolved protocol,
/// so waiting on that value alone would be satisfied by an implementation that
/// never reads one. Requiring a resolved model id as well means the worker at
/// least completed introspection against a fixture that answers both
/// endpoints.
async fn wait_for_protocol(
registry: &Arc<WorkerRegistry>,
worker_id: &str,
expected: WireProtocol,
) -> bool {
let id = WorkerId(worker_id.into());
wait_for(Duration::from_secs(2), || {
registry
.get(&id)
.is_some_and(|w| !w.model_ids.is_empty() && w.protocol() == expected)
})
.await
}
/// Poll `cond` every 20 ms until it returns true or `budget` elapses.
async fn wait_for(budget: Duration, mut cond: impl FnMut() -> bool) -> bool {
let deadline = Instant::now() + budget;
while Instant::now() < deadline {
if cond() {
return true;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
cond()
}
#[tokio::test] #[tokio::test]
async fn mode_changed_preserves_active_requests_and_breaker() { async fn mode_changed_preserves_active_requests_and_breaker() {
let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await; let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;