From c610c403990255e6ffe41d7314f6fb748a53085f Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Sun, 20 Sep 2026 04:40:04 -0700 Subject: [PATCH] [sgl-router] refactor - config and organize CLI options (#39867) --- experimental/sgl-router/src/config/cli.rs | 1516 +++++++++-------- experimental/sgl-router/src/config/mod.rs | 274 ++- .../sgl-router/src/config/sampling.rs | 297 +--- experimental/sgl-router/src/config/types.rs | 410 +---- 4 files changed, 1021 insertions(+), 1476 deletions(-) diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 9692ae9f7..581072783 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -1,15 +1,13 @@ // 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`]. +//! Grouped CLI options and conversion into a validated [`Config`]. -use anyhow::{anyhow, Result}; +use anyhow::{anyhow, ensure, Result}; use clap::Parser; use std::num::NonZeroU32; -use crate::config::sampling::{parse_sampling_overrides, ConflictPolicy, SamplingOverrides}; +use crate::config::sampling::{parse_sampling_overrides, ConflictPolicy}; use crate::config::{ default_cb_cool_down, default_host, default_port, default_proxy_request_timeout_secs, default_shutdown_drain_secs, default_stale_request_timeout_secs, resolve_mode, @@ -23,88 +21,66 @@ use crate::config::{ const DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS: u64 = 100; const DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT: usize = sgl_kv_indexer::DEFAULT_QUERY_MAX_INFLIGHT; -/// `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" + about = "Slim KV-aware OpenAI-compatible router for SGLang workers", + after_help = "Examples:\n sgl-router --model-id Qwen/Qwen3-0.6B --worker-urls http://localhost:30001\n sgl-router --model-id Qwen/Qwen3-0.6B --service-discovery --selector app=sglang\n\nChoose exactly one discovery backend. Policy-specific options name their required policy in the descriptions." )] pub struct Cli { - // ---- server ---- - /// Address to bind the HTTP server to. - #[arg(long, default_value_t = default_host())] - pub host: String, - /// Port to bind the HTTP server to. - #[arg(long, default_value_t = default_port())] - pub port: u16, - /// Seconds to keep serving after SIGTERM, with `/readyz` returning 503, - /// before the server stops accepting — so the endpoint removal reaches - /// kube-proxy first. Leave room under terminationGracePeriodSeconds for the - /// in-flight drain that follows. If you rely on a readiness probe (rather - /// than pod deletion) to deregister, size this above your - /// failureThreshold * periodSeconds. The default equals the k8s default - /// terminationGracePeriodSeconds, so on a pod that has not raised its grace - /// period startup warns until you do — and declare it with - /// --termination-grace-secs so the check uses the real budget. - /// 0 disables the pause. - #[arg(long, default_value_t = default_shutdown_drain_secs())] - pub shutdown_drain_secs: u64, - /// The pod's terminationGracePeriodSeconds, if you have raised it from the - /// k8s default of 30. Only used to check --shutdown-drain-secs leaves room - /// for the in-flight drain at startup: the router cannot read its own pod - /// spec, so without this it warns against the default and a deliberately - /// long drain has no way to say it is safe. - #[arg(long)] - pub termination_grace_secs: Option, + #[command(flatten, next_help_heading = "Model, tokenizer and sampling")] + pub model: ModelArgs, + #[command(flatten, next_help_heading = "Server, timeouts and logging")] + pub server: ServerArgs, + #[command( + flatten, + next_help_heading = "Worker discovery (static URLs or Kubernetes)" + )] + pub discovery: DiscoveryArgs, + #[command( + flatten, + next_help_heading = "Routing policies, admission and circuit breaker" + )] + pub routing: RoutingArgs, + #[command( + flatten, + next_help_heading = "Cache-aware routing (--policy cache_aware)" + )] + pub cache: CacheArgs, + #[command( + flatten, + next_help_heading = "Session affinity, sticky routing and pressure guards" + )] + pub affinity: AffinityArgs, +} - // ---- model (exactly one) ---- +#[derive(clap::Args, Debug)] +pub struct ModelArgs { /// 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`). + + /// Local tokenizer.json or HuggingFace repo id. Defaults to --model-id; honors HF_TOKEN / HF_HOME. #[arg(long)] pub tokenizer_path: Option, - /// Disable router-generated input_ids for this model. Workers tokenize messages - /// themselves; cache-aware routing still renders locally. Use for worker-only - /// thinking/effort defaults, parser/template overrides, or template stop strings. + + /// Disable generated input_ids; workers tokenize messages, while routing still renders locally. + /// Use for worker-only thinking defaults, parser/template overrides, or template stop strings. #[arg(long)] pub disable_input_ids_forwarding: bool, - /// Routing policy. - #[arg(long, value_enum, default_value = "round_robin")] - pub policy: PolicyKind, - /// Policy used to select decode workers for PD requests. - #[arg(long, value_enum, default_value = "power_of_two")] - pub decode_policy: DecodePolicyKind, - /// Static P/D bucket configuration. Omit to use the global candidate domain. - #[arg(long)] - pub bucket_config: Option, - // ---- fleet-wide sampling contract (opt-in) ---- - /// Sampling parameters fixed fleet-wide, as one JSON object keyed by the - /// request-body field names — e.g. `{"temperature": 1, "top_p": 0.95}`. - /// Keys: temperature, top_p, top_k, min_p, repetition_penalty, - /// frequency_penalty, presence_penalty, n. Each value is a number, or an - /// inclusive band `{"min": LO, "max": HI}`. - /// - /// A configured value is injected whenever the request omits that field; - /// `--sampling-param-conflict` decides what a request that sends one gets. - /// Unknown or repeated keys, out-of-domain values and a band under `allow` - /// all fail the launch, naming the offending key. Full contract — domains, - /// `null` handling, cost — in the router README. + /// Fleet sampling defaults as JSON, e.g. {"temperature": 1, "top_p": 0.95}. + /// Accepts temperature, top_p, top_k, min_p, repetition_penalty, + /// frequency_penalty, presence_penalty, and n. Numeric values fill absent + /// fields; {"min": LO, "max": HI} bands only constrain supplied values + /// and require reject mode. See README for domains and null handling. #[arg(long, value_name = "JSON")] pub override_sampling_params: Option, - /// What a request that sends a value differing from - /// `--override-sampling-params` gets: `reject` (the default) 400s it - /// before admission, quoting the configured value; `allow` forwards the - /// client's value untouched. Only accepted alongside - /// `--override-sampling-params`. + + /// How to handle sampling values that differ from configured defaults: + /// reject returns 400 before admission; allow forwards the client value. + /// Defaults to reject. Requires --override-sampling-params. #[arg( long, value_enum, @@ -112,625 +88,295 @@ pub struct Cli { requires = "override_sampling_params" )] pub sampling_param_conflict: Option, +} - // ---- 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, +#[derive(clap::Args, Debug)] +pub struct ServerArgs { + /// Address to bind the HTTP server to. + #[arg(long, default_value_t = default_host())] + pub host: String, - /// External KV indexer gRPC endpoint used as the authoritative cache signal. - /// Needs an explicit scheme, e.g. `http://10.0.0.1:50051`. - #[arg(long)] - pub kv_indexer_endpoint: Option, - /// KV Indexer query timeout in milliseconds. Requires - /// `--kv-indexer-endpoint`; defaults to 100. - #[arg(long)] - pub kv_indexer_query_timeout_ms: Option, - /// Maximum concurrent KV Indexer queries issued by this Router. Requires - /// `--kv-indexer-endpoint`; defaults to 32. - #[arg(long)] - pub kv_indexer_query_max_inflight: Option, - /// Prefix-match source for native Cache-Aware. - #[arg(long, value_enum)] - pub cache_prefix_provider: Option, + /// Port to bind the HTTP server to. + #[arg(long, default_value_t = default_port())] + pub port: u16, - // ---- session-affinity tuning ---- - /// Header carrying the session ID for `--policy session_aware`. - #[arg(long)] - pub session_id_header: Option, - /// Idle timeout for a session assignment, in seconds. - #[arg(long)] - pub session_idle_secs: Option, - /// Session-assignment eviction cadence, in seconds. - #[arg(long)] - pub session_eviction_interval_secs: Option, - /// Use a deterministic backup for the affinity key and candidate range. - #[arg(long)] - pub stable_pair: bool, - /// Session-affinity admission mode. - #[arg(long, value_enum)] - pub affinity_mode: Option, - /// Session-affinity primary lookup and fallback behavior. - #[arg(long, value_enum)] - pub session_affinity_mode: Option, - /// Disables the Session/Cache-Aware pressure guard. - #[arg(long)] - pub disable_pressure_guard: bool, - /// Absolute waiting-uncached-token gap required by the pressure guard. - #[arg(long)] - pub pressure_abs_threshold_tokens: Option, - /// Absolute millisecond gap when a Prefill queue estimate is available. - #[arg(long)] - pub pressure_abs_threshold_ms: Option, - /// Relative waiting-uncached-token multiplier required by the pressure guard. - #[arg(long)] - pub pressure_rel_threshold: Option, - /// Minimum cache-hit tokens for a cache-aware candidate. - #[arg(long)] - pub cache_affinity_min_matched_tokens: Option, - /// Minimum cache-hit ratio for a cache-aware candidate. - #[arg(long)] - pub cache_affinity_min_match_ratio: Option, - /// Minimum number of cache-aware candidates to try. - #[arg(long)] - pub cache_candidate_min_workers: Option, - /// Fraction of healthy prefill workers considered as cache-aware candidates. - #[arg(long)] - pub cache_candidate_ratio: Option, - /// Maximum number of cache-aware candidates to try. - #[arg(long)] - pub cache_candidate_max_workers: Option, - /// Maximum uncached-work difference that pressure may override. - #[arg(long)] - pub cache_switch_margin_tokens: Option, - /// Queue gate for cache affinity: a worker whose engine reports at least - /// this many waiting (queued) requests cannot win a selection on cache - /// affinity. The request goes to another worker holding the same prefix, - /// or failing that to the least-loaded worker that is not queueing; when - /// every worker is queueing the least-loaded worker overall keeps the - /// fleet routable. Unset disables the gate. Requires - /// `--policy cache_aware`; scale with the engine's `--dp-size` because - /// the published queue sums across a worker's DP ranks. - #[arg(long)] - pub worker_queue_limit: Option, - /// Saturation floor for `--worker-queue-limit` diversions: when no - /// cache candidate survives both the queue limit and hard admission, - /// at least one was over the limit, AND no worker in the routable - /// fleet has a fresh queue reading strictly below this floor, the - /// diverted request would wait wherever it lands, so it stays with the - /// least-pressured prefix owner instead — same wait, but prefilled - /// from cache instead of a full cold prefill that evicts other - /// prefixes and manufactures the next round of misses. Unset disables - /// the pin. Requires `--worker-queue-limit` (there is no diversion to - /// cancel without it) and must be at most the limit; scale with - /// `--dp-size` like the limit. - #[arg(long)] - pub saturation_queue_floor: Option, - /// Number of random candidates sampled for the min-load fallback; the - /// least-pressured of the sample wins. The default 2 keeps today's - /// power-of-2 behavior unchanged. `k >= pool` skips the shuffle and - /// returns the exact minimum, with ties broken randomly — an idle - /// fleet ties on every comparison, so a fixed order would pin every - /// fallback dispatch to one worker; `k = 1` is a uniform draw within - /// the tier, and because a one-member sample has no runner-up the - /// proposal carries no backup, which disables the backup-admission - /// and pressure-guard paths. Note - /// the division of labor with `--cache-candidate-min-workers`, - /// `--cache-candidate-ratio`, and `--cache-candidate-max-workers`: - /// those bound the cache-affinity OWNER candidate set; this flag - /// bounds the min-load FALLBACK sample used when no owner is usable. - /// Requires `--policy cache_aware`. - #[arg(long)] - pub min_load_choices: Option, + /// Keep serving after SIGTERM with /readyz returning 503 before stopping accepts. + /// Leave time in the pod grace period for in-flight requests; cover readiness + /// probe failureThreshold * periodSeconds when probe-driven. 0 disables the pause. + #[arg(long, default_value_t = default_shutdown_drain_secs())] + pub shutdown_drain_secs: u64, - // ---- score composition ---- - /// Policies to sum, spelled exactly as `--policy` spells them and each - /// optionally weighted: `--fuse prefix_cache=2.0,load_based=0.3`. An - /// omitted weight keeps that policy's own default. Requires `--policy - /// score_policy` or `fused_score`; when either policy is set and this flag - /// is omitted, the terms default to `prefix_cache,load_based`. + /// Pod terminationGracePeriodSeconds for the startup drain-budget check. + /// Omitting this assumes 30 seconds; this flag does not change the pod spec. + #[arg(long)] + pub termination_grace_secs: Option, + + /// 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, + + /// 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, +} + +#[derive(clap::Args, Debug)] +pub struct DiscoveryArgs { + /// Static worker URLs, space-separated or repeated. Conflicts with --service-discovery. + #[arg(long, num_args = 1..)] + pub worker_urls: Vec, + + /// 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, AND-joined; Kubernetes selector grammar. + /// Mutually exclusive with --prefill-selector and --decode-selector. + #[arg(long, num_args = 1..)] + pub selector: Vec, + + /// Prefill equality selector terms (key=value or key==value). Requires --decode-selector. + #[arg(long, num_args = 1..)] + pub prefill_selector: Vec, + + /// Decode equality selector terms (key=value or key==value). Requires --prefill-selector. + #[arg(long, num_args = 1..)] + pub decode_selector: Vec, +} + +#[derive(clap::Args, Debug)] +pub struct RoutingArgs { + /// Routing policy. + #[arg(long, value_enum, default_value = "round_robin")] + pub policy: PolicyKind, + + /// Policy used to select decode workers for PD requests. + #[arg(long, value_enum, default_value = "power_of_two")] + pub decode_policy: DecodePolicyKind, + + /// Static P/D bucket configuration. Omit to use the global candidate domain. + #[arg(long)] + pub bucket_config: Option, + + /// Weighted scoring terms, e.g. prefix_cache=2.0,load_based=0.3. + /// Defaults to prefix_cache,load_based for score_policy or fused_score. + /// Requires --policy score_policy or fused_score. Omitted weights use each term's default. #[arg(long, value_delimiter = ',')] pub fuse: Vec, /// Ordered hard constraints applied before policy selection. #[arg(long, value_delimiter = ',')] pub filter: Vec, + /// Router-local in-flight limit for `--filter overloaded`. #[arg(long)] pub max_in_flight: Option, + /// Minimum cached prompt share for `--filter prefix_cache`. #[arg(long)] pub prefix_cache_min_share: Option, - // ---- sticky-session policy (only used by `--policy sticky`) ---- - /// Request header carrying the routing key for sticky-session routing. - /// Defaults to `x-sgl-routing-key` when `--policy sticky` is set. + /// Consecutive upstream failures before opening the breaker. Must be positive; enables the breaker. + #[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, +} + +#[derive(clap::Args, Debug)] +pub struct CacheArgs { + /// Prefix-match source: indexer when --kv-indexer-endpoint is set, otherwise radix_tree. + #[arg(long, value_enum)] + pub cache_prefix_provider: Option, + + /// External KV indexer gRPC endpoint used as the authoritative cache signal. + /// Needs an explicit scheme, e.g. `http://10.0.0.1:50051`. + #[arg(long)] + pub kv_indexer_endpoint: Option, + + /// KV Indexer query timeout in milliseconds. Requires + /// `--kv-indexer-endpoint`; defaults to 100. + #[arg(long)] + pub kv_indexer_query_timeout_ms: Option, + + /// Maximum concurrent KV Indexer queries issued by this Router. Requires + /// `--kv-indexer-endpoint`; defaults to 32. + #[arg(long)] + pub kv_indexer_query_max_inflight: Option, + + /// Minimum cache-hit tokens for a candidate. Defaults to 1024. + #[arg(long)] + pub cache_affinity_min_matched_tokens: Option, + + /// Minimum cache-hit ratio for a candidate. Unset by default. + #[arg(long)] + pub cache_affinity_min_match_ratio: Option, + + /// Minimum number of cache candidates to try. Defaults to 8. + #[arg(long)] + pub cache_candidate_min_workers: Option, + + /// Fraction of healthy prefill workers considered as cache candidates. Defaults to 0.05. + #[arg(long)] + pub cache_candidate_ratio: Option, + + /// Maximum number of cache candidates to try. Defaults to 32. + #[arg(long)] + pub cache_candidate_max_workers: Option, + + /// Maximum uncached-work difference that pressure may override. Defaults to 1024 tokens. + #[arg(long)] + pub cache_switch_margin_tokens: Option, + + /// Divert cache-affine requests when an engine queue reaches this limit. + /// Prefer another prefix owner, then the least-loaded worker. Missing fresh + /// queue data leaves affinity intact. Unset disables; scale with engine --dp-size. + #[arg(long)] + pub worker_queue_limit: Option, + + /// Keep the least-pressured prefix owner when the queue gate rejects all + /// admitted cache candidates and no worker has a fresh queue below this floor. + /// Requires --worker-queue-limit; must be positive and at most that limit. + /// Unset disables; scale with engine --dp-size. + #[arg(long)] + pub saturation_queue_floor: Option, + + /// Min-load fallback sample size. Defaults to 2; requires --policy cache_aware. + /// Values at least the pool size choose the exact minimum with random ties. + /// A value of 1 draws uniformly with no backup for admission or pressure guards. + /// Unlike --cache-candidate-* (prefix owners), this bounds the fallback sample. + #[arg(long)] + pub min_load_choices: Option, +} + +#[derive(clap::Args, Debug)] +pub struct AffinityArgs { + /// Header carrying the session ID for `--policy session_aware`. + #[arg(long)] + pub session_id_header: Option, + + /// Session idle timeout in seconds (--policy session_aware). Defaults to 600. + #[arg(long)] + pub session_idle_secs: Option, + + /// Session eviction interval in seconds (--policy session_aware). Defaults to 60. + #[arg(long)] + pub session_eviction_interval_secs: Option, + + /// Use a deterministic session backup (--policy session_aware). + #[arg(long)] + pub stable_pair: bool, + + /// Session admission mode (--policy session_aware). Defaults to soft (allow backup selection). + #[arg(long, value_enum)] + pub affinity_mode: Option, + + /// Session lookup mode (--policy session_aware). Defaults to bucket (search the target bucket). + #[arg(long, value_enum)] + pub session_affinity_mode: Option, + + /// Routing-key header (--policy sticky). Defaults to x-sgl-routing-key. #[arg(long)] pub routing_key_header: Option, - /// Policy used to select a worker for requests with no routing key, and - /// to pick the initial worker when a new key is first seen. One of - /// `round_robin` / `random` / `power_of_two` / `load_based`. Defaults - /// to `round_robin`. + + /// Policy for new or missing routing keys (--policy sticky). Defaults to round_robin. #[arg(long, value_enum)] pub sticky_fallback_policy: Option, - /// Evict a sticky assignment after it has been idle (unreferenced) this - /// many seconds. Defaults to 600. + + /// Idle timeout in seconds (--policy sticky). Defaults to 600. #[arg(long)] pub sticky_idle_secs: Option, - /// Wall-clock cadence of the sticky idle-eviction sweep, in seconds. + + /// Eviction sweep interval in seconds (--policy sticky). /// Defaults to 60. #[arg(long)] pub sticky_eviction_interval_secs: 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. + /// Disable the pressure guard (--policy session_aware or cache_aware). #[arg(long)] - pub service_discovery: bool, - /// Namespace to watch. Unset/empty watches all namespaces (requires - /// cluster-wide RBAC). + pub disable_pressure_guard: bool, + + /// Pressure-guard token gap (session_aware or cache_aware). Defaults to 1024. #[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, + pub pressure_abs_threshold_tokens: Option, - // ---- 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, + /// Pressure-guard gap in ms (session_aware or cache_aware), when a queue estimate exists. Unset by default. + #[arg(long)] + pub pressure_abs_threshold_ms: Option, - // ---- 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, + /// Pressure-guard token multiplier (session_aware or cache_aware). Defaults to 1.5. + #[arg(long)] + pub pressure_rel_threshold: Option, } 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). + /// Resolve CLI options and validate the resulting configuration. pub fn into_config(self) -> Result { - let discovery = self.build_discovery()?; + let affinity = self + .affinity + .build_config(&self.cache, self.routing.policy)?; + let discovery = self.discovery.into_config()?; let bucket_config = self + .routing .bucket_config .as_deref() .map(load_bucket_config) .transpose()?; - - // 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 cache_prefix_provider = self.cache_prefix_provider.unwrap_or_else(|| { - if self.kv_indexer_endpoint.is_some() { - CachePrefixProvider::Indexer - } else { - CachePrefixProvider::RadixTree - } - }); - if self.cache_prefix_provider.is_some() && self.policy != PolicyKind::CacheAware { - return Err(anyhow!( - "--cache-prefix-provider requires --policy cache_aware" - )); - } - if self.kv_indexer_query_timeout_ms == Some(0) { - return Err(anyhow!( - "--kv-indexer-query-timeout-ms must be greater than zero" - )); - } - if self.kv_indexer_query_timeout_ms.is_some() && self.kv_indexer_endpoint.is_none() { - return Err(anyhow!( - "--kv-indexer-query-timeout-ms requires --kv-indexer-endpoint" - )); - } - if self.kv_indexer_query_max_inflight == Some(0) { - return Err(anyhow!( - "--kv-indexer-query-max-inflight must be greater than zero" - )); - } - if self.kv_indexer_query_max_inflight.is_some() && self.kv_indexer_endpoint.is_none() { - return Err(anyhow!( - "--kv-indexer-query-max-inflight requires --kv-indexer-endpoint" - )); - } - let cache_aware_uses_indexer = self.policy == PolicyKind::CacheAware - && cache_prefix_provider == CachePrefixProvider::Indexer; - if self.kv_indexer_endpoint.is_some() && !cache_aware_uses_indexer { - if self.policy == PolicyKind::CacheAware { - return Err(anyhow!( - "--kv-indexer-endpoint requires --cache-prefix-provider indexer" - )); - } - return Err(anyhow!( - "--kv-indexer-endpoint requires --policy cache_aware" - )); - } - if cache_aware_uses_indexer && self.kv_indexer_endpoint.is_none() { - return Err(anyhow!( - "--cache-prefix-provider indexer requires --kv-indexer-endpoint" - )); - } - let tuned_cache_aware = self.policy == PolicyKind::CacheAware; - let affinity_policy = matches!( - self.policy, - PolicyKind::SessionAware | PolicyKind::CacheAware - ); - let tuned_session_affinity = self.session_id_header.is_some() - || self.session_idle_secs.is_some() - || self.session_eviction_interval_secs.is_some() - || self.stable_pair - || self.affinity_mode.is_some() - || self.session_affinity_mode.is_some(); - if tuned_session_affinity && self.policy != PolicyKind::SessionAware { - return Err(anyhow!( - "--session-id-header, --session-*-secs, --stable-pair, --affinity-mode, and \ - --session-affinity-mode require --policy session_aware" - )); - } - if self.disable_pressure_guard && !affinity_policy { - return Err(anyhow!( - "--disable-pressure-guard requires --policy session_aware or cache_aware" - )); - } - let tuned_cache_candidates = self.cache_affinity_min_matched_tokens.is_some() - || self.cache_affinity_min_match_ratio.is_some() - || self.cache_candidate_min_workers.is_some() - || self.cache_candidate_ratio.is_some() - || self.cache_candidate_max_workers.is_some() - || self.cache_switch_margin_tokens.is_some() - || self.worker_queue_limit.is_some() - || self.saturation_queue_floor.is_some() - || self.min_load_choices.is_some(); - // Value checks before the policy check: a value that is wrong under - // every policy should say so, rather than pointing at --policy. - if self.worker_queue_limit == Some(0) { - return Err(anyhow!("--worker-queue-limit must be at least 1")); - } - if let Some(floor) = self.saturation_queue_floor { - // The floor modifies the gate's diversion; without the gate - // there is no diversion to cancel and the knob would sit dead. - let Some(limit) = self.worker_queue_limit else { - return Err(anyhow!( - "--saturation-queue-floor requires --worker-queue-limit (there is no \ - diversion to cancel without it)" - )); - }; - if floor == 0 { - return Err(anyhow!("--saturation-queue-floor must be at least 1")); - } - // floor <= limit keeps the saturation label readable: a floor - // above the limit would declare the fleet saturated while - // workers the gate still admits exist. - if floor > limit { - return Err(anyhow!( - "--saturation-queue-floor ({floor}) must be at most --worker-queue-limit \ - ({limit})" - )); - } - } - if self.min_load_choices == Some(0) { - return Err(anyhow!("--min-load-choices must be at least 1")); - } - if tuned_cache_candidates && self.policy != PolicyKind::CacheAware { - return Err(anyhow!( - "cache candidate tuning flags require --policy cache_aware" - )); - } - if (self.pressure_abs_threshold_tokens.is_some() - || self.pressure_abs_threshold_ms.is_some() - || self.pressure_rel_threshold.is_some()) - && !affinity_policy - { - return Err(anyhow!( - "pressure guard tuning requires --policy session_aware or cache_aware" - )); - } - let is_score_composition = matches!( - self.policy, - PolicyKind::FusedScore | PolicyKind::ScorePolicy - ); - if !self.fuse.is_empty() && !is_score_composition { - return Err(anyhow!( - "--fuse requires --policy score_policy or fused_score" - )); - } - let fused = if is_score_composition { - let terms = if self.fuse.is_empty() { - DEFAULT_FUSE - .iter() - .map(|&kind| FusedTerm { kind, weight: None }) - .collect() - } else { - self.fuse.clone() - }; - for (i, t) in terms.iter().enumerate() { - if terms[..i].iter().any(|p| p.kind == t.kind) { - return Err(anyhow!("--fuse: `{}` is listed more than once", t.kind)); - } - } - Some(terms) - } else { - None - }; - - for (i, kind) in self.filter.iter().enumerate() { - if self.filter[..i].contains(kind) { - return Err(anyhow!("--filter: `{kind}` is listed more than once")); - } - } - let has = |k: FilterKind| self.filter.contains(&k); - if self.max_in_flight.is_some() != has(FilterKind::Overloaded) { - return Err(anyhow!( - "--max-in-flight and `--filter overloaded` require each other" - )); - } - if self.max_in_flight == Some(0) { - return Err(anyhow!("--max-in-flight must be greater than 0")); - } - if self.prefix_cache_min_share.is_some() != has(FilterKind::PrefixCache) { - return Err(anyhow!( - "--prefix-cache-min-share and `--filter prefix_cache` require each other" - )); - } - if self - .prefix_cache_min_share - .is_some_and(|s| !(s > 0.0 && s <= 1.0)) - { - return Err(anyhow!("--prefix-cache-min-share must be in (0, 1]")); - } - if self.policy == PolicyKind::Sticky && !self.filter.is_empty() { - return Err(anyhow!("--filter cannot be combined with --policy sticky")); - } - let eligibility = (!self.filter.is_empty()).then(|| EligibilityConfig { - filters: self.filter.clone(), - max_in_flight: self.max_in_flight, - min_prefix_share: self.prefix_cache_min_share, - }); - - let tuned_sticky = self.routing_key_header.is_some() - || self.sticky_fallback_policy.is_some() - || self.sticky_idle_secs.is_some() - || self.sticky_eviction_interval_secs.is_some(); - if tuned_sticky && self.policy != PolicyKind::Sticky { - return Err(anyhow!( - "--routing-key-header / --sticky-fallback-policy / --sticky-idle-secs / \ - --sticky-eviction-interval-secs require --policy sticky" - )); - } - - // Build and validate the sticky config exactly when the sticky - // policy is selected. The header name must parse as an HTTP header - // name so a typo fails at startup rather than silently never - // matching any request header. - let sticky = if self.policy == PolicyKind::Sticky { - let d = StickyConfig::default(); - let header_name = self.routing_key_header.unwrap_or(d.header_name); - axum::http::HeaderName::try_from(header_name.as_str()).map_err(|e| { - anyhow!("--routing-key-header {header_name:?} is not a valid HTTP header name: {e}") - })?; - let fallback_policy = self.sticky_fallback_policy.unwrap_or(d.fallback_policy); - let idle_secs = self.sticky_idle_secs.unwrap_or(d.idle_secs); - let eviction_interval_secs = self - .sticky_eviction_interval_secs - .unwrap_or(d.eviction_interval_secs); - // Reject zero durations: `--sticky-eviction-interval-secs 0` would - // panic `tokio::time::interval` at startup, and `--sticky-idle-secs - // 0` would evict every assignment on the next sweep (defeating - // stickiness entirely). Fail fast with a clear message instead. - if eviction_interval_secs == 0 { - return Err(anyhow!( - "--sticky-eviction-interval-secs must be greater than 0" - )); - } - if idle_secs == 0 { - return Err(anyhow!( - "--sticky-idle-secs must be greater than 0 (0 would evict every \ - assignment immediately, defeating sticky routing)" - )); - } - Some(StickyConfig { - header_name, - fallback_policy, - idle_secs, - eviction_interval_secs, - }) - } else { - None - }; - - let affinity = if affinity_policy { - let d = AffinityConfig::default(); - let session_id_header = self.session_id_header.unwrap_or(d.session_id_header); - axum::http::HeaderName::try_from(session_id_header.as_str()).map_err(|e| { - anyhow!( - "--session-id-header {session_id_header:?} is not a valid HTTP header name: {e}" + let circuit_breaker = self.routing.build_circuit_breaker()?; + let cache_aware = self.cache.into_config(self.routing.policy)?; + let fused = self.routing.build_fused()?; + let eligibility = self.routing.build_eligibility()?; + let sticky = self.affinity.into_sticky_config(self.routing.policy)?; + let sampling_overrides = self + .model + .override_sampling_params + .as_deref() + .map(|raw| { + parse_sampling_overrides( + raw, + self.model.sampling_param_conflict.unwrap_or_default(), ) - })?; - let pressure_rel_threshold = self - .pressure_rel_threshold - .unwrap_or(d.pressure_rel_threshold); - if !pressure_rel_threshold.is_finite() || pressure_rel_threshold <= 1.0 { - return Err(anyhow!( - "--pressure-rel-threshold must be finite and greater than 1" - )); - } - if self - .pressure_abs_threshold_ms - .is_some_and(|threshold| !threshold.is_finite() || threshold < 0.0) - { - return Err(anyhow!( - "--pressure-abs-threshold-ms must be finite and non-negative" - )); - } - let cache_affinity_min_match_ratio = self - .cache_affinity_min_match_ratio - .or(d.cache_affinity_min_match_ratio); - if cache_affinity_min_match_ratio - .is_some_and(|ratio| !ratio.is_finite() || !(0.0..=1.0).contains(&ratio)) - { - return Err(anyhow!( - "--cache-affinity-min-match-ratio must be finite and in [0, 1]" - )); - } - let cache_candidate_ratio = self - .cache_candidate_ratio - .unwrap_or(d.cache_candidate_ratio); - if !cache_candidate_ratio.is_finite() || !(0.0..=1.0).contains(&cache_candidate_ratio) { - return Err(anyhow!( - "--cache-candidate-ratio must be finite and in [0, 1]" - )); - } - let cache_candidate_min_workers = self - .cache_candidate_min_workers - .unwrap_or(d.cache_candidate_min_workers); - let cache_candidate_max_workers = self - .cache_candidate_max_workers - .unwrap_or(d.cache_candidate_max_workers); - if cache_candidate_min_workers == 0 - || cache_candidate_max_workers == 0 - || cache_candidate_min_workers > cache_candidate_max_workers - { - return Err(anyhow!( - "--cache-candidate-min-workers and --cache-candidate-max-workers must be \ - positive and min must not exceed max" - )); - } - let session_idle_secs = self.session_idle_secs.unwrap_or(d.session_idle_secs); - let session_eviction_interval_secs = self - .session_eviction_interval_secs - .unwrap_or(d.session_eviction_interval_secs); - if session_idle_secs == 0 { - return Err(anyhow!("--session-idle-secs must be greater than 0")); - } - if session_eviction_interval_secs == 0 { - return Err(anyhow!( - "--session-eviction-interval-secs must be greater than 0" - )); - } - Some(AffinityConfig { - session_id_header, - session_idle_secs, - session_eviction_interval_secs, - stable_pair: self.stable_pair, - mode: self.affinity_mode.unwrap_or(d.mode), - session_affinity_mode: self - .session_affinity_mode - .unwrap_or(d.session_affinity_mode), - pressure_guard: !self.disable_pressure_guard && d.pressure_guard, - pressure_abs_threshold_tokens: self - .pressure_abs_threshold_tokens - .unwrap_or(d.pressure_abs_threshold_tokens), - pressure_abs_threshold_ms: self - .pressure_abs_threshold_ms - .or(d.pressure_abs_threshold_ms), - pressure_rel_threshold, - cache_affinity_min_matched_tokens: self - .cache_affinity_min_matched_tokens - .or(d.cache_affinity_min_matched_tokens), - cache_affinity_min_match_ratio, - cache_candidate_min_workers, - cache_candidate_ratio, - cache_candidate_max_workers, - cache_switch_margin_tokens: self - .cache_switch_margin_tokens - .unwrap_or(d.cache_switch_margin_tokens), - worker_queue_limit: self.worker_queue_limit.or(d.worker_queue_limit), - saturation_queue_floor: self.saturation_queue_floor.or(d.saturation_queue_floor), - min_load_choices: self.min_load_choices.unwrap_or(d.min_load_choices), }) - } else { - None - }; - - 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), - }); - - // Keep the selected prefix provider and optional Indexer settings - // together with the native Cache-Aware policy. - let kv_indexer_query_timeout_ms = self - .kv_indexer_query_timeout_ms - .unwrap_or(DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS); - let kv_indexer_query_max_inflight = self - .kv_indexer_query_max_inflight - .unwrap_or(DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT); - let cache_aware = if tuned_cache_aware { - let kv_indexer_endpoint = self.kv_indexer_endpoint.map(|url| KvIndexerEndpointConfig { - url, - query_timeout_ms: kv_indexer_query_timeout_ms, - query_max_inflight: kv_indexer_query_max_inflight, - }); - Some(CacheAwareConfig { - prefix_provider: cache_prefix_provider, - kv_indexer_endpoint, - }) - } else { - None - }; - - let sampling_overrides = match &self.override_sampling_params { - None => SamplingOverrides::default(), - Some(raw) => { - parse_sampling_overrides(raw, self.sampling_param_conflict.unwrap_or_default())? - } - }; + .transpose()? + .unwrap_or_default(); let config = Config { server: ServerConfig { - host: self.host, - port: self.port, - shutdown_drain_secs: self.shutdown_drain_secs, - termination_grace_secs: self.termination_grace_secs, + host: self.server.host, + port: self.server.port, + shutdown_drain_secs: self.server.shutdown_drain_secs, + termination_grace_secs: self.server.termination_grace_secs, }, observability: ObservabilityConfig { - log_level: self.log_level, - log_format: self.log_format, + log_level: self.server.log_level, + log_format: self.server.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, - disable_input_ids_forwarding: self.disable_input_ids_forwarding, - policy: self.policy, - decode_policy: self.decode_policy, + tokenizer_path: self + .model + .tokenizer_path + .unwrap_or_else(|| self.model.model_id.clone()), + id: self.model.model_id, + disable_input_ids_forwarding: self.model.disable_input_ids_forwarding, + policy: self.routing.policy, + decode_policy: self.routing.decode_policy, bucket_config, circuit_breaker, cache_aware, @@ -742,26 +388,19 @@ impl Cli { }, discovery, proxy: ProxyConfig { - request_timeout_secs: self.request_timeout_secs, + request_timeout_secs: self.server.request_timeout_secs, }, active_load: ActiveLoadConfig { - stale_request_timeout_secs: self.stale_request_timeout_secs, + stale_request_timeout_secs: self.server.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 { +impl DiscoveryArgs { + fn into_config(self) -> Result { let has_static = !self.worker_urls.is_empty(); let backend = match (has_static, self.service_discovery) { (true, true) => { @@ -776,32 +415,26 @@ impl Cli { )); } (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 / \ + ensure!( + self.service_discovery_namespace.is_none() + && self.selector.is_empty() + && self.prefill_selector.is_empty() + && self.decode_selector.is_empty(), + "--service-discovery-namespace / --selector / --prefill-selector / \ --decode-selector require --service-discovery" - )); - } + ); DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig { - urls: self.worker_urls.clone(), + urls: self.worker_urls, }) } (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(), + namespace: self.service_discovery_namespace.unwrap_or_default(), mode, }) } @@ -810,6 +443,365 @@ impl Cli { } } +impl RoutingArgs { + fn build_circuit_breaker(&self) -> Result> { + ensure!( + self.cb_cool_down_secs.is_none() || self.cb_threshold.is_some(), + "--cb-cool-down-secs requires --cb-threshold (the circuit breaker is \ + enabled by --cb-threshold)" + ); + 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), + }); + + Ok(circuit_breaker) + } + + fn build_fused(&self) -> Result>> { + let is_score_composition = matches!( + self.policy, + PolicyKind::FusedScore | PolicyKind::ScorePolicy + ); + ensure!( + self.fuse.is_empty() || is_score_composition, + "--fuse requires --policy score_policy or fused_score" + ); + if !is_score_composition { + return Ok(None); + } + let terms = if self.fuse.is_empty() { + DEFAULT_FUSE + .iter() + .map(|&kind| FusedTerm { kind, weight: None }) + .collect() + } else { + self.fuse.clone() + }; + for (i, t) in terms.iter().enumerate() { + ensure!( + !terms[..i].iter().any(|p| p.kind == t.kind), + "--fuse: `{}` is listed more than once", + t.kind + ); + } + Ok(Some(terms)) + } + + fn build_eligibility(&self) -> Result> { + for (i, kind) in self.filter.iter().enumerate() { + ensure!( + !self.filter[..i].contains(kind), + "--filter: `{kind}` is listed more than once" + ); + } + let has = |k: FilterKind| self.filter.contains(&k); + ensure!( + (self.max_in_flight.is_some() == has(FilterKind::Overloaded)), + "--max-in-flight and `--filter overloaded` require each other" + ); + ensure!( + self.max_in_flight != Some(0), + "--max-in-flight must be greater than 0" + ); + ensure!( + (self.prefix_cache_min_share.is_some() == has(FilterKind::PrefixCache)), + "--prefix-cache-min-share and `--filter prefix_cache` require each other" + ); + ensure!( + self.prefix_cache_min_share + .is_none_or(|s| s > 0.0 && s <= 1.0), + "--prefix-cache-min-share must be in (0, 1]" + ); + ensure!( + self.policy != PolicyKind::Sticky || self.filter.is_empty(), + "--filter cannot be combined with --policy sticky" + ); + let eligibility = (!self.filter.is_empty()).then_some(EligibilityConfig { + filters: self.filter.clone(), + max_in_flight: self.max_in_flight, + min_prefix_share: self.prefix_cache_min_share, + }); + + Ok(eligibility) + } +} + +impl CacheArgs { + fn into_config(self, policy: PolicyKind) -> Result> { + let cache_prefix_provider = self.cache_prefix_provider.unwrap_or_else(|| { + if self.kv_indexer_endpoint.is_some() { + CachePrefixProvider::Indexer + } else { + CachePrefixProvider::RadixTree + } + }); + ensure!( + self.cache_prefix_provider.is_none() || policy == PolicyKind::CacheAware, + "--cache-prefix-provider requires --policy cache_aware" + ); + ensure!( + self.kv_indexer_query_timeout_ms != Some(0), + "--kv-indexer-query-timeout-ms must be greater than zero" + ); + ensure!( + self.kv_indexer_query_timeout_ms.is_none() || self.kv_indexer_endpoint.is_some(), + "--kv-indexer-query-timeout-ms requires --kv-indexer-endpoint" + ); + ensure!( + self.kv_indexer_query_max_inflight != Some(0), + "--kv-indexer-query-max-inflight must be greater than zero" + ); + ensure!( + self.kv_indexer_query_max_inflight.is_none() || self.kv_indexer_endpoint.is_some(), + "--kv-indexer-query-max-inflight requires --kv-indexer-endpoint" + ); + let cache_aware_uses_indexer = policy == PolicyKind::CacheAware + && cache_prefix_provider == CachePrefixProvider::Indexer; + if self.kv_indexer_endpoint.is_some() && !cache_aware_uses_indexer { + return Err(if policy == PolicyKind::CacheAware { + anyhow!("--kv-indexer-endpoint requires --cache-prefix-provider indexer") + } else { + anyhow!("--kv-indexer-endpoint requires --policy cache_aware") + }); + } + ensure!( + !cache_aware_uses_indexer || self.kv_indexer_endpoint.is_some(), + "--cache-prefix-provider indexer requires --kv-indexer-endpoint" + ); + let kv_indexer_query_timeout_ms = self + .kv_indexer_query_timeout_ms + .unwrap_or(DEFAULT_KV_INDEXER_QUERY_TIMEOUT_MS); + let kv_indexer_query_max_inflight = self + .kv_indexer_query_max_inflight + .unwrap_or(DEFAULT_KV_INDEXER_QUERY_MAX_INFLIGHT); + if policy != PolicyKind::CacheAware { + return Ok(None); + } + let kv_indexer_endpoint = self.kv_indexer_endpoint.map(|url| KvIndexerEndpointConfig { + url, + query_timeout_ms: kv_indexer_query_timeout_ms, + query_max_inflight: kv_indexer_query_max_inflight, + }); + Ok(Some(CacheAwareConfig { + prefix_provider: cache_prefix_provider, + kv_indexer_endpoint, + })) + } +} + +impl AffinityArgs { + fn build_config( + &self, + cache: &CacheArgs, + policy: PolicyKind, + ) -> Result> { + let affinity_policy = matches!(policy, PolicyKind::SessionAware | PolicyKind::CacheAware); + let tuned_session_affinity = self.session_id_header.is_some() + || self.session_idle_secs.is_some() + || self.session_eviction_interval_secs.is_some() + || self.stable_pair + || self.affinity_mode.is_some() + || self.session_affinity_mode.is_some(); + ensure!( + !tuned_session_affinity || policy == PolicyKind::SessionAware, + "--session-id-header, --session-*-secs, --stable-pair, --affinity-mode, and \ + --session-affinity-mode require --policy session_aware" + ); + ensure!( + !self.disable_pressure_guard || affinity_policy, + "--disable-pressure-guard requires --policy session_aware or cache_aware" + ); + let tuned_cache_candidates = cache.cache_affinity_min_matched_tokens.is_some() + || cache.cache_affinity_min_match_ratio.is_some() + || cache.cache_candidate_min_workers.is_some() + || cache.cache_candidate_ratio.is_some() + || cache.cache_candidate_max_workers.is_some() + || cache.cache_switch_margin_tokens.is_some() + || cache.worker_queue_limit.is_some() + || cache.saturation_queue_floor.is_some() + || cache.min_load_choices.is_some(); + // Value checks before the policy check: a value that is wrong under + // every policy should say so, rather than pointing at --policy. + ensure!( + cache.worker_queue_limit != Some(0), + "--worker-queue-limit must be at least 1" + ); + if let Some(floor) = cache.saturation_queue_floor { + // The floor modifies the gate's diversion; without the gate + // there is no diversion to cancel and the knob would sit dead. + let Some(limit) = cache.worker_queue_limit else { + return Err(anyhow!( + "--saturation-queue-floor requires --worker-queue-limit (there is no \ + diversion to cancel without it)" + )); + }; + ensure!(floor != 0, "--saturation-queue-floor must be at least 1"); + ensure!( + floor <= limit, + "--saturation-queue-floor ({floor}) must be at most --worker-queue-limit \ + ({limit})" + ); + } + ensure!( + cache.min_load_choices != Some(0), + "--min-load-choices must be at least 1" + ); + ensure!( + !tuned_cache_candidates || policy == PolicyKind::CacheAware, + "cache candidate tuning flags require --policy cache_aware" + ); + ensure!( + affinity_policy + || (self.pressure_abs_threshold_tokens.is_none() + && self.pressure_abs_threshold_ms.is_none() + && self.pressure_rel_threshold.is_none()), + "pressure guard tuning requires --policy session_aware or cache_aware" + ); + if !affinity_policy { + return Ok(None); + } + let defaults = AffinityConfig::default(); + let session_id_header = self + .session_id_header + .clone() + .unwrap_or(defaults.session_id_header); + axum::http::HeaderName::try_from(session_id_header.as_str()).map_err(|e| { + anyhow!( + "--session-id-header {session_id_header:?} is not a valid HTTP header name: {e}" + ) + })?; + let pressure_rel_threshold = self + .pressure_rel_threshold + .unwrap_or(defaults.pressure_rel_threshold); + ensure!( + pressure_rel_threshold.is_finite() && pressure_rel_threshold > 1.0, + "--pressure-rel-threshold must be finite and greater than 1" + ); + ensure!( + self.pressure_abs_threshold_ms + .is_none_or(|threshold| threshold.is_finite() && threshold >= 0.0), + "--pressure-abs-threshold-ms must be finite and non-negative" + ); + let cache_affinity_min_match_ratio = cache + .cache_affinity_min_match_ratio + .or(defaults.cache_affinity_min_match_ratio); + ensure!( + cache_affinity_min_match_ratio + .is_none_or(|ratio| ratio.is_finite() && (0.0..=1.0).contains(&ratio)), + "--cache-affinity-min-match-ratio must be finite and in [0, 1]" + ); + let cache_candidate_ratio = cache + .cache_candidate_ratio + .unwrap_or(defaults.cache_candidate_ratio); + ensure!( + cache_candidate_ratio.is_finite() && (0.0..=1.0).contains(&cache_candidate_ratio), + "--cache-candidate-ratio must be finite and in [0, 1]" + ); + let cache_candidate_min_workers = cache + .cache_candidate_min_workers + .unwrap_or(defaults.cache_candidate_min_workers); + let cache_candidate_max_workers = cache + .cache_candidate_max_workers + .unwrap_or(defaults.cache_candidate_max_workers); + ensure!( + cache_candidate_min_workers > 0 + && cache_candidate_max_workers > 0 + && cache_candidate_min_workers <= cache_candidate_max_workers, + "--cache-candidate-min-workers and --cache-candidate-max-workers must be \ + positive and min must not exceed max" + ); + let session_idle_secs = self.session_idle_secs.unwrap_or(defaults.session_idle_secs); + let session_eviction_interval_secs = self + .session_eviction_interval_secs + .unwrap_or(defaults.session_eviction_interval_secs); + ensure!( + session_idle_secs != 0, + "--session-idle-secs must be greater than 0" + ); + ensure!( + session_eviction_interval_secs != 0, + "--session-eviction-interval-secs must be greater than 0" + ); + Ok(Some(AffinityConfig { + session_id_header, + session_idle_secs, + session_eviction_interval_secs, + stable_pair: self.stable_pair, + mode: self.affinity_mode.unwrap_or(defaults.mode), + session_affinity_mode: self + .session_affinity_mode + .unwrap_or(defaults.session_affinity_mode), + pressure_guard: !self.disable_pressure_guard && defaults.pressure_guard, + pressure_abs_threshold_tokens: self + .pressure_abs_threshold_tokens + .unwrap_or(defaults.pressure_abs_threshold_tokens), + pressure_abs_threshold_ms: self + .pressure_abs_threshold_ms + .or(defaults.pressure_abs_threshold_ms), + pressure_rel_threshold, + cache_affinity_min_matched_tokens: cache + .cache_affinity_min_matched_tokens + .or(defaults.cache_affinity_min_matched_tokens), + cache_affinity_min_match_ratio, + cache_candidate_min_workers, + cache_candidate_ratio, + cache_candidate_max_workers, + cache_switch_margin_tokens: cache + .cache_switch_margin_tokens + .unwrap_or(defaults.cache_switch_margin_tokens), + worker_queue_limit: cache.worker_queue_limit.or(defaults.worker_queue_limit), + saturation_queue_floor: cache + .saturation_queue_floor + .or(defaults.saturation_queue_floor), + min_load_choices: cache.min_load_choices.unwrap_or(defaults.min_load_choices), + })) + } + + fn into_sticky_config(self, policy: PolicyKind) -> Result> { + let tuned_sticky = self.routing_key_header.is_some() + || self.sticky_fallback_policy.is_some() + || self.sticky_idle_secs.is_some() + || self.sticky_eviction_interval_secs.is_some(); + ensure!( + !tuned_sticky || policy == PolicyKind::Sticky, + "--routing-key-header / --sticky-fallback-policy / --sticky-idle-secs / \ + --sticky-eviction-interval-secs require --policy sticky" + ); + + if policy != PolicyKind::Sticky { + return Ok(None); + } + let defaults = StickyConfig::default(); + let header_name = self.routing_key_header.unwrap_or(defaults.header_name); + axum::http::HeaderName::try_from(header_name.as_str()).map_err(|e| { + anyhow!("--routing-key-header {header_name:?} is not a valid HTTP header name: {e}") + })?; + let fallback_policy = self + .sticky_fallback_policy + .unwrap_or(defaults.fallback_policy); + let idle_secs = self.sticky_idle_secs.unwrap_or(defaults.idle_secs); + let eviction_interval_secs = self + .sticky_eviction_interval_secs + .unwrap_or(defaults.eviction_interval_secs); + ensure!( + eviction_interval_secs != 0, + "--sticky-eviction-interval-secs must be greater than 0" + ); + ensure!( + idle_secs != 0, + "--sticky-idle-secs must be greater than 0 (0 would evict every \ + assignment immediately, defeating sticky routing)" + ); + Ok(Some(StickyConfig { + header_name, + fallback_policy, + idle_secs, + eviction_interval_secs, + })) + } +} + fn load_bucket_config(path: &str) -> Result { let raw = std::fs::read_to_string(path) .map_err(|error| anyhow!("--bucket-config cannot read {path:?}: {error}"))?; @@ -817,16 +809,8 @@ fn load_bucket_config(path: &str) -> Result { .map_err(|error| anyhow!("--bucket-config {path:?} is not valid JSON: {error}")) } -/// 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(",")) - } + (!terms.is_empty()).then(|| terms.join(",")) } #[cfg(test)] @@ -861,6 +845,59 @@ mod tests { into_config(&refs) } + #[test] + fn help_groups_options_by_purpose() { + use clap::CommandFactory; + + Cli::command().debug_assert(); + for long in [false, true] { + let mut command = Cli::command(); + let help = if long { + command.render_long_help() + } else { + command.render_help() + } + .to_string(); + let headings: std::collections::HashSet<_> = command + .get_arguments() + .filter_map(|arg| arg.get_help_heading()) + .collect(); + assert_eq!( + headings.len(), + 6, + "keep related options in six broad groups" + ); + for arg in command + .get_arguments() + .filter(|arg| !matches!(arg.get_id().as_str(), "help" | "version")) + { + let heading = arg.get_help_heading().expect("every option has a group"); + assert!( + arg.get_help().is_some(), + "missing help for {}", + arg.get_id() + ); + let section = help + .split_once(&format!("{heading}:\n")) + .unwrap_or_else(|| panic!("missing help section: {heading}")) + .1 + .split("\n\n") + .next() + .unwrap(); + // Long help separates individual options with blank lines. + if !long { + assert!( + section.contains(&format!("--{}", arg.get_long().unwrap())), + "{} is outside {heading}", + arg.get_id() + ); + } + } + assert!(help.contains("Examples:")); + assert!(help.contains("Choose exactly one discovery backend")); + } + } + #[test] fn defaults_host_port_and_policy() { let c = into_config_owned(with_model(&["--worker-urls", "http://10.0.0.1:30000"])).unwrap(); @@ -898,9 +935,6 @@ mod tests { } } - /// The ceiling is enforced on the CLI path, not only on a hand-built - /// `Config`: a drain carrying a fat-fingered extra digit must fail at - /// startup rather than turn every later termination into a SIGKILL. #[test] fn shutdown_drain_secs_past_the_ceiling_is_rejected() { let error = into_config_owned(with_model(&[ @@ -917,10 +951,6 @@ mod tests { ); } - /// `--termination-grace-secs` exists only to feed the startup advisory, so - /// the one thing that matters is that it reaches the config — and that - /// omitting it stays `None` (assume the k8s default) rather than - /// defaulting to a number that would silently become the compared budget. #[test] fn termination_grace_secs_maps_into_config_and_defaults_to_none() { let c = into_config_owned(with_model(&["--worker-urls", "http://10.0.0.1:30000"])).unwrap(); @@ -965,19 +995,6 @@ mod tests { assert_eq!(c.model.tokenizer_path, "/models/qwen3/tokenizer.json"); } - #[test] - fn input_ids_forwarding_can_be_disabled_for_the_model() { - let defaults = into_config_owned(with_model(&["--worker-urls", "http://x:30000"])).unwrap(); - assert!(!defaults.model.disable_input_ids_forwarding); - let disabled = into_config_owned(with_model(&[ - "--worker-urls", - "http://x:30000", - "--disable-input-ids-forwarding", - ])) - .unwrap(); - assert!(disabled.model.disable_input_ids_forwarding); - } - #[test] fn static_urls_backend() { let c = into_config_owned(with_model(&[ @@ -1156,10 +1173,6 @@ mod tests { 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(&[ @@ -1448,9 +1461,6 @@ mod tests { 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(&[ @@ -1495,8 +1505,11 @@ mod tests { .split_once("--sticky-fallback-policy ") .expect("sticky fallback option is documented"); let choices = after - .split_once("--sticky-idle-secs") - .expect("sticky fallback precedes its tuning") + .split_once("[possible values:") + .expect("sticky fallback lists its choices") + .1 + .split_once(']') + .unwrap() .0; for value in ["round_robin", "random", "power_of_two", "load_based"] { @@ -1786,11 +1799,6 @@ mod tests { assert!(fused_of("").is_none(), "round_robin builds no term list"); } - /// Non-finite and negative weights are refused, naming the term. - /// - /// `nan`/`inf` matter more than they look: `str::parse::` accepts - /// both, and a NaN weight makes every worker's fused total NaN, so argmax - /// discards them all and the router silently degrades to least-load. #[test] fn fuse_rejects_non_finite_and_negative_weights() { for bad in ["nan", "NaN", "inf", "-inf", "-0.5", "banana"] { @@ -2117,43 +2125,6 @@ mod tests { ); } - #[test] - fn min_load_choices_is_plumbed_and_validated() { - let config = cfg_of("--policy cache_aware --min-load-choices 5").unwrap(); - assert_eq!( - config - .model - .affinity - .expect("cache-aware needs affinity config") - .min_load_choices, - 5 - ); - - // Unset keeps the pre-existing power-of-2 behavior. - let defaults = cfg_of("--policy cache_aware").unwrap(); - assert_eq!( - defaults - .model - .affinity - .expect("default affinity config") - .min_load_choices, - 2 - ); - - let err = cfg_of("--policy cache_aware --min-load-choices 0") - .expect_err("a zero sample size would select nothing") - .to_string(); - assert!(err.contains("--min-load-choices"), "got: {err}"); - - let err = cfg_of("--policy power_of_two --min-load-choices 3") - .expect_err("the knob only tunes the cache-aware fallback") - .to_string(); - assert!( - err.contains("cache candidate tuning flags require --policy cache_aware"), - "got: {err}" - ); - } - #[test] fn cache_candidate_cli_rejects_invalid_bounds() { for (args, expected) in [ @@ -2318,4 +2289,53 @@ mod tests { .to_string(); assert!(err.contains("unknown parameter"), "got: {err}"); } + #[test] + fn input_ids_forwarding_can_be_disabled_for_the_model() { + let defaults = into_config_owned(with_model(&["--worker-urls", "http://x:30000"])).unwrap(); + assert!(!defaults.model.disable_input_ids_forwarding); + let disabled = into_config_owned(with_model(&[ + "--worker-urls", + "http://x:30000", + "--disable-input-ids-forwarding", + ])) + .unwrap(); + assert!(disabled.model.disable_input_ids_forwarding); + } + + #[test] + fn min_load_choices_is_plumbed_and_validated() { + let config = cfg_of("--policy cache_aware --min-load-choices 5").unwrap(); + assert_eq!( + config + .model + .affinity + .expect("cache-aware needs affinity config") + .min_load_choices, + 5 + ); + + // Unset keeps the pre-existing power-of-2 behavior. + let defaults = cfg_of("--policy cache_aware").unwrap(); + assert_eq!( + defaults + .model + .affinity + .expect("default affinity config") + .min_load_choices, + 2 + ); + + let err = cfg_of("--policy cache_aware --min-load-choices 0") + .expect_err("a zero sample size would select nothing") + .to_string(); + assert!(err.contains("--min-load-choices"), "got: {err}"); + + let err = cfg_of("--policy power_of_two --min-load-choices 3") + .expect_err("the knob only tunes the cache-aware fallback") + .to_string(); + assert!( + err.contains("cache candidate tuning flags require --policy cache_aware"), + "got: {err}" + ); + } } diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index 9d20f82e4..f5b84fd62 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -5,51 +5,29 @@ pub use cli::Cli; pub use sampling::*; pub use types::*; -use anyhow::{anyhow, Result}; +use anyhow::{anyhow, ensure, Result}; -/// The k8s default `terminationGracePeriodSeconds`, assumed when the operator -/// has not declared the pod's real one. A `shutdown_drain_secs` at or above the -/// grace period leaves no time for the in-flight drain, so the pod is SIGKILLed -/// before it finishes — the opposite of what the drain is for. +/// Default pod termination grace period when none is declared. pub const K8S_DEFAULT_GRACE_SECS: u64 = 30; -/// Ceiling on `shutdown_drain_secs`, enforced by [`Config::validate`]. Sized -/// for the workload rather than for the k8s default grace period: a single -/// streaming completion can hold the router for many minutes, so a deployment -/// that does not want terminations cutting one off runs a -/// `terminationGracePeriodSeconds` in the tens of minutes and a drain to match. -/// Deciding whether a particular drain fits a particular grace period is -/// [`shutdown_drain_advisory`]'s job — advice, because the operator can raise -/// the budget. This constant is the separate, harder gate: it rejects a value -/// that is not a drain at all, an extra digit or seconds confused with -/// milliseconds, which no grace period could ever service. +/// Maximum shutdown pause; deployments must also allow time to drain in-flight requests. +/// This is the hard typo gate (an extra digit, seconds confused with milliseconds); +/// whether a legal drain fits a particular grace period is [`shutdown_drain_advisory`]'s +/// job, because the operator can raise the budget. pub const MAX_SHUTDOWN_DRAIN_SECS: u64 = 1800; -/// A `shutdown_drain_secs` that leaves no room under the grace period for the -/// in-flight drain that follows the pause. Carries the compared values as -/// fields so the caller logs a static message with structured data rather than -/// interpolating the numbers into the message text. +/// A shutdown pause that exhausts the declared or assumed pod grace period. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ShutdownDrainAdvisory { pub shutdown_drain_secs: u64, /// The budget the drain was compared against. pub termination_grace_secs: u64, - /// Whether that budget came from the operator or from - /// [`K8S_DEFAULT_GRACE_SECS`]. An assumed budget makes the advisory a - /// guess; a declared one makes it a fact. + /// Whether the operator declared the grace period. pub grace_declared: bool, } -/// Advisory (not a hard error: the drain may well be right and the grace period -/// raised to match) for a drain that leaves no room for the in-flight drain. -/// The bound is `>=`, not `>`: a drain of exactly the grace period already -/// consumes all of it. -/// -/// `termination_grace_secs` is the pod's real `terminationGracePeriodSeconds` -/// when the operator declared it. The router cannot read its own pod spec, so -/// `None` falls back to the k8s default — which is why declaring the real value -/// is the way to silence this on a deployment that raised the grace period -/// deliberately, rather than lowering a drain that was correct. +/// Warn when the pause leaves no time for in-flight draining. +/// An undeclared grace period uses the Kubernetes default. pub fn shutdown_drain_advisory( shutdown_drain_secs: u64, termination_grace_secs: Option, @@ -63,51 +41,35 @@ pub fn shutdown_drain_advisory( } impl Config { - /// 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. + /// Validate invariants not enforced by the CLI parser. pub(crate) fn validate(&self) -> Result<()> { - if self.model.id.is_empty() { - return Err(anyhow!("model id must be non-empty")); - } + ensure!(!self.model.id.is_empty(), "model id must be non-empty"); if let Some(bucket_config) = self.model.bucket_config.as_ref() { validate_bucket_config(bucket_config)?; } self.model.sampling_overrides.validate()?; - if self.server.shutdown_drain_secs > MAX_SHUTDOWN_DRAIN_SECS { - return Err(anyhow!( - "shutdown_drain_secs must be at most {MAX_SHUTDOWN_DRAIN_SECS} (got {}); \ - past the ceiling a value is a typo rather than a drain, and the pod \ - would be SIGKILLed long before the pause elapsed. A long but deliberate \ - drain is fine — declare --termination-grace-secs so startup can check \ - it against the pod's real budget", - self.server.shutdown_drain_secs, - )); - } + ensure!( + self.server.shutdown_drain_secs <= MAX_SHUTDOWN_DRAIN_SECS, + "shutdown_drain_secs must be at most {MAX_SHUTDOWN_DRAIN_SECS} (got {}); \ + past the ceiling a value is a typo rather than a drain. A long but \ + deliberate drain is fine — declare --termination-grace-secs so startup \ + can check it against the pod's real budget", + self.server.shutdown_drain_secs, + ); match &self.discovery { DiscoveryBackend::StaticUrls(s) => { - if s.urls.is_empty() { - return Err(anyhow!( - "discovery.static_urls.urls must be a non-empty list" - )); - } - // Validate every entry up front so typos surface at - // 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 - // form (trimmed + trailing `/` stripped) so - // `"http://x:30000"` and `"http://x:30000/"` collide. + ensure!( + !s.urls.is_empty(), + "discovery.static_urls.urls must be a non-empty list" + ); + // Normalize URLs before deduplication so trailing slashes cannot register a worker twice. let mut seen = std::collections::HashSet::new(); for raw in &s.urls { let trimmed = raw.trim(); - if trimmed.is_empty() { - return Err(anyhow!( - "discovery.static_urls.urls contains an empty or whitespace-only entry" - )); - } + ensure!( + !trimmed.is_empty(), + "discovery.static_urls.urls contains an empty or whitespace-only entry" + ); let parsed = url::Url::parse(trimmed).map_err(|e| { anyhow!("discovery.static_urls.urls entry {raw:?} is not a valid URL: {e}") })?; @@ -120,17 +82,13 @@ impl Config { } } let normalized = parsed.as_str().trim_end_matches('/').to_string(); - if !seen.insert(normalized.clone()) { - return Err(anyhow!( - "discovery.static_urls.urls contains duplicate entry {raw:?} (normalized: {normalized:?})" - )); - } + ensure!( + seen.insert(normalized.clone()), + "discovery.static_urls.urls contains duplicate entry {raw:?} (normalized: {normalized:?})" + ); } } - // 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. + // Kubernetes selector combinations are validated by `resolve_mode` during construction. DiscoveryBackend::K8s(_) => {} } Ok(()) @@ -138,50 +96,44 @@ impl Config { } fn validate_bucket_config(bucket_config: &BucketConfig) -> Result<()> { - if bucket_config.buckets.is_empty() { - return Err(anyhow!( - "bucket_config.buckets must be non-empty when configured" - )); - } + ensure!( + !bucket_config.buckets.is_empty(), + "bucket_config.buckets must be non-empty when configured" + ); let mut ids = std::collections::HashSet::new(); let mut ranks = std::collections::HashSet::new(); let mut stage_workers = std::collections::HashSet::new(); let mut has_prefill_bucket = false; for bucket in &bucket_config.buckets { has_prefill_bucket |= bucket.stage == BucketStage::Prefill; - if bucket.id.is_empty() || !ids.insert(bucket.id.as_str()) { - return Err(anyhow!( - "bucket_config bucket id must be non-empty and unique: {:?}", - bucket.id - )); - } - if !ranks.insert((bucket.stage, bucket.rank)) { - return Err(anyhow!( - "bucket_config rank must be unique within each stage: {}", - bucket.rank - )); - } - if bucket.worker_ids.is_empty() { - return Err(anyhow!( - "bucket_config bucket {:?} has no worker_ids", - bucket.id - )); - } + ensure!( + !bucket.id.is_empty() && ids.insert(bucket.id.as_str()), + "bucket_config bucket id must be non-empty and unique: {:?}", + bucket.id + ); + ensure!( + ranks.insert((bucket.stage, bucket.rank)), + "bucket_config rank must be unique within each stage: {}", + bucket.rank + ); + ensure!( + !bucket.worker_ids.is_empty(), + "bucket_config bucket {:?} has no worker_ids", + bucket.id + ); let mut worker_ids = std::collections::HashSet::new(); for worker_id in &bucket.worker_ids { - if worker_id.is_empty() || !worker_ids.insert(worker_id.as_str()) { - return Err(anyhow!( - "bucket_config bucket {:?} has an empty or duplicate worker id", - bucket.id - )); - } - if !stage_workers.insert((bucket.stage, worker_id.as_str())) { - return Err(anyhow!( - "bucket_config worker {:?} belongs to more than one {:?} bucket", - worker_id, - bucket.stage - )); - } + ensure!( + !worker_id.is_empty() && worker_ids.insert(worker_id.as_str()), + "bucket_config bucket {:?} has an empty or duplicate worker id", + bucket.id + ); + ensure!( + stage_workers.insert((bucket.stage, worker_id.as_str())), + "bucket_config worker {:?} belongs to more than one {:?} bucket", + worker_id, + bucket.stage + ); } validate_range( bucket.min_extend_tokens, @@ -195,33 +147,28 @@ fn validate_bucket_config(bucket_config: &BucketConfig) -> Result<()> { &bucket.id, "sequence", )?; - if bucket.max_context_tokens == Some(0) { - return Err(anyhow!( - "bucket_config bucket {:?} max_context_tokens must be > 0", - bucket.id - )); - } - if bucket.ttft_p95_at_capacity_ms == Some(0) { - return Err(anyhow!( - "bucket_config bucket {:?} TTFT p95 must be > 0", - bucket.id - )); - } - if bucket - .tps_p05_at_capacity - .is_some_and(|value| !value.is_finite() || value <= 0.0) - { - return Err(anyhow!( - "bucket_config bucket {:?} TPS p05 must be finite and > 0", - bucket.id - )); - } - if bucket.max_pending_prefill_tokens == Some(0) { - return Err(anyhow!( - "bucket_config bucket {:?} max_pending_prefill_tokens must be > 0", - bucket.id - )); - } + ensure!( + bucket.max_context_tokens != Some(0), + "bucket_config bucket {:?} max_context_tokens must be > 0", + bucket.id + ); + ensure!( + bucket.ttft_p95_at_capacity_ms != Some(0), + "bucket_config bucket {:?} TTFT p95 must be > 0", + bucket.id + ); + ensure!( + bucket + .tps_p05_at_capacity + .is_none_or(|value| value.is_finite() && value > 0.0), + "bucket_config bucket {:?} TPS p05 must be finite and > 0", + bucket.id + ); + ensure!( + bucket.max_pending_prefill_tokens != Some(0), + "bucket_config bucket {:?} max_pending_prefill_tokens must be > 0", + bucket.id + ); match bucket.stage { BucketStage::Prefill if bucket.min_sequence_tokens.is_some() @@ -247,20 +194,19 @@ fn validate_bucket_config(bucket_config: &BucketConfig) -> Result<()> { _ => {} } } - if !has_prefill_bucket { - return Err(anyhow!( - "bucket_config must contain at least one Prefill bucket; enabling Bucket routing otherwise leaves every request without a Prefill domain" - )); - } + ensure!( + has_prefill_bucket, + "bucket_config must contain at least one Prefill bucket; enabling Bucket \ + routing otherwise leaves every request without a Prefill domain" + ); Ok(()) } fn validate_range(min: Option, max: Option, id: &str, name: &str) -> Result<()> { - if min.zip(max).is_some_and(|(min, max)| min > max) { - return Err(anyhow!( - "bucket_config bucket {id:?} has invalid {name} range: min > max" - )); - } + ensure!( + min.zip(max).is_none_or(|(min, max)| min <= max), + "bucket_config bucket {id:?} has invalid {name} range: min > max" + ); Ok(()) } @@ -268,10 +214,6 @@ fn validate_range(min: Option, max: Option, id: &str, name: &str) -> R mod tests { use super::*; - /// 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::default(), @@ -549,19 +491,13 @@ mod tests { #[test] fn shutdown_drain_advisory_is_silent_below_the_k8s_default_grace() { - // Anything strictly under the k8s default terminationGracePeriodSeconds - // (30 s) still leaves room for the in-flight drain, so it is safe - // without operator action. assert!(shutdown_drain_advisory(29, None).is_none()); assert!(shutdown_drain_advisory(0, None).is_none()); } - /// The default drain is exactly the assumed grace period, so out of the box - /// it leaves nothing for the in-flight drain and says so on every startup. - /// Asserted rather than left implicit because it is the one case an - /// operator meets without choosing it: a later edit to either constant that - /// silenced the warning would be changing the default deployment's - /// behaviour, and should have to say so here. + /// Pinned because this is the one case an operator meets without choosing it: + /// an edit to either constant that silenced the warning would change the default + /// deployment's behaviour, and should have to say so here. #[test] fn the_default_drain_warns_until_the_grace_period_is_raised() { let advisory = shutdown_drain_advisory(default_shutdown_drain_secs(), None) @@ -585,11 +521,6 @@ mod tests { #[test] fn shutdown_drain_advisory_warns_once_the_drain_consumes_the_whole_grace() { - // A drain of exactly the 30 s k8s default leaves zero seconds for the - // in-flight drain, so the pod is SIGKILLed mid-drain — the boundary - // itself must warn, not just values past it. The ceiling is in the list - // because it is startable: `validate` accepts it, so the advisory is - // the only thing left to say it does not fit the default grace period. for drain in [K8S_DEFAULT_GRACE_SECS, 120, MAX_SHUTDOWN_DRAIN_SECS] { let advisory = shutdown_drain_advisory(drain, None) .unwrap_or_else(|| panic!("{drain}s must warn")); @@ -602,9 +533,6 @@ mod tests { } } - /// The advisory's whole purpose is to be silenceable by declaring the real - /// budget: a 60 s drain under a 120 s grace period is a correct - /// configuration, and warning about it trains operators to ignore the line. #[test] fn shutdown_drain_advisory_respects_a_declared_grace_period() { assert!( @@ -624,18 +552,12 @@ mod tests { shutdown_drain_advisory(10, Some(10)).is_some(), "a short declared grace period must still be compared against", ); - // The configuration the ceiling was raised for: a completion streaming - // for minutes wants a drain of minutes, under a grace period declared - // to match. That is correct, not merely tolerated, so it must be silent. assert!( shutdown_drain_advisory(MAX_SHUTDOWN_DRAIN_SECS, Some(3600)).is_none(), "a long drain under a grace period declared to cover it must not warn", ); } - /// `validate` is the hard gate the advisory deliberately is not: past the - /// ceiling the value can only be a typo, and starting on it would make - /// every later termination a SIGKILL. #[test] fn validate_rejects_a_shutdown_drain_past_the_ceiling() { let mut config = cfg("qwen3-0.6b", &["http://10.0.0.1:30000"]); diff --git a/experimental/sgl-router/src/config/sampling.rs b/experimental/sgl-router/src/config/sampling.rs index 3a71d1944..c039ddf6d 100644 --- a/experimental/sgl-router/src/config/sampling.rs +++ b/experimental/sgl-router/src/config/sampling.rs @@ -1,58 +1,30 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 -//! Fleet-wide sampling contract (`--override-sampling-params` / -//! `--sampling-param-conflict`): the parameters an operator fixes for every -//! request this router serves, and what a request that disagrees gets. +//! Fleet sampling defaults and constraints. Custom JSON visitors preserve duplicate +//! keys so validation can reject them and report the offending parameter. //! -//! WHY this is parsed by hand rather than with `serde(deny_unknown_fields)`: -//! the flag is read once, at startup, on a router that crash-loops if it is -//! wrong, so the message an operator reads out of `kubectl logs` is the whole -//! debugging session. Every rejection here names the offending key, the value -//! it saw, and the domain it violated. The same reasoning is why both the -//! outer object and a band are decoded as ordered ENTRIES instead of a -//! `serde_json::Map`: a map keeps only the last of a repeated key, so -//! `{"temperature": 0, "temperature": 1}` would start cleanly and enforce a -//! value the operator did not write. +//! The flag is read once, at startup, on a router that crash-loops if it is wrong, +//! so the message an operator reads out of `kubectl logs` is the whole debugging +//! session: every rejection names the offending key, the value it saw, and the +//! domain it violated. (A `serde_json::Map` would keep only the last of a repeated +//! key and silently enforce a value the operator did not write.) -use anyhow::{anyhow, Result}; +use anyhow::{anyhow, ensure, Result}; use std::collections::BTreeMap; -/// Sampling parameters fixed fleet-wide, and what to do with a request that -/// disagrees. -/// -/// A configured parameter is always injected into the forwarded body when the -/// request OMITS it, so the engine's own defaults cannot drift from what the -/// operator declared. What differs between the two [`ConflictPolicy`] modes is -/// only the request that DOES send the field: `Reject` makes the value an -/// immutability contract (400 before admission — never a silent rewrite), -/// while `Allow` lets the client value through untouched, degrading the -/// configured value to a fleet-wide default. +/// Fleet sampling defaults. Exact values fill absent fields; [`ConflictPolicy`] +/// determines whether differing client values are rejected or forwarded. #[derive(Debug, Clone, Default, PartialEq)] pub struct SamplingOverrides { - /// Configured parameters, keyed so enforcement and injection are one loop - /// over whatever the operator set instead of a per-field ladder repeated - /// at each site. Iterating a `BTreeMap` keyed by the field enum is what - /// fixes the order values are injected in, so a forwarded body is - /// byte-identical across runs. + /// Parameters in deterministic injection order. pub params: BTreeMap, - /// Applies to every configured parameter: there is deliberately no - /// per-parameter mode, so an operator reads one knob off one manifest. + /// Conflict behavior shared by all configured parameters. pub conflict: ConflictPolicy, } impl SamplingOverrides { - /// Re-check every invariant [`parse_sampling_overrides`] enforces, on an - /// already-built value. - /// - /// WHY this is separate from the parser: the parser turns a raw JSON - /// string into this struct and is reachable only from the CLI, but the - /// struct itself is reachable from anywhere — a test fixture, a future - /// config file, an admin API. [`crate::config::Config::validate`] calls - /// this so no such path can hold a spec the flag would have refused to - /// start with (an out-of-domain exact value, an inverted band whose - /// `contains` rejects every value, or a band under `allow`, which names - /// nothing to inject and rejects nothing). + /// Validate parsed and programmatically constructed overrides alike. pub(crate) fn validate(&self) -> Result<()> { for (&field, spec) in &self.params { validate_spec(field, spec, self.conflict)?; @@ -65,10 +37,7 @@ impl SamplingOverrides { /// (`--sampling-param-conflict`). #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)] pub enum ConflictPolicy { - /// 400 before admission, quoting the configured value. The default: the - /// point of declaring a fleet-wide sampling contract is usually that it - /// holds, and silently serving something other than what the client asked - /// for is the one behavior no client can detect. + /// Reject differing client values with 400 before admission. #[default] Reject, /// Forward the client's value to the engine untouched. The configured @@ -80,30 +49,16 @@ pub enum ConflictPolicy { /// accepted ones. #[derive(Debug, Clone, PartialEq)] pub enum ParamSpec { - /// A single value: injected when the request omits the field, and under - /// [`ConflictPolicy::Reject`] the only value a request may send. - /// - /// Held as the parsed JSON number rather than an `f64` so injection - /// re-emits the operator's literal — `"n": 1` stays `1` and does not - /// become `1.0` on the wire for the integer-typed fields. + /// Injected when absent; under [`ConflictPolicy::Reject`], the only accepted value. + /// JSON numbers preserve integer wire types. Exact(serde_json::Number), - /// An inclusive `[lo, hi]` band of accepted values, for a contract that - /// fixes most sampling knobs but leaves one tunable inside a range. A band - /// names no single value, so it never injects; it only rejects - /// out-of-band values, which is why a band under [`ConflictPolicy::Allow`] - /// is a startup error rather than a no-op. - /// - /// A band therefore constrains only the requests that NAME the parameter. - /// A request that omits it gets the model's own default (the engine reads - /// `generation_config`), which the router cannot see and which may itself - /// lie outside the band. An operator who needs the omitting majority - /// pinned too wants an exact value, not a band. + /// Inclusive bounds for supplied values; never injects a default. + /// Requires [`ConflictPolicy::Reject`]. Omitted fields use the engine default, + /// which may lie outside the band. Range { lo: f64, hi: f64 }, } -/// A sampling parameter that can be fixed fleet-wide. The enum is what makes a -/// typo in the `--override-sampling-params` JSON a startup error instead of a -/// key that silently never matches a request field. +/// Supported fleet sampling parameters. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum SamplingField { Temperature, @@ -129,15 +84,7 @@ impl SamplingField { Self::N, ]; - /// This field's slot in [`Self::ALL`], and in the request probe's - /// fixed-size array of probed values. - /// - /// Declaration order IS the slot order, so this cannot assign a wrong - /// one. What still needs checking is that [`Self::ALL`] agrees — see the - /// assertion below; `from_wire_name` and `supported_fields` both iterate - /// `ALL`, so a field missing from it is rejected at startup as an unknown - /// key, the failure that is invisible to any test that also iterates - /// `ALL`. + /// Slot in [`Self::ALL`] and in the request probe array. pub const fn index(self) -> usize { self as usize } @@ -190,12 +137,11 @@ pub(crate) fn parse_sampling_overrides( '{{\"temperature\": 1, \"top_p\": 0.95}}': {e}" ) })?; - if entries.is_empty() { - return Err(anyhow!( - "--override-sampling-params is empty: pass at least one of {}, or omit the flag", - supported_fields() - )); - } + ensure!( + !entries.is_empty(), + "--override-sampling-params is empty: pass at least one of {}, or omit the flag", + supported_fields() + ); let mut params = BTreeMap::new(); for (key, value) in entries { let field = SamplingField::from_wire_name(&key).ok_or_else(|| { @@ -216,26 +162,20 @@ pub(crate) fn parse_sampling_overrides( )) } }; - if params.insert(field, spec).is_some() { - return Err(anyhow!( - "--override-sampling-params: {} is set more than once", - field.wire_name() - )); - } + ensure!( + params.insert(field, spec).is_none(), + "--override-sampling-params: {} is set more than once", + field.wire_name() + ); } let overrides = SamplingOverrides { params, conflict }; - // Re-checks the domains `checked_value` already covered above. That first - // pass is not redundant: it is what quotes the operator's own literal - // (`1e2`, not `100`) and what guards `canonical_number`'s saturating - // `as i64` cast before it runs. This pass is what a hand-built - // `SamplingOverrides` gets, and is the only check a band's bounds see. + // Validate complete specs too, including bands and programmatically built overrides. + // The earlier exact-value check protects the integer cast in `canonical_number`. overrides.validate()?; Ok(overrides) } -/// Check one already-built spec. Shared by [`parse_sampling_overrides`] and -/// [`SamplingOverrides::validate`] so a hand-built `SamplingOverrides` is held -/// to exactly the domain the flag is. +/// Validate a spec independently of how it was constructed. fn validate_spec(field: SamplingField, spec: &ParamSpec, conflict: ConflictPolicy) -> Result<()> { let key = field.wire_name(); match spec { @@ -245,31 +185,24 @@ fn validate_spec(field: SamplingField, spec: &ParamSpec, conflict: ConflictPolic &ParamSpec::Range { lo, hi } => { check_domain(field, lo, &lo.to_string())?; check_domain(field, hi, &hi.to_string())?; - if lo > hi { - return Err(anyhow!( - "--override-sampling-params: {key} band needs min <= max, got min {lo} > max {hi}" - )); - } - // Bounds are checked one at a time, which is only sufficient for a - // contiguous domain. `top_k`'s is not ({-1} U [1, inf)): `{"min": -1, - // "max": 100}` has two individually legal bounds and would admit - // `top_k: 0`, which is rejected as an exact value. -1 is a sentinel, - // not a range endpoint. - if field == SamplingField::TopK && lo < 1.0 { - return Err(anyhow!( - "--override-sampling-params: top_k band bounds must both be >= 1 \ + ensure!( + lo <= hi, + "--override-sampling-params: {key} band needs min <= max, got min {lo} > max {hi}" + ); + // `top_k = -1` disables filtering; it cannot bound a band that would admit zero. + ensure!( + field != SamplingField::TopK || lo >= 1.0, + "--override-sampling-params: top_k band bounds must both be >= 1 \ (-1 disables top_k entirely and cannot bound a range)" - )); - } + ); // A band only ever rejects, so under `allow` it would be dead config // that silently accepts everything. - if conflict == ConflictPolicy::Allow { - return Err(anyhow!( - "--override-sampling-params: the {key} band requires \ + ensure!( + conflict == ConflictPolicy::Reject, + "--override-sampling-params: the {key} band requires \ --sampling-param-conflict reject — under `allow` nothing is rejected \ and a band names no value to inject" - )); - } + ); } } Ok(()) @@ -294,11 +227,10 @@ fn parse_band( )) } }; - if slot.is_some() { - return Err(anyhow!( - "--override-sampling-params: {key} band sets \"{bound}\" more than once" - )); - } + ensure!( + slot.is_none(), + "--override-sampling-params: {key} band sets \"{bound}\" more than once" + ); let serde_json::Value::Number(n) = v else { return Err(anyhow!( "--override-sampling-params: {key} band needs numeric bounds, got {bound}: {v}" @@ -312,96 +244,61 @@ fn parse_band( {{\"min\": LO, \"max\": HI}} with numeric bounds" )); }; - // `lo <= hi`, `top_k`'s discontiguous domain and the band-under-`allow` - // rule are all properties of the finished spec, so they live in - // `validate_spec` and hold for a hand-built `SamplingOverrides` too. + // Finished-spec constraints are checked by `validate_spec` for all construction paths. Ok(ParamSpec::Range { lo, hi }) } -/// Check one configured value against its parameter's domain, at startup -/// instead of per request. Written as positive containment so a NaN bound -/// fails too. -/// -/// These are the OpenAI API's domains, which are NARROWER than what the engine -/// itself accepts (`SamplingParams.verify` requires only that `temperature` be -/// non-negative and finite, so it would take `temperature: 5`). Narrower is -/// deliberate: the values here are injected into request bodies, and a fleet -/// contract outside the range every OpenAI client library validates against is -/// far more likely a typo than an intent. The one exception is `top_k`, where -/// `-1` is the engine's own "disable / whole vocabulary" spelling and its -/// default — a legitimate thing to fix fleet-wide. Note `top_k: 1` is greedy -/// decoding, NOT "disabled". +/// Validate before normalization so diagnostics retain the configured number. +/// Domains follow the OpenAI contract plus engine-specific parameters — deliberately +/// NARROWER than what the engine accepts: these values are injected into request +/// bodies, and a fleet contract outside the range every OpenAI client library +/// validates against is far more likely a typo than an intent. fn checked_value(field: SamplingField, n: &serde_json::Number) -> Result { let name = field.wire_name(); - // `as_f64` is infallible for a JSON number unless serde_json's - // `arbitrary_precision` is on (it is not); kept total rather than - // `expect`-ing, so enabling that feature can't turn config into a panic. + // Handle conversion failure even if serde_json arbitrary precision is enabled later. let v = n.as_f64().ok_or_else(|| { anyhow!("--override-sampling-params: {name} ({n}) is not a finite number") })?; - // The operator's own literal is what the message quotes, not the parsed - // f64: `1e2` should read back as `1e2`. check_domain(field, v, &n.to_string())?; Ok(v) } -/// The domain half of [`checked_value`], over an f64 that may not have come -/// from a literal (a band's bounds are stored as f64). `shown` is what the -/// error quotes back to the operator. +/// Validate a numeric domain; `shown` is the value quoted in diagnostics. fn check_domain(field: SamplingField, v: f64, shown: &str) -> Result<()> { let name = field.wire_name(); let (ok, domain) = match field { SamplingField::Temperature => ((0.0..=2.0).contains(&v), "in [0, 2]"), SamplingField::TopP => (v > 0.0 && v <= 1.0, "in (0, 1]"), SamplingField::TopK => (v >= 1.0 || v == -1.0, ">= 1, or -1 to disable"), - // Not an OpenAI parameter: `min_p` is the engine's own nucleus floor, - // and 0 is its default (disabled), so the whole [0, 1] range is - // legitimate to fix fleet-wide. + // Engine-specific: zero disables `min_p`. SamplingField::MinP => ((0.0..=1.0).contains(&v), "in [0, 1]"), - // Also engine-only. 1.0 is "no penalty"; the engine requires > 0, and - // values above ~2 degrade output badly enough that a fleet-wide pin - // there is far more likely a typo than an intent. + // Engine-specific: one disables the penalty; cap fleet defaults at two. SamplingField::RepetitionPenalty => (v > 0.0 && v <= 2.0, "in (0, 2]"), SamplingField::FrequencyPenalty | SamplingField::PresencePenalty => { ((-2.0..=2.0).contains(&v), "in [-2, 2]") } - // OpenAI caps `n` at 128. Unbounded here, a typo'd digit would be - // injected into every request that omits `n` and fan each one out to - // that many sequences at the engine — the exact per-request failure - // this startup check exists to convert into a launch failure. + // Bound sequence fan-out when `n` is injected into requests. SamplingField::N => ((1.0..=128.0).contains(&v), "in [1, 128]"), }; - if !ok { - return Err(anyhow!( - "--override-sampling-params: {name} ({shown}) must be {domain}" - )); - } + ensure!( + ok, + "--override-sampling-params: {name} ({shown}) must be {domain}" + ); if field.is_integral() { - if v.fract() != 0.0 { - return Err(anyhow!( - "--override-sampling-params: {name} ({shown}) must be a whole number" - )); - } - // `canonical_number` casts to `i64`, and a Rust float-to-int cast - // SATURATES rather than failing, so a literal past the i64 range would - // silently become `i64::MAX` in every forwarded body. The exactly - // convertible f64s are [-2^63, 2^63), which is this half-open range - // written as positive containment — `i64::MAX as f64` rounds UP to - // 2^63, so an inclusive `<=` against it would admit 2^63 itself and - // saturate exactly as described. - if !(i64::MIN as f64..i64::MAX as f64).contains(&v) { - return Err(anyhow!( - "--override-sampling-params: {name} ({shown}) is too large to forward" - )); - } + ensure!( + v.fract() == 0.0, + "--override-sampling-params: {name} ({shown}) must be a whole number" + ); + // Float-to-i64 casts saturate. Use [-2^63, 2^63): `i64::MAX as f64` rounds up. + ensure!( + (i64::MIN as f64..i64::MAX as f64).contains(&v), + "--override-sampling-params: {name} ({shown}) is too large to forward" + ); } Ok(()) } -/// Normalize an integer-typed parameter's literal so injection writes `1` -/// rather than `1.0` for a config that spelled it `1.0` — the engine types -/// these fields as `int`, and the forwarded body should look like what a -/// client would have sent. Non-integral fields keep the operator's literal. +/// Emit integer-typed parameters as integers; preserve other JSON numbers. fn canonical_number( field: SamplingField, value: f64, @@ -423,8 +320,7 @@ fn supported_fields() -> String { .join(", ") } -/// A JSON object decoded to its entries IN ORDER, keeping a repeated key -/// instead of collapsing it. See the module WHY note. +/// Ordered JSON entries preserve duplicate keys for validation. struct ObjectEntries(Vec<(String, ParamValue)>); impl<'de> serde::Deserialize<'de> for ObjectEntries { @@ -454,10 +350,7 @@ impl<'de> serde::Deserialize<'de> for ObjectEntries { } } -/// One parameter's raw value: a number, a band's entries, or anything else. -/// `Other` keeps the offending value so the caller can name it, rather than -/// degrading a wrong-type message into a serde type error behind the outer -/// object's context. +/// Raw parameter value; `Other` retains invalid values for precise diagnostics. enum ParamValue { Number(serde_json::Number), Band(Vec<(String, serde_json::Value)>), @@ -510,9 +403,7 @@ impl<'de> serde::Deserialize<'de> for ParamValue { Ok(ParamValue::Other(v.into())) } - /// JSON `null`. There is deliberately no `visit_none`: this type - /// is only ever reached through `deserialize_any`, which routes - /// null here and never to the `Option` hook. + /// `deserialize_any` routes JSON null to `visit_unit`. fn visit_unit(self) -> Result { Ok(ParamValue::Other(serde_json::Value::Null)) } @@ -654,18 +545,12 @@ mod tests { assert_eq!(exact_of(&o, SamplingField::TopK), Some(1000.0)); } - /// `top_k: -1` is the engine's own "disable / whole vocabulary" spelling - /// (and its default), so a fleet may legitimately fix `top_k` to it — - /// unlike every other parameter, whose domain is the OpenAI one. #[test] fn top_k_accepts_the_engines_disable_sentinel() { let o = parse(r#"{"top_k": -1}"#).unwrap(); assert_eq!(exact_of(&o, SamplingField::TopK), Some(-1.0)); } - /// An integer-typed parameter spelled as a float is normalized, so the - /// forwarded body carries `1` and not `1.0` for a field the engine types - /// as `int`. #[test] fn integral_params_are_normalized_to_integers() { let o = parse(r#"{"n": 1.0, "top_k": 20.0}"#).unwrap(); @@ -689,11 +574,7 @@ mod tests { } assert_eq!(SamplingField::from_wire_name("max_tokens"), None); } - /// `canonical_number` casts to `i64` and a Rust float-to-int cast - /// SATURATES, so a literal past the i64 range must fail the launch rather - /// than be injected as `i64::MAX`. The boundary case is the trap: `i64::MAX - /// as f64` rounds UP to 2^63, so a `>` comparison against it admits 2^63 - /// itself. + /// Float-to-i64 casts saturate. Use [-2^63, 2^63): `i64::MAX as f64` rounds up. #[test] fn integral_literals_beyond_i64_fail_the_launch() { for raw in [ @@ -716,12 +597,6 @@ mod tests { ); } - /// `min_p` and `repetition_penalty` are the only two parameters besides - /// `temperature`/`top_p`/`top_k` that the engine resolves from the model's - /// own `generation_config`, so they are exactly the ones a fleet-wide pin - /// exists to stop drifting when an image is swapped. Rejecting them as - /// unknown keys would crash-loop the router for the operator who needs the - /// flag most. #[test] fn governs_the_engine_defaulted_parameters() { let o = parse(r#"{"min_p": 0.05, "repetition_penalty": 1.1}"#).unwrap(); @@ -744,16 +619,8 @@ mod tests { } } - /// `ALL` is what `from_wire_name` and `supported_fields` iterate, so a - /// field missing from it is silently rejected at startup as an unknown - /// key — invisible to any test that also iterates `ALL`. Pin the length - /// and the slot mapping against the wire names instead. #[test] fn all_covers_every_field_exactly_once() { - // The slot mapping itself is asserted at compile time (see the - // `const _` block above `parse_sampling_overrides`); what only a test - // can catch is a field missing from `ALL` entirely, which is why the - // names below are written out rather than derived from it. let names: std::collections::BTreeSet<_> = SamplingField::ALL.iter().map(|f| f.wire_name()).collect(); assert_eq!(names.len(), SamplingField::ALL.len(), "duplicate wire name"); @@ -771,10 +638,6 @@ mod tests { } } - /// The struct-level invariants must hold for a `SamplingOverrides` that - /// never went through the parser — a test fixture, a future config file or - /// admin API. An inverted band is the nastiest of these: `(lo..=hi)` - /// contains nothing, so it would 400 every request naming the parameter. #[test] fn validate_rejects_hand_built_specs_the_parser_would_refuse() { let bad = [ diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 6aeebc0b4..ed704ec28 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -2,32 +2,22 @@ use crate::config::sampling::SamplingOverrides; use serde::Deserialize; use std::num::NonZeroU32; -/// 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. +/// Single-model configuration built and validated by [`crate::config::Cli::into_config`]. #[derive(Debug, Clone)] pub struct Config { pub server: ServerConfig, pub observability: ObservabilityConfig, 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`]. + /// Discovery mode resolved from CLI options; static URLs are checked by [`Config::validate`]. pub discovery: DiscoveryBackend, pub proxy: ProxyConfig, 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. +/// Outbound request timeout settings. #[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. + /// Timeout for upstream response headers and body. Counts as a circuit-breaker failure. pub request_timeout_secs: u64, } @@ -43,15 +33,10 @@ impl Default for ProxyConfig { } } -/// Active-load (per-request) tracking. Production default (10 min) -/// 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. +/// Request-tracking timeout; defaults above the proxy timeout. #[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. + /// Maximum request-entry lifetime before cancellation with 504 `stale_request_expired`. pub stale_request_timeout_secs: u64, } @@ -67,13 +52,7 @@ impl Default for ActiveLoadConfig { } } -/// Routing policy selector — the enum form lets `clap` reject unknown -/// values at parse time and removes the runtime string match in the -/// policy factory. -/// -/// Accepted on the CLI (`--policy`) as `round_robin` / `random` / -/// `power_of_two` / `load_based` / `fused_score` / `score_policy` / -/// `session_aware` / `cache_aware` / `sticky`. +/// Routing strategies accepted by `--policy`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)] pub enum PolicyKind { #[default] @@ -98,11 +77,7 @@ pub enum PolicyKind { /// Selects cache-affine prefill candidates from the configured prefix provider. #[value(name = "cache_aware")] CacheAware, - /// Sticky-session routing: pins a routing key (read from a - /// configurable request header) to a worker via an in-memory map, so - /// stateful sessions land on the same backend. Tuning — header name, - /// keyless-fallback policy, and TTL eviction — lives on - /// `ModelConfig::sticky`. + /// Pin a request-header routing key to a worker. #[value(name = "sticky")] Sticky, } @@ -248,35 +223,16 @@ impl std::fmt::Display for StickyFallbackKind { pub struct ServerConfig { pub host: String, pub port: u16, - /// Seconds to keep serving after SIGTERM — with `/readyz` flipped to 503 — - /// before the HTTP server stops accepting. The default covers both - /// deregistration paths: endpoint removal reaching kube-proxy after the - /// pod's `deletionTimestamp` is stamped, and a probe-driven load balancer, - /// which cannot act until `failureThreshold * periodSeconds` of `/readyz` - /// failures have accumulated. - /// - /// Note that it equals the k8s default `terminationGracePeriodSeconds`, so - /// a pod that has not raised its grace period is left with nothing for the - /// in-flight drain that follows the pause, and - /// [`shutdown_drain_advisory`](crate::config::shutdown_drain_advisory) - /// warns at every startup. That is the intended reading rather than a - /// misconfigured default: a router whose completions stream for minutes - /// cannot terminate cleanly inside 30 s at all, and the grace period is the - /// thing to raise. 0 disables the pause. + /// Pause after SIGTERM with `/readyz` returning 503 before stopping accepts. + /// Allows endpoint removal or readiness-probe failures to reach load balancers. + /// Leave time in the pod grace period for in-flight draining; 0 disables the pause. pub shutdown_drain_secs: u64, - /// The pod's actual `terminationGracePeriodSeconds`, when the operator - /// declares it. The router cannot read its own pod spec, so without this - /// the startup advisory can only compare the drain against the k8s - /// default — and warns, wrongly, about a deployment that raised the grace - /// period on purpose. `None` means "assume the default". + /// Declared pod termination grace period; `None` uses the Kubernetes default for advisories. pub termination_grace_secs: Option, } impl ServerConfig { - /// [`Self::shutdown_drain_secs`] as a `Duration`. Keeps the seconds-to- - /// `Duration` conversion in the library, where a test can pin it, rather - /// than in `main.rs` where a `from_secs`/`from_millis` slip would silently - /// shorten every drain by a factor of 1000. + /// Shutdown pause as a duration. pub fn shutdown_drain(&self) -> std::time::Duration { std::time::Duration::from_secs(self.shutdown_drain_secs) } @@ -294,10 +250,7 @@ pub fn default_shutdown_drain_secs() -> u64 { 30 } -/// Exists so test fixtures can spell out only the fields they care about -/// (`tests/` is a separate crate, so a `#[cfg(test)]` constructor cannot reach -/// the integration fixtures). Keep `Cli::into_config` exhaustive so adding a -/// field still forces a decision on the production path. +/// Defaults for config construction; the CLI mapping remains exhaustive. impl Default for ServerConfig { fn default() -> Self { Self { @@ -312,10 +265,7 @@ impl Default for ServerConfig { #[derive(Debug, Clone)] pub struct ObservabilityConfig { pub log_level: String, - /// 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). + /// Tracing output format. pub log_format: LogFormat, } @@ -346,9 +296,8 @@ impl Default for ObservabilityConfig { #[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`]. + /// Local tokenizer.json or HuggingFace repo id; defaults to `id`. + /// Resolved by [`crate::tokenizer::adapter::load`]. pub tokenizer_path: String, /// Disable router-generated input IDs for this model; keep routing tokenization. /// Use when workers have rendering defaults or template stops the router cannot see. @@ -361,25 +310,15 @@ pub struct ModelConfig { pub circuit_breaker: Option, /// Cache-Aware prefix configuration. pub cache_aware: Option, - /// Tuning for the sticky-session policy. `Some` exactly when - /// `policy = "sticky"` (built by [`crate::config::cli::Cli::into_config`]). - /// The chat handler reads `sticky.header_name` to populate - /// [`crate::policies::SelectionContext::routing_key`]. + /// Present only for the sticky policy; the header supplies the request routing key. pub sticky: Option, /// Session and cache-affinity tuning. pub affinity: Option, - /// Terms the score-composition policy sums. `Some` exactly when - /// `policy = "fused_score"` or `policy = "score_policy"` (built by - /// [`crate::config::cli::Cli::into_config`]), defaulting to - /// [`DEFAULT_FUSE`] when `--fuse` is omitted. + /// Terms for `fused_score` or `score_policy`; defaults to [`DEFAULT_FUSE`]. pub fused: Option>, /// Hard constraints applied before policy selection. pub eligibility: Option, - /// Sampling parameters fixed fleet-wide for this model, and what happens - /// to a request that sends a different value: a 400 before admission, or - /// the client value forwarded untouched. Either way the configured value - /// is injected when the request omits the field — see - /// [`SamplingOverrides`]. Empty (default) preserves today's behavior. + /// Fleet sampling defaults and conflict behavior. See [`SamplingOverrides`]. pub sampling_overrides: SamplingOverrides, } @@ -458,9 +397,7 @@ pub struct CacheAwareConfig { pub kv_indexer_endpoint: Option, } -/// Default routing-key header for the sticky policy. The `x-sgl-` prefix -/// matches the router's other emitted/consumed metadata headers -/// (`x-sgl-decode-url`, `x-sgl-router-error-code`). +/// Default request header for sticky routing. pub const DEFAULT_STICKY_HEADER: &str = "x-sgl-routing-key"; /// Default request header for session-aware routing. @@ -520,66 +457,23 @@ pub struct AffinityConfig { pub cache_candidate_ratio: f64, pub cache_candidate_max_workers: usize, pub cache_switch_margin_tokens: u64, - /// Queue gate (`--worker-queue-limit`): a worker whose engine reports at - /// least this many *waiting* requests cannot win a selection on cache - /// affinity — the request goes to another worker holding the same - /// prefix, or failing that to the least-loaded worker that is not - /// queueing. `None` disables the gate. - /// - /// Gating on the queue rather than on total depth is what makes this - /// targeted: `num_waiting_reqs` IS the question the request cares about - /// — will I sit behind other work before my prefill starts — whereas - /// depth only proxies it, and proxies it badly (an engine can queue at - /// 7-8 running on long-prompt traffic, far below its running cap, so a - /// depth threshold either fires on healthy busy workers or misses the - /// workers actually making requests wait). - /// - /// The gate reads the engine-published load sample and fails OPEN on a - /// worker with no fresh sample: the router-side in-flight counter cannot - /// separate a running request from a waiting one, so there is no honest - /// substitute to compare the limit against. - /// - /// Note the firing point scales with `dp_size`: the sample sums `waiting` - /// across a worker's DP ranks while a request lands on one of them, so - /// scale the limit with `--dp-size` on DP-attention deployments. - /// - /// The companion `saturation_queue_floor` cancels the gate's diversions - /// when they have no payoff (nothing in the fleet reads below the - /// floor). + /// Waiting-request limit for cache affinity; `None` disables. + /// Gates on the engine-published *waiting* count rather than total depth because + /// waiting is the question the request cares about — will it sit behind other work — + /// while depth proxies it badly (an engine can queue far below its running cap on + /// long-prompt traffic). Fails open without a fresh sample: the router-side + /// in-flight counter cannot separate running from waiting requests. + /// Counts sum across DP ranks, so scale the limit with `dp_size`. pub worker_queue_limit: Option, - /// Saturation pin (`--saturation-queue-floor`): cancels queue-gate - /// diversions that have no payoff. When no cache candidate survives - /// both `worker_queue_limit` and hard admission, at least one was over - /// the limit, AND no worker in the routable fleet has a fresh queue - /// reading strictly below this floor, the diverted request would wait - /// wherever it lands — so it pins to the least-pressured prefix owner - /// instead of cold-prefilling on a non-owner (which evicts other - /// prefixes and manufactures the next round of misses). `None` — the - /// default — preserves the pure gate behavior. - /// - /// Polarity note: a worker with no fresh sample does NOT count as idle - /// — the opposite of the gate's fail-open, and deliberately so. The - /// gate keeps affinity because that is the safe default action; the - /// pin asks whether a *provably better* destination exists, and an - /// unknown queue is not proof. Both polarities leave the request with - /// its prefix owner when the signal is missing. - /// - /// The CLI enforces `floor <= worker_queue_limit` and requires the - /// gate; like the limit, scale the floor with `dp_size`. + /// Keep the least-pressured prefix owner when the queue gate rejects all admitted + /// cache candidates and no fresh fleet queue is below this floor. Unknown queues + /// do not count as idle — the opposite of the gate's fail-open, deliberately: the + /// pin asks whether a provably better destination exists, and an unknown queue is + /// not proof. Requires `floor <= worker_queue_limit`; scale with `dp_size`. pub saturation_queue_floor: Option, - /// Number of random candidates sampled for the min-load fallback - /// (`--min-load-choices`); the least-pressured of the sample wins. - /// [`DEFAULT_MIN_LOAD_CHOICES`] is the pre-existing power-of-2 - /// behavior, so upgrading changes nothing. `k >= pool` skips the - /// shuffle and returns the exact minimum, with ties broken randomly - /// (an idle fleet ties on every comparison, so a fixed order would pin - /// every fallback dispatch to one worker); `k = 1` is a uniform draw - /// within the tier, and its sample has no second member, so the - /// proposal carries no backup and admission loses its backup-admission - /// and pressure-guard paths. The - /// `--cache-candidate-*` knobs bound the cache-affinity OWNER candidate - /// set; this bounds the min-load FALLBACK sample used when no owner is - /// usable. + /// Min-load fallback sample size; defaults to power-of-two. At least the pool + /// size chooses the exact minimum with random ties; 1 draws uniformly without + /// a backup for admission or pressure guards. Separate from cache-owner limits. pub min_load_choices: usize, } @@ -610,18 +504,13 @@ impl Default for AffinityConfig { } } -/// Per-model sticky-session tuning. Built from the `--routing-key-header` -/// / `--sticky-*` flags by [`crate::config::cli::Cli::into_config`], which -/// also validates that `header_name` parses as an HTTP header name. +/// Sticky routing settings; the CLI validates the header name and positive durations. #[derive(Debug, Clone)] pub struct StickyConfig { /// Request header carrying the routing key. Validated to parse as a /// `http::HeaderName` at config-build time. pub header_name: String, - /// Policy used to pick a worker when a request has no routing key, and - /// to pick the initial worker when a new key is first seen. One of - /// `round_robin` / `random` / `power_of_two` / `load_based` — the - /// dependency-free policies the factory can build standalone. + /// Fallback for new or missing routing keys. pub fallback_policy: StickyFallbackKind, /// Evict an assignment after it has been idle (unreferenced) this many /// seconds. Bounds the map against unbounded routing-key cardinality. @@ -650,10 +539,7 @@ impl Default for StickyConfig { #[derive(Debug, Clone)] pub struct CircuitBreakerConfig { - /// 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". + /// Consecutive failures before opening the breaker; zero is invalid. pub threshold: NonZeroU32, pub cool_down_secs: u64, } @@ -670,43 +556,15 @@ pub enum DiscoveryBackend { K8s(K8sDiscoveryConfig), } -/// Fixed list of worker URLs. Each URL is registered once at startup; -/// `mode`, `model_ids`, and `bootstrap_port` are resolved per-worker -/// from `/server_info` (see [`crate::workers::introspect`]). -/// -/// No file watcher, no hot-reload: topology change requires a restart. +/// Workers registered at startup. Roles, models, and bootstrap ports come from +/// `/server_info`; topology changes require a restart. #[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 flags are set: -/// -/// 1. **Plain** — all matched workers share the same role: -/// `--service-discovery-namespace default --selector app=sglang` -/// -/// 2. **PD disaggregation** — prefill and decode workers are separated by -/// different selectors: -/// `--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 -/// `WorkerMode` and `bootstrap_port` for each worker are filled in by -/// the worker manager from each worker's `/server_info` introspection, -/// so PD works without any pod-level annotations — see -/// [`crate::workers::introspect`] for the `disaggregation_mode` and -/// `disaggregation_bootstrap_port` extraction. -/// -/// [`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. +/// Kubernetes EndpointSlice discovery. Selectors classify slices; worker roles +/// and bootstrap ports come from `/server_info` introspection. #[derive(Debug, Clone)] pub struct K8sDiscoveryConfig { pub namespace: String, @@ -714,14 +572,8 @@ pub struct K8sDiscoveryConfig { pub mode: K8sDiscoveryMode, } -/// 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; -/// PD: empty, with classification done client-side per slice), and -/// * assign each `EndpointSlice` a [`crate::discovery::WorkerMode`] in -/// `extract_workers`. +/// Validated selector mode. Plain selectors run server-side; PD selectors +/// classify EndpointSlices client-side. #[derive(Debug, Clone, PartialEq, Eq)] pub enum K8sDiscoveryMode { /// One global label selector; every matched EndpointSlice becomes a @@ -774,52 +626,31 @@ pub enum ConfigError { IdenticalPdSelectors, } -/// Returns `true` when `selector` has zero non-empty terms after -/// trimming and splitting on `,`. `labels_match_selector` then returns -/// `true` for every label set, which is the "matches everything" -/// degenerate case PD mode must reject. +/// An empty selector matches every slice, which is invalid for a PD role. fn is_selector_empty(selector: &str) -> bool { selector.split(',').all(|t| t.trim().is_empty()) } -/// Canonicalize a comma-separated equality selector to a sorted list of -/// parsed `(key, value)` tuples. Comparison happens at the parsed-term -/// level — *not* the raw string level — because `labels_match_selector` -/// already strips whitespace and treats `key=value` and `key==value` as -/// the same equality test. Comparing raw strings would let -/// `"app=sglang"` vs `"app==sglang"` (and `"app = sglang"` vs -/// `"app=sglang"`) past the identical-selector check, even though -/// `classify_mode` would treat them identically at runtime — exactly -/// the silent decode-pool-empty failure mode this check exists to -/// prevent. -/// -/// Returns an empty `Vec` for selectors with no parseable terms -/// (whitespace-only, comma-only, or any term that doesn't match the -/// `key=value` / `key==value` grammar). Callers must run -/// [`is_equality_selector`] before this to surface malformed -/// selectors as `UnsupportedSelectorGrammar`. -fn canonical_selector(selector: &str) -> Vec<(String, String)> { - let mut terms: Vec<(String, String)> = selector +/// Normalize validated equality terms for comparison, matching runtime whitespace +/// and `=`/`==` handling. Term order does not affect matching. +fn canonical_selector(selector: &str) -> Vec<(&str, &str)> { + let mut terms: Vec<_> = selector .split(',') .filter_map(|raw| { let term = raw.trim(); if term.is_empty() { return None; } - // Mirror `labels_match_selector`: prefer the `==` alias so a - // term like `key==value` parses to `(key, value)` instead of - // `(key, =value)`. + // Prefer `==` so its second equals sign does not become part of the value. let (k, v) = term.split_once("==").or_else(|| term.split_once('='))?; - Some((k.trim().to_string(), v.trim().to_string())) + Some((k.trim(), v.trim())) }) .collect(); - terms.sort(); + terms.sort_unstable(); terms } -/// Returns `true` when `selector` parses as a comma-separated equality -/// selector — every term has the shape `key=value` or `key==value`. -/// See [`ConfigError::UnsupportedSelectorGrammar`] for rationale. +/// Check the equality-only grammar supported by client-side PD matching. fn is_equality_selector(selector: &str) -> bool { for term in selector.split(',') { let term = term.trim(); @@ -835,9 +666,7 @@ fn is_equality_selector(selector: &str) -> bool { continue; } if let Some((k, _value)) = term.split_once('=') { - // Reject `!=` (rendered as `key!` + `=value` by split_once). - // Empty value is legal in K8s — `label_selector = "tier="` - // matches pods with `tier=""` — so we don't constrain it. + // Reject `!=`; empty label values are valid. if k.trim().is_empty() || k.trim().ends_with('!') { return false; } @@ -849,10 +678,7 @@ fn is_equality_selector(selector: &str) -> bool { true } -/// 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`]. +/// Resolve and validate the plain or prefill/decode selector combination. pub fn resolve_mode( label_selector: Option<&str>, prefill_selector: Option<&str>, @@ -860,57 +686,30 @@ pub fn resolve_mode( ) -> 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. + // Plain selectors run on the Kubernetes API, which supports the full grammar. + // PD selectors are checked client-side and only support equality. 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(), - }); + // Validate both grammars before checking for empty or identical selectors. + let selectors = [("prefill", prefill), ("decode", decode)]; + for (selector, value) in selectors { + if !is_equality_selector(value) { + return Err(ConfigError::UnsupportedSelectorGrammar { + selector, + value: value.to_string(), + }); + } } - if !is_equality_selector(decode) { - return Err(ConfigError::UnsupportedSelectorGrammar { - selector: "decode", - value: decode.to_string(), - }); + // Empty PD selectors match everything and starve the opposite role. + for (selector, value) in selectors { + if is_selector_empty(value) { + return Err(ConfigError::EmptyPdSelector { selector }); + } } - // 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. + // Prefill wins when both selectors match, leaving decode empty. if canonical_selector(prefill) == canonical_selector(decode) { return Err(ConfigError::IdenticalPdSelectors); } @@ -931,10 +730,6 @@ mod k8s_discovery_config_tests { #[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 = resolve_mode(None, Some("app=sglang,role=p"), Some("app=sglang,role=d")) .expect("PD mode is now valid"); assert_eq!( @@ -948,9 +743,6 @@ mod k8s_discovery_config_tests { #[test] fn mode_pd_rejects_set_based_prefill_selector() { - // 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 = resolve_mode(None, Some("app in (sglang, vllm)"), Some("app=sglang")).unwrap_err(); assert!( @@ -992,12 +784,8 @@ 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 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. + /// Plain selectors run on the Kubernetes API, which supports the full grammar. + /// PD selectors are checked client-side and only support equality. #[test] fn mode_accepts_set_based_selector_in_plain_mode() { let m = resolve_mode(Some("app in (sglang,sglang-small)"), None, None) @@ -1010,9 +798,6 @@ mod k8s_discovery_config_tests { ); } - /// `notin`, presence (`key`), absence (`!key`), and inequality (`!=`) - /// are all valid K8s server-side selector grammar — plain mode must - /// pass them through. #[test] fn mode_accepts_other_set_based_forms_in_plain_mode() { for raw in [ @@ -1033,15 +818,6 @@ mod k8s_discovery_config_tests { } } - /// PD mode evaluates selectors *client-side* via - /// `labels_match_selector`, which only handles equality. A set-based - /// PD selector would silently match zero pods → fail-fast at load. - /// Pins the plain-server-side / PD-client-side asymmetry: relaxing - /// the grammar check for plain (see `mode_accepts_set_based_*` - /// above) must not accidentally relax it for PD selectors. Uses - /// `notin` so this test covers a different set-based form than - /// `mode_pd_rejects_set_based_prefill_selector` (which uses `in`) - /// — both must keep failing. #[test] fn mode_pd_rejects_notin_prefill_selector() { let err = @@ -1101,9 +877,6 @@ mod k8s_discovery_config_tests { ); } - /// Empty plain `label_selector` is valid — matches every - /// EndpointSlice in the namespace (documented K8s behavior; the - /// operator opts in by setting plain mode at all). #[test] fn mode_accepts_empty_plain_label_selector() { let m = resolve_mode(Some(""), None, None).unwrap(); @@ -1115,12 +888,6 @@ mod k8s_discovery_config_tests { ); } - /// PD mode is the *opposite* of plain: an empty selector would match - /// every EndpointSlice, and since `classify_mode` checks prefill - /// before decode, an empty `prefill_selector` would classify - /// everything as Prefill — decode pool stays empty and the resolver - /// surfaces the wrong `no_decode_workers_available` error. Fail-fast - /// at config load. #[test] fn mode_pd_rejects_empty_prefill_selector() { let err = resolve_mode(None, Some(""), Some("role=decode")).unwrap_err(); @@ -1144,9 +911,6 @@ mod k8s_discovery_config_tests { ); } - /// Whitespace-only / comma-only PD selector parses to zero terms in - /// `labels_match_selector` and matches every slice at runtime — same - /// failure mode as a literal empty string. #[test] fn mode_pd_rejects_whitespace_only_prefill_selector() { let err = resolve_mode(None, Some(" , "), Some("role=decode")).unwrap_err(); @@ -1161,9 +925,6 @@ mod k8s_discovery_config_tests { ); } - /// Identical prefill and decode selectors degrade silently: every - /// slice matches both, but `classify_mode` returns `Prefill` first, - /// so the decode pool stays empty. #[test] fn mode_pd_rejects_identical_prefill_and_decode_selectors() { let err = resolve_mode(None, Some("app=sglang"), Some("app=sglang")).unwrap_err(); @@ -1184,14 +945,6 @@ mod k8s_discovery_config_tests { ); } - /// `labels_match_selector` accepts both `key=value` and `key==value` - /// for equality and parses them to the same `(key, value)` tuple. - /// Two selectors that differ only in this alias choice are runtime- - /// equivalent — they'd match the same EndpointSlices, then - /// `classify_mode`'s prefill-first ordering would funnel every slice - /// into Prefill, leaving decode empty. The check must canonicalize - /// at the term level (parsed `(key, value)` tuples), not the raw - /// string level. #[test] fn mode_pd_rejects_identical_selectors_under_eq_alias() { let err = resolve_mode(None, Some("app=sglang"), Some("app==sglang")).unwrap_err(); @@ -1201,11 +954,6 @@ mod k8s_discovery_config_tests { ); } - /// Inner whitespace inside a term (`"app = sglang"`) is the same - /// label as no whitespace (`"app=sglang"`) — the runtime - /// `labels_match_selector` trims key and value independently - /// (see `key.trim()` / `expected.trim()` in `k8s.rs`). Canonical - /// form must agree. #[test] fn mode_pd_rejects_identical_selectors_under_inner_whitespace() { let err = resolve_mode(None, Some("app=sglang"), Some("app = sglang")).unwrap_err(); @@ -1215,11 +963,6 @@ mod k8s_discovery_config_tests { ); } - /// Term order doesn't matter for label matching, so `"a=1,b=2"` and - /// `"b=2,a=1"` must be treated as identical. (Implied by the sort - /// in `canonical_selector`, but pinned explicitly so a future - /// "preserve user order for diagnostics" refactor can't silently - /// reintroduce the silent-failure bug.) #[test] fn mode_pd_rejects_identical_selectors_under_term_order_permutation() { let err = @@ -1230,9 +973,6 @@ mod k8s_discovery_config_tests { ); } - /// Sanity: two selectors that genuinely differ at the term level - /// must still pass validation — the canonicalizer must not be so - /// aggressive that it false-positives on legitimate PD configs. #[test] fn mode_pd_accepts_truly_distinct_selectors() { let m = resolve_mode(