[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:
co-authored by
Kangyan Zhou
Claude Opus 5
parent
0dabef3d30
commit
63845a1bb2
@@ -6,10 +6,11 @@
|
||||
//! Two concurrent requests, because the worker answers two different
|
||||
//! questions on two different endpoints: `/model_info` reports the identity
|
||||
//! the worker currently serves under (a weight update moves it), while
|
||||
//! `/server_info` reports its launch configuration — kv-event publisher and
|
||||
//! disaggregation role. The result is dispatched by the manager: registry
|
||||
//! consumes `served_model_name`, the optional `KvEventIndex` consumes the
|
||||
//! resolved `EventConfig`.
|
||||
//! `/server_info` reports its launch configuration — kv-event publisher,
|
||||
//! disaggregation role, and whether the engine serves cleartext h2c. The
|
||||
//! result is dispatched by the manager: registry consumes
|
||||
//! `served_model_name` and `enable_http2`, the optional `KvEventIndex`
|
||||
//! consumes the resolved `EventConfig`.
|
||||
//!
|
||||
//! `served_model_name` is taken from `/model_info`, falling back to
|
||||
//! `/server_info` for workers that predate the field there.
|
||||
@@ -56,6 +57,14 @@ pub struct ServerInfo {
|
||||
pub served_model_name: Option<String>,
|
||||
pub event_config: Option<EventConfig>,
|
||||
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.
|
||||
@@ -157,6 +166,7 @@ impl WorkerIntrospector {
|
||||
served_model_name,
|
||||
event_config,
|
||||
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
|
||||
/// 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)]
|
||||
struct ModelInfoBody {
|
||||
#[serde(default)]
|
||||
@@ -360,6 +371,12 @@ struct ServerInfoBody {
|
||||
/// bootstrap server binds to exactly this port (no internal offset).
|
||||
#[serde(default)]
|
||||
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)]
|
||||
@@ -706,6 +723,41 @@ mod tests {
|
||||
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
|
||||
/// `None` so the manager keeps the discovery backend's
|
||||
/// classification. The alternative — forcing Plain — would silently
|
||||
|
||||
@@ -7,7 +7,7 @@ use crate::health::circuit_breaker::CircuitBreakerConfig;
|
||||
use crate::policies::active_load::ActiveLoadRegistry;
|
||||
use crate::policies::kv_events::KvEventIndex;
|
||||
use crate::workers::introspect::{DisaggregationRole, WorkerIntrospector};
|
||||
use crate::workers::WorkerRegistry;
|
||||
use crate::workers::{WireProtocol, WorkerRegistry};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
@@ -15,14 +15,13 @@ use tokio::sync::mpsc;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
/// Production reconcile cadence. Workers that register without resolving
|
||||
/// their model IDs (a `/server_info` introspection that failed at `Added`
|
||||
/// time — e.g. the EndpointSlice flipped `ready=true` before the engine's
|
||||
/// scheduler-backed `/server_info` could answer) are re-introspected on
|
||||
/// this interval until they join their model pool. The worst-case
|
||||
/// "registered but invisible" window is about one interval plus the
|
||||
/// introspection round-trip; steady state costs one cheap registry scan
|
||||
/// per interval. See `reconcile_unresolved_workers` for the (benign)
|
||||
/// case of a worker that answers but never advertises a model name.
|
||||
/// their model IDs (an introspection that failed at `Added` time — e.g. the
|
||||
/// EndpointSlice flipped `ready=true` before the engine's HTTP server could
|
||||
/// answer) are re-introspected on this interval until they join their model
|
||||
/// pool. The worst-case "registered but invisible" window is about one
|
||||
/// interval plus the introspection round-trip; steady state costs one cheap
|
||||
/// registry scan per interval. See `reconcile_unresolved_workers` for the
|
||||
/// (benign) case of a worker that answers but never advertises a model name.
|
||||
const RECONCILE_INTERVAL: Duration = Duration::from_secs(30);
|
||||
|
||||
/// 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
|
||||
}
|
||||
|
||||
/// 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>) {
|
||||
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
|
||||
/// 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
|
||||
/// `active_load` is `None` the active-load bookkeeping is not pruned
|
||||
/// 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
|
||||
/// (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
|
||||
/// [`run_with_introspector`] directly.
|
||||
pub async fn run_with_config(
|
||||
@@ -303,17 +360,15 @@ async fn handle_discovery_event(
|
||||
}
|
||||
}
|
||||
|
||||
/// Re-introspect workers that registered without resolving their model
|
||||
/// IDs.
|
||||
/// Re-introspect workers that registered without resolving their model IDs.
|
||||
///
|
||||
/// A worker lands in the registry with empty `model_ids` when its
|
||||
/// `/server_info` introspection failed at `Added` time (e.g. the
|
||||
/// EndpointSlice flips `ready=true` before the engine can answer the
|
||||
/// scheduler-backed `/server_info` round-trip, and the introspector's
|
||||
/// bounded retry budget is exhausted). Such a worker is present in
|
||||
/// `by_id` but absent from every `by_model` pool, so it gets zero traffic
|
||||
/// — and it never recovers on its own: discovery re-lists carry the same
|
||||
/// empty `model_ids` the backend always emits, so they produce no new
|
||||
/// introspection failed at `Added` time (e.g. the EndpointSlice flips
|
||||
/// `ready=true` before the engine's HTTP server can answer, and the
|
||||
/// introspector's bounded retry budget is exhausted). Such a worker is
|
||||
/// present in `by_id` but absent from every `by_model` pool, so it gets zero
|
||||
/// traffic — and it never recovers on its own: discovery re-lists carry the
|
||||
/// same empty `model_ids` the backend always emits, so they produce no new
|
||||
/// `Added` event.
|
||||
///
|
||||
/// 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
|
||||
/// registry, so a worker that genuinely left is not resurrected.
|
||||
///
|
||||
/// A worker that answers `/server_info` but never advertises a
|
||||
/// `served_model_name` also stays in this set and is re-introspected
|
||||
/// every interval. That is a benign, bounded poll (one cheap round-trip
|
||||
/// per worker per interval), not a leak — but it is also never escalated,
|
||||
/// so the per-attempt failure logging stays at `debug!`; the introspector
|
||||
/// itself emits the `warn!` that surfaces a persistently failing worker.
|
||||
/// A worker that answers but never advertises a `served_model_name` also
|
||||
/// stays in this set and is re-introspected every interval — a benign,
|
||||
/// bounded poll, not a leak. It is never escalated, so the per-attempt
|
||||
/// logging stays at `debug!`; the introspector emits the `warn!` that
|
||||
/// surfaces a persistently failing worker.
|
||||
fn reconcile_unresolved_workers(
|
||||
registry: &Arc<WorkerRegistry>,
|
||||
cfg: &Option<Arc<Config>>,
|
||||
@@ -348,31 +402,33 @@ fn reconcile_unresolved_workers(
|
||||
// finish rather than racing a second introspection.
|
||||
continue;
|
||||
}
|
||||
// Rebuild a discovery-shaped spec: empty `model_ids` so
|
||||
// `register_one` re-resolves them from `/server_info`; current
|
||||
// mode + bootstrap_port as the seed (`register_one` re-applies
|
||||
// any `/server_info` override).
|
||||
let registry_t = registry.clone();
|
||||
let introspector_t = introspector.clone();
|
||||
let worker_url = worker.url.clone();
|
||||
// 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 {
|
||||
id: id.clone(),
|
||||
url: worker.url.clone(),
|
||||
url: worker_url.clone(),
|
||||
mode: worker.mode(),
|
||||
model_ids: Vec::new(),
|
||||
bootstrap_port: worker.bootstrap_port(),
|
||||
};
|
||||
// `debug!` not `info!`: this fires every interval for each
|
||||
// still-unresolved worker, so info-level would spam for a worker
|
||||
// that is permanently model-less. The introspector logs the
|
||||
// underlying `/server_info` failure at `warn!` on each attempt,
|
||||
// which is the operator-facing signal.
|
||||
// still-unresolved worker, so info-level would spam for one that is
|
||||
// permanently model-less. The introspector logs the underlying failure
|
||||
// at `warn!` on each attempt, which is the operator-facing signal.
|
||||
tracing::debug!(
|
||||
worker_id = %id,
|
||||
worker_url = %worker.url,
|
||||
worker_url = %worker_url,
|
||||
"reconcile: re-introspecting worker that registered without model_ids",
|
||||
);
|
||||
let registry_t = registry.clone();
|
||||
let cfg_t = cfg.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 {
|
||||
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
|
||||
/// 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
|
||||
@@ -427,8 +527,19 @@ async fn register_one(
|
||||
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));
|
||||
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
|
||||
// time. Log loudly so the operator notices the conflicting
|
||||
// worker — the alternative (silently dropping into either pool)
|
||||
@@ -442,6 +553,11 @@ async fn register_one(
|
||||
);
|
||||
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 {
|
||||
// Pass the pre-resolved EventConfig so the KvEventIndex does
|
||||
// not issue a second `/server_info` round-trip.
|
||||
@@ -509,19 +625,54 @@ mod tests {
|
||||
assert_eq!(cb.cool_down, Duration::from_secs(60));
|
||||
}
|
||||
|
||||
/// Helper: spawn a tiny fake worker that returns the supplied JSON body
|
||||
/// on `GET /server_info`. Returns the worker URL + a shutdown channel.
|
||||
async fn spawn_fake_server_info_worker(body: Value) -> (String, oneshot::Sender<()>) {
|
||||
let body = Arc::new(body);
|
||||
#[test]
|
||||
fn resolve_protocol_upgrades_to_h2c_only_for_cleartext_http2_engine() {
|
||||
// The one case that gets h2c: engine self-reports --enable-http2 and
|
||||
// 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 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::<()>();
|
||||
tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app)
|
||||
@@ -533,6 +684,19 @@ mod tests {
|
||||
(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
|
||||
/// connection attempts during the test fail fast with
|
||||
/// ConnectionRefused.
|
||||
@@ -941,43 +1105,122 @@ mod tests {
|
||||
let _ = manager_handle.await;
|
||||
}
|
||||
|
||||
/// Fake `/server_info` worker whose readiness is switchable at
|
||||
/// runtime. While `ready` is false it answers `503` (mimicking an
|
||||
/// engine whose EndpointSlice flipped `ready=true` before its
|
||||
/// scheduler-backed `/server_info` could answer); flip `ready` to
|
||||
/// true and it serves `body`.
|
||||
async fn spawn_switchable_server_info_worker(
|
||||
/// Fake worker whose readiness is switchable at runtime. While `ready` is
|
||||
/// false both introspection endpoints answer `503` (mimicking an engine
|
||||
/// whose EndpointSlice flipped `ready=true` before its HTTP server could
|
||||
/// answer anything); flip `ready` to true and both serve `body`.
|
||||
async fn spawn_switchable_worker(
|
||||
body: Value,
|
||||
ready: Arc<std::sync::atomic::AtomicBool>,
|
||||
) -> (String, oneshot::Sender<()>) {
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::IntoResponse;
|
||||
let body = Arc::new(body);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
let app = Router::new().route(
|
||||
"/server_info",
|
||||
get(move || {
|
||||
let body = body.clone();
|
||||
let ready = ready.clone();
|
||||
async move {
|
||||
if ready.load(std::sync::atomic::Ordering::SeqCst) {
|
||||
Json((*body).clone()).into_response()
|
||||
} else {
|
||||
StatusCode::SERVICE_UNAVAILABLE.into_response()
|
||||
let when_ready = move || {
|
||||
let body = body.clone();
|
||||
let ready = ready.clone();
|
||||
async move {
|
||||
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 {
|
||||
let _ = axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async move {
|
||||
let _ = rx.await;
|
||||
})
|
||||
.await;
|
||||
});
|
||||
(format!("http://127.0.0.1:{port}"), tx)
|
||||
|
||||
drop(tx);
|
||||
let _ = manager_handle.await;
|
||||
}
|
||||
|
||||
/// A worker that registers with empty `model_ids` because
|
||||
@@ -993,8 +1236,7 @@ mod tests {
|
||||
|
||||
let ready = Arc::new(AtomicBool::new(false));
|
||||
let (worker_url, _shutdown) =
|
||||
spawn_switchable_server_info_worker(json!({"served_model_name": "m"}), ready.clone())
|
||||
.await;
|
||||
spawn_switchable_worker(json!({"served_model_name": "m"}), ready.clone()).await;
|
||||
|
||||
let registry = Arc::new(WorkerRegistry::default());
|
||||
let (tx, rx) = mpsc::channel::<DiscoveryEvent>(8);
|
||||
@@ -1066,6 +1308,94 @@ mod tests {
|
||||
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
|
||||
/// re-introspection for the same id is in-flight must NOT resurrect
|
||||
/// 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 registry::WorkerRegistry;
|
||||
pub use worker::LoadGuard;
|
||||
pub use worker::WireProtocol;
|
||||
pub use worker::Worker;
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
|
||||
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use crate::health::circuit_breaker::CircuitBreakerConfig;
|
||||
use crate::workers::worker::Worker;
|
||||
use crate::workers::worker::{WireProtocol, Worker};
|
||||
use dashmap::DashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::{Arc, Mutex};
|
||||
@@ -51,11 +51,12 @@ pub struct WorkerRegistry {
|
||||
|
||||
impl WorkerRegistry {
|
||||
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.
|
||||
/// Pass `None` to use the circuit-breaker default (threshold = 3).
|
||||
/// Add a worker, optionally supplying a circuit-breaker config, and with
|
||||
/// 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
|
||||
/// `by_model` memberships are cleared first so a model that the new
|
||||
@@ -83,6 +84,7 @@ impl WorkerRegistry {
|
||||
&self,
|
||||
spec: WorkerSpec,
|
||||
cb: Option<CircuitBreakerConfig>,
|
||||
protocol: WireProtocol,
|
||||
) -> Result<(), AddWorkerError> {
|
||||
let incoming_mode = spec.mode;
|
||||
// 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();
|
||||
self.remove_locked(&id);
|
||||
for m in &w.model_ids {
|
||||
@@ -310,7 +312,11 @@ mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
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
|
||||
// flips it to Open.
|
||||
let _ = r.add_with_cb(
|
||||
@@ -319,6 +325,7 @@ mod tests {
|
||||
threshold: NonZeroU32::new(1).unwrap(),
|
||||
cool_down: Duration::from_secs(30),
|
||||
}),
|
||||
WireProtocol::default(),
|
||||
);
|
||||
let bad = r.get(&WorkerId("bad".into())).expect("bad worker present");
|
||||
bad.breaker.record_failure();
|
||||
|
||||
@@ -8,6 +8,31 @@ use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
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`
|
||||
/// fallback chain: parse as-is, retry with `http://` prefix if missing,
|
||||
/// 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
|
||||
/// dropping the Worker (which would reset `active_requests` + breaker).
|
||||
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 breaker: Arc<CircuitBreaker>,
|
||||
pub active_requests: Arc<AtomicUsize>,
|
||||
@@ -155,14 +183,16 @@ pub struct Worker {
|
||||
|
||||
impl Worker {
|
||||
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.
|
||||
/// Pass `None` to use the default config (threshold = 3, cool_down = 30 s).
|
||||
/// Construct a worker with an explicit circuit-breaker configuration and
|
||||
/// forwarding protocol. Pass `None` for the default breaker config
|
||||
/// (threshold = 3, cool_down = 30 s).
|
||||
pub fn with_cb_config(
|
||||
spec: crate::discovery::WorkerSpec,
|
||||
cb: Option<CircuitBreakerConfig>,
|
||||
protocol: WireProtocol,
|
||||
) -> Self {
|
||||
let breaker = match cb {
|
||||
Some(cfg) => Arc::new(CircuitBreaker::with_config(cfg)),
|
||||
@@ -175,6 +205,7 @@ impl Worker {
|
||||
id: spec.id,
|
||||
url: spec.url,
|
||||
mode: AtomicU8::new(spec.mode.as_u8()),
|
||||
protocol,
|
||||
model_ids: spec.model_ids,
|
||||
breaker,
|
||||
active_requests,
|
||||
@@ -210,6 +241,11 @@ impl Worker {
|
||||
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 {
|
||||
self.active_requests.load(Ordering::Relaxed)
|
||||
}
|
||||
@@ -247,6 +283,7 @@ impl std::fmt::Debug for Worker {
|
||||
.field("id", &self.id)
|
||||
.field("url", &self.url)
|
||||
.field("mode", &self.mode())
|
||||
.field("protocol", &self.protocol)
|
||||
.field("active_load", &self.active_load())
|
||||
.finish()
|
||||
}
|
||||
@@ -335,6 +372,24 @@ mod tests {
|
||||
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]
|
||||
fn bootstrap_port_returns_spec_value_for_prefill() {
|
||||
let w = Worker::new(WorkerSpec {
|
||||
|
||||
@@ -4,26 +4,43 @@
|
||||
use axum::{routing::get, Json, Router};
|
||||
use serde_json::{json, Value};
|
||||
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::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{mpsc, oneshot, Barrier};
|
||||
|
||||
/// Spin up a tiny fake worker that returns `body` on `GET /server_info`.
|
||||
/// Returns the worker base URL and a shutdown channel.
|
||||
/// Spin up a tiny fake worker that returns `body` on both introspection
|
||||
/// 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<()>) {
|
||||
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 listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
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 serve_body = move || {
|
||||
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::<()>();
|
||||
tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app)
|
||||
@@ -176,6 +193,208 @@ async fn manager_handles_mode_changed() {
|
||||
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(®istry, "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(®istry, "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(®istry, "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(®istry, "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]
|
||||
async fn mode_changed_preserves_active_requests_and_breaker() {
|
||||
let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
|
||||
|
||||
Reference in New Issue
Block a user