[router] Configure experimental sgl-router via CLI flags instead of a config file (#27073)

Signed-off-by: Kangyan Zhou <zky314343421@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
2026-06-05 10:02:10 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 631db6c757
commit bcf89928b4
43 changed files with 2355 additions and 1237 deletions
+5 -6
View File
@@ -37,10 +37,11 @@ reqwest = { version = "0.12", features = ["stream", "json", "rustls-tls"], defau
# Serialization
serde = { version = "1", features = ["derive"] }
serde_json = { version = "1", features = ["preserve_order"] }
# `humantime-serde` lets `WorkerConfig.request_timeout` accept human-readable
# durations like `"60s"` / `"500ms"` / `"2m"` in YAML / TOML, rather than
# forcing operators to write raw milliseconds.
humantime-serde = "1"
# Tokenizer auto-download from HuggingFace when --tokenizer-path is omitted.
# Sync (`ureq`) API only — it runs once at startup; `ureq` is on rustls, so
# this pulls no openssl/native-tls (matching reqwest's rustls-tls above).
hf-hub = { version = "0.4", default-features = false, features = ["ureq"] }
# Utilities
anyhow = "1"
@@ -53,8 +54,6 @@ bytes = "1"
rand = "0.8"
tokio-stream = "0.1"
dashmap = "6"
serde_yaml = "0.9"
toml = "0.8"
kube = { version = "0.96", features = ["runtime", "derive"] }
k8s-openapi = { version = "0.23", features = ["v1_31"] }
tokio-util = "0.7"
+40 -5
View File
@@ -2,11 +2,11 @@
Slim, KV-aware, OpenAI-compatible router for SGLang workers.
**Status:** functional single-worker HTTP proxy. Exposes `/v1/tokenize`,
`/v1/detokenize`, `/v1/models`, `/v1/chat/completions` (buffered and SSE),
plus `/healthz` / `/readyz`. Forwards to one configured worker via reqwest;
parity-tested against `transformers.AutoTokenizer`. Multi-worker routing,
service discovery, and observability still pending.
Serves a single model and routes across its workers. Exposes
`/v1/tokenize`, `/v1/detokenize`, `/v1/models`, `/v1/chat/completions`
(buffered and SSE), plus `/healthz` / `/readyz` and `/metrics`. Worker
pools come from either a static URL list or Kubernetes EndpointSlice
discovery.
## Building
@@ -15,6 +15,41 @@ cd experimental/sgl-router
cargo build --release
```
## Running
The router is configured entirely through CLI flags (run
`sgl-router --help` for the full list). It serves exactly one model, so
`--model-id` is required, along with exactly one discovery backend.
`--tokenizer-path` is optional: give it a local `tokenizer.json` path or a
HuggingFace repo id, and when omitted the router downloads the tokenizer
for `--model-id` from HuggingFace (honoring `HF_TOKEN` / `HF_HOME`).
Static worker list:
```bash
sgl-router \
--host 0.0.0.0 --port 30000 \
--model-id qwen3 \
--tokenizer-path /models/qwen3/tokenizer.json \
--worker-urls http://10.0.0.1:30000 http://10.0.0.2:30000
```
Kubernetes EndpointSlice discovery:
```bash
sgl-router \
--host 0.0.0.0 --port 30000 \
--model-id qwen3 \
--tokenizer-path /models/qwen3/tokenizer.json \
--service-discovery \
--service-discovery-namespace prod \
--selector app=engines-qwen3
```
Omit `--service-discovery-namespace` to watch all namespaces (requires
cluster-wide RBAC). For prefill/decode disaggregation, replace `--selector`
with `--prefill-selector` and `--decode-selector`.
## License
Apache-2.0.
+711
View File
@@ -0,0 +1,711 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0
//! Command-line interface. The router is configured entirely through
//! flags — there is no config file. [`Cli::into_config`] resolves the
//! flags into a validated [`Config`].
use anyhow::{anyhow, Result};
use clap::Parser;
use std::num::NonZeroU32;
use crate::config::{
default_cb_cool_down, default_proxy_request_timeout_secs, default_stale_request_timeout_secs,
resolve_mode, ActiveLoadConfig, CacheAwareConfig, CircuitBreakerConfig, Config,
DiscoveryBackend, K8sDiscoveryConfig, LogFormat, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
/// `sgl-router` — slim KV-aware OpenAI-compatible router for SGLang workers.
///
/// Discovery is mutually exclusive: pass `--worker-urls` for a static
/// worker list, or `--service-discovery` for Kubernetes EndpointSlice
/// discovery — exactly one is required.
#[derive(Parser, Debug)]
#[command(
name = "sgl-router",
version,
about = "Slim KV-aware OpenAI-compatible router for SGLang workers"
)]
pub struct Cli {
// ---- server ----
/// Address to bind the HTTP server to.
#[arg(long, default_value = "127.0.0.1")]
pub host: String,
/// Port to bind the HTTP server to.
#[arg(long, default_value_t = 30000)]
pub port: u16,
// ---- model (exactly one) ----
/// Model id this router serves (the OpenAI `model` field).
#[arg(long)]
pub model_id: String,
/// Tokenizer source: a local `tokenizer.json` path, or a HuggingFace
/// repo id to download from. When omitted, falls back to `--model-id`
/// as the repo id (download honors `HF_TOKEN` / `HF_HOME`).
#[arg(long)]
pub tokenizer_path: Option<String>,
/// Routing policy.
#[arg(long, value_enum, default_value = "round_robin")]
pub policy: PolicyKind,
// ---- circuit breaker (opt-in via --cb-threshold) ----
/// Consecutive upstream failures before the circuit breaker opens.
/// Setting this enables the circuit breaker; `0` is rejected.
#[arg(long)]
pub cb_threshold: Option<NonZeroU32>,
/// Circuit-breaker cool-down in seconds. Only meaningful with
/// `--cb-threshold`; defaults to 30 when the breaker is enabled.
#[arg(long)]
pub cb_cool_down_secs: Option<u64>,
// ---- cache-aware-zmq tuning (only used by that policy) ----
/// Min `matched_blocks / total_blocks` for a cache match to win.
#[arg(long)]
pub cache_threshold: Option<f32>,
/// Absolute load spread above which the cache check is skipped.
#[arg(long)]
pub balance_abs_threshold: Option<usize>,
/// Multiplicative load spread gating the absolute balance check.
#[arg(long)]
pub balance_rel_threshold: Option<f32>,
// ---- discovery: static ----
/// Static worker URLs (space-separated or repeated). Mutually
/// exclusive with `--service-discovery`.
#[arg(long, num_args = 1..)]
pub worker_urls: Vec<String>,
// ---- discovery: kubernetes ----
/// Enable Kubernetes EndpointSlice discovery.
#[arg(long)]
pub service_discovery: bool,
/// Namespace to watch. Unset/empty watches all namespaces (requires
/// cluster-wide RBAC).
#[arg(long)]
pub service_discovery_namespace: Option<String>,
/// Plain-mode label selector terms, e.g. `app=engines-qwen3`
/// (space-separated or repeated `key=value`, AND-joined). Mutually
/// exclusive with the prefill/decode selectors.
#[arg(long, num_args = 1..)]
pub selector: Vec<String>,
/// PD-mode prefill label selector terms. Requires `--decode-selector`.
#[arg(long, num_args = 1..)]
pub prefill_selector: Vec<String>,
/// PD-mode decode label selector terms. Requires `--prefill-selector`.
#[arg(long, num_args = 1..)]
pub decode_selector: Vec<String>,
// ---- proxy / active-load ----
/// Per-request upstream timeout in seconds.
#[arg(long, default_value_t = default_proxy_request_timeout_secs())]
pub request_timeout_secs: u64,
/// Max lifetime of an in-flight request entry before the janitor
/// reaps it (returns 504 `stale_request_expired`).
#[arg(long, default_value_t = default_stale_request_timeout_secs())]
pub stale_request_timeout_secs: u64,
// ---- observability ----
/// Default tracing level (overridden by `RUST_LOG`).
#[arg(long, default_value = "info")]
pub log_level: String,
/// Log output format.
#[arg(long, value_enum, default_value = "text")]
pub log_format: LogFormat,
}
impl Cli {
/// Resolve parsed flags into a validated [`Config`].
///
/// Builds the [`DiscoveryBackend`] (enforcing static-vs-k8s mutual
/// exclusivity and resolving the k8s selector grammar via
/// [`resolve_mode`]), assembles the single [`ModelConfig`], then runs
/// [`Config::validate`] for the remaining value-level invariants
/// (model id, static worker URLs).
pub fn into_config(self) -> Result<Config> {
let discovery = self.build_discovery()?;
// Reject knobs that only take effect alongside another flag, rather
// than silently dropping them — mirrors the discovery mutual-exclusion
// checks. Otherwise an operator believes they tuned something that has
// no effect.
if self.cb_cool_down_secs.is_some() && self.cb_threshold.is_none() {
return Err(anyhow!(
"--cb-cool-down-secs requires --cb-threshold (the circuit breaker is \
enabled by --cb-threshold)"
));
}
let tuned_cache_aware = self.cache_threshold.is_some()
|| self.balance_abs_threshold.is_some()
|| self.balance_rel_threshold.is_some();
if tuned_cache_aware && self.policy != PolicyKind::CacheAwareZmq {
return Err(anyhow!(
"--cache-threshold / --balance-abs-threshold / --balance-rel-threshold \
require --policy cache_aware_zmq"
));
}
let circuit_breaker = self.cb_threshold.map(|threshold| CircuitBreakerConfig {
threshold,
cool_down_secs: self.cb_cool_down_secs.unwrap_or_else(default_cb_cool_down),
});
// Only build a CacheAwareConfig when the operator tuned at least
// one knob; otherwise leave it None so the policy uses its own
// defaults. Unset knobs fall back to the per-field defaults.
let cache_aware = if tuned_cache_aware {
let d = CacheAwareConfig::default();
Some(CacheAwareConfig {
cache_threshold: self.cache_threshold.unwrap_or(d.cache_threshold),
balance_abs_threshold: self
.balance_abs_threshold
.unwrap_or(d.balance_abs_threshold),
balance_rel_threshold: self
.balance_rel_threshold
.unwrap_or(d.balance_rel_threshold),
})
} else {
None
};
let config = Config {
server: ServerConfig {
host: self.host,
port: self.port,
},
observability: ObservabilityConfig {
log_level: self.log_level,
log_format: self.log_format,
},
model: ModelConfig {
// Default the tokenizer source to the model id (treated as a
// HuggingFace repo id) when --tokenizer-path is omitted.
tokenizer_path: self.tokenizer_path.unwrap_or_else(|| self.model_id.clone()),
id: self.model_id,
policy: self.policy,
circuit_breaker,
cache_aware,
},
discovery,
proxy: ProxyConfig {
request_timeout_secs: self.request_timeout_secs,
},
active_load: ActiveLoadConfig {
stale_request_timeout_secs: self.stale_request_timeout_secs,
},
};
config.validate()?;
Ok(config)
}
/// Resolve the discovery flags into a [`DiscoveryBackend`].
///
/// `--worker-urls` (static) and `--service-discovery` (k8s) are
/// mutually exclusive and exactly one is required. K8s-only flags
/// passed without `--service-discovery` are rejected so a typo can't
/// silently fall back to the static (empty) path. The k8s selector
/// grammar (plain vs PD) is validated eagerly here by [`resolve_mode`]
/// before the `K8sDiscoveryConfig` is constructed, so an invalid
/// combination is never stored.
fn build_discovery(&self) -> Result<DiscoveryBackend> {
let has_static = !self.worker_urls.is_empty();
let backend = match (has_static, self.service_discovery) {
(true, true) => {
return Err(anyhow!(
"--worker-urls and --service-discovery are mutually exclusive; pass exactly one"
))
}
(false, false) => {
return Err(anyhow!(
"no discovery backend selected; pass --worker-urls <URL...> (static) \
or --service-discovery (kubernetes)"
))
}
(true, false) => {
if self.service_discovery_namespace.is_some()
|| !self.selector.is_empty()
|| !self.prefill_selector.is_empty()
|| !self.decode_selector.is_empty()
{
return Err(anyhow!(
"--service-discovery-namespace / --selector / --prefill-selector / \
--decode-selector require --service-discovery"
));
}
DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: self.worker_urls.clone(),
})
}
(false, true) => {
// Resolve (and validate) the selector flags into a
// K8sDiscoveryMode here, so an invalid combination can't be
// stored. Surfaces ConfigError as anyhow for the CLI.
let mode = resolve_mode(
join_selector(&self.selector).as_deref(),
join_selector(&self.prefill_selector).as_deref(),
join_selector(&self.decode_selector).as_deref(),
)
.map_err(|e| anyhow!("{e}"))?;
DiscoveryBackend::K8s(K8sDiscoveryConfig {
namespace: self.service_discovery_namespace.clone().unwrap_or_default(),
mode,
})
}
};
Ok(backend)
}
}
/// Join space/repeated `key=value` selector terms into the single
/// comma-joined string the k8s backend's `labels_match_selector`
/// expects. `None` for an empty term list so [`resolve_mode`] can apply
/// its plain-vs-PD rules (and surface `NoSelector`).
fn join_selector(terms: &[String]) -> Option<String> {
if terms.is_empty() {
None
} else {
Some(terms.join(","))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::{DiscoveryBackend, K8sDiscoveryMode};
/// Parse argv (without the leading binary name) into a `Config`.
fn into_config(args: &[&str]) -> Result<Config> {
let argv = std::iter::once("sgl-router").chain(args.iter().copied());
let cli = Cli::try_parse_from(argv).map_err(|e| anyhow!("{e}"))?;
cli.into_config()
}
const MODEL_ARGS: &[&str] = &[
"--model-id",
"qwen3-0.6b",
"--tokenizer-path",
"/tmp/qwen.json",
];
fn with_model(extra: &[&str]) -> Vec<String> {
MODEL_ARGS
.iter()
.chain(extra.iter())
.map(|s| s.to_string())
.collect()
}
fn into_config_owned(args: Vec<String>) -> Result<Config> {
let refs: Vec<&str> = args.iter().map(String::as_str).collect();
into_config(&refs)
}
#[test]
fn defaults_host_port_and_policy() {
let c = into_config_owned(with_model(&["--worker-urls", "http://10.0.0.1:30000"])).unwrap();
assert_eq!(c.server.host, "127.0.0.1");
assert_eq!(c.server.port, 30000);
assert_eq!(c.model.policy, PolicyKind::RoundRobin);
assert_eq!(c.model.id, "qwen3-0.6b");
assert_eq!(c.proxy.request_timeout_secs, 300);
assert_eq!(c.active_load.stale_request_timeout_secs, 600);
}
/// With `--tokenizer-path` omitted, the tokenizer source defaults to the
/// model id (treated as an HF repo id at load time).
#[test]
fn tokenizer_path_defaults_to_model_id_when_omitted() {
let c = into_config(&[
"--model-id",
"Qwen/Qwen3-0.6B",
"--worker-urls",
"http://x:30000",
])
.unwrap();
assert_eq!(c.model.id, "Qwen/Qwen3-0.6B");
assert_eq!(c.model.tokenizer_path, "Qwen/Qwen3-0.6B");
}
#[test]
fn explicit_tokenizer_path_is_used() {
let c = into_config(&[
"--model-id",
"qwen3",
"--tokenizer-path",
"/models/qwen3/tokenizer.json",
"--worker-urls",
"http://x:30000",
])
.unwrap();
assert_eq!(c.model.tokenizer_path, "/models/qwen3/tokenizer.json");
}
#[test]
fn static_urls_backend() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://10.0.0.1:30000",
"http://10.0.0.2:30000",
]))
.unwrap();
match &c.discovery {
DiscoveryBackend::StaticUrls(s) => assert_eq!(
s.urls,
vec![
"http://10.0.0.1:30000".to_string(),
"http://10.0.0.2:30000".to_string()
]
),
_ => panic!("expected static_urls backend"),
}
}
#[test]
fn rejects_no_discovery_backend() {
let err = into_config_owned(with_model(&[])).unwrap_err().to_string();
assert!(err.contains("no discovery backend"), "got: {err}");
}
#[test]
fn rejects_both_discovery_backends() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--service-discovery",
]))
.unwrap_err()
.to_string();
assert!(err.contains("mutually exclusive"), "got: {err}");
}
#[test]
fn rejects_k8s_flags_without_service_discovery() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--selector",
"app=sglang",
]))
.unwrap_err()
.to_string();
assert!(err.contains("require --service-discovery"), "got: {err}");
}
#[test]
fn rejects_static_urls_duplicate() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"http://x:30000",
]))
.unwrap_err()
.to_string();
assert!(err.contains("duplicate"), "got: {err}");
}
#[test]
fn rejects_static_urls_schemeless() {
let err = into_config_owned(with_model(&["--worker-urls", "10.0.0.1:30000"]))
.unwrap_err()
.to_string();
assert!(
err.contains("not a valid URL") || err.contains("unsupported scheme"),
"got: {err}"
);
}
#[test]
fn rejects_static_urls_non_http_scheme() {
let err = into_config_owned(with_model(&["--worker-urls", "ws://x:30000"]))
.unwrap_err()
.to_string();
assert!(err.contains("unsupported scheme"), "got: {err}");
}
#[test]
fn k8s_plain_backend() {
let c = into_config_owned(with_model(&[
"--service-discovery",
"--service-discovery-namespace",
"prod",
"--selector",
"app=engines-qwen3",
]))
.unwrap();
match &c.discovery {
DiscoveryBackend::K8s(k) => {
assert_eq!(k.namespace, "prod");
assert_eq!(
k.mode,
K8sDiscoveryMode::Plain {
label_selector: "app=engines-qwen3".to_string()
}
);
}
_ => panic!("expected k8s backend"),
}
}
/// Multiple `--selector` terms AND-join into one comma-separated
/// label selector (matches the Python router's space-separated form).
#[test]
fn k8s_plain_selector_joins_multiple_terms() {
let c = into_config_owned(with_model(&[
"--service-discovery",
"--selector",
"app=sglang",
"zone=us-east",
]))
.unwrap();
match &c.discovery {
DiscoveryBackend::K8s(k) => assert_eq!(
k.mode,
K8sDiscoveryMode::Plain {
label_selector: "app=sglang,zone=us-east".to_string()
}
),
_ => panic!("expected k8s backend"),
}
}
/// Empty namespace is intentional — it triggers a cluster-wide watch.
#[test]
fn k8s_empty_namespace_watches_all() {
let c = into_config_owned(with_model(&[
"--service-discovery",
"--selector",
"app=sglang",
]))
.unwrap();
match &c.discovery {
DiscoveryBackend::K8s(k) => assert_eq!(k.namespace, ""),
_ => panic!("expected k8s backend"),
}
}
#[test]
fn k8s_pd_backend() {
let c = into_config_owned(with_model(&[
"--service-discovery",
"--service-discovery-namespace",
"default",
"--prefill-selector",
"app=sglang,role=prefill",
"--decode-selector",
"app=sglang,role=decode",
]))
.unwrap();
match &c.discovery {
DiscoveryBackend::K8s(k) => assert_eq!(
k.mode,
K8sDiscoveryMode::PdDisaggregation {
prefill_selector: "app=sglang,role=prefill".to_string(),
decode_selector: "app=sglang,role=decode".to_string(),
}
),
_ => panic!("expected k8s backend"),
}
}
/// `--service-discovery` with no selector at all fails `resolve_mode`
/// validation with the `NoSelector` wording.
#[test]
fn rejects_k8s_without_selector() {
let err = into_config_owned(with_model(&["--service-discovery"]))
.unwrap_err()
.to_string()
.to_lowercase();
assert!(err.contains("none were set"), "got: {err}");
}
/// `--prefill-selector` without `--decode-selector` is rejected through
/// the full CLI path — pins that `build_discovery` feeds the right
/// selectors into `resolve_mode` (a positional mix-up would surface a
/// different error or none).
#[test]
fn rejects_k8s_partial_pd_selectors() {
let err = into_config_owned(with_model(&[
"--service-discovery",
"--prefill-selector",
"app=sglang,role=prefill",
]))
.unwrap_err()
.to_string();
assert!(
err.contains("PD mode requires BOTH"),
"expected PartialPdSelectors wording, got: {err}"
);
}
/// Identical prefill/decode selectors are rejected through the full CLI
/// path (would silently leave the decode pool empty at runtime).
#[test]
fn rejects_k8s_identical_pd_selectors() {
let err = into_config_owned(with_model(&[
"--service-discovery",
"--prefill-selector",
"app=sglang",
"--decode-selector",
"app=sglang",
]))
.unwrap_err()
.to_string();
assert!(
err.contains("must differ"),
"expected IdenticalPdSelectors wording, got: {err}"
);
}
/// clap rejects an unknown `--policy` value at parse time.
#[test]
fn rejects_unknown_policy() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"bogus_policy",
]))
.unwrap_err()
.to_string();
assert!(
err.contains("bogus_policy") || err.contains("policy"),
"got: {err}"
);
}
/// clap rejects `--cb-threshold 0` because the field is `NonZeroU32`.
#[test]
fn rejects_zero_cb_threshold() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--cb-threshold",
"0",
]))
.unwrap_err()
.to_string();
assert!(err.contains("cb-threshold"), "got: {err}");
}
#[test]
fn cb_threshold_enables_circuit_breaker_with_default_cool_down() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--cb-threshold",
"5",
]))
.unwrap();
let cb = c.model.circuit_breaker.expect("cb enabled");
assert_eq!(cb.threshold.get(), 5);
assert_eq!(cb.cool_down_secs, 30);
}
#[test]
fn cb_cool_down_honors_explicit_override() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--cb-threshold",
"3",
"--cb-cool-down-secs",
"10",
]))
.unwrap();
let cb = c.model.circuit_breaker.expect("cb enabled");
assert_eq!(cb.cool_down_secs, 10);
}
#[test]
fn rejects_cb_cool_down_without_threshold() {
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--cb-cool-down-secs",
"10",
]))
.unwrap_err()
.to_string();
assert!(
err.contains("--cb-cool-down-secs requires --cb-threshold"),
"got: {err}"
);
}
#[test]
fn cache_aware_knob_builds_partial_config() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"cache_aware_zmq",
"--cache-threshold",
"0.7",
]))
.unwrap();
let ca = c.model.cache_aware.expect("cache_aware set");
assert_eq!(ca.cache_threshold, 0.7);
// Untouched knobs fall back to defaults.
assert_eq!(ca.balance_abs_threshold, 32);
}
#[test]
fn no_cache_aware_flags_leaves_none() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--policy",
"cache_aware_zmq",
]))
.unwrap();
assert!(c.model.cache_aware.is_none());
}
#[test]
fn rejects_cache_aware_knob_without_cache_aware_policy() {
// Default policy is round_robin, so a cache knob has no effect —
// reject rather than silently ignore it.
let err = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--cache-threshold",
"0.7",
]))
.unwrap_err()
.to_string();
assert!(
err.contains("require --policy cache_aware_zmq"),
"got: {err}"
);
}
#[test]
fn log_format_parses_json() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--log-format",
"json",
]))
.unwrap();
assert_eq!(c.observability.log_format, LogFormat::Json);
}
/// Pins that the two timeout overrides land in the right fields — they
/// are adjacent `u64`s with similar names, so a copy-paste swap would
/// otherwise go unnoticed (and `stale` must sit above `proxy`).
#[test]
fn timeout_overrides_land_in_distinct_fields() {
let c = into_config_owned(with_model(&[
"--worker-urls",
"http://x:30000",
"--request-timeout-secs",
"120",
"--stale-request-timeout-secs",
"240",
]))
.unwrap();
assert_eq!(c.proxy.request_timeout_secs, 120);
assert_eq!(c.active_load.stale_request_timeout_secs, 240);
}
}
+72 -458
View File
@@ -1,42 +1,21 @@
pub mod cli;
pub mod types;
pub use cli::Cli;
pub use types::*;
use anyhow::Context as _;
use anyhow::{anyhow, Result};
use std::path::Path;
impl Config {
pub fn from_path(p: &Path) -> Result<Self> {
let raw =
std::fs::read_to_string(p).with_context(|| format!("read config {}", p.display()))?;
let ext = p.extension().and_then(|s| s.to_str()).unwrap_or("");
let cfg: Config = match ext {
"yaml" | "yml" => serde_yaml::from_str(&raw)
.map_err(|e| anyhow!("parse yaml {}: {e}", p.display()))?,
"toml" => {
toml::from_str(&raw).map_err(|e| anyhow!("parse toml {}: {e}", p.display()))?
}
other => {
return Err(anyhow!(
"unsupported config extension {other:?}; want yaml/yml/toml"
))
}
};
cfg.validate()?;
Ok(cfg)
}
fn validate(&self) -> Result<()> {
// Unknown policy names are rejected by serde via `PolicyKind`'s
// `rename_all = "snake_case"`; threshold = 0 is rejected by
// `NonZeroU32`. Only fields without a type-system constraint are
// checked here.
for m in &self.models {
if m.id.is_empty() {
return Err(anyhow!("model.id must be non-empty"));
}
/// Check invariants the type system and `clap` don't already enforce.
/// Called by [`cli::Cli::into_config`] after assembling the `Config`
/// from flags. Unknown policy names and `--cb-threshold 0` are
/// rejected at parse time (`ValueEnum` / `NonZeroU32`); only the
/// remaining value-level invariants are checked here.
pub(crate) fn validate(&self) -> Result<()> {
if self.model.id.is_empty() {
return Err(anyhow!("model id must be non-empty"));
}
match &self.discovery.backend {
match &self.discovery {
DiscoveryBackend::StaticUrls(s) => {
if s.urls.is_empty() {
return Err(anyhow!(
@@ -44,7 +23,7 @@ impl Config {
));
}
// Validate every entry up front so typos surface at
// config-load with a precise diagnostic instead of as
// startup with a precise diagnostic instead of as
// per-worker introspect failures or as two registry
// entries pointing at the same SGLang (trailing-slash
// near-duplicates). Dedupe runs against a normalized
@@ -77,14 +56,11 @@ impl Config {
}
}
}
DiscoveryBackend::K8s(k) => {
// Empty namespace is intentional: triggers `Api::all(client)`
// for cluster-wide EndpointSlice watch (see
// `discovery::k8s::spawn`). Only validate the selector
// combination here.
let _ = &k.namespace;
k.mode().map_err(|e| anyhow!("{e}"))?;
}
// K8s selector validity is resolved at construction time
// (`resolve_mode` in `Cli::build_discovery`), so the stored
// `K8sDiscoveryMode` is already valid here. Any namespace
// (including empty, for a cluster-wide watch) is accepted.
DiscoveryBackend::K8s(_) => {}
}
Ok(())
}
@@ -94,447 +70,85 @@ impl Config {
mod tests {
use super::*;
/// Write `body` to a temp file with the given extension and load it
/// through `Config::from_path`. Failures still surface the offending
/// config because each call site passes its body inline.
fn load(ext: &str, body: &str) -> Result<Config> {
let dir = tempfile::tempdir().unwrap();
let p = dir.path().join(format!("c.{ext}"));
std::fs::write(&p, body).unwrap();
Config::from_path(&p)
}
#[test]
fn loads_minimal_yaml() {
let c = load(
"yaml",
r#"
server:
host: "0.0.0.0"
port: 8090
models:
- id: "qwen3-0.6b"
tokenizer_path: "/tmp/qwen.json"
discovery:
backend: static_urls
static_urls:
urls:
- "http://10.0.0.1:30000"
"#,
)
.unwrap();
assert_eq!(c.server.port, 8090);
assert_eq!(c.models[0].id, "qwen3-0.6b");
match &c.discovery.backend {
DiscoveryBackend::StaticUrls(s) => {
assert_eq!(s.urls, vec!["http://10.0.0.1:30000".to_string()])
}
_ => panic!("expected static_urls backend"),
/// Build a minimal valid-shape `Config` with the given static worker
/// URLs and model id, so the `validate()` branches can be exercised
/// directly. CLI parsing and the static-vs-k8s mapping are covered in
/// the `cli` module tests; the k8s selector grammar in `types`.
fn cfg(model_id: &str, urls: &[&str]) -> Config {
Config {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 30000,
},
observability: ObservabilityConfig::default(),
model: ModelConfig {
id: model_id.into(),
tokenizer_path: "/tmp/tok.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: urls.iter().map(|s| s.to_string()).collect(),
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
}
#[test]
fn loads_minimal_toml() {
let c = load(
"toml",
r#"
[server]
host = "0.0.0.0"
port = 8090
[[models]]
id = "qwen3-0.6b"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["http://10.0.0.1:30000"]
"#,
)
.unwrap();
assert_eq!(c.server.port, 8090);
match &c.discovery.backend {
DiscoveryBackend::StaticUrls(s) => {
assert_eq!(s.urls, vec!["http://10.0.0.1:30000".to_string()])
}
_ => panic!("expected static_urls backend"),
}
fn accepts_minimal_static_config() {
cfg("qwen3", &["http://10.0.0.1:30000"]).validate().unwrap();
}
#[test]
fn rejects_missing_discovery_section() {
let err = load(
"yaml",
"server:\n host: \"0.0.0.0\"\n port: 8090\nmodels: []\n",
)
.unwrap_err();
let msg = err.to_string().to_lowercase();
assert!(
msg.contains("discovery") || msg.contains("missing"),
"got: {err}"
);
fn rejects_empty_model_id() {
let err = cfg("", &["http://10.0.0.1:30000"])
.validate()
.unwrap_err()
.to_string();
assert!(err.contains("model id"), "got: {err}");
}
#[test]
fn rejects_unknown_extension() {
let err = load("txt", "").unwrap_err();
assert!(err.to_string().contains("yaml") && err.to_string().contains("toml"));
}
#[test]
fn loads_static_urls_discovery() {
let c = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "qwen3-0.6b"
tokenizer_path = "/tmp/qwen.json"
policy = "round_robin"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["http://10.0.0.1:30000", "http://10.0.0.2:30000"]
"#,
)
.unwrap();
match &c.discovery.backend {
DiscoveryBackend::StaticUrls(s) => {
assert_eq!(
s.urls,
vec![
"http://10.0.0.1:30000".to_string(),
"http://10.0.0.2:30000".to_string(),
],
);
}
_ => panic!("expected static_urls backend"),
}
assert_eq!(c.models[0].policy, PolicyKind::RoundRobin);
}
#[test]
fn rejects_static_urls_with_empty_list() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "m"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = []
"#,
)
.unwrap_err()
.to_string();
fn rejects_empty_static_urls_list() {
let err = cfg("qwen3", &[]).validate().unwrap_err().to_string();
assert!(err.contains("non-empty"), "got: {err}");
}
#[test]
fn rejects_static_urls_with_duplicate_entry() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "m"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["http://x:30000", "http://x:30000"]
"#,
)
.unwrap_err()
.to_string();
assert!(err.contains("duplicate"), "got: {err}");
}
#[test]
fn rejects_static_urls_with_empty_entry() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "m"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["http://x:30000", ""]
"#,
)
.unwrap_err()
.to_string();
fn rejects_static_urls_empty_entry() {
let err = cfg("qwen3", &["http://x:30000", ""])
.validate()
.unwrap_err()
.to_string();
assert!(err.contains("empty"), "got: {err}");
}
/// Whitespace-only entries are user typos that previously slipped
/// through `is_empty()` checks and surfaced as "introspect against
/// ` /server_info` failed" at runtime. Catch at load.
#[test]
fn rejects_static_urls_with_whitespace_only_entry() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "m"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["http://x:30000", " "]
"#,
)
.unwrap_err()
.to_string();
assert!(err.contains("whitespace"), "got: {err}");
fn rejects_static_urls_whitespace_only_entry() {
let err = cfg("qwen3", &["http://x:30000", " "])
.validate()
.unwrap_err()
.to_string();
assert!(err.contains("empty or whitespace"), "got: {err}");
}
/// `"10.0.0.1:30000"` (missing scheme) used to pass validation; the
/// scheme/`http://` would only fail (or worse, silently degrade
/// because of the `parse_bootstrap_host` localhost fallback) at
/// introspect time. Reject at load.
#[test]
fn rejects_static_urls_with_schemeless_entry() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "m"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["10.0.0.1:30000"]
"#,
)
.unwrap_err()
.to_string();
assert!(
err.contains("not a valid URL") || err.contains("unsupported scheme"),
"got: {err}"
);
}
/// Non-http(s) schemes are rejected. The router speaks HTTP to
/// workers; a `tcp://` or `ws://` entry is almost certainly an
/// operator typo.
#[test]
fn rejects_static_urls_with_non_http_scheme() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "m"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["ws://x:30000"]
"#,
)
.unwrap_err()
.to_string();
assert!(err.contains("unsupported scheme"), "got: {err}");
}
/// Trailing-slash near-duplicates collide in the registry but used
/// to pass byte-equality dedupe. Normalize before checking so two
/// pointers at the same SGLang surface as a config error.
#[test]
fn rejects_static_urls_with_trailing_slash_near_duplicate() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "m"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["http://x:30000", "http://x:30000/"]
"#,
)
.unwrap_err()
.to_string();
fn rejects_static_urls_trailing_slash_near_duplicate() {
let err = cfg("qwen3", &["http://x:30000", "http://x:30000/"])
.validate()
.unwrap_err()
.to_string();
assert!(err.contains("duplicate"), "got: {err}");
}
#[test]
fn loads_k8s_discovery() {
let c = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "qwen3-0.6b"
tokenizer_path = "/tmp/qwen.json"
policy = "round_robin"
[discovery]
backend = "k8s"
[discovery.k8s]
namespace = "default"
label_selector = "app=sglang"
"#,
)
.unwrap();
match &c.discovery.backend {
DiscoveryBackend::K8s(k) => {
assert_eq!(k.namespace, "default");
assert_eq!(k.label_selector.as_deref(), Some("app=sglang"));
assert!(k.prefill_selector.is_none());
assert!(k.decode_selector.is_none());
}
_ => panic!("expected k8s backend"),
}
}
/// K8s PD selectors drive slice-classification only; per-worker
/// bootstrap_port comes from `/server_info` post-discovery
/// (`crate::workers::introspect`). This test pins the wire-shape;
/// the selector grammar itself is covered in `types.rs`.
#[test]
fn loads_k8s_pd_discovery_with_prefill_and_decode_selectors() {
let c = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "qwen3-0.6b"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "k8s"
[discovery.k8s]
namespace = "default"
prefill_selector = "app=sglang,role=prefill"
decode_selector = "app=sglang,role=decode"
"#,
)
.expect("k8s PD config must load");
match &c.discovery.backend {
DiscoveryBackend::K8s(k) => {
assert_eq!(k.namespace, "default");
assert_eq!(
k.prefill_selector.as_deref(),
Some("app=sglang,role=prefill")
);
assert_eq!(k.decode_selector.as_deref(), Some("app=sglang,role=decode"));
assert!(k.label_selector.is_none());
}
_ => panic!("expected k8s backend"),
}
}
#[test]
fn rejects_k8s_config_with_no_selector() {
let err = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "qwen"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "k8s"
[discovery.k8s]
namespace = "default"
"#,
)
.unwrap_err();
// Pin the specific variant: `ConfigError::NoSelector` ("none were
// set"). A bare `contains("selector")` would also pass for
// EmptyPdSelector / PartialPdSelectors / IdenticalPdSelectors /
// UnsupportedSelectorGrammar — variants that have semantically
// different error wording but all mention "selector". A future
// regression that returned, say, `PartialPdSelectors` for the
// all-None input would be caught here.
let msg = err.to_string().to_lowercase();
assert!(
msg.contains("none were set"),
"expected NoSelector wording (\"none were set\"); got: {err}",
);
}
// Direct `K8sDiscoveryConfig::mode()` unit tests live alongside the
// type in `src/config/types.rs::k8s_discovery_config_tests`.
// The tests in this module exercise the `Config::from_path` ↔ K8s
// selector wiring, not the selector grammar itself.
#[test]
fn rejects_unknown_policy_name() {
let err = load(
"yaml",
"
server:
host: 0.0.0.0
port: 8090
discovery:
backend: static_urls
static_urls:
urls:
- http://x:30000
models:
- id: qwen
tokenizer_path: /tmp/qwen.json
policy: bogus_policy
",
)
.unwrap_err();
let msg = err.to_string().to_lowercase();
assert!(
msg.contains("bogus_policy") || msg.contains("policy"),
"got: {err}"
);
}
#[test]
fn defaults_policy_to_round_robin() {
let c = load(
"toml",
r#"
[server]
host = "127.0.0.1"
port = 8090
[[models]]
id = "qwen"
tokenizer_path = "/tmp/qwen.json"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["http://x:30000"]
"#,
)
.unwrap();
assert_eq!(c.models[0].policy, PolicyKind::RoundRobin);
fn rejects_static_urls_non_http_scheme() {
let err = cfg("qwen3", &["ws://x:30000"])
.validate()
.unwrap_err()
.to_string();
assert!(err.contains("unsupported scheme"), "got: {err}");
}
}
+168 -312
View File
@@ -1,32 +1,35 @@
use serde::{Deserialize, Serialize};
use std::num::NonZeroU32;
#[derive(Debug, Clone, Serialize, Deserialize)]
/// In-memory router configuration, built from CLI flags by
/// [`crate::config::cli::Cli::into_config`] and validated by
/// [`Config::validate`]. The router serves exactly one model.
#[derive(Debug, Clone)]
pub struct Config {
pub server: ServerConfig,
#[serde(default)]
pub observability: ObservabilityConfig,
pub models: Vec<ModelConfig>,
pub discovery: DiscoveryConfig,
#[serde(default)]
pub model: ModelConfig,
/// Selected discovery backend. Built from CLI flags by
/// [`crate::config::cli::Cli::into_config`]: the static-vs-k8s choice
/// and the k8s selector grammar are resolved there (the latter via
/// [`resolve_mode`]); static worker-URL validity is checked by
/// [`Config::validate`].
pub discovery: DiscoveryBackend,
pub proxy: ProxyConfig,
#[serde(default)]
pub active_load: ActiveLoadConfig,
}
/// Outbound proxy tuning. Default mirrors SGLang's typical prefill /
/// decode latency budget; e2e tests lower it so per-request failures
/// trip the circuit breaker within the test's wall-time.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[derive(Debug, Clone, Copy)]
pub struct ProxyConfig {
/// Maximum time to wait for a single upstream HTTP request to
/// return headers + body. Default 300 s. The circuit breaker
/// records a failure when this fires.
#[serde(default = "default_proxy_request_timeout_secs")]
pub request_timeout_secs: u64,
}
fn default_proxy_request_timeout_secs() -> u64 {
pub fn default_proxy_request_timeout_secs() -> u64 {
300
}
@@ -42,16 +45,15 @@ impl Default for ProxyConfig {
/// sits above `proxy.request_timeout_secs` so the proxy timeout is the
/// one users hit first for normal slow upstreams; tests lower it to
/// let the janitor fire within their wall-time budget.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[derive(Debug, Clone, Copy)]
pub struct ActiveLoadConfig {
/// How long a request entry can live in the registry before the
/// janitor fires its `cancel_token` and the chat handler returns
/// 504 `stale_request_expired`. Default 600 s.
#[serde(default = "default_stale_request_timeout_secs")]
pub stale_request_timeout_secs: u64,
}
fn default_stale_request_timeout_secs() -> u64 {
pub fn default_stale_request_timeout_secs() -> u64 {
600
}
@@ -63,51 +65,52 @@ impl Default for ActiveLoadConfig {
}
}
/// Routing policy selector — the enum form lets serde reject unknown
/// values at deserialization time and removes the runtime string match in
/// the policy factory.
/// Routing policy selector — the enum form lets `clap` reject unknown
/// values at parse time and removes the runtime string match in the
/// policy factory.
///
/// Serialised as `"round_robin"` / `"random"` / `"power_of_two"` /
/// `"cache_aware_zmq"`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
/// Accepted on the CLI (`--policy`) as `round_robin` / `random` /
/// `power_of_two` / `cache_aware_zmq`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)]
pub enum PolicyKind {
#[default]
#[value(name = "round_robin")]
RoundRobin,
#[value(name = "random")]
Random,
#[value(name = "power_of_two")]
PowerOfTwo,
/// Cache-aware routing fed by SGLang's ZMQ KV-cache event publisher.
/// Requires the model to have a tokenizer loaded; cache_aware tuning
/// lives on `ModelConfig::cache_aware`.
#[value(name = "cache_aware_zmq")]
CacheAwareZmq,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone)]
pub struct ServerConfig {
pub host: String,
pub port: u16,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone)]
pub struct ObservabilityConfig {
#[serde(default = "default_log_level")]
pub log_level: String,
/// Selects the tracing-subscriber output format. Serde rejects
/// unrecognized values at config-load (`"jsonl"` and similar
/// plausible typos surface as an error instead of silently
/// degrading to text), matching the discoverability pattern used
/// by `policy` and `discovery.backend`.
#[serde(default)]
/// Selects the tracing-subscriber output format. `clap` rejects
/// unrecognized values at parse time (`--log-format jsonl` and
/// similar typos surface as an error instead of silently degrading
/// to text).
pub log_format: LogFormat,
}
/// `text` for human-readable dev output, `json` for one-line-per-record
/// JSON suitable for k8s log aggregators (fluent-bit / vector / Loki).
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)]
pub enum LogFormat {
#[default]
#[value(name = "text")]
Text,
#[value(name = "json")]
Json,
}
@@ -124,40 +127,37 @@ impl Default for ObservabilityConfig {
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone)]
pub struct ModelConfig {
pub id: String,
/// Tokenizer source: a local `tokenizer.json` path or a HuggingFace repo
/// id (downloaded on demand). Defaults to `id` when `--tokenizer-path`
/// is omitted. Resolved by [`crate::tokenizer::adapter::load`].
pub tokenizer_path: String,
#[serde(default)]
pub policy: PolicyKind,
#[serde(default)]
pub circuit_breaker: Option<CircuitBreakerConfig>,
/// Tuning for the cache-aware ZMQ policy. Ignored unless
/// `policy = "cache_aware_zmq"`. `None` falls back to defaults at
/// policy construction time.
#[serde(default)]
pub cache_aware: Option<CacheAwareConfig>,
}
/// Per-model cache-aware-ZMQ tuning.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[derive(Debug, Clone, Copy)]
pub struct CacheAwareConfig {
/// Lower bound on `matched_blocks / total_blocks` for the tree match
/// to win the selection. Below this, the policy falls back to
/// min-load. Default 0.5 — a half-cached prompt is still a strong
/// signal but not so weak that random hash collisions could trigger
/// affinity to an arbitrary worker.
#[serde(default = "default_cache_threshold")]
pub cache_threshold: f32,
/// Absolute load spread (`max - min`) above which the cache check is
/// skipped in favour of min-load. Default 32 — picked to dominate
/// over typical batch-of-8 effect.
#[serde(default = "default_balance_abs")]
pub balance_abs_threshold: usize,
/// Multiplicative load spread (`max > min * balance_rel_threshold`)
/// that the absolute check is gated on. Default 1.1 — 10 % relative
/// difference triggers re-balancing.
#[serde(default = "default_balance_rel")]
pub balance_rel_threshold: f32,
}
@@ -181,125 +181,22 @@ fn default_balance_rel() -> f32 {
1.1
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
/// Consecutive failures required before the breaker opens. Encoded
/// as `NonZeroU32` so a config setting `threshold = 0` (which would
/// open the breaker before any failure) is rejected at deserialization
/// rather than silently behaving as "always open".
#[serde(default = "default_cb_threshold")]
/// Consecutive failures required before the breaker opens. Encoded
/// as `NonZeroU32` so `--cb-threshold 0` (which would open the
/// breaker before any failure) is rejected at CLI-parse time rather
/// than silently behaving as "always open".
pub threshold: NonZeroU32,
#[serde(default = "default_cb_cool_down")]
pub cool_down_secs: u64,
}
fn default_cb_threshold() -> NonZeroU32 {
NonZeroU32::new(3).unwrap()
}
fn default_cb_cool_down() -> u64 {
/// Default circuit-breaker cool-down, applied when `--cb-threshold` is
/// set without an explicit `--cb-cool-down-secs`.
pub fn default_cb_cool_down() -> u64 {
30
}
/// Config-level discovery section. Deserialized from:
///
/// TOML:
/// ```toml
/// [discovery]
/// backend = "static_urls"
/// [discovery.static_urls]
/// urls = ["http://10.0.0.1:30000", "http://10.0.0.2:30000"]
/// ```
///
/// YAML:
/// ```yaml
/// discovery:
/// backend: static_urls
/// static_urls:
/// urls:
/// - http://10.0.0.1:30000
/// - http://10.0.0.2:30000
/// ```
///
/// The custom `Deserialize` impl on [`DiscoveryConfig`] converts the
/// raw fields into the resolved `DiscoveryBackend` enum via `try_from`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DiscoveryConfigRaw {
pub backend: String,
pub static_urls: Option<StaticUrlsDiscoveryConfig>,
pub k8s: Option<K8sDiscoveryConfig>,
}
/// Post-validation discovery config with a resolved `DiscoveryBackend` enum.
/// Constructed by `Config::from_path` after `validate()`.
#[derive(Debug, Clone)]
pub struct DiscoveryConfig {
pub backend: DiscoveryBackend,
}
impl<'de> Deserialize<'de> for DiscoveryConfig {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let raw = DiscoveryConfigRaw::deserialize(deserializer)?;
raw.try_into().map_err(serde::de::Error::custom)
}
}
impl Serialize for DiscoveryConfig {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let raw: DiscoveryConfigRaw = self.clone().into();
raw.serialize(serializer)
}
}
impl TryFrom<DiscoveryConfigRaw> for DiscoveryConfig {
type Error = String;
fn try_from(raw: DiscoveryConfigRaw) -> Result<Self, Self::Error> {
let backend = match raw.backend.as_str() {
"static_urls" => {
let s = raw.static_urls.ok_or(
"discovery.backend = \"static_urls\" requires [discovery.static_urls] section",
)?;
DiscoveryBackend::StaticUrls(s)
}
"k8s" => {
let k = raw
.k8s
.ok_or("discovery.backend = \"k8s\" requires [discovery.k8s] section")?;
DiscoveryBackend::K8s(k)
}
other => {
return Err(format!(
"unknown discovery.backend = {other:?}; valid: \"static_urls\", \"k8s\""
))
}
};
Ok(DiscoveryConfig { backend })
}
}
impl From<DiscoveryConfig> for DiscoveryConfigRaw {
fn from(cfg: DiscoveryConfig) -> Self {
match cfg.backend {
DiscoveryBackend::StaticUrls(s) => DiscoveryConfigRaw {
backend: "static_urls".to_string(),
static_urls: Some(s),
k8s: None,
},
DiscoveryBackend::K8s(k) => DiscoveryConfigRaw {
backend: "k8s".to_string(),
static_urls: None,
k8s: Some(k),
},
}
}
}
#[derive(Debug, Clone)]
pub enum DiscoveryBackend {
StaticUrls(StaticUrlsDiscoveryConfig),
@@ -311,30 +208,25 @@ pub enum DiscoveryBackend {
/// from `/server_info` (see [`crate::workers::introspect`]).
///
/// No file watcher, no hot-reload: topology change requires a restart.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone)]
pub struct StaticUrlsDiscoveryConfig {
pub urls: Vec<String>,
}
/// Configuration for the Kubernetes `EndpointSlice` discovery backend.
/// Built from the `--service-discovery*` / `--selector` / `--prefill-selector`
/// / `--decode-selector` flags by [`crate::config::cli::Cli::build_discovery`].
///
/// Two operating modes, distinguished by which selector fields are set:
/// Two operating modes, distinguished by which selector flags are set:
///
/// 1. **Plain** — all matched workers share the same role:
/// ```toml
/// [discovery.k8s]
/// namespace = "default"
/// label_selector = "app=sglang"
/// ```
/// `--service-discovery-namespace default --selector app=sglang`
///
/// 2. **PD disaggregation** — prefill and decode workers are separated by
/// different selectors:
/// ```toml
/// [discovery.k8s]
/// namespace = "default"
/// prefill_selector = "app=sglang,role=prefill"
/// decode_selector = "app=sglang,role=decode"
/// ```
/// `--service-discovery-namespace default
/// --prefill-selector app=sglang,role=prefill
/// --decode-selector app=sglang,role=decode`
///
/// In PD mode, the selectors drive **slice-classification** (which
/// EndpointSlices feed the prefill pool vs the decode pool). The actual
@@ -344,20 +236,19 @@ pub struct StaticUrlsDiscoveryConfig {
/// [`crate::workers::introspect`] for the `disaggregation_mode` and
/// `disaggregation_bootstrap_port` extraction.
///
/// `mode()` validates the combination and returns the resolved
/// [`K8sDiscoveryMode`]; any other selector combination is rejected.
#[derive(Debug, Clone, Serialize, Deserialize)]
/// [`resolve_mode`] validates the selector flags and produces the
/// resolved [`K8sDiscoveryMode`] once, at construction in
/// [`crate::config::cli::Cli::build_discovery`] — so an invalid selector
/// combination is unrepresentable here.
#[derive(Debug, Clone)]
pub struct K8sDiscoveryConfig {
pub namespace: String,
#[serde(default)]
pub label_selector: Option<String>,
#[serde(default)]
pub prefill_selector: Option<String>,
#[serde(default)]
pub decode_selector: Option<String>,
/// Resolved + validated selector mode (plain vs PD).
pub mode: K8sDiscoveryMode,
}
/// Resolved discovery mode derived from a [`K8sDiscoveryConfig`].
/// Resolved discovery mode, produced by [`resolve_mode`] from the CLI
/// selector flags and stored on [`K8sDiscoveryConfig`].
///
/// The discovery backend uses this to:
/// * pick the server-side `LIST` label selector (Plain: the single selector;
@@ -377,8 +268,8 @@ pub enum K8sDiscoveryMode {
},
}
/// Error returned by [`K8sDiscoveryConfig::mode`] when the selector
/// combination is invalid.
/// Error returned by [`resolve_mode`] when the selector combination is
/// invalid.
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
#[error("discovery.k8s requires either `label_selector` (plain) or both `prefill_selector` and `decode_selector` (PD); none were set")]
@@ -391,7 +282,7 @@ pub enum ConfigError {
"discovery.k8s: {selector}_selector `{value}` uses unsupported syntax — \
only equality terms (`key=value` or `key==value`) joined by `,` are accepted. \
Set-based operators (`in`, `notin`), presence tests, and `!=` silently match \
zero endpoints at runtime and are rejected at config-load time."
zero endpoints at runtime and are rejected at startup."
)]
UnsupportedSelectorGrammar {
selector: &'static str,
@@ -487,80 +378,79 @@ fn is_equality_selector(selector: &str) -> bool {
true
}
impl K8sDiscoveryConfig {
/// Validate the selector combination and return the resolved mode.
pub fn mode(&self) -> Result<K8sDiscoveryMode, ConfigError> {
let plain = self.label_selector.as_deref();
let prefill = self.prefill_selector.as_deref();
let decode = self.decode_selector.as_deref();
match (plain, prefill, decode) {
(Some(label), None, None) => {
// Plain mode pushes `label` to the K8s API as the
// server-side `labelSelector` of the EndpointSlice
// watcher (`watcher::Config::default().labels(&label)`
// in `discovery::k8s::spawn`). K8s itself parses the
// full label-selector grammar — equality, set-based
// (`in` / `notin`), presence (`key` / `!key`), and
// `!=` — and rejects malformed selectors at
// watch-start time. So at config-load we don't
// grammar-check `label` and let the K8s API be the
// syntax authority (README.md:25 and the multi-model
// e2e in tests/e2e/k8s_integration/test_multi_model.py
// depend on this). PD mode, in contrast, evaluates
// selectors client-side via `labels_match_selector`
// which only understands equality — so PD selectors
// are still grammar-checked below.
Ok(K8sDiscoveryMode::Plain {
label_selector: label.to_string(),
})
}
(None, Some(prefill), Some(decode)) => {
// Both selectors validated individually so the operator
// sees which one is malformed. WorkerMode + bootstrap_port
// for each prefill pod are filled in by the worker
// manager from each worker's `/server_info` — these
// selectors only drive client-side classification per
// EndpointSlice (see `classify_mode` in discovery/k8s.rs).
if !is_equality_selector(prefill) {
return Err(ConfigError::UnsupportedSelectorGrammar {
selector: "prefill",
value: prefill.to_string(),
});
}
if !is_equality_selector(decode) {
return Err(ConfigError::UnsupportedSelectorGrammar {
selector: "decode",
value: decode.to_string(),
});
}
// Empty PD selector matches every EndpointSlice at
// runtime; combined with classify_mode's prefill-first
// ordering, an empty selector would silently funnel all
// workers into one role. Reject up front.
if is_selector_empty(prefill) {
return Err(ConfigError::EmptyPdSelector {
selector: "prefill",
});
}
if is_selector_empty(decode) {
return Err(ConfigError::EmptyPdSelector { selector: "decode" });
}
// Identical selectors degrade the same way as an empty
// one: every slice matches both, prefill wins, decode
// stays empty.
if canonical_selector(prefill) == canonical_selector(decode) {
return Err(ConfigError::IdenticalPdSelectors);
}
Ok(K8sDiscoveryMode::PdDisaggregation {
prefill_selector: prefill.to_string(),
decode_selector: decode.to_string(),
})
}
(None, None, None) => Err(ConfigError::NoSelector),
(None, Some(_), None) | (None, None, Some(_)) => Err(ConfigError::PartialPdSelectors),
(Some(_), _, _) => Err(ConfigError::MixedModes),
/// Validate the selector combination and return the resolved
/// [`K8sDiscoveryMode`]. Called once at construction by
/// [`crate::config::cli::Cli::build_discovery`], so an invalid
/// combination can never be stored on a [`K8sDiscoveryConfig`].
pub fn resolve_mode(
label_selector: Option<&str>,
prefill_selector: Option<&str>,
decode_selector: Option<&str>,
) -> Result<K8sDiscoveryMode, ConfigError> {
match (label_selector, prefill_selector, decode_selector) {
(Some(label), None, None) => {
// Plain mode pushes `label` to the K8s API as the
// server-side `labelSelector` of the EndpointSlice
// watcher (`watcher::Config::default().labels(&label)`
// in `discovery::k8s::spawn`). K8s itself parses the
// full label-selector grammar — equality, set-based
// (`in` / `notin`), presence (`key` / `!key`), and
// `!=` — and rejects malformed selectors at
// watch-start time. So we don't grammar-check `label`
// here and let the K8s API be the syntax authority. PD
// mode, in contrast, evaluates selectors client-side via
// `labels_match_selector` which only understands
// equality — so PD selectors are still grammar-checked
// below.
Ok(K8sDiscoveryMode::Plain {
label_selector: label.to_string(),
})
}
(None, Some(prefill), Some(decode)) => {
// Both selectors validated individually so the operator
// sees which one is malformed. WorkerMode + bootstrap_port
// for each prefill pod are filled in by the worker
// manager from each worker's `/server_info` — these
// selectors only drive client-side classification per
// EndpointSlice (see `classify_mode` in discovery/k8s.rs).
if !is_equality_selector(prefill) {
return Err(ConfigError::UnsupportedSelectorGrammar {
selector: "prefill",
value: prefill.to_string(),
});
}
if !is_equality_selector(decode) {
return Err(ConfigError::UnsupportedSelectorGrammar {
selector: "decode",
value: decode.to_string(),
});
}
// Empty PD selector matches every EndpointSlice at
// runtime; combined with classify_mode's prefill-first
// ordering, an empty selector would silently funnel all
// workers into one role. Reject up front.
if is_selector_empty(prefill) {
return Err(ConfigError::EmptyPdSelector {
selector: "prefill",
});
}
if is_selector_empty(decode) {
return Err(ConfigError::EmptyPdSelector { selector: "decode" });
}
// Identical selectors degrade the same way as an empty
// one: every slice matches both, prefill wins, decode
// stays empty.
if canonical_selector(prefill) == canonical_selector(decode) {
return Err(ConfigError::IdenticalPdSelectors);
}
Ok(K8sDiscoveryMode::PdDisaggregation {
prefill_selector: prefill.to_string(),
decode_selector: decode.to_string(),
})
}
(None, None, None) => Err(ConfigError::NoSelector),
(None, Some(_), None) | (None, None, Some(_)) => Err(ConfigError::PartialPdSelectors),
(Some(_), _, _) => Err(ConfigError::MixedModes),
}
}
@@ -568,23 +458,13 @@ impl K8sDiscoveryConfig {
mod k8s_discovery_config_tests {
use super::*;
fn cfg(plain: Option<&str>, prefill: Option<&str>, decode: Option<&str>) -> K8sDiscoveryConfig {
K8sDiscoveryConfig {
namespace: "ns".to_string(),
label_selector: plain.map(str::to_string),
prefill_selector: prefill.map(str::to_string),
decode_selector: decode.map(str::to_string),
}
}
#[test]
fn mode_constructs_pd_disaggregation_from_prefill_and_decode_selectors() {
// K8s PD now works without per-pod annotations: each worker's
// `/server_info` carries `disaggregation_bootstrap_port`, and the
// worker manager applies it post-discovery. The K8s config layer's
// job is just to validate the selector combination.
let m = cfg(None, Some("app=sglang,role=p"), Some("app=sglang,role=d"))
.mode()
let m = resolve_mode(None, Some("app=sglang,role=p"), Some("app=sglang,role=d"))
.expect("PD mode is now valid");
assert_eq!(
m,
@@ -600,9 +480,8 @@ mod k8s_discovery_config_tests {
// Both PD selectors get the same equality-only grammar check as
// the plain label_selector. A set-based prefill selector would
// silently match zero pods at runtime → fail-fast at load.
let err = cfg(None, Some("app in (sglang, vllm)"), Some("app=sglang"))
.mode()
.unwrap_err();
let err =
resolve_mode(None, Some("app in (sglang, vllm)"), Some("app=sglang")).unwrap_err();
assert!(
matches!(
err,
@@ -617,9 +496,8 @@ mod k8s_discovery_config_tests {
#[test]
fn mode_pd_rejects_set_based_decode_selector() {
let err = cfg(None, Some("app=sglang"), Some("app in (sglang, vllm)"))
.mode()
.unwrap_err();
let err =
resolve_mode(None, Some("app=sglang"), Some("app in (sglang, vllm)")).unwrap_err();
assert!(
matches!(
err,
@@ -634,7 +512,7 @@ mod k8s_discovery_config_tests {
#[test]
fn mode_accepts_plain_with_equality_selector() {
let m = cfg(Some("app=sglang"), None, None).mode().unwrap();
let m = resolve_mode(Some("app=sglang"), None, None).unwrap();
assert_eq!(
m,
K8sDiscoveryMode::Plain {
@@ -646,15 +524,12 @@ mod k8s_discovery_config_tests {
/// Plain mode pushes its selector to the K8s API server-side
/// (`watcher::Config::default().labels(&selector)` in
/// `discovery::k8s::spawn`), so the full K8s label-selector grammar
/// — including set-based operators — is supported. README.md:25
/// advertises this, and `tests/e2e/k8s_integration/test_multi_model.py`
/// relies on it (`label_selector = "app in (sglang,sglang-small)"`).
/// Rejecting set-based selectors at config-load broke the documented
/// multi-model k8s path.
/// — including set-based operators like `app in (a,b)` — is
/// supported and must not be grammar-checked at startup. PD mode
/// (checked client-side) is the opposite; see the PD tests below.
#[test]
fn mode_accepts_set_based_selector_in_plain_mode() {
let m = cfg(Some("app in (sglang,sglang-small)"), None, None)
.mode()
let m = resolve_mode(Some("app in (sglang,sglang-small)"), None, None)
.expect("plain mode must accept set-based selectors");
assert_eq!(
m,
@@ -675,8 +550,7 @@ mod k8s_discovery_config_tests {
"!deprecated",
"tier!=canary",
] {
let m = cfg(Some(raw), None, None)
.mode()
let m = resolve_mode(Some(raw), None, None)
.unwrap_or_else(|e| panic!("plain mode must accept `{raw}`, got {e:?}"));
assert_eq!(
m,
@@ -699,9 +573,8 @@ mod k8s_discovery_config_tests {
/// — both must keep failing.
#[test]
fn mode_pd_rejects_notin_prefill_selector() {
let err = cfg(None, Some("app notin (vllm, trtllm)"), Some("app=sglang"))
.mode()
.unwrap_err();
let err =
resolve_mode(None, Some("app notin (vllm, trtllm)"), Some("app=sglang")).unwrap_err();
assert!(
matches!(
err,
@@ -717,9 +590,7 @@ mod k8s_discovery_config_tests {
#[test]
fn mode_accepts_comma_separated_equality_terms() {
// The canonical Plain-mode selector form: `key1=v1,key2=v2`.
let m = cfg(Some("app=sglang,zone=us-east"), None, None)
.mode()
.unwrap();
let m = resolve_mode(Some("app=sglang,zone=us-east"), None, None).unwrap();
assert_eq!(
m,
K8sDiscoveryMode::Plain {
@@ -730,30 +601,29 @@ mod k8s_discovery_config_tests {
#[test]
fn mode_rejects_when_no_selector_is_set() {
let err = cfg(None, None, None).mode().unwrap_err();
let err = resolve_mode(None, None, None).unwrap_err();
assert!(matches!(err, ConfigError::NoSelector), "got {err:?}");
}
#[test]
fn mode_rejects_mixed_plain_and_pd_selectors() {
let err = cfg(
let err = resolve_mode(
Some("app=sglang"),
Some("role=prefill"),
Some("role=decode"),
)
.mode()
.unwrap_err();
assert!(matches!(err, ConfigError::MixedModes), "got {err:?}");
}
#[test]
fn mode_rejects_partial_pd_selectors() {
let err = cfg(None, Some("role=prefill"), None).mode().unwrap_err();
let err = resolve_mode(None, Some("role=prefill"), None).unwrap_err();
assert!(
matches!(err, ConfigError::PartialPdSelectors),
"got {err:?}"
);
let err = cfg(None, None, Some("role=decode")).mode().unwrap_err();
let err = resolve_mode(None, None, Some("role=decode")).unwrap_err();
assert!(
matches!(err, ConfigError::PartialPdSelectors),
"got {err:?}"
@@ -765,7 +635,7 @@ mod k8s_discovery_config_tests {
/// operator opts in by setting plain mode at all).
#[test]
fn mode_accepts_empty_plain_label_selector() {
let m = cfg(Some(""), None, None).mode().unwrap();
let m = resolve_mode(Some(""), None, None).unwrap();
assert_eq!(
m,
K8sDiscoveryMode::Plain {
@@ -782,7 +652,7 @@ mod k8s_discovery_config_tests {
/// at config load.
#[test]
fn mode_pd_rejects_empty_prefill_selector() {
let err = cfg(None, Some(""), Some("role=decode")).mode().unwrap_err();
let err = resolve_mode(None, Some(""), Some("role=decode")).unwrap_err();
assert!(
matches!(
err,
@@ -796,9 +666,7 @@ mod k8s_discovery_config_tests {
#[test]
fn mode_pd_rejects_empty_decode_selector() {
let err = cfg(None, Some("role=prefill"), Some(""))
.mode()
.unwrap_err();
let err = resolve_mode(None, Some("role=prefill"), Some("")).unwrap_err();
assert!(
matches!(err, ConfigError::EmptyPdSelector { selector: "decode" },),
"expected EmptyPdSelector(decode), got {err:?}",
@@ -810,9 +678,7 @@ mod k8s_discovery_config_tests {
/// failure mode as a literal empty string.
#[test]
fn mode_pd_rejects_whitespace_only_prefill_selector() {
let err = cfg(None, Some(" , "), Some("role=decode"))
.mode()
.unwrap_err();
let err = resolve_mode(None, Some(" , "), Some("role=decode")).unwrap_err();
assert!(
matches!(
err,
@@ -829,9 +695,7 @@ mod k8s_discovery_config_tests {
/// so the decode pool stays empty.
#[test]
fn mode_pd_rejects_identical_prefill_and_decode_selectors() {
let err = cfg(None, Some("app=sglang"), Some("app=sglang"))
.mode()
.unwrap_err();
let err = resolve_mode(None, Some("app=sglang"), Some("app=sglang")).unwrap_err();
assert!(
matches!(err, ConfigError::IdenticalPdSelectors),
"expected IdenticalPdSelectors, got {err:?}",
@@ -842,9 +706,7 @@ mod k8s_discovery_config_tests {
/// identical-selector check.
#[test]
fn mode_pd_rejects_identical_selectors_under_whitespace_normalization() {
let err = cfg(None, Some("app=sglang"), Some(" app=sglang "))
.mode()
.unwrap_err();
let err = resolve_mode(None, Some("app=sglang"), Some(" app=sglang ")).unwrap_err();
assert!(
matches!(err, ConfigError::IdenticalPdSelectors),
"expected IdenticalPdSelectors, got {err:?}",
@@ -861,9 +723,7 @@ mod k8s_discovery_config_tests {
/// string level.
#[test]
fn mode_pd_rejects_identical_selectors_under_eq_alias() {
let err = cfg(None, Some("app=sglang"), Some("app==sglang"))
.mode()
.unwrap_err();
let err = resolve_mode(None, Some("app=sglang"), Some("app==sglang")).unwrap_err();
assert!(
matches!(err, ConfigError::IdenticalPdSelectors),
"expected IdenticalPdSelectors, got {err:?}",
@@ -877,9 +737,7 @@ mod k8s_discovery_config_tests {
/// form must agree.
#[test]
fn mode_pd_rejects_identical_selectors_under_inner_whitespace() {
let err = cfg(None, Some("app=sglang"), Some("app = sglang"))
.mode()
.unwrap_err();
let err = resolve_mode(None, Some("app=sglang"), Some("app = sglang")).unwrap_err();
assert!(
matches!(err, ConfigError::IdenticalPdSelectors),
"expected IdenticalPdSelectors, got {err:?}",
@@ -893,9 +751,8 @@ mod k8s_discovery_config_tests {
/// reintroduce the silent-failure bug.)
#[test]
fn mode_pd_rejects_identical_selectors_under_term_order_permutation() {
let err = cfg(None, Some("role=p,app=sglang"), Some("app=sglang,role=p"))
.mode()
.unwrap_err();
let err =
resolve_mode(None, Some("role=p,app=sglang"), Some("app=sglang,role=p")).unwrap_err();
assert!(
matches!(err, ConfigError::IdenticalPdSelectors),
"expected IdenticalPdSelectors, got {err:?}",
@@ -907,12 +764,11 @@ mod k8s_discovery_config_tests {
/// aggressive that it false-positives on legitimate PD configs.
#[test]
fn mode_pd_accepts_truly_distinct_selectors() {
let m = cfg(
let m = resolve_mode(
None,
Some("app=sglang,role=prefill"),
Some("app=sglang,role=decode"),
)
.mode()
.expect("distinct selectors must validate");
assert!(matches!(m, K8sDiscoveryMode::PdDisaggregation { .. }));
}
+35 -3
View File
@@ -328,16 +328,18 @@ pub async fn spawn(
cfg: K8sDiscoveryConfig,
tx: mpsc::Sender<DiscoveryEvent>,
) -> Result<tokio::task::JoinHandle<()>> {
let mode = cfg.mode().context("validate k8s discovery selectors")?;
// The mode was resolved + validated at construction (`resolve_mode` in
// `Cli::build_discovery`); just destructure it here.
let K8sDiscoveryConfig { namespace, mode } = cfg;
let client = Client::try_default()
.await
.context("kube client default config")?;
let api: Api<EndpointSlice> = if cfg.namespace.is_empty() {
let api: Api<EndpointSlice> = if namespace.is_empty() {
Api::all(client)
} else {
Api::namespaced(client, &cfg.namespace)
Api::namespaced(client, &namespace)
};
// Plain mode pushes the single selector to the server side so the LIST
@@ -351,6 +353,36 @@ pub async fn spawn(
};
let watcher_cfg = watcher::Config::default().labels(&server_side_selector);
// Log the resolved namespace + selector(s) at startup. We can't
// verify the namespace exists (the router's RBAC covers
// endpointslices/services/pods, not namespaces, and a correct
// namespace legitimately has zero matching workers until they come
// up), so a typo'd `--service-discovery-namespace` silently watches
// an empty namespace. Surfacing the watch target here lets an
// operator spot the typo in the first log lines instead of only
// discovering it via later `no workers available` request failures.
let namespace_display: &str = if namespace.is_empty() {
"<all namespaces>"
} else {
&namespace
};
match &mode {
K8sDiscoveryMode::Plain { label_selector } => tracing::info!(
namespace = %namespace_display,
label_selector = %label_selector,
"k8s discovery starting (plain mode); a wrong namespace or selector matches zero EndpointSlices"
),
K8sDiscoveryMode::PdDisaggregation {
prefill_selector,
decode_selector,
} => tracing::info!(
namespace = %namespace_display,
prefill_selector = %prefill_selector,
decode_selector = %decode_selector,
"k8s discovery starting (PD mode); a wrong namespace or selector matches zero EndpointSlices"
),
}
let handle = tokio::spawn(async move {
let stream = watcher(api, watcher_cfg);
tokio::pin!(stream);
+1 -1
View File
@@ -25,7 +25,7 @@ pub async fn spawn_discovery(
cfg: &Config,
) -> Result<(mpsc::Receiver<DiscoveryEvent>, tokio::task::JoinHandle<()>)> {
let (tx, rx) = mpsc::channel(DISCOVERY_CHANNEL_CAP);
let handle = match &cfg.discovery.backend {
let handle = match &cfg.discovery {
DiscoveryBackend::StaticUrls(s) => static_urls::spawn(s.clone(), tx).await?,
DiscoveryBackend::K8s(k) => k8s::spawn(k.clone(), tx).await?,
};
+9 -16
View File
@@ -3,18 +3,10 @@
use anyhow::{Context, Result};
use clap::Parser;
use sgl_router::config::LogFormat;
use std::path::PathBuf;
use sgl_router::config::{Cli, LogFormat};
use std::sync::Arc;
use tokio::signal::unix::{signal, Signal, SignalKind};
#[derive(Parser, Debug)]
#[command(name = "sgl-router", version)]
struct Cli {
#[arg(long, env = "SGL_ROUTER_CONFIG")]
config: PathBuf,
}
/// Install the global tracing subscriber.
///
/// Idempotent: a second call returns `Ok` without panicking. When
@@ -54,13 +46,13 @@ fn init_tracing(default_level: &str, format: LogFormat) -> Result<()> {
Ok(())
}
/// Install a minimal text-format subscriber BEFORE config parsing so a
/// config-load error has somewhere to surface. The real subscriber
/// Install a minimal text-format subscriber BEFORE config resolution so a
/// config-resolution error has somewhere to surface. The real subscriber
/// (driven by `Config.observability`) is installed after; the second
/// `try_init` is a no-op because a subscriber is already present.
/// The bootstrap subscriber respects `RUST_LOG` so an operator can
/// debug startup with `RUST_LOG=debug` even when the config file is
/// missing or malformed.
/// debug startup with `RUST_LOG=debug` even when configuration resolution
/// fails.
fn install_bootstrap_subscriber() {
let filter = tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
@@ -83,12 +75,13 @@ fn install_signal_handlers() -> Result<(Signal, Signal)> {
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
// Bootstrap subscriber so a Config::from_path error has structured
// Bootstrap subscriber so a config-resolution error has structured
// output. The configured-format subscriber installs after this and
// becomes a no-op via try_init's idempotency.
install_bootstrap_subscriber();
let cfg = sgl_router::config::Config::from_path(&cli.config)
.with_context(|| format!("load config from {}", cli.config.display()))?;
let cfg = cli
.into_config()
.context("resolve configuration from CLI flags")?;
init_tracing(&cfg.observability.log_level, cfg.observability.log_format)?;
@@ -36,11 +36,14 @@
use crate::config::CacheAwareConfig;
use crate::discovery::ModelId;
use crate::policies::kv_events::{compute_block_hashes, BlockSizeOracle, HashTree};
use crate::policies::kv_events::{
compute_block_hashes, compute_block_hashes_bigram, BlockSizeOracle, HashTree,
};
use crate::policies::{Policy, SelectionContext};
use crate::server::metrics::MetricsRegistry;
use crate::tokenizer::{adapter, TokenizerRegistry};
use crate::workers::Worker;
use std::sync::Arc;
use std::sync::{Arc, OnceLock};
/// Selection policy that scores candidates by tree-overlap with the
/// request's prefix and falls back to load-based picking when the tree
@@ -59,6 +62,14 @@ pub struct CacheAwareZmqPolicy {
/// degrades to min-load — the router cannot hash a prompt without
/// a block size that matches what the worker publishes.
block_size_oracle: Arc<BlockSizeOracle>,
/// Optional metrics sink. Set via [`Self::with_metrics`] by the policy
/// factory for the production policy; `None` in unit tests and
/// non-cache-aware call sites. When set, each cache-aware selection
/// records the prefix-overlap block count into
/// `sgl_router_overlap_blocks`. Set once via [`Self::with_metrics`]
/// (tests) or the `Policy::attach_metrics` hook (production, called by
/// `PolicyRegistry::attach_metrics` after the registry is built).
metrics: OnceLock<Arc<MetricsRegistry>>,
}
impl std::fmt::Debug for CacheAwareZmqPolicy {
@@ -82,9 +93,19 @@ impl CacheAwareZmqPolicy {
tree,
tokenizers,
block_size_oracle,
metrics: OnceLock::new(),
}
}
/// Attach a metrics sink so each cache-aware selection records the
/// prefix-overlap block count into `sgl_router_overlap_blocks`. Builder
/// form used by tests; production wiring goes through the
/// `Policy::attach_metrics` hook.
pub fn with_metrics(self, metrics: Arc<MetricsRegistry>) -> Self {
let _ = self.metrics.set(metrics);
self
}
/// Lowest-load worker — ties broken by stable iteration order (which
/// is the order the registry returned, i.e. dashmap-undefined). For
/// production traffic the ties are rare; tests pin the load skew.
@@ -218,9 +239,22 @@ impl Policy for CacheAwareZmqPolicy {
// has registered yet (oracle empty), cache-aware routing has no
// ground truth to score against; fall back to min-load.
let Some(block_size) = self.block_size_oracle.get() else {
tracing::debug!(
model = %ctx.model(),
"cache-aware-zmq: block size unknown (no worker page_size yet), falling back to min-load",
);
return Self::pick_min_load(workers);
};
let block_hashes = compute_block_hashes(&tokens, block_size as usize);
// EAGLE-family workers hash KV blocks over token bigrams; the query
// hashes must match the worker's stored hashes or the tree lookup
// always misses (overlap stays 0). The oracle carries the worker-
// reported flag.
let is_bigram = self.block_size_oracle.is_bigram();
let block_hashes = if is_bigram {
compute_block_hashes_bigram(&tokens, block_size as usize)
} else {
compute_block_hashes(&tokens, block_size as usize)
};
if block_hashes.is_empty() {
return Self::pick_min_load(workers);
}
@@ -228,13 +262,28 @@ impl Policy for CacheAwareZmqPolicy {
let match_rate = matched.matched_blocks as f32 / block_hashes.len() as f32;
tracing::debug!(
model = %ctx.model(),
hashing = if is_bigram { "bigram" } else { "unigram" },
n_blocks = block_hashes.len(),
matched_blocks = matched.matched_blocks,
match_rate,
cache_threshold = self.config.cache_threshold,
"cache-aware-zmq match_prefix",
);
// Record the matched overlap into `sgl_router_overlap_blocks` before
// the threshold branch, so the histogram captures the full
// distribution — including low-overlap selections that fall back to
// min-load. This is the quantitative signal that cache-aware routing
// is matching prefixes at all.
if let Some(m) = self.metrics.get() {
m.observe_overlap_blocks(ctx.model().0.as_str(), matched.matched_blocks as u64);
}
if match_rate <= self.config.cache_threshold || matched.workers.is_empty() {
tracing::debug!(
model = %ctx.model(),
match_rate,
cache_threshold = self.config.cache_threshold,
"cache-aware-zmq: overlap below threshold, falling back to min-load",
);
return Self::pick_min_load(workers);
}
// Among workers in the matched set, pick the lowest-load one.
@@ -245,7 +294,20 @@ impl Policy for CacheAwareZmqPolicy {
.filter(|w| matched_urls.contains(w.url.as_str()))
.min_by_key(|w| w.active_load())
.map(Arc::clone);
best_matched.or_else(|| Self::pick_min_load(workers))
let chosen = best_matched.or_else(|| Self::pick_min_load(workers));
if let Some(w) = &chosen {
tracing::debug!(
model = %ctx.model(),
worker = %w.url,
matched_blocks = matched.matched_blocks,
"cache-aware-zmq: selected worker by cache overlap",
);
}
chosen
}
fn attach_metrics(&self, metrics: Arc<MetricsRegistry>) {
let _ = self.metrics.set(metrics);
}
}
@@ -292,20 +354,18 @@ mod tests {
port: 0,
},
observability: Default::default(),
models: vec![crate::config::ModelConfig {
model: crate::config::ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: crate::config::PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: crate::config::DiscoveryConfig {
backend: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
},
discovery: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
proxy: crate::config::ProxyConfig::default(),
active_load: crate::config::ActiveLoadConfig::default(),
};
@@ -392,6 +452,266 @@ mod tests {
assert_eq!(chosen.url, "http://w0:30000");
}
/// The cache-aware path records the matched prefix-overlap block count
/// into `sgl_router_overlap_blocks`. Regression: the metric was defined
/// but never observed in production, so the histogram stayed empty and
/// gave no signal that cache-aware routing was matching anything.
#[test]
fn records_overlap_blocks_metric() {
let tree = Arc::new(HashTree::new());
let registry = tokenizer_registry_with_tiny();
let text = "hello world hello world hello world";
let tok = registry.get("tiny").unwrap();
let ids = adapter::encode(&tok, text).unwrap();
let block_size = 4u32;
let hashes = compute_block_hashes(&ids, block_size as usize);
assert!(!hashes.is_empty());
tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes);
let metrics = MetricsRegistry::new();
let policy = CacheAwareZmqPolicy::new(
CacheAwareConfig {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
},
tree,
registry,
oracle_for_tests(4),
)
.with_metrics(Arc::clone(&metrics));
let workers = vec![
worker("http://w0:30000", "tiny"),
worker("http://w1:30000", "tiny"),
];
let model = ModelId("tiny".into());
let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap();
let ctx = SelectionContext::new(&model, Some(&body));
let _ = policy.select(&workers, &ctx).expect("must pick");
let rendered = metrics.render();
assert!(
rendered.contains("sgl_router_overlap_blocks_count{model_id=\"tiny\"}"),
"overlap_blocks histogram must be observed on a cache-aware selection; got:\n{rendered}"
);
}
/// Production wiring path: the policy is stored as `Arc<dyn Policy>` in a
/// `PolicyRegistry`, then `PolicyRegistry::attach_metrics` injects the
/// registry — exactly what `AppContext::with_active_load` does at startup.
/// Exercises trait dispatch (the default no-op vs the `CacheAwareZmqPolicy`
/// override) and the registry fan-out, neither of which the `with_metrics`
/// builder test covers.
#[test]
fn attach_metrics_via_registry_records_overlap() {
let tree = Arc::new(HashTree::new());
let toks = tokenizer_registry_with_tiny();
let text = "hello world hello world hello world";
let tok = toks.get("tiny").unwrap();
let ids = adapter::encode(&tok, text).unwrap();
let hashes = compute_block_hashes(&ids, 4);
assert!(!hashes.is_empty());
tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes);
let policy = CacheAwareZmqPolicy::new(
CacheAwareConfig {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
},
tree,
toks,
oracle_for_tests(4),
);
let model = ModelId("tiny".into());
let registry = crate::policies::PolicyRegistry::default();
registry.insert(model.clone(), Arc::new(policy));
// The production injection point — not the `with_metrics` builder.
let metrics = MetricsRegistry::new();
registry.attach_metrics(Arc::clone(&metrics));
let chosen_policy = registry.get(&model).unwrap();
let workers = vec![
worker("http://w0:30000", "tiny"),
worker("http://w1:30000", "tiny"),
];
let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap();
let ctx = SelectionContext::new(&model, Some(&body));
let _ = chosen_policy.select(&workers, &ctx).expect("must pick");
let rendered = metrics.render();
assert!(
rendered.contains("sgl_router_overlap_blocks_count{model_id=\"tiny\"}"),
"PolicyRegistry::attach_metrics must wire overlap recording through the trait; got:\n{rendered}"
);
}
/// The overlap observation is recorded *before* the cache-threshold branch,
/// so low-overlap selections that fall back to min-load are still counted.
/// `cache_threshold: 1.0` forces the fallback (match_rate is always <= 1.0)
/// even on a full prefix match; assert the histogram is still observed AND
/// the pick came from min-load (w1), not the cache-overlap worker (w0).
#[test]
fn overlap_recorded_even_when_selection_falls_back() {
let tree = Arc::new(HashTree::new());
let toks = tokenizer_registry_with_tiny();
let text = "hello world hello world hello world";
let tok = toks.get("tiny").unwrap();
let ids = adapter::encode(&tok, text).unwrap();
let hashes = compute_block_hashes(&ids, 4);
assert!(!hashes.is_empty());
tree.insert(&KvWorkerId::new("http://w0:30000".into(), 0), None, &hashes);
let metrics = MetricsRegistry::new();
let policy = CacheAwareZmqPolicy::new(
CacheAwareConfig {
cache_threshold: 1.0, // match_rate <= 1.0 always -> always fall back
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
},
tree,
toks,
oracle_for_tests(4),
)
.with_metrics(Arc::clone(&metrics));
// Bump w0's load so min-load picks w1 — distinguishing a min-load
// fallback from the cache-overlap pick (which would be w0). Two guards
// mirror `empty_tree_falls_back_to_min_load` (below the imbalance
// threshold, so the cache-aware path is still reached).
let w0 = worker("http://w0:30000", "tiny");
let w1 = worker("http://w1:30000", "tiny");
let _g = w0.load_guard();
let _g2 = w0.load_guard();
let workers = vec![Arc::clone(&w0), Arc::clone(&w1)];
let model = ModelId("tiny".into());
let body = serde_json::to_vec(&serde_json::json!({"prompt": text})).unwrap();
let ctx = SelectionContext::new(&model, Some(&body));
let chosen = policy.select(&workers, &ctx).expect("must pick");
assert_eq!(
chosen.url, "http://w1:30000",
"cache_threshold 1.0 must force a min-load fallback (w1), not the overlap worker (w0)"
);
let rendered = metrics.render();
assert!(
rendered.contains("sgl_router_overlap_blocks_count{model_id=\"tiny\"}"),
"overlap must be recorded even on the below-threshold fallback; got:\n{rendered}"
);
}
/// End-to-end bigram wiring (the fix that takes `overlap_blocks_sum` from
/// 0 to non-zero for EAGLE models): an EAGLE worker publishes its blocks
/// under BIGRAM hashes. Only a router whose oracle reports `is_bigram` —
/// and thus hashes its query with the bigram hasher — matches them, so
/// overlap is non-zero and it picks the cached worker. A unigram-hashing
/// router against the SAME tree matches nothing (overlap recorded as 0).
#[test]
fn bigram_routing_matches_only_with_bigram_hashing() {
fn overlap_sum(rendered: &str) -> f64 {
rendered
.lines()
.find(|l| l.starts_with("sgl_router_overlap_blocks_sum{model_id=\"tiny\"}"))
.and_then(|l| l.split_whitespace().last())
.and_then(|v| v.parse::<f64>().ok())
.unwrap_or(-1.0)
}
let registry = tokenizer_registry_with_tiny();
let text = "hello world hello world hello world";
let tok = registry.get("tiny").unwrap();
let ids = adapter::encode(&tok, text).unwrap();
let block_size = 4u32;
// The EAGLE worker publishes BIGRAM block hashes.
let bigram_hashes = compute_block_hashes_bigram(&ids, block_size as usize);
assert!(!bigram_hashes.is_empty());
assert_ne!(
bigram_hashes,
compute_block_hashes(&ids, block_size as usize),
"bigram and unigram hashes must differ for this prefix"
);
let model = ModelId("tiny".into());
let body = serde_json::to_vec(&serde_json::json!({ "prompt": text })).unwrap();
// Bigram-aware router (oracle.is_bigram == true): query hashes match
// the bigram tree -> overlap > 0 and it picks the matched worker w0.
{
let tree = Arc::new(HashTree::new());
tree.insert(
&KvWorkerId::new("http://w0:30000".into(), 0),
None,
&bigram_hashes,
);
let oracle = BlockSizeOracle::new();
oracle.try_set(block_size).unwrap();
oracle.set_bigram(true);
let metrics = MetricsRegistry::new();
let policy = CacheAwareZmqPolicy::new(
CacheAwareConfig {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
},
tree,
Arc::clone(&registry),
oracle,
)
.with_metrics(Arc::clone(&metrics));
let workers = vec![
worker("http://w0:30000", "tiny"),
worker("http://w1:30000", "tiny"),
];
let ctx = SelectionContext::new(&model, Some(&body));
let chosen = policy.select(&workers, &ctx).expect("must pick");
assert_eq!(
chosen.url, "http://w0:30000",
"bigram-aware router must match w0's bigram-hashed prefix"
);
assert!(
overlap_sum(&metrics.render()) > 0.0,
"overlap_blocks_sum must be > 0 once the router hashes with bigram"
);
}
// Unigram router (default is_bigram == false) vs the SAME bigram tree:
// query hashes never match -> overlap recorded as 0.
{
let tree = Arc::new(HashTree::new());
tree.insert(
&KvWorkerId::new("http://w0:30000".into(), 0),
None,
&bigram_hashes,
);
let oracle = BlockSizeOracle::new();
oracle.try_set(block_size).unwrap();
let metrics = MetricsRegistry::new();
let policy = CacheAwareZmqPolicy::new(
CacheAwareConfig {
cache_threshold: 0.0,
balance_abs_threshold: 32,
balance_rel_threshold: 1.1,
},
tree,
Arc::clone(&registry),
oracle,
)
.with_metrics(Arc::clone(&metrics));
let workers = vec![
worker("http://w0:30000", "tiny"),
worker("http://w1:30000", "tiny"),
];
let ctx = SelectionContext::new(&model, Some(&body));
let _ = policy.select(&workers, &ctx).expect("must pick");
assert_eq!(
overlap_sum(&metrics.render()),
0.0,
"unigram hashing matches nothing in a bigram tree -> overlap_sum == 0"
);
}
}
/// Two workers both hold the prefix; the lower-load one wins.
#[test]
fn tie_break_by_lowest_active_load() {
+25 -35
View File
@@ -76,17 +76,16 @@ pub fn build_registry(
block_size_oracle: Arc<BlockSizeOracle>,
) -> Result<PolicyRegistry> {
let reg = PolicyRegistry::default();
for m in &cfg.models {
reg.insert(
ModelId(m.id.clone()),
build_policy(
m,
Arc::clone(&tree),
Arc::clone(&tokenizers),
Arc::clone(&block_size_oracle),
),
);
}
let m = &cfg.model;
reg.insert(
ModelId(m.id.clone()),
build_policy(
m,
Arc::clone(&tree),
Arc::clone(&tokenizers),
Arc::clone(&block_size_oracle),
),
);
Ok(reg)
}
@@ -111,34 +110,29 @@ pub fn build_registry_with_defaults(cfg: &Config) -> Result<PolicyRegistry> {
mod tests {
use super::*;
use crate::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ProxyConfig,
ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ProxyConfig, ServerConfig,
StaticUrlsDiscoveryConfig,
};
use crate::config::PolicyKind;
fn cfg_with_models(policies: &[(&str, PolicyKind)]) -> Config {
fn cfg_with_model(id: &str, policy: PolicyKind) -> Config {
Config {
server: ServerConfig {
host: "0".into(),
port: 0,
},
observability: Default::default(),
models: policies
.iter()
.map(|(id, p)| ModelConfig {
id: (*id).into(),
tokenizer_path: "/tmp/x".into(),
policy: *p,
circuit_breaker: None,
cache_aware: None,
})
.collect(),
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
model: ModelConfig {
id: id.into(),
tokenizer_path: "/tmp/x".into(),
policy,
circuit_breaker: None,
cache_aware: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
@@ -154,22 +148,18 @@ mod tests {
}
#[test]
fn registry_assigns_per_model() {
let cfg = cfg_with_models(&[
("qwen", PolicyKind::RoundRobin),
("deepseek", PolicyKind::Random),
]);
fn registry_assigns_configured_model() {
let cfg = cfg_with_model("qwen", PolicyKind::RoundRobin);
let tree = Arc::new(HashTree::new());
let tokenizers = Arc::new(TokenizerRegistry::default());
let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap();
assert!(reg.get(&ModelId("qwen".into())).is_some());
assert!(reg.get(&ModelId("deepseek".into())).is_some());
assert!(reg.get(&ModelId("missing".into())).is_none());
}
#[test]
fn cache_aware_zmq_builds_via_factory() {
let cfg = cfg_with_models(&[("modelA", PolicyKind::CacheAwareZmq)]);
let cfg = cfg_with_model("modelA", PolicyKind::CacheAwareZmq);
let tree = Arc::new(HashTree::new());
let tokenizers = Arc::new(TokenizerRegistry::default());
let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap();
@@ -31,17 +31,32 @@
//! through `KvEventIndex::add_worker`; that refactor can land later
//! without changing the oracle's public surface.
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::atomic::{AtomicU32, AtomicU8, Ordering};
use std::sync::Arc;
/// Tri-state for the bigram flag: distinguishes "not yet reported" from an
/// established `false`, so [`BlockSizeOracle::set_bigram`] can be first-wins
/// (matching `try_set`) rather than last-writer-wins.
const BIGRAM_UNKNOWN: u8 = 0;
const BIGRAM_UNIGRAM: u8 = 1;
const BIGRAM_BIGRAM: u8 = 2;
/// First-wins, idempotent block-size publisher.
///
/// Internally an `AtomicU32` where 0 means "not yet known". Use
/// [`Self::try_set`] to publish a worker-reported value and
/// [`Self::get`] to read at routing time.
///
/// Also carries a `bigram` flag — EAGLE-family workers hash KV blocks over
/// token bigrams, so the policy must pick the bigram hasher. Like `value` it
/// is a per-cluster property (all workers run the same model) and is
/// established first-wins with a loud warning on disagreement, mirroring
/// `try_set` — a heterogeneous EAGLE/non-EAGLE cluster would otherwise let the
/// last registrant silently flip the global hashing mode.
#[derive(Debug, Default)]
pub struct BlockSizeOracle {
value: AtomicU32,
bigram: AtomicU8,
}
/// Returned by [`BlockSizeOracle::try_set`] when the candidate disagrees
@@ -70,6 +85,46 @@ impl BlockSizeOracle {
}
}
/// Publish whether the cluster's workers use bigram (EAGLE-family) KV-block
/// hashing. Called from `KvEventIndex::add_worker` alongside `try_set`.
/// First-wins: the first worker establishes the mode; a later worker that
/// disagrees is logged (not silently honored), since the query-hashing mode
/// is process-wide and one mismatched worker would zero out cache-aware
/// routing for the cluster.
pub fn set_bigram(&self, is_bigram: bool) {
let candidate = if is_bigram {
BIGRAM_BIGRAM
} else {
BIGRAM_UNIGRAM
};
match self.bigram.compare_exchange(
BIGRAM_UNKNOWN,
candidate,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => {}
Err(existing) if existing == candidate => {}
Err(existing) => {
tracing::warn!(
established_bigram = existing == BIGRAM_BIGRAM,
worker_bigram = is_bigram,
"kv-events: worker hashing mode (bigram/EAGLE) disagrees with the \
established cluster value; keeping the first. A heterogeneous \
EAGLE/non-EAGLE cluster will silently never match cache for the \
minority workers — check that all workers run the same model.",
);
}
}
}
/// Whether query hashing should use the bigram variant
/// ([`super::hash::compute_block_hashes_bigram`]). Defaults to `false`
/// until a worker reports an EAGLE-family `speculative_algorithm`.
pub fn is_bigram(&self) -> bool {
self.bigram.load(Ordering::Relaxed) == BIGRAM_BIGRAM
}
/// Publish a candidate block size. Returns the established value on
/// success (idempotent: same candidate as already set is `Ok`);
/// returns `Err(BlockSizeMismatch)` when the candidate disagrees.
@@ -123,6 +178,36 @@ mod tests {
assert_eq!(oracle.get(), Some(64));
}
#[test]
fn bigram_flag_defaults_false_first_wins_and_is_idempotent() {
let oracle = BlockSizeOracle::new();
assert!(
!oracle.is_bigram(),
"unknown (no worker reported yet) reads as non-bigram"
);
oracle.set_bigram(true);
assert!(oracle.is_bigram(), "first worker establishes the mode");
oracle.set_bigram(true); // idempotent agreement
assert!(oracle.is_bigram());
// Independent of block_size establishment.
assert_eq!(oracle.get(), None);
// First-wins: a conflicting later worker is logged, not honored.
oracle.set_bigram(false);
assert!(
oracle.is_bigram(),
"a disagreeing worker must not flip the established mode"
);
}
#[test]
fn bigram_flag_establishes_false_first_wins() {
let oracle = BlockSizeOracle::new();
oracle.set_bigram(false);
assert!(!oracle.is_bigram(), "established as unigram");
oracle.set_bigram(true); // conflicting; first (unigram) wins
assert!(!oracle.is_bigram());
}
#[test]
fn mismatching_set_fails_without_changing_state() {
let oracle = BlockSizeOracle::new();
@@ -48,6 +48,14 @@ pub struct EventConfig {
/// many SUB connections (one per rank), skipping any rank whose
/// `port_base + dp_rank` overflows `u16`.
pub dp_size: u32,
/// Whether the worker uses EAGLE-family speculative decoding (EAGLE /
/// EAGLE3 / FROZEN_KV_MTP), reported via `/server_info`'s top-level
/// `speculative_algorithm`. When true the worker hashes KV blocks over
/// overlapping token *bigrams* (`is_bigram = is_eagle`), so the router must
/// use [`super::hash::compute_block_hashes_bigram`] for its query hashes to
/// match the worker's stored hashes — otherwise cache-aware routing
/// silently never matches and degrades to min-load.
pub is_bigram: bool,
}
/// Default timeout for the `/server_info` introspection request. The
@@ -88,6 +96,10 @@ pub async fn fetch_event_config(
let body = fetch_with_retry(&server_info_url, worker_url, client).await?;
// EAGLE-family speculative decoding ⇒ the worker hashes KV blocks over
// token bigrams; the router must mirror that on the selection side.
let is_bigram = classify_bigram(body.speculative_algorithm.as_deref(), worker_url);
let block = match body.kv_events {
Some(b) => b,
None => {
@@ -117,6 +129,7 @@ pub async fn fetch_event_config(
topic: block.topic,
block_size: block.block_size,
dp_size: block.dp_size,
is_bigram,
}))
}
@@ -185,6 +198,43 @@ async fn fetch_with_retry(
struct ServerInfoResponse {
#[serde(default)]
kv_events: Option<KvEventsBlock>,
/// Top-level `/server_info` field. EAGLE-family values
/// (EAGLE / EAGLE3 / FROZEN_KV_MTP) mean the worker hashes KV blocks over
/// token bigrams — see [`EventConfig::is_bigram`].
#[serde(default)]
speculative_algorithm: Option<String>,
}
/// Whether a worker's `/server_info` `speculative_algorithm` means it hashes KV
/// blocks over token bigrams. Recognizes the engine's `is_eagle()` set
/// (`EAGLE`, `EAGLE3`, `FROZEN_KV_MTP`, case-insensitive).
///
/// An *unrecognized* value that looks EAGLE-family (contains `EAGLE` or `MTP`)
/// is logged loudly and treated as non-bigram — it most likely means a new
/// EAGLE variant the router doesn't know yet, which would otherwise silently
/// zero out cache-aware routing (the exact failure this whole path fixes).
/// Recognized non-EAGLE algorithms (and the absent field) map to `false`
/// silently.
pub(crate) fn classify_bigram(speculative_algorithm: Option<&str>, worker_url: &str) -> bool {
let Some(algo) = speculative_algorithm else {
return false;
};
let upper = algo.to_ascii_uppercase();
match upper.as_str() {
"EAGLE" | "EAGLE3" | "FROZEN_KV_MTP" => true,
_ => {
if upper.contains("EAGLE") || upper.contains("MTP") {
tracing::warn!(
worker_url = %worker_url,
speculative_algorithm = %algo,
"kv-events: unrecognized EAGLE-like speculative_algorithm; treating as \
non-bigram (unigram) hashing. If this is an EAGLE-family algorithm, \
cache-aware routing will silently never match — add it to classify_bigram",
);
}
false
}
}
}
#[derive(Deserialize)]
@@ -270,10 +320,47 @@ mod tests {
topic: "kv".to_string(),
block_size: 64,
dp_size: 2,
is_bigram: false,
})
);
}
/// EAGLE-family `speculative_algorithm` (and only those) must set
/// `is_bigram`, so the router selects the bigram hasher and its query
/// hashes match the worker's bigram-stored block hashes.
#[tokio::test]
async fn fetch_sets_is_bigram_for_eagle_family_only() {
for (algo, expected) in [
(Some("EAGLE"), true),
(Some("EAGLE3"), true),
(Some("FROZEN_KV_MTP"), true),
(Some("eagle"), true), // case-insensitive
(Some("NONE"), false),
(Some("NEXTN"), false), // non-eagle speculative algorithm
(None, false), // no speculative decoding
] {
let mut obj = json!({
"kv_events": {
"publisher": "zmq",
"endpoint_host": "*",
"endpoint_port_base": 5557,
"topic": "",
"block_size": 64,
"dp_size": 1,
}
});
if let Some(a) = algo {
obj["speculative_algorithm"] = json!(a);
}
let (url, _shutdown) = spawn_fake_worker(Arc::new(obj)).await;
let got = fetch_event_config(&url, &client()).await.unwrap().unwrap();
assert_eq!(
got.is_bigram, expected,
"speculative_algorithm={algo:?} should map to is_bigram={expected}"
);
}
}
/// Worker reports a specific bind host (not wildcard): gateway must
/// honour it instead of overwriting from the URL.
#[tokio::test]
@@ -34,10 +34,11 @@
//!
//! ### Bigram mode
//!
//! Not supported in v1. SGLang's bigram mode interleaves overlapping
//! `(t_i, t_{i+1})` pairs into the hash. The gateway does not need this
//! today; if/when it does, add a separate `compute_block_hashes_bigram` rather
//! than complicating the non-bigram fast path.
//! EAGLE-family workers (`is_bigram = is_eagle`) hash KV blocks over
//! overlapping `(t_i, t_{i+1})` token pairs. That path is implemented as a
//! separate [`compute_block_hashes_bigram`] (below) rather than branching
//! inside the non-bigram fast path; `CacheAwareZmqPolicy::select` chooses
//! between the two from the worker-reported bigram flag.
use sha2::{Digest, Sha256};
@@ -107,10 +108,157 @@ pub fn sha256_to_i64(digest: &[u8; 32]) -> i64 {
i64::from_be_bytes(top)
}
/// Bigram variant of [`compute_block_hashes`], matching SGLang's `radix_cache`
/// worker when the model runs **EAGLE speculative decoding** (`is_bigram =
/// is_eagle`). Mirrors `RadixKey.hash_page` (Python:
/// `mem_cache/radix_cache.py`) on the bigram path:
///
/// - The logical sequence is the `N-1` overlapping bigrams of `N` raw tokens,
/// so the page count is `ceil((len-1) / block_size)`. Fewer than 2 tokens
/// yields no blocks.
/// - Each page `[start, end)` (in bigram-index space) feeds **both** tokens of
/// every bigram into the SHA256 hasher — `t[j]` then `t[j+1]`, each as 4
/// little-endian bytes — vs. the unigram path's single token per unit.
/// - Pages chain on the prior page's full 32-byte digest and truncate to i64
/// exactly as the unigram path does.
///
/// Use this (instead of [`compute_block_hashes`]) when the worker advertises an
/// EAGLE speculative algorithm via `/server_info`; otherwise the router's query
/// hashes won't match the worker's stored bigram block hashes and cache-aware
/// routing silently degrades to min-load.
pub fn compute_block_hashes_bigram(token_ids: &[u32], block_size: usize) -> Vec<i64> {
assert!(block_size > 0, "block_size must be positive");
// N raw tokens -> N-1 overlapping bigrams; fewer than 2 tokens -> no blocks.
let logical_len = token_ids.len().saturating_sub(1);
if logical_len == 0 {
return Vec::new();
}
let num_blocks = logical_len.div_ceil(block_size);
let mut out = Vec::with_capacity(num_blocks);
let mut prior: Option<[u8; 32]> = None;
let mut start = 0;
while start < logical_len {
let end = (start + block_size).min(logical_len);
let digest = chain_block_bigram(prior.as_ref(), token_ids, start, end);
out.push(sha256_to_i64(&digest));
prior = Some(digest);
start = end;
}
out
}
/// Hash a single bigram page: for each unit `j` in `[start, end)`, feed
/// `tokens[j]` then `tokens[j + 1]` (4 little-endian bytes each), chained on the
/// parent block's full 32-byte digest. The caller guarantees `end <= len - 1`,
/// so `tokens[j + 1]` is always in bounds. Mirrors the engine's `hash_page`
/// bigram branch.
#[inline]
fn chain_block_bigram(
parent_digest: Option<&[u8; 32]>,
tokens: &[u32],
start: usize,
end: usize,
) -> [u8; 32] {
let mut hasher = Sha256::new();
if let Some(parent) = parent_digest {
hasher.update(parent);
}
for j in start..end {
hasher.update(tokens[j].to_le_bytes());
hasher.update(tokens[j + 1].to_le_bytes());
}
hasher.finalize().into()
}
#[cfg(test)]
mod tests {
use super::*;
// ---- Bigram cross-language goldens ----
// Values produced by SGLang's REAL `RadixKey(..., is_bigram=True).hash_page`
// + `compute_node_hash_values` chunking + `hash_str_to_int64`, run against
// the deployed DeepSeek-V4-Flash engine. These lock byte-exact equivalence
// with the worker's stored block hashes — the contract that makes
// cache-aware routing actually match for EAGLE/bigram models.
#[test]
fn bigram_golden_single_block_full() {
// engine: chain_bigram([10,20,30,40], 4) -> [-2735951481331064195]
assert_eq!(
compute_block_hashes_bigram(&[10, 20, 30, 40], 4),
vec![-2735951481331064195_i64]
);
}
#[test]
fn bigram_golden_multi_block() {
// engine: chain_bigram([10,20,30,40,50], 2) -> [-8847804484166691499, 4989791362144317498]
assert_eq!(
compute_block_hashes_bigram(&[10, 20, 30, 40, 50], 2),
vec![-8847804484166691499_i64, 4989791362144317498_i64]
);
}
#[test]
fn bigram_golden_partial_last_block() {
// engine: chain_bigram([1,2,3,4,5,6], 4) -> [-638950109823820341, 3604587133525381017]
assert_eq!(
compute_block_hashes_bigram(&[1, 2, 3, 4, 5, 6], 4),
vec![-638950109823820341_i64, 3604587133525381017_i64]
);
}
#[test]
fn bigram_golden_longer_multi_block() {
// engine: chain_bigram([5,6,7,8,9,10,11,12,13], 4) -> [-2900568514773989563, -322435596280658912]
assert_eq!(
compute_block_hashes_bigram(&[5, 6, 7, 8, 9, 10, 11, 12, 13], 4),
vec![-2900568514773989563_i64, -322435596280658912_i64]
);
}
#[test]
fn bigram_single_bigram_equals_unigram_pair() {
// One bigram (10,20) feeds bytes 10,20 — identical to a unigram block
// [10,20]. engine: chain_bigram([10,20], 4) -> [978178666101069530],
// which equals the unigram block hash of [10,20].
assert_eq!(
compute_block_hashes_bigram(&[10, 20], 4),
vec![978178666101069530_i64]
);
assert_eq!(
compute_block_hashes_bigram(&[10, 20], 4),
compute_block_hashes(&[10, 20], 4)
);
}
#[test]
fn bigram_fewer_than_two_tokens_yields_no_blocks() {
// N tokens -> N-1 bigrams; <2 tokens -> 0 bigrams -> empty.
assert!(compute_block_hashes_bigram(&[10], 4).is_empty());
assert!(compute_block_hashes_bigram(&[], 4).is_empty());
}
#[test]
fn bigram_differs_from_unigram_for_multi_token_blocks() {
// Sanity: for >2 tokens the bigram hash must NOT equal the unigram hash
// (different byte stream) — this is exactly why a unigram-hashing
// router gets zero overlap against a bigram worker.
let toks = [10u32, 20, 30, 40];
assert_ne!(
compute_block_hashes_bigram(&toks, 4),
compute_block_hashes(&toks, 4)
);
}
#[test]
#[should_panic(expected = "block_size must be positive")]
fn bigram_zero_block_size_panics() {
let _ = compute_block_hashes_bigram(&[1, 2, 3], 0);
}
/// Helper for tests: derive the expected i64 from a list of tokens
/// chained against an optional parent digest. This mirrors `chain_block`
/// but is duplicated here so a regression in the production helper
@@ -204,11 +204,16 @@ impl KvEventIndex {
);
return;
}
// Establish the bigram flag alongside block_size. EAGLE-family workers
// hash KV blocks over token bigrams, so the policy must use the bigram
// hasher for its query hashes to match the worker's stored hashes.
self.block_size_oracle.set_bigram(cfg.is_bigram);
info!(
worker_url = %worker_url,
dp_size = cfg.dp_size,
port_base = cfg.port_base,
block_size = cfg.block_size,
is_bigram = cfg.is_bigram,
"kv-events: subscribing",
);
// Compute the DP ranks that will actually be subscribed (skip
@@ -646,6 +651,7 @@ mod tests {
topic: String::new(),
block_size: 128,
dp_size: 1,
is_bigram: false,
};
index
.add_worker("http://127.0.0.1:30100", Some(bad_cfg))
@@ -674,9 +680,34 @@ mod tests {
topic: String::new(),
block_size: 64,
dp_size: 0,
is_bigram: false,
};
index.add_worker("http://127.0.0.1:30200", Some(cfg)).await;
assert_eq!(index.block_size_oracle().get(), Some(64));
index.shutdown().await;
}
#[tokio::test]
async fn add_worker_seeds_bigram_flag_from_event_config() {
// The discovery->routing seam: add_worker must publish
// EventConfig.is_bigram into the oracle (alongside block_size) so
// select() picks the bigram hasher for EAGLE workers.
let index = KvEventIndex::new();
assert!(!index.block_size_oracle().is_bigram());
// dp_size=0 short-circuits the subscriber spawn but still runs the seed.
let cfg = EventConfig {
host: "127.0.0.1".into(),
port_base: 30300,
topic: String::new(),
block_size: 64,
dp_size: 0,
is_bigram: true,
};
index.add_worker("http://127.0.0.1:30300", Some(cfg)).await;
assert!(
index.block_size_oracle().is_bigram(),
"add_worker must seed the bigram flag from EventConfig"
);
index.shutdown().await;
}
}
@@ -23,8 +23,9 @@ pub mod tree;
pub mod wire;
pub use block_size_oracle::BlockSizeOracle;
pub(crate) use discovery::classify_bigram;
pub use discovery::{fetch_event_config, EventConfig};
pub use hash::{compute_block_hashes, sha256_to_i64};
pub use hash::{compute_block_hashes, compute_block_hashes_bigram, sha256_to_i64};
pub use index::KvEventIndex;
pub use subscriber::{KvEventSubscriberRegistry, WorkerEvent};
pub use tree::{HashTree, KvWorkerId, MatchResult};
@@ -593,6 +593,7 @@ mod tests {
topic: String::new(),
block_size: 64,
dp_size,
is_bigram: false,
}
}
@@ -213,7 +213,64 @@ impl<'de> Deserialize<'de> for BoundedI64Vec {
}
}
/// `BoundedI64Vec`'s `u32` twin. Same shape, different cap.
/// One element of a `token_ids` array. SGLang emits a flat `u32` per token for
/// unigram pages, but a 2-element `[t_i, t_{i+1}]` array per token for *bigram*
/// pages (`mem_cache/events.py`, `is_bigram` branch — DeepSeek-V4-class models).
/// `token_ids` is purely informational for the gateway (routing keys off the
/// engine-provided `block_hashes`), so we accept either shape and flatten the
/// ints rather than model the bigram pairing.
enum TokenCell {
One(u32),
Many(Vec<u32>),
}
impl<'de> Deserialize<'de> for TokenCell {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
struct V;
impl<'de> Visitor<'de> for V {
type Value = TokenCell;
fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.write_str("a token id (u32) or an array of token ids")
}
// serde's default visit_u8/u16/u32 forward to visit_u64, and
// visit_i8/i16/i32 forward to visit_i64, so these two cover every
// integer width msgpack might use for a scalar token id.
fn visit_u64<E: de::Error>(self, v: u64) -> Result<TokenCell, E> {
Ok(TokenCell::One(v as u32))
}
fn visit_i64<E: de::Error>(self, v: i64) -> Result<TokenCell, E> {
Ok(TokenCell::One(v as u32))
}
fn visit_seq<A>(self, mut seq: A) -> Result<TokenCell, A::Error>
where
A: SeqAccess<'de>,
{
let mut ts: Vec<u32> = match seq.size_hint() {
Some(h) => Vec::with_capacity(h.min(8)),
None => Vec::new(),
};
while let Some(t) = seq.next_element::<u32>()? {
if ts.len() >= MAX_TOKENS_PER_EVENT {
return Err(de::Error::custom(format!(
"{PAYLOAD_TOO_LARGE_TAG}:token_ids:{}:{MAX_TOKENS_PER_EVENT}",
ts.len() + 1
)));
}
ts.push(t);
}
Ok(TokenCell::Many(ts))
}
}
deserializer.deserialize_any(V)
}
}
/// `BoundedI64Vec`'s `u32` twin. Same shape, different cap. Accepts both flat
/// (unigram) token ids and bigram `[t_i, t_{i+1}]` pairs via [`TokenCell`],
/// flattening the latter.
#[derive(Debug, Clone, PartialEq)]
struct BoundedU32Vec(Vec<u32>);
@@ -243,14 +300,27 @@ impl<'de> Deserialize<'de> for BoundedU32Vec {
Some(h) => Vec::with_capacity(h),
None => Vec::new(),
};
while let Some(v) = seq.next_element::<u32>()? {
if out.len() >= MAX_TOKENS_PER_EVENT {
return Err(de::Error::custom(format!(
"{PAYLOAD_TOO_LARGE_TAG}:token_ids:{}:{MAX_TOKENS_PER_EVENT}",
out.len() + 1
)));
// Each element is either a scalar token id (unigram) or a
// `[t_i, t_{i+1}]` pair (bigram); flatten both into `out`.
while let Some(cell) = seq.next_element::<TokenCell>()? {
let push = |t: u32, out: &mut Vec<u32>| -> Result<(), A::Error> {
if out.len() >= MAX_TOKENS_PER_EVENT {
return Err(de::Error::custom(format!(
"{PAYLOAD_TOO_LARGE_TAG}:token_ids:{}:{MAX_TOKENS_PER_EVENT}",
out.len() + 1
)));
}
out.push(t);
Ok(())
};
match cell {
TokenCell::One(t) => push(t, &mut out)?,
TokenCell::Many(ts) => {
for t in ts {
push(t, &mut out)?;
}
}
}
out.push(v);
}
Ok(out)
}
@@ -417,6 +487,84 @@ mod tests {
}
}
/// Encode `token_ids` the way SGLang's *bigram* pages do: a sequence of
/// 2-element `[t_i, t_{i+1}]` arrays instead of flat ints. See
/// `mem_cache/events.py` (`is_bigram` branch).
fn write_bigram_token_array(buf: &mut Vec<u8>, pairs: &[(u32, u32)]) {
mp::write_array_len(buf, pairs.len() as u32).unwrap();
for (a, b) in pairs {
mp::write_array_len(buf, 2).unwrap();
mp::write_uint(buf, *a as u64).unwrap();
mp::write_uint(buf, *b as u64).unwrap();
}
}
/// Like `build_block_stored_bytes`, but `token_ids` is the bigram
/// list-of-pairs shape that DeepSeek-V4-class models emit.
fn build_block_stored_bigram_bytes(
block_hashes: &[i64],
parent: Option<i64>,
token_pairs: &[(u32, u32)],
block_size: u32,
lora_id: Option<i64>,
medium: Option<&str>,
) -> Vec<u8> {
let mut buf = Vec::new();
write_event_array(&mut buf, "BlockStored", 7);
write_i64_array(&mut buf, block_hashes);
match parent {
Some(v) => {
mp::write_sint(&mut buf, v).unwrap();
}
None => mp::write_nil(&mut buf).unwrap(),
}
write_bigram_token_array(&mut buf, token_pairs);
mp::write_uint(&mut buf, block_size as u64).unwrap();
match lora_id {
Some(v) => {
mp::write_sint(&mut buf, v).unwrap();
}
None => mp::write_nil(&mut buf).unwrap(),
}
match medium {
Some(s) => mp::write_str(&mut buf, s).unwrap(),
None => mp::write_nil(&mut buf).unwrap(),
}
buf
}
/// Regression: bigram models (e.g. DeepSeek-V4-Flash) emit `token_ids` as
/// `[[t_i, t_{i+1}], ...]`. The decoder previously read `token_ids` as a
/// flat `u32` array and failed the entire batch with
/// "wrong msgpack marker FixArray(2)", silently disabling cache-aware
/// routing. It must instead accept the bigram shape (flattening the ints).
#[test]
fn decodes_block_stored_with_bigram_token_ids() {
let event = build_block_stored_bigram_bytes(
&[111_i64],
None,
&[(10, 20), (20, 30)],
2,
None,
Some("GPU"),
);
let bytes = build_batch_bytes(1.5, &[event], Some(0), true);
let batch = decode_event_batch(&bytes).expect("decode bigram token_ids");
assert_eq!(batch.events.len(), 1);
match &batch.events[0] {
KvCacheEvent::BlockStored(b) => {
// routing-relevant fields decode unchanged
assert_eq!(b.block_hashes, vec![111]);
assert_eq!(b.parent_block_hash, None);
assert_eq!(b.block_size, 2);
// bigram pairs are flattened into the (informational) token vec
assert_eq!(b.token_ids, vec![10, 20, 20, 30]);
}
other => panic!("expected BlockStored, got {other:?}"),
}
}
/// Build a full BlockStored event as msgspec would emit it (all 7
/// elements: tag + 6 fields). `medium` may be Some/None.
fn build_block_stored_bytes(
@@ -11,6 +11,7 @@ pub mod registry;
pub mod round_robin;
use crate::discovery::ModelId;
use crate::server::metrics::MetricsRegistry;
use crate::workers::Worker;
use dashmap::DashMap;
use std::sync::Arc;
@@ -46,6 +47,13 @@ impl<'a> SelectionContext<'a> {
pub trait Policy: Send + Sync + std::fmt::Debug {
fn select(&self, workers: &[Arc<Worker>], ctx: &SelectionContext<'_>) -> Option<Arc<Worker>>;
/// Attach the process metrics registry after construction. Default is a
/// no-op — only policies that emit metrics (cache-aware-zmq's
/// `sgl_router_overlap_blocks`) override it. Mirrors
/// `ActiveLoadRegistry::attach_metrics`: the registry is built after the
/// policies, so it is injected here rather than passed to the constructor.
fn attach_metrics(&self, _metrics: Arc<MetricsRegistry>) {}
}
#[derive(Debug, Default)]
@@ -61,4 +69,13 @@ impl PolicyRegistry {
pub fn get(&self, model: &ModelId) -> Option<Arc<dyn Policy>> {
self.by_model.get(model).map(|p| p.clone())
}
/// Inject the metrics registry into every registered policy. Called once
/// at startup (after the registry is built) so metrics-emitting policies
/// can record into the shared registry.
pub fn attach_metrics(&self, metrics: Arc<MetricsRegistry>) {
for entry in self.by_model.iter() {
entry.value().attach_metrics(Arc::clone(&metrics));
}
}
}
@@ -69,6 +69,11 @@ impl AppContext {
// Without this, the metric is permanently 0 in production even
// though the chat handler is faithfully calling `register`.
active_load.attach_metrics(Arc::clone(&metrics));
// Same rationale for the cache-aware-zmq policy's
// `sgl_router_overlap_blocks`: the metrics registry is built here,
// after the policy registry, so inject it now. No-op for policies
// that don't emit metrics.
policies.attach_metrics(Arc::clone(&metrics));
Self {
config,
tokenizers,
@@ -100,14 +105,18 @@ impl AppContext {
port: 0,
},
observability: Default::default(),
models: vec![],
discovery: crate::config::DiscoveryConfig {
backend: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
model: crate::config::ModelConfig {
id: "stub-model".into(),
tokenizer_path: "stub".into(),
policy: crate::config::PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
},
discovery: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
proxy: crate::config::ProxyConfig::default(),
active_load: crate::config::ActiveLoadConfig::default(),
},
@@ -144,6 +144,13 @@ impl ApiError {
ApiError::Internal(_) => (StatusCode::INTERNAL_SERVER_ERROR, "internal_error"),
}
}
/// The HTTP status this error maps to — same value the client receives via
/// `into_response`. Exposed so the access log records the real status
/// (e.g. 502/503/504) instead of a sentinel.
pub fn status_code(&self) -> StatusCode {
self.status_and_code().0
}
}
#[derive(Serialize)]
@@ -75,6 +75,7 @@ pub async fn chat_completions(
headers: HeaderMap,
body: Bytes,
) -> Result<Response<Body>, ApiError> {
let start = std::time::Instant::now();
let probe = parse_probe(&body)?;
let streaming = probe.stream.unwrap_or(false);
let model_str = probe
@@ -397,6 +398,37 @@ pub async fn chat_completions(
ctx.metrics
.record_request(&metrics_worker_url, &metrics_model, metrics_mode, outcome);
// Per-request access log — always on at INFO so incoming traffic and its
// status are visible without DEBUG. `request_id` is the client/gateway
// X-Request-Id (echoed end-to-end); `worker` is the engine the policy
// selected. The cache-aware routing rationale is logged separately at
// DEBUG by the policy.
let request_id = headers
.get("x-request-id")
.and_then(|v| v.to_str().ok())
.unwrap_or("-");
let http_status = match &result {
Ok(resp) => resp.status().as_u16(),
Err(e) => e.status_code().as_u16(),
};
let outcome_str = match outcome {
RequestOutcome::Success => "success",
RequestOutcome::Error => "error",
RequestOutcome::Cancelled => "cancelled",
};
tracing::info!(
request_id = %request_id,
method = "POST",
path = "/v1/chat/completions",
model = %metrics_model,
worker = %metrics_worker_url,
outcome = outcome_str,
http_status,
stream = streaming,
latency_ms = start.elapsed().as_millis() as u64,
"chat_completions",
);
// Mirror the upstream `x-sgl-decode-url` hint onto the response so
// external tests / sidecars can observe PD decode affinity without
// sniffing the proxy hop. The request-side header was set above for
@@ -21,16 +21,14 @@ pub struct ModelEntry {
}
pub async fn list_models(State(ctx): State<Arc<AppContext>>) -> Json<ModelsList> {
let data = ctx
.config
.models
.iter()
.map(|m| ModelEntry {
id: m.id.clone(),
object: "model",
owned_by: "sglang",
})
.collect();
// The router serves a single configured model; OpenAI clients still
// expect a list shape, so return a one-element `data` array.
let m = &ctx.config.model;
let data = vec![ModelEntry {
id: m.id.clone(),
object: "model",
owned_by: "sglang",
}];
Json(ModelsList {
object: "list",
data,
@@ -47,24 +45,15 @@ mod tests {
use crate::config::PolicyKind;
#[tokio::test]
async fn lists_configured_models() {
async fn lists_configured_model() {
let mut ctx = crate::server::app_context::AppContext::stub();
ctx.config.models = vec![
crate::config::ModelConfig {
id: "qwen3".into(),
tokenizer_path: "x".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
},
crate::config::ModelConfig {
id: "deepseek".into(),
tokenizer_path: "y".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
},
];
ctx.config.model = crate::config::ModelConfig {
id: "qwen3".into(),
tokenizer_path: "x".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
};
let app = crate::server::app::build_router(std::sync::Arc::new(ctx));
let res = app
.oneshot(
@@ -85,13 +74,12 @@ mod tests {
.iter()
.map(|m| m["id"].as_str().unwrap())
.collect();
assert_eq!(ids, vec!["qwen3", "deepseek"]);
assert_eq!(ids, vec!["qwen3"]);
assert_eq!(v["data"][0]["object"], "model");
// Pin `owned_by` so a refactor that flips the hardcoded value to
// "openai" / "" / a typo would fail loudly here. OpenAI clients
// expect this field and some (e.g. langchain-openai) treat
// `owned_by != "system"` as a meaningful signal.
assert_eq!(v["data"][0]["owned_by"], "sglang");
assert_eq!(v["data"][1]["owned_by"], "sglang");
}
}
@@ -114,20 +114,18 @@ mod tests {
port: 0,
},
observability: Default::default(),
models: vec![crate::config::ModelConfig {
model: crate::config::ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: crate::config::DiscoveryConfig {
backend: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
},
discovery: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
proxy: crate::config::ProxyConfig::default(),
active_load: crate::config::ActiveLoadConfig::default(),
};
@@ -3,12 +3,64 @@
use anyhow::{Context, Result};
use dynamo_tokenizers::{traits::DecodeResult, Tokenizer};
use std::path::Path;
use std::sync::Arc;
pub fn load(path: &str) -> Result<Arc<Tokenizer>> {
/// Load a tokenizer from `source`, which is either a local `tokenizer.json`
/// path or a HuggingFace repo id.
///
/// An existing local file (or anything with a filesystem-path shape) is
/// loaded directly via `Tokenizer::from_file`. Otherwise `source` is treated
/// as a HuggingFace repo id and its `tokenizer.json` is downloaded (once, at
/// startup) into the HF cache, honoring `HF_TOKEN` / `HF_HOME` /
/// `HF_HUB_OFFLINE`. `dynamo_tokenizers` itself has no HF-download path, so
/// the fetch is done here via `hf-hub`.
pub fn load(source: &str) -> Result<Arc<Tokenizer>> {
if Path::new(source).is_file() || looks_like_path(source) {
return Tokenizer::from_file(source)
.map(Arc::new)
.with_context(|| format!("load tokenizer from {source}"));
}
let downloaded = download_tokenizer_json(source)?;
let path = downloaded
.to_str()
.context("downloaded tokenizer path is not valid UTF-8")?;
Tokenizer::from_file(path)
.map(Arc::new)
.with_context(|| format!("load tokenizer from {path}"))
.with_context(|| format!("load downloaded tokenizer for {source}"))
}
/// Treat `source` as a filesystem path (rather than a HuggingFace repo id)
/// when it has a path-like shape — an absolute/relative prefix or a `.json`
/// suffix. HF repo ids are `namespace/name` with none of these markers, so a
/// missing local file like `/models/tok.json` reports a load error instead of
/// silently attempting a (doomed) network fetch.
fn looks_like_path(source: &str) -> bool {
source.starts_with('/')
|| source.starts_with("./")
|| source.starts_with("../")
|| source.starts_with('~')
|| source.ends_with(".json")
}
/// Download `tokenizer.json` for a HuggingFace repo id and return the cached
/// local path. Uses the blocking `ureq` API (this runs once at startup,
/// before the server begins serving) and `from_env` so `HF_TOKEN` /
/// `HF_HOME` / endpoint overrides are honored.
fn download_tokenizer_json(repo_id: &str) -> Result<std::path::PathBuf> {
use hf_hub::api::sync::ApiBuilder;
let api = ApiBuilder::from_env()
.build()
.context("initialize HuggingFace Hub client")?;
api.model(repo_id.to_string())
.get("tokenizer.json")
.with_context(|| {
format!(
"download tokenizer.json for HuggingFace repo {repo_id:?} \
(pass --tokenizer-path with a local tokenizer.json, or set HF_TOKEN \
for a gated/private repo)"
)
})
}
pub fn encode(t: &Tokenizer, text: &str) -> Result<Vec<u32>> {
+10 -13
View File
@@ -24,10 +24,9 @@ impl std::fmt::Debug for TokenizerRegistry {
impl TokenizerRegistry {
pub fn load_from_config(cfg: &crate::config::Config) -> Result<Self> {
let me = TokenizerRegistry::default();
for m in &cfg.models {
let t = adapter::load(&m.tokenizer_path)?;
me.inner.insert(m.id.clone(), t);
}
let m = &cfg.model;
let t = adapter::load(&m.tokenizer_path)?;
me.inner.insert(m.id.clone(), t);
Ok(me)
}
@@ -54,20 +53,18 @@ mod tests {
port: 0,
},
observability: Default::default(),
models: vec![crate::config::ModelConfig {
model: crate::config::ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: crate::config::DiscoveryConfig {
backend: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
},
discovery: crate::config::DiscoveryBackend::StaticUrls(
crate::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
proxy: crate::config::ProxyConfig::default(),
active_load: crate::config::ActiveLoadConfig::default(),
}
@@ -192,7 +189,7 @@ mod tests {
#[test]
fn missing_file_errors() {
let mut c = cfg();
c.models[0].tokenizer_path = "/nonexistent.json".into();
c.model.tokenizer_path = "/nonexistent.json".into();
let err = TokenizerRegistry::load_from_config(&c).unwrap_err();
assert!(err.to_string().to_lowercase().contains("tokenizer"));
}
@@ -125,9 +125,15 @@ impl WorkerIntrospector {
None => None,
};
// EAGLE-family speculative decoding ⇒ the worker hashes KV blocks over
// token bigrams; the router must mirror that on the selection side.
let is_bigram = crate::policies::kv_events::classify_bigram(
parsed.speculative_algorithm.as_deref(),
worker_url,
);
let event_config = parsed
.kv_events
.map(|block| resolve_event_config(block, worker_url));
.map(|block| resolve_event_config(block, worker_url, is_bigram));
let disaggregation_role = resolve_disaggregation_role(
parsed.disaggregation_mode.as_deref(),
@@ -265,7 +271,11 @@ impl Default for WorkerIntrospector {
/// unchanged: the subsequent ZMQ connect will fail visibly with the
/// wildcard literal, which is the same observable failure mode that
/// would occur today if the bind/connect were skipped.
pub(crate) fn resolve_event_config(block: KvEventsBlock, worker_url: &str) -> EventConfig {
pub(crate) fn resolve_event_config(
block: KvEventsBlock,
worker_url: &str,
is_bigram: bool,
) -> EventConfig {
let host = if matches!(
block.endpoint_host.as_str(),
"*" | "0.0.0.0" | "::" | "[::]"
@@ -292,6 +302,7 @@ pub(crate) fn resolve_event_config(block: KvEventsBlock, worker_url: &str) -> Ev
topic: block.topic,
block_size: block.block_size,
dp_size: block.dp_size,
is_bigram,
}
}
@@ -304,6 +315,11 @@ struct ServerInfoBody {
served_model_name: Option<String>,
#[serde(default)]
kv_events: Option<KvEventsBlock>,
/// Top-level `speculative_algorithm`. EAGLE-family values
/// (EAGLE / EAGLE3 / FROZEN_KV_MTP) ⇒ the worker hashes KV blocks over
/// token bigrams. Absent on workers without speculative decoding.
#[serde(default)]
speculative_algorithm: Option<String>,
/// Carries the value of `ServerArgs.disaggregation_mode`
/// (`"null"` | `"prefill"` | `"decode"`). Absent on older SGLang
/// versions that predate the field.
@@ -369,6 +385,59 @@ mod tests {
WorkerIntrospector::new(Duration::from_millis(500))
}
/// The PRIMARY `/server_info` path (the introspector, not the discovery.rs
/// fallback) must flag `is_bigram` for an EAGLE worker so the policy picks
/// the bigram hasher. Regression guard for the duplicated parse + the
/// `resolve_event_config(.., is_bigram)` threading.
#[tokio::test]
async fn fetch_sets_is_bigram_for_eagle_worker() {
let (url, _shutdown) = spawn_fake_worker(json!({
"served_model_name": "m",
"speculative_algorithm": "EAGLE",
"kv_events": {
"publisher": "zmq",
"endpoint_host": "*",
"endpoint_port_base": 5557,
"topic": "",
"block_size": 64,
"dp_size": 1,
}
}))
.await;
let cfg = fast_introspector()
.fetch(&url)
.await
.event_config
.expect("kv_events present");
assert!(
cfg.is_bigram,
"EAGLE worker via the introspector must set is_bigram"
);
}
/// A non-speculative worker (no `speculative_algorithm`) must NOT be bigram.
#[tokio::test]
async fn fetch_no_bigram_without_speculative_algorithm() {
let (url, _shutdown) = spawn_fake_worker(json!({
"served_model_name": "m",
"kv_events": {
"publisher": "zmq",
"endpoint_host": "*",
"endpoint_port_base": 5557,
"topic": "",
"block_size": 64,
"dp_size": 1,
}
}))
.await;
let cfg = fast_introspector()
.fetch(&url)
.await
.event_config
.expect("kv_events present");
assert!(!cfg.is_bigram, "non-speculative worker must not be bigram");
}
#[tokio::test]
async fn fetch_returns_both_served_model_name_and_event_config() {
let (url, _shutdown) = spawn_fake_worker(json!({
+16 -19
View File
@@ -16,18 +16,17 @@ use tokio::task::JoinHandle;
/// Resolve the circuit-breaker config for all model IDs carried by a spec.
///
/// Workers may serve multiple models; we use the config of the **first** model
/// that has an explicit CB config, falling back to `None` (default config).
/// The router serves a single configured model; apply its circuit-breaker
/// config when this worker advertises that model id. Falls back to `None`
/// (default config) otherwise.
fn cb_config_for_spec(spec: &WorkerSpec, cfg: &Config) -> Option<CircuitBreakerConfig> {
for model_id in &spec.model_ids {
if let Some(mc) = cfg.models.iter().find(|m| m.id == model_id.0) {
if let Some(cbc) = &mc.circuit_breaker {
return Some(CircuitBreakerConfig {
threshold: cbc.threshold,
cool_down: Duration::from_secs(cbc.cool_down_secs),
});
}
}
let model = &cfg.model;
let cbc = model.circuit_breaker.as_ref()?;
if spec.model_ids.iter().any(|id| id.0 == model.id) {
return Some(CircuitBreakerConfig {
threshold: cbc.threshold,
cool_down: Duration::from_secs(cbc.cool_down_secs),
});
}
None
}
@@ -279,8 +278,8 @@ async fn register_one(
mod tests {
use super::*;
use crate::config::{
ActiveLoadConfig, CircuitBreakerConfig as RawCbConfig, DiscoveryBackend, DiscoveryConfig,
ModelConfig, PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, CircuitBreakerConfig as RawCbConfig, DiscoveryBackend, ModelConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use crate::discovery::{WorkerId, WorkerMode};
use axum::{routing::get, Json, Router};
@@ -296,7 +295,7 @@ mod tests {
port: 0,
},
observability: Default::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: id.into(),
tokenizer_path: "/tmp/x".into(),
policy: PolicyKind::RoundRobin,
@@ -305,12 +304,10 @@ mod tests {
cool_down_secs,
}),
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://test:30000".into()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://test:30000".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
@@ -90,8 +90,7 @@ async fn static_urls_pd_role_resolved_end_to_end() {
use axum::{routing::get, Json, Router};
use serde_json::json;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ObservabilityConfig,
ProxyConfig, ServerConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ObservabilityConfig, ProxyConfig, ServerConfig,
};
use sgl_router::discovery::{spawn_discovery, WorkerId};
use sgl_router::workers::{manager, WorkerRegistry};
@@ -127,12 +126,16 @@ async fn static_urls_pd_role_resolved_end_to_end() {
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec![url.clone()],
}),
model: sgl_router::config::ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: sgl_router::config::PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec![url.clone()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
@@ -65,20 +65,18 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
port: 0,
},
observability: Default::default(),
models: vec![sgl_router::config::ModelConfig {
model: sgl_router::config::ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: sgl_router::config::PolicyKind::CacheAwareZmq,
circuit_breaker: None,
cache_aware: None,
}],
discovery: sgl_router::config::DiscoveryConfig {
backend: sgl_router::config::DiscoveryBackend::StaticUrls(
sgl_router::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
},
discovery: sgl_router::config::DiscoveryBackend::StaticUrls(
sgl_router::config::StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
},
),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
@@ -127,6 +125,7 @@ async fn zmq_indexer_routes_to_publishing_worker_e2e() {
topic: String::new(),
block_size,
dp_size: 1,
is_bigram: false,
};
kv_index.add_worker(url_a, Some(preresolved.clone())).await;
kv_index.add_worker(url_b, Some(preresolved)).await;
@@ -41,6 +41,7 @@ async fn two_independent_subscribers_converge_to_same_tree_state() {
topic: String::new(),
block_size,
dp_size: 1,
is_bigram: false,
};
// 2. Two independent router-process surrogates, each with its own
@@ -173,6 +174,7 @@ async fn two_subscribers_merge_events_from_two_publishers() {
topic: String::new(),
block_size,
dp_size: 1,
is_bigram: false,
};
let cfg_y = EventConfig {
host: "127.0.0.1".into(),
@@ -180,6 +182,7 @@ async fn two_subscribers_merge_events_from_two_publishers() {
topic: String::new(),
block_size,
dp_size: 1,
is_bigram: false,
};
// Both routers subscribe to BOTH workers — the production fan-out.
+34 -51
View File
@@ -166,81 +166,64 @@ def _find_tokenizer_path(model: str) -> str:
return model
def build_smoke_router_config(
def build_smoke_router_args(
*,
host: str,
port: int,
model: str,
tokenizer_path: str,
sglang_url: str,
) -> str:
"""Build the TOML the smoke `router` fixture writes to disk.
) -> list[str]:
"""Build the sgl-router CLI flags the smoke ``router`` fixture launches.
Returns ``main_config_text`` carrying ``[server]``, ``[[models]]``,
and ``[discovery] backend = "static_urls"`` with the worker URL
inline. The Rust ``Config`` struct requires a ``[discovery]``
section (``DiscoveryConfig`` has no ``#[serde(default)]``) and has
no top-level ``workers`` field. The previous ``static_file``
backend was replaced by ``static_urls`` (which holds the URL list
inline rather than via a side-car file).
Static single-worker discovery (``--worker-urls``) pointed at the one
SGLang worker, serving exactly one model.
"""
return f"""\
[server]
host = "{host}"
port = {port}
[[models]]
id = "{model}"
tokenizer_path = "{tokenizer_path}"
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = ["{sglang_url}"]
"""
return [
"--host",
host,
"--port",
str(port),
"--model-id",
model,
"--tokenizer-path",
tokenizer_path,
"--worker-urls",
sglang_url,
]
@pytest.fixture(scope="session")
def router(sglang_server): # noqa: ARG001 (sglang_server must start first)
"""Launch sgl-router on port 8090 pointed at the SGLang worker."""
tok_path = _find_tokenizer_path(MODEL)
cfg_handle = tempfile.NamedTemporaryFile(mode="w", suffix=".toml", delete=False)
cfg_path = Path(cfg_handle.name)
main_text = build_smoke_router_config(
args = build_smoke_router_args(
host="0.0.0.0",
port=ROUTER_PORT,
model=MODEL,
tokenizer_path=tok_path,
sglang_url=f"http://localhost:{SGLANG_PORT}",
)
cfg_handle.write(main_text)
cfg_handle.close()
proc = subprocess.Popen(
[str(_BINARY), *args],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
# try/finally so the router is always reaped — on a readiness-probe
# failure, a test-body error, or a session-teardown exception alike.
try:
proc = subprocess.Popen(
[str(_BINARY), "--config", str(cfg_path)],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
try:
_wait_http(f"http://localhost:{ROUTER_PORT}/readyz", timeout=60)
except Exception:
proc.send_signal(signal.SIGTERM)
proc.wait(timeout=30)
raise
_wait_http(f"http://localhost:{ROUTER_PORT}/readyz", timeout=60)
yield f"http://localhost:{ROUTER_PORT}"
proc.send_signal(signal.SIGTERM)
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait()
finally:
cfg_path.unlink(missing_ok=True)
if proc.poll() is None:
proc.send_signal(signal.SIGTERM)
try:
proc.wait(timeout=30)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait()
# ---------------------------------------------------------------------------
@@ -1,23 +1,22 @@
"""Minimal sgl-router Gateway class — adapted from SMG's e2e_test/infra/gateway.py.
"""Minimal sgl-router Gateway class for e2e tests.
Differences from SMG:
- SMG drives a Python launcher (`python3 -m sglang_router.launch_router`)
with worker URLs on the CLI.
- sgl-router uses a Rust binary (`experimental/sgl-router/target/release/sgl-router`)
with a TOML config file. Worker discovery is config-file-based; this
Gateway writes a TOML to a tempfile and execs the binary with
`--config <tempfile>`.
sgl-router is a Rust binary
(`experimental/sgl-router/target/release/sgl-router`) configured entirely
through CLI flags. This Gateway execs the binary with `--worker-urls <...>`
(static discovery) plus the model + policy flags.
Supported lifecycles:
- Regular mode: one model, N worker URLs, single policy.
- PD mode: one model, prefill_workers + decode_workers (lists of URLs),
discovery emits separate `WorkerMode::Prefill` / `WorkerMode::Decode`
entries. The router resolves PD pool isolation at request time.
- PD mode: one model; prefill + decode URLs all go into one
`--worker-urls` static list. Each worker is seeded as
`WorkerMode::Plain` and its actual prefill/decode role + bootstrap
port are resolved from `/server_info` introspection, after which the
router isolates the PD pools at request time.
Use as a context manager:
with Gateway() as gw:
gw.start_regular(model_path="...", worker_urls=[...])
gw.start_regular(model_id="...", tokenizer_path="...", worker_urls=[...])
resp = httpx.post(f"{gw.base_url}/v1/chat/completions", json=...)
or pytest fixture style (see e2e_test/conftest.py).
@@ -30,7 +29,6 @@ import os
import signal
import socket
import subprocess
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
@@ -93,8 +91,11 @@ def _resolve_tokenizer_path(tokenizer_path: str) -> str:
cached = try_to_load_from_cache(tokenizer_path, "tokenizer.json")
if cached and Path(cached).is_file():
return str(cached)
except Exception: # noqa: BLE001
pass
except Exception as exc: # noqa: BLE001
# A cache miss is normal; log other failures (corrupt cache,
# signature change) so a later tokenizer-load error is traceable
# rather than mysterious.
logger.debug("HF tokenizer cache lookup failed for %r: %s", tokenizer_path, exc)
return tokenizer_path
@@ -150,7 +151,6 @@ class Gateway:
self.stale_request_timeout_secs = stale_request_timeout_secs
self.process: subprocess.Popen | None = None
self._config_path: Path | None = None
self._started: bool = False
# Track child workers we spawned so __exit__ can tear them down.
self._owned_workers: list[subprocess.Popen] = []
@@ -172,7 +172,6 @@ class Gateway:
tokenizer_path: str,
worker_urls: list[str],
policy: str = "round_robin",
extra_models: list[dict] | None = None,
timeout: float = 60.0,
) -> None:
"""Start the router in regular (non-PD) mode.
@@ -190,12 +189,11 @@ class Gateway:
timeout: How long to wait for ``/readyz`` before giving up.
"""
self._launch(
self._build_config(
self._build_args(
model_id=model_id,
tokenizer_path=tokenizer_path,
urls=list(worker_urls),
policy=policy,
extra_models=extra_models or [],
),
timeout=timeout,
)
@@ -222,12 +220,11 @@ class Gateway:
assumed.
"""
self._launch(
self._build_config(
self._build_args(
model_id=model_id,
tokenizer_path=tokenizer_path,
urls=list(prefill_urls) + list(decode_urls),
policy=policy,
extra_models=[],
),
timeout=timeout,
)
@@ -247,9 +244,6 @@ class Gateway:
except ProcessLookupError:
pass
self.process = None
if self._config_path and self._config_path.exists():
self._config_path.unlink(missing_ok=True)
self._config_path = None
self._started = False
# Tear down any owned upstream workers.
for w in self._owned_workers:
@@ -292,77 +286,53 @@ class Gateway:
# ----- internals ------------------------------------------------------
def _build_config(
def _build_args(
self,
*,
model_id: str,
tokenizer_path: str,
urls: list[str],
policy: str,
extra_models: list[dict],
) -> str:
) -> list[str]:
resolved_tokenizer = _resolve_tokenizer_path(tokenizer_path)
extra_model_toml = ""
for em in extra_models:
extra_model_toml += (
f'\n[[models]]\nid = "{em["id"]}"\n'
f'tokenizer_path = "{_resolve_tokenizer_path(em["tokenizer_path"])}"\n'
f'policy = "{em.get("policy", policy)}"\n'
)
# Optional tunables — only emit the [proxy] and [active_load]
# sections if a test has overridden them, so production defaults
# apply otherwise.
proxy_section = ""
args = [
"--host",
self.host,
"--port",
str(self.port),
"--model-id",
model_id,
"--tokenizer-path",
resolved_tokenizer,
"--policy",
policy,
]
# Optional tunables — only pass them if a test overrode them, so
# the router's production defaults apply otherwise.
if self.proxy_request_timeout_secs is not None:
proxy_section = (
f"\n[proxy]\nrequest_timeout_secs = {self.proxy_request_timeout_secs}\n"
)
active_load_section = ""
args += ["--request-timeout-secs", str(self.proxy_request_timeout_secs)]
if self.stale_request_timeout_secs is not None:
active_load_section = (
f"\n[active_load]\nstale_request_timeout_secs = "
f"{self.stale_request_timeout_secs}\n"
)
args += [
"--stale-request-timeout-secs",
str(self.stale_request_timeout_secs),
]
# `--worker-urls` is multi-valued; keep it last so clap doesn't
# absorb a following flag as a URL.
args += ["--worker-urls", *urls]
return args
urls_toml = ", ".join(f'"{u}"' for u in urls)
return f"""\
[server]
host = "{self.host}"
port = {self.port}
[[models]]
id = "{model_id}"
tokenizer_path = "{resolved_tokenizer}"
policy = "{policy}"
{extra_model_toml}
[discovery]
backend = "static_urls"
[discovery.static_urls]
urls = [{urls_toml}]
{proxy_section}{active_load_section}"""
def _launch(self, config_text: str, *, timeout: float) -> None:
def _launch(self, args: list[str], *, timeout: float) -> None:
if not self.binary.exists():
raise RuntimeError(
f"sgl-router binary not found at {self.binary}. "
"Build it first: `cd experimental/sgl-router && cargo build --release` "
"or set SGL_ROUTER_BINARY to the binary path."
)
# Write the main config.
fd, path = tempfile.mkstemp(suffix=".toml", prefix="sgl-router-")
os.close(fd)
self._config_path = Path(path)
self._config_path.write_text(config_text, encoding="utf-8")
logger.info("sgl-router config: %s", self._config_path)
logger.debug("sgl-router config text:\n%s", config_text)
logger.info("sgl-router args: %s", args)
self.process = subprocess.Popen(
[str(self.binary), "--config", str(self._config_path)],
[str(self.binary), *args],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
start_new_session=True,
@@ -380,16 +350,19 @@ urls = [{urls_toml}]
last_exc: Exception | None = None
while time.time() < deadline:
if self.process is not None and self.process.poll() is not None:
# Process exited early — surface stdout/stderr.
out = b""
# Process exited early — surface stdout/stderr. This is the
# primary startup-failure diagnostic, so if the read itself
# fails, report that instead of blanking the output.
try:
out = b""
if self.process.stdout is not None:
out = self.process.stdout.read() or b""
except Exception: # noqa: BLE001
pass
output = out.decode(errors="replace")
except Exception as read_exc: # noqa: BLE001
output = f"<failed to read router stdout: {read_exc}>"
raise RuntimeError(
f"sgl-router exited during startup with code "
f"{self.process.returncode}. output:\n{out.decode(errors='replace')}",
f"{self.process.returncode}. output:\n{output}",
)
try:
resp = httpx.get(f"{self.base_url}/readyz", timeout=2.0)
@@ -20,9 +20,23 @@ spec:
- name: router
image: sgl-router:e2e
imagePullPolicy: Never
# Configured entirely via CLI flags. No --service-discovery-namespace
# means a cluster-wide EndpointSlice watch (all namespaces); the
# `cross-ns-test=true` selector term scopes it to this test's workers.
args:
- "--config"
- "/etc/config/router-cluster.toml"
- "--host"
- "0.0.0.0"
- "--port"
- "8091"
- "--model-id"
- "tiny"
- "--tokenizer-path"
- "/etc/tokenizer/tiny.json"
- "--policy"
- "round_robin"
- "--service-discovery"
- "--selector"
- "app=sglang,cross-ns-test=true"
ports:
- containerPort: 8091
name: http
@@ -38,13 +52,6 @@ spec:
port: 8091
initialDelaySeconds: 5
periodSeconds: 10
volumeMounts:
- name: config
mountPath: /etc/config
volumes:
- name: config
configMap:
name: sgl-router-cluster-config
---
apiVersion: v1
kind: Service
@@ -18,9 +18,32 @@ spec:
- name: router
image: sgl-router:e2e
imagePullPolicy: Never
# Configured entirely via CLI flags. K8s EndpointSlice discovery
# watches `app=sglang` pods in this namespace. The aggressive
# circuit breaker (threshold 1, 5s cool-down) lets a terminating
# pod's connection-refused immediately drop it from the candidate
# set — the reconciliation tests scale workers rapidly and depend
# on fast eviction to absorb the churn.
args:
- "--config"
- "/etc/config/router.toml"
- "--host"
- "0.0.0.0"
- "--port"
- "8090"
- "--model-id"
- "tiny"
- "--tokenizer-path"
- "/etc/tokenizer/tiny.json"
- "--policy"
- "round_robin"
- "--cb-threshold"
- "1"
- "--cb-cool-down-secs"
- "5"
- "--service-discovery"
- "--service-discovery-namespace"
- "sgl-router-test"
- "--selector"
- "app=sglang"
ports:
- containerPort: 8090
name: http
@@ -36,13 +59,6 @@ spec:
port: 8090
initialDelaySeconds: 5
periodSeconds: 10
volumeMounts:
- name: config
mountPath: /etc/config
volumes:
- name: config
configMap:
name: sgl-router-config
---
apiVersion: v1
kind: Service
@@ -141,38 +141,9 @@ log "Waiting for fake-worker rollout..."
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" rollout status deployment/fake-worker --timeout=120s
# ---------------------------------------------------------------------------
# Step 6: Create sgl-router ConfigMap with k8s discovery pointing at the
# namespace where fake-worker pods live.
# ---------------------------------------------------------------------------
log "Creating sgl-router-config ConfigMap..."
ROUTER_CONFIG="[server]
host = \"0.0.0.0\"
port = 8090
[[models]]
id = \"tiny\"
tokenizer_path = \"/etc/tokenizer/tiny.json\"
policy = \"round_robin\"
# Aggressive breaker so a terminating pod's connection-refused
# immediately excludes it from the next request's candidate set —
# the reconciliation tests scale workers rapidly and depend on
# fast worker eviction to absorb the churn.
circuit_breaker = { threshold = 1, cool_down_secs = 5 }
[discovery]
backend = \"k8s\"
[discovery.k8s]
namespace = \"${NAMESPACE}\"
label_selector = \"app=sglang\""
kubectl --context "${CONTEXT}" -n "${NAMESPACE}" create configmap sgl-router-config \
--from-literal=router.toml="${ROUTER_CONFIG}" \
--dry-run=client -o yaml \
| kubectl --context "${CONTEXT}" apply -f -
# ---------------------------------------------------------------------------
# Step 7: Deploy sgl-router
# Step 6: Deploy sgl-router. It is configured entirely via CLI flags in
# router.yaml — k8s EndpointSlice discovery watches `app=sglang`
# pods in the sgl-router-test namespace (where fake-worker lives).
# ---------------------------------------------------------------------------
log "Deploying sgl-router..."
kubectl --context "${CONTEXT}" apply -f "${MANIFESTS_DIR}/router.yaml"
@@ -149,49 +149,9 @@ def cluster_scoped_router(k8s_cluster):
_ensure_namespace(EXTRA_NAMESPACE)
_ensure_service_in_ns(EXTRA_NAMESPACE)
# ConfigMap for the cluster-scoped router: empty namespace = watch all
cluster_config = """[server]
host = "0.0.0.0"
port = 8091
[[models]]
id = "tiny"
tokenizer_path = "/etc/tokenizer/tiny.json"
policy = "round_robin"
[discovery]
backend = "k8s"
[discovery.k8s]
namespace = ""
label_selector = "app=sglang,cross-ns-test=true"
"""
_kubectl(
"create",
"configmap",
"sgl-router-cluster-config",
f"--from-literal=router-cluster.toml={cluster_config}",
"-n",
NAMESPACE,
"--dry-run=client",
"-o",
"yaml",
check=True,
)
# pipe through apply
proc = _kubectl(
"create",
"configmap",
"sgl-router-cluster-config",
f"--from-literal=router-cluster.toml={cluster_config}",
"-n",
NAMESPACE,
"--dry-run=client",
"-o",
"yaml",
)
_apply_from_stdin(proc.stdout)
# The cluster-scoped router is configured via CLI flags in
# router-cluster-scoped.yaml: no --service-discovery-namespace (watch
# all namespaces) and --selector app=sglang,cross-ns-test=true.
_kubectl("apply", "-f", str(router_manifest))
# The cluster-scoped router's /readyz blocks on registry-not-empty, so
@@ -2,8 +2,8 @@
// SPDX-License-Identifier: Apache-2.0
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
@@ -29,18 +29,16 @@ fn config_for(_worker_url: &str) -> Config {
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
@@ -619,7 +617,7 @@ async fn no_healthy_workers_returns_503() {
);
}
/// A worker is registered for a model that is NOT in `cfg.models` (so the
/// A worker is registered for a model that is NOT the configured `cfg.model` (so the
/// policy registry has no entry for it). The handler returns 404
/// `model_not_found` rather than 500 — clients can recover by sending a
/// different model name; an internal_error would mask the misconfiguration.
@@ -31,7 +31,7 @@ async fn failover_when_one_worker_dies() {
port: 0,
},
observability: Default::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
@@ -40,12 +40,10 @@ async fn failover_when_one_worker_dies() {
cool_down_secs: 30,
}),
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec![w1.url.clone(), w2.url.clone(), w3.url.clone()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
@@ -16,8 +16,8 @@
use bytes::Bytes;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults;
@@ -40,18 +40,16 @@ fn build_ctx_with_worker(worker_url: &str) -> Arc<AppContext> {
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
@@ -4,8 +4,8 @@
use axum::body::Body;
use axum::http::Request;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
@@ -27,18 +27,16 @@ async fn forwards_whitelisted_headers_strips_others() {
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
};
@@ -21,8 +21,8 @@ use axum::http::{Request, StatusCode};
use bytes::Bytes;
use serde_json::{json, Value};
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults;
@@ -42,18 +42,16 @@ fn config() -> Config {
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
@@ -20,8 +20,8 @@ use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults;
@@ -41,18 +41,16 @@ fn config() -> Config {
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}
@@ -13,8 +13,8 @@ use axum::body::Body;
use axum::http::{Request, StatusCode};
use http_body_util::BodyExt;
use sgl_router::config::{
ActiveLoadConfig, Config, DiscoveryBackend, DiscoveryConfig, ModelConfig, ObservabilityConfig,
PolicyKind, ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
ActiveLoadConfig, Config, DiscoveryBackend, ModelConfig, ObservabilityConfig, PolicyKind,
ProxyConfig, ServerConfig, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults as build_policy_registry;
@@ -34,18 +34,16 @@ fn config(_worker_url: &str) -> Config {
port: 0,
},
observability: ObservabilityConfig::default(),
models: vec![ModelConfig {
model: ModelConfig {
id: "tiny".into(),
tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
policy: PolicyKind::RoundRobin,
circuit_breaker: None,
cache_aware: None,
}],
discovery: DiscoveryConfig {
backend: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
},
discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
urls: vec!["http://placeholder:0".into()],
}),
proxy: ProxyConfig::default(),
active_load: ActiveLoadConfig::default(),
}