From bcf89928b46dc9b24ed00f58bd548f126824d60d Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Fri, 5 Jun 2026 10:02:10 +0800 Subject: [PATCH] [router] Configure experimental sgl-router via CLI flags instead of a config file (#27073) Signed-off-by: Kangyan Zhou Co-authored-by: Claude Opus 4.8 --- experimental/sgl-router/Cargo.toml | 11 +- experimental/sgl-router/README.md | 45 +- experimental/sgl-router/src/config/cli.rs | 711 ++++++++++++++++++ experimental/sgl-router/src/config/mod.rs | 530 ++----------- experimental/sgl-router/src/config/types.rs | 480 +++++------- experimental/sgl-router/src/discovery/k8s.rs | 38 +- experimental/sgl-router/src/discovery/mod.rs | 2 +- experimental/sgl-router/src/main.rs | 25 +- .../src/policies/cache_aware_zmq.rs | 344 ++++++++- .../sgl-router/src/policies/factory.rs | 60 +- .../policies/kv_events/block_size_oracle.rs | 87 ++- .../src/policies/kv_events/discovery.rs | 87 +++ .../sgl-router/src/policies/kv_events/hash.rs | 156 +++- .../src/policies/kv_events/index.rs | 31 + .../sgl-router/src/policies/kv_events/mod.rs | 3 +- .../src/policies/kv_events/subscriber.rs | 1 + .../sgl-router/src/policies/kv_events/wire.rs | 164 +++- experimental/sgl-router/src/policies/mod.rs | 17 + .../sgl-router/src/server/app_context.rs | 23 +- experimental/sgl-router/src/server/error.rs | 7 + .../sgl-router/src/server/routes/chat.rs | 32 + .../sgl-router/src/server/routes/models.rs | 46 +- .../sgl-router/src/server/routes/tokenize.rs | 14 +- .../sgl-router/src/tokenizer/adapter.rs | 56 +- experimental/sgl-router/src/tokenizer/mod.rs | 23 +- .../sgl-router/src/workers/introspect.rs | 73 +- .../sgl-router/src/workers/manager.rs | 35 +- .../tests/component/discovery/static_urls.rs | 17 +- .../component/policies/cache_aware_zmq.rs | 15 +- .../policies/kv_events_two_subscribers.rs | 3 + experimental/sgl-router/tests/e2e/conftest.py | 85 +-- .../sgl-router/tests/e2e/infra/gateway.py | 135 ++-- .../manifests/router-cluster-scoped.yaml | 25 +- .../e2e/k8s_integration/manifests/router.yaml | 34 +- .../tests/e2e/k8s_integration/setup.sh | 35 +- .../k8s_integration/test_cross_namespace.py | 46 +- .../sgl-router/tests/proxy/chat_routing.rs | 16 +- .../sgl-router/tests/proxy/failover.rs | 10 +- .../tests/proxy/graceful_shutdown.rs | 14 +- .../tests/proxy/header_forwarding.rs | 14 +- .../tests/proxy/pd_bootstrap_injection.rs | 14 +- .../tests/proxy/pd_pool_isolation.rs | 14 +- .../sgl-router/tests/proxy/timeout.rs | 14 +- 43 files changed, 2355 insertions(+), 1237 deletions(-) create mode 100644 experimental/sgl-router/src/config/cli.rs diff --git a/experimental/sgl-router/Cargo.toml b/experimental/sgl-router/Cargo.toml index 74368ab7e..de419ab8b 100644 --- a/experimental/sgl-router/Cargo.toml +++ b/experimental/sgl-router/Cargo.toml @@ -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" diff --git a/experimental/sgl-router/README.md b/experimental/sgl-router/README.md index 0231d7952..fa1e05b01 100644 --- a/experimental/sgl-router/README.md +++ b/experimental/sgl-router/README.md @@ -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. diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs new file mode 100644 index 000000000..6957e383b --- /dev/null +++ b/experimental/sgl-router/src/config/cli.rs @@ -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, + /// 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, + /// 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, + + // ---- 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, + /// Absolute load spread above which the cache check is skipped. + #[arg(long)] + pub balance_abs_threshold: Option, + /// Multiplicative load spread gating the absolute balance check. + #[arg(long)] + pub balance_rel_threshold: Option, + + // ---- discovery: static ---- + /// Static worker URLs (space-separated or repeated). Mutually + /// exclusive with `--service-discovery`. + #[arg(long, num_args = 1..)] + pub worker_urls: Vec, + + // ---- 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, + /// 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, + /// PD-mode prefill label selector terms. Requires `--decode-selector`. + #[arg(long, num_args = 1..)] + pub prefill_selector: Vec, + /// PD-mode decode label selector terms. Requires `--prefill-selector`. + #[arg(long, num_args = 1..)] + pub decode_selector: Vec, + + // ---- 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 { + 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 { + 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 (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 { + 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 { + 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 { + MODEL_ARGS + .iter() + .chain(extra.iter()) + .map(|s| s.to_string()) + .collect() + } + + fn into_config_owned(args: Vec) -> Result { + 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); + } +} diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index 79f5ffb02..baef6ddc8 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -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 { - 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 { - 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}"); } } diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 0b55df66c..5cff09325 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -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, - 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, /// 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, } /// 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, - pub k8s: Option, -} - -/// 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(deserializer: D) -> Result - 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(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - let raw: DiscoveryConfigRaw = self.clone().into(); - raw.serialize(serializer) - } -} - -impl TryFrom for DiscoveryConfig { - type Error = String; - - fn try_from(raw: DiscoveryConfigRaw) -> Result { - 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 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, } /// 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, - #[serde(default)] - pub prefill_selector: Option, - #[serde(default)] - pub decode_selector: Option, + /// 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 { - 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 { + 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 { .. })); } diff --git a/experimental/sgl-router/src/discovery/k8s.rs b/experimental/sgl-router/src/discovery/k8s.rs index bf2798954..1a7d27e63 100644 --- a/experimental/sgl-router/src/discovery/k8s.rs +++ b/experimental/sgl-router/src/discovery/k8s.rs @@ -328,16 +328,18 @@ pub async fn spawn( cfg: K8sDiscoveryConfig, tx: mpsc::Sender, ) -> Result> { - 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 = if cfg.namespace.is_empty() { + let api: Api = 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() { + "" + } 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); diff --git a/experimental/sgl-router/src/discovery/mod.rs b/experimental/sgl-router/src/discovery/mod.rs index f9544f836..d5986a814 100644 --- a/experimental/sgl-router/src/discovery/mod.rs +++ b/experimental/sgl-router/src/discovery/mod.rs @@ -25,7 +25,7 @@ pub async fn spawn_discovery( cfg: &Config, ) -> Result<(mpsc::Receiver, 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?, }; diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index 00e1b4e7b..08ea5472e 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -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)?; diff --git a/experimental/sgl-router/src/policies/cache_aware_zmq.rs b/experimental/sgl-router/src/policies/cache_aware_zmq.rs index 804a415e1..9c6ecb7e2 100644 --- a/experimental/sgl-router/src/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/src/policies/cache_aware_zmq.rs @@ -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, + /// 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>, } 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) -> 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) { + 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` 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::().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(®istry), + 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(®istry), + 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() { diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index cd0a34faf..6bd2bcb00 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -76,17 +76,16 @@ pub fn build_registry( block_size_oracle: Arc, ) -> Result { 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 { 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(); diff --git a/experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs b/experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs index a2ec7604b..6077a2df3 100644 --- a/experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs +++ b/experimental/sgl-router/src/policies/kv_events/block_size_oracle.rs @@ -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(); diff --git a/experimental/sgl-router/src/policies/kv_events/discovery.rs b/experimental/sgl-router/src/policies/kv_events/discovery.rs index 0afb038d2..217af1dbb 100644 --- a/experimental/sgl-router/src/policies/kv_events/discovery.rs +++ b/experimental/sgl-router/src/policies/kv_events/discovery.rs @@ -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, + /// 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, +} + +/// 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] diff --git a/experimental/sgl-router/src/policies/kv_events/hash.rs b/experimental/sgl-router/src/policies/kv_events/hash.rs index 900a49171..0d37fec66 100644 --- a/experimental/sgl-router/src/policies/kv_events/hash.rs +++ b/experimental/sgl-router/src/policies/kv_events/hash.rs @@ -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 { + 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 diff --git a/experimental/sgl-router/src/policies/kv_events/index.rs b/experimental/sgl-router/src/policies/kv_events/index.rs index c29fa5b71..de979b6f3 100644 --- a/experimental/sgl-router/src/policies/kv_events/index.rs +++ b/experimental/sgl-router/src/policies/kv_events/index.rs @@ -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; + } } diff --git a/experimental/sgl-router/src/policies/kv_events/mod.rs b/experimental/sgl-router/src/policies/kv_events/mod.rs index 4f76dd89a..1923816f8 100644 --- a/experimental/sgl-router/src/policies/kv_events/mod.rs +++ b/experimental/sgl-router/src/policies/kv_events/mod.rs @@ -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}; diff --git a/experimental/sgl-router/src/policies/kv_events/subscriber.rs b/experimental/sgl-router/src/policies/kv_events/subscriber.rs index 1c1de438e..d564f9087 100644 --- a/experimental/sgl-router/src/policies/kv_events/subscriber.rs +++ b/experimental/sgl-router/src/policies/kv_events/subscriber.rs @@ -593,6 +593,7 @@ mod tests { topic: String::new(), block_size: 64, dp_size, + is_bigram: false, } } diff --git a/experimental/sgl-router/src/policies/kv_events/wire.rs b/experimental/sgl-router/src/policies/kv_events/wire.rs index a71606962..47714fa2a 100644 --- a/experimental/sgl-router/src/policies/kv_events/wire.rs +++ b/experimental/sgl-router/src/policies/kv_events/wire.rs @@ -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), +} + +impl<'de> Deserialize<'de> for TokenCell { + fn deserialize(deserializer: D) -> Result + 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(self, v: u64) -> Result { + Ok(TokenCell::One(v as u32)) + } + fn visit_i64(self, v: i64) -> Result { + Ok(TokenCell::One(v as u32)) + } + fn visit_seq(self, mut seq: A) -> Result + where + A: SeqAccess<'de>, + { + let mut ts: Vec = match seq.size_hint() { + Some(h) => Vec::with_capacity(h.min(8)), + None => Vec::new(), + }; + while let Some(t) = seq.next_element::()? { + 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); @@ -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::()? { - 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::()? { + let push = |t: u32, out: &mut Vec| -> 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, 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, + token_pairs: &[(u32, u32)], + block_size: u32, + lora_id: Option, + medium: Option<&str>, + ) -> Vec { + 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( diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index 4bad0fb1d..f30652c81 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -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], ctx: &SelectionContext<'_>) -> Option>; + + /// 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) {} } #[derive(Debug, Default)] @@ -61,4 +69,13 @@ impl PolicyRegistry { pub fn get(&self, model: &ModelId) -> Option> { 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) { + for entry in self.by_model.iter() { + entry.value().attach_metrics(Arc::clone(&metrics)); + } + } } diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 15bc289c4..2ff31bf1a 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -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(), }, diff --git a/experimental/sgl-router/src/server/error.rs b/experimental/sgl-router/src/server/error.rs index bce6d631f..de4676916 100644 --- a/experimental/sgl-router/src/server/error.rs +++ b/experimental/sgl-router/src/server/error.rs @@ -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)] diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index 868f84850..9fb5aa036 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -75,6 +75,7 @@ pub async fn chat_completions( headers: HeaderMap, body: Bytes, ) -> Result, 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 diff --git a/experimental/sgl-router/src/server/routes/models.rs b/experimental/sgl-router/src/server/routes/models.rs index eed5ab19f..5d66cb9ed 100644 --- a/experimental/sgl-router/src/server/routes/models.rs +++ b/experimental/sgl-router/src/server/routes/models.rs @@ -21,16 +21,14 @@ pub struct ModelEntry { } pub async fn list_models(State(ctx): State>) -> Json { - 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"); } } diff --git a/experimental/sgl-router/src/server/routes/tokenize.rs b/experimental/sgl-router/src/server/routes/tokenize.rs index 932f276d7..8be123bbd 100644 --- a/experimental/sgl-router/src/server/routes/tokenize.rs +++ b/experimental/sgl-router/src/server/routes/tokenize.rs @@ -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(), }; diff --git a/experimental/sgl-router/src/tokenizer/adapter.rs b/experimental/sgl-router/src/tokenizer/adapter.rs index 02fe72b63..efec68053 100644 --- a/experimental/sgl-router/src/tokenizer/adapter.rs +++ b/experimental/sgl-router/src/tokenizer/adapter.rs @@ -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> { +/// 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> { + 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 { + 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> { diff --git a/experimental/sgl-router/src/tokenizer/mod.rs b/experimental/sgl-router/src/tokenizer/mod.rs index 36d3a5cee..fcb7bdeac 100644 --- a/experimental/sgl-router/src/tokenizer/mod.rs +++ b/experimental/sgl-router/src/tokenizer/mod.rs @@ -24,10 +24,9 @@ impl std::fmt::Debug for TokenizerRegistry { impl TokenizerRegistry { pub fn load_from_config(cfg: &crate::config::Config) -> Result { 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")); } diff --git a/experimental/sgl-router/src/workers/introspect.rs b/experimental/sgl-router/src/workers/introspect.rs index fa829674b..f6330bf69 100644 --- a/experimental/sgl-router/src/workers/introspect.rs +++ b/experimental/sgl-router/src/workers/introspect.rs @@ -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, #[serde(default)] kv_events: Option, + /// 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, /// 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!({ diff --git a/experimental/sgl-router/src/workers/manager.rs b/experimental/sgl-router/src/workers/manager.rs index 4a583de66..1e8670057 100644 --- a/experimental/sgl-router/src/workers/manager.rs +++ b/experimental/sgl-router/src/workers/manager.rs @@ -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 { - 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(), } diff --git a/experimental/sgl-router/tests/component/discovery/static_urls.rs b/experimental/sgl-router/tests/component/discovery/static_urls.rs index 8097eac70..8c7ae5fc5 100644 --- a/experimental/sgl-router/tests/component/discovery/static_urls.rs +++ b/experimental/sgl-router/tests/component/discovery/static_urls.rs @@ -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(), }; diff --git a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs index ba5b0fd34..a5585f819 100644 --- a/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs +++ b/experimental/sgl-router/tests/component/policies/cache_aware_zmq.rs @@ -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; diff --git a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs index e603978fc..87f4afe1c 100644 --- a/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs +++ b/experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs @@ -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. diff --git a/experimental/sgl-router/tests/e2e/conftest.py b/experimental/sgl-router/tests/e2e/conftest.py index ae5fb2728..07bba336e 100644 --- a/experimental/sgl-router/tests/e2e/conftest.py +++ b/experimental/sgl-router/tests/e2e/conftest.py @@ -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() # --------------------------------------------------------------------------- diff --git a/experimental/sgl-router/tests/e2e/infra/gateway.py b/experimental/sgl-router/tests/e2e/infra/gateway.py index 389267cff..b580975b4 100644 --- a/experimental/sgl-router/tests/e2e/infra/gateway.py +++ b/experimental/sgl-router/tests/e2e/infra/gateway.py @@ -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 `. +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"" 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) diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml index 407173bed..d04cab9ac 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml +++ b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router-cluster-scoped.yaml @@ -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 diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml index 0d9b3900c..021d8bfcb 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml +++ b/experimental/sgl-router/tests/e2e/k8s_integration/manifests/router.yaml @@ -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 diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/setup.sh b/experimental/sgl-router/tests/e2e/k8s_integration/setup.sh index 882a90a60..3719e10e2 100755 --- a/experimental/sgl-router/tests/e2e/k8s_integration/setup.sh +++ b/experimental/sgl-router/tests/e2e/k8s_integration/setup.sh @@ -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" diff --git a/experimental/sgl-router/tests/e2e/k8s_integration/test_cross_namespace.py b/experimental/sgl-router/tests/e2e/k8s_integration/test_cross_namespace.py index 7dae30881..831c18f16 100644 --- a/experimental/sgl-router/tests/e2e/k8s_integration/test_cross_namespace.py +++ b/experimental/sgl-router/tests/e2e/k8s_integration/test_cross_namespace.py @@ -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 diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index c55fdc48d..7347ee18b 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -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. diff --git a/experimental/sgl-router/tests/proxy/failover.rs b/experimental/sgl-router/tests/proxy/failover.rs index 4e048497f..07920d806 100644 --- a/experimental/sgl-router/tests/proxy/failover.rs +++ b/experimental/sgl-router/tests/proxy/failover.rs @@ -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(), }; diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 6d87029c9..8c5145299 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs @@ -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 { 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(), }; diff --git a/experimental/sgl-router/tests/proxy/header_forwarding.rs b/experimental/sgl-router/tests/proxy/header_forwarding.rs index 410ec6e89..79b86bc02 100644 --- a/experimental/sgl-router/tests/proxy/header_forwarding.rs +++ b/experimental/sgl-router/tests/proxy/header_forwarding.rs @@ -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(), }; diff --git a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs index b48bf5fe4..059494ec4 100644 --- a/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs +++ b/experimental/sgl-router/tests/proxy/pd_bootstrap_injection.rs @@ -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(), } diff --git a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs index ba484ddd3..5d6fdec99 100644 --- a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs +++ b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs @@ -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(), } diff --git a/experimental/sgl-router/tests/proxy/timeout.rs b/experimental/sgl-router/tests/proxy/timeout.rs index 24c71f160..c5426d871 100644 --- a/experimental/sgl-router/tests/proxy/timeout.rs +++ b/experimental/sgl-router/tests/proxy/timeout.rs @@ -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(), }