From 7fce436f030869ca9ff11122c3420920d5b684e4 Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 13:36:45 -0700 Subject: [PATCH 01/27] feat(rl): weight version ordering and the version-policy grammar Signed-off-by: key4ng --- crates/rl/src/lib.rs | 4 + crates/rl/src/policy.rs | 177 +++++++++++++++++++++++++++++++++++++++ crates/rl/src/version.rs | 133 +++++++++++++++++++++++++++++ 3 files changed, 314 insertions(+) create mode 100644 crates/rl/src/policy.rs create mode 100644 crates/rl/src/version.rs diff --git a/crates/rl/src/lib.rs b/crates/rl/src/lib.rs index a4b1895e53..acfb58c00b 100644 --- a/crates/rl/src/lib.rs +++ b/crates/rl/src/lib.rs @@ -9,11 +9,13 @@ pub mod error; pub mod fanout; pub mod metrics; pub mod path; +pub mod policy; pub mod proxy; pub mod selector; pub mod state; #[cfg(test)] pub(crate) mod testing; +pub mod version; pub mod view; use std::sync::Arc; @@ -22,7 +24,9 @@ use axum::{routing::get, Router}; pub use config::RlConfig; pub use error::RlError; pub use metrics::init_rl_metrics; +pub use policy::{VersionPolicy, VersionPolicyError, VERSION_POLICY_HEADER}; pub use state::RlState; +pub use version::Version; pub use view::{RlWorkerInfo, RlWorkerView}; /// Build the `/v1/rl` router. `with_state` returns `Router` for any `S`, diff --git a/crates/rl/src/policy.rs b/crates/rl/src/policy.rs new file mode 100644 index 0000000000..90363a0cb3 --- /dev/null +++ b/crates/rl/src/policy.rs @@ -0,0 +1,177 @@ +//! The version-policy grammar shared by `--rl-version-policy` and the +//! `x-smg-version-policy` request header. + +use std::{fmt, str::FromStr}; + +use http::HeaderMap; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use thiserror::Error; + +use crate::version::Version; + +/// Per-request override of the configured policy, same grammar. +pub const VERSION_POLICY_HEADER: &str = "x-smg-version-policy"; + +/// Which engines a request may be routed to, judged against the model's +/// fleet maximum: `any | latest-only | min-version: | max-staleness:`. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum VersionPolicy { + #[default] + Any, + LatestOnly, + MinVersion(Version), + MaxStaleness(u64), +} + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +#[error("invalid version policy `{input}`: {reason}")] +pub struct VersionPolicyError { + pub input: String, + pub reason: &'static str, +} + +impl VersionPolicy { + /// The header override, if the request carries one. `Err` when it is + /// present but malformed (including non-UTF-8 bytes). + pub fn from_headers(headers: &HeaderMap) -> Result, VersionPolicyError> { + let Some(value) = headers.get(VERSION_POLICY_HEADER) else { + return Ok(None); + }; + let text = value.to_str().map_err(|_| VersionPolicyError { + input: String::from_utf8_lossy(value.as_bytes()).into_owned(), + reason: "header value is not UTF-8", + })?; + text.trim().parse().map(Some) + } +} + +impl FromStr for VersionPolicy { + type Err = VersionPolicyError; + + fn from_str(s: &str) -> Result { + let err = |reason| VersionPolicyError { + input: s.to_string(), + reason, + }; + match s { + "any" => Ok(Self::Any), + "latest-only" => Ok(Self::LatestOnly), + _ => { + if let Some(v) = s.strip_prefix("min-version:") { + if v.is_empty() || v.contains(char::is_whitespace) { + return Err(err("min-version needs a version with no whitespace")); + } + return Ok(Self::MinVersion(Version::parse(v))); + } + if let Some(k) = s.strip_prefix("max-staleness:") { + return k + .parse::() + .map(Self::MaxStaleness) + .map_err(|_| err("max-staleness needs a decimal integer")); + } + Err(err( + "expected any | latest-only | min-version: | max-staleness:", + )) + } + } + } +} + +impl fmt::Display for VersionPolicy { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Any => f.write_str("any"), + Self::LatestOnly => f.write_str("latest-only"), + Self::MinVersion(v) => write!(f, "min-version:{v}"), + Self::MaxStaleness(k) => write!(f, "max-staleness:{k}"), + } + } +} + +impl Serialize for VersionPolicy { + fn serialize(&self, serializer: S) -> Result { + serializer.collect_str(self) + } +} + +impl<'de> Deserialize<'de> for VersionPolicy { + fn deserialize>(deserializer: D) -> Result { + String::deserialize(deserializer)? + .parse() + .map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use http::HeaderValue; + + use super::*; + + #[test] + fn grammar_round_trips() { + for text in [ + "any", + "latest-only", + "min-version:42", + "min-version:step-3", + "max-staleness:2", + ] { + let policy: VersionPolicy = text.parse().unwrap(); + assert_eq!(policy.to_string(), text); + let json = serde_json::to_string(&policy).unwrap(); + assert_eq!(json, format!("\"{text}\"")); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + policy + ); + } + assert_eq!(VersionPolicy::default(), VersionPolicy::Any); + assert_eq!( + "max-staleness:1".parse::().unwrap(), + VersionPolicy::MaxStaleness(1) + ); + } + + #[test] + fn rejects_bad_inputs() { + for bad in [ + "", + "Any", + "latest", + "min-version:", + "min-version:a b", + "max-staleness:x", + "max-staleness:-1", + "max-staleness:", + ] { + let err = bad.parse::().unwrap_err(); + assert_eq!(err.input, bad); + } + assert!(serde_json::from_str::("\"nope\"").is_err()); + } + + #[test] + fn header_is_optional_and_validated() { + let mut headers = HeaderMap::new(); + assert_eq!(VersionPolicy::from_headers(&headers).unwrap(), None); + headers.insert( + VERSION_POLICY_HEADER, + HeaderValue::from_static(" latest-only "), + ); + assert_eq!( + VersionPolicy::from_headers(&headers).unwrap(), + Some(VersionPolicy::LatestOnly) + ); + headers.insert(VERSION_POLICY_HEADER, HeaderValue::from_static("bogus")); + assert!(VersionPolicy::from_headers(&headers).is_err()); + headers.insert( + VERSION_POLICY_HEADER, + HeaderValue::from_bytes(b"\xff").unwrap(), + ); + assert_eq!( + VersionPolicy::from_headers(&headers).unwrap_err().reason, + "header value is not UTF-8" + ); + } +} diff --git a/crates/rl/src/version.rs b/crates/rl/src/version.rs new file mode 100644 index 0000000000..7ff562c562 --- /dev/null +++ b/crates/rl/src/version.rs @@ -0,0 +1,133 @@ +//! Weight versions as engines report them, with the ordering the pre-filter +//! and the per-model fleet maximum use. + +use std::{ + cmp::Ordering, + fmt, + sync::{ + atomic::{AtomicBool, Ordering as AtomicOrdering}, + Arc, + }, +}; + +use tracing::warn; + +/// The label an engine reports before any refit; seeds as "no version". +pub const UNVERSIONED_LABEL: &str = "default"; + +/// A weight version: numeric when the text parses as `u64`, opaque text +/// otherwise. Numeric versions compare by value, text versions lexically, +/// and a numeric/text pair lexically with one warning per process. +#[derive(Clone, Debug)] +pub struct Version { + raw: Arc, + numeric: Option, +} + +static MIXED_ORDERING_WARNED: AtomicBool = AtomicBool::new(false); + +impl Version { + pub fn parse(raw: &str) -> Self { + let raw = raw.trim(); + Self { + raw: Arc::from(raw), + numeric: raw.parse::().ok(), + } + } + + /// The version a registration label carries: `None` for an absent or + /// empty label and for [`UNVERSIONED_LABEL`]. + pub fn from_label(label: Option<&str>) -> Option { + let label = label?.trim(); + if label.is_empty() || label == UNVERSIONED_LABEL { + return None; + } + Some(Self::parse(label)) + } + + pub fn as_str(&self) -> &str { + &self.raw + } + + pub fn numeric(&self) -> Option { + self.numeric + } +} + +impl fmt::Display for Version { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.raw) + } +} + +impl PartialEq for Version { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Eq for Version {} + +impl PartialOrd for Version { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Version { + fn cmp(&self, other: &Self) -> Ordering { + match (self.numeric, other.numeric) { + (Some(a), Some(b)) => a.cmp(&b), + (None, None) => self.raw.cmp(&other.raw), + _ => { + if !MIXED_ORDERING_WARNED.swap(true, AtomicOrdering::Relaxed) { + warn!( + target: "smg_rl", + a = %self.raw, b = %other.raw, + "comparing a numeric weight version with a text one; using lexical order" + ); + } + self.raw.cmp(&other.raw) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn numeric_versions_compare_by_value() { + assert!(Version::parse("10") > Version::parse("9")); + assert_eq!(Version::parse("042"), Version::parse("42")); + assert_eq!(Version::parse(" 7 ").as_str(), "7"); + assert_eq!(Version::parse("7").numeric(), Some(7)); + } + + #[test] + fn text_versions_compare_lexically() { + assert!(Version::parse("step-9") > Version::parse("step-10")); + assert_eq!(Version::parse("abc").numeric(), None); + assert_ne!(Version::parse("abc"), Version::parse("abd")); + } + + #[test] + fn mixed_versions_fall_back_to_lexical_order() { + // '1' (0x31) sorts before 'd' (0x64). + assert!(Version::parse("1") < Version::parse("default")); + assert!(Version::parse("1") < Version::parse("1a")); + } + + #[test] + fn labels_seed_versions_except_the_placeholder() { + assert_eq!(Version::from_label(None), None); + assert_eq!(Version::from_label(Some("")), None); + assert_eq!(Version::from_label(Some("default")), None); + assert_eq!( + Version::from_label(Some("42")).map(|v| v.numeric()), + Some(Some(42)) + ); + assert_eq!(format!("{}", Version::parse("v3")), "v3"); + } +} From 89d08d6b32ec17e2536ef43435328fa6d8cda057 Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 13:48:13 -0700 Subject: [PATCH 02/27] feat(rl): version and control-state table with eviction sink and metrics Signed-off-by: key4ng --- Cargo.lock | 1 + crates/rl/Cargo.toml | 1 + crates/rl/src/config.rs | 23 +++ crates/rl/src/fanout.rs | 1 + crates/rl/src/lib.rs | 2 + crates/rl/src/metrics.rs | 56 ++++- crates/rl/src/proxy.rs | 1 + crates/rl/src/state.rs | 102 ++++++++- crates/rl/src/table.rs | 436 +++++++++++++++++++++++++++++++++++++++ crates/rl/src/view.rs | 7 + 10 files changed, 618 insertions(+), 12 deletions(-) create mode 100644 crates/rl/src/table.rs diff --git a/Cargo.lock b/Cargo.lock index 08a9784201..02b8791a3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7730,6 +7730,7 @@ name = "smg-rl" version = "0.1.0" dependencies = [ "axum", + "dashmap", "futures", "http", "http-body-util", diff --git a/crates/rl/Cargo.toml b/crates/rl/Cargo.toml index fb4bfa0bba..26059294a7 100644 --- a/crates/rl/Cargo.toml +++ b/crates/rl/Cargo.toml @@ -17,6 +17,7 @@ name = "smg_rl" [dependencies] openai-protocol.workspace = true axum = { workspace = true, features = ["macros"] } +dashmap.workspace = true http.workspace = true reqwest = { workspace = true, features = ["rustls", "json", "http2"] } serde = { workspace = true, features = ["derive"] } diff --git a/crates/rl/src/config.rs b/crates/rl/src/config.rs index d817cc58cc..a0a6b9ef3e 100644 --- a/crates/rl/src/config.rs +++ b/crates/rl/src/config.rs @@ -2,6 +2,8 @@ use serde::{Deserialize, Serialize}; +use crate::policy::VersionPolicy; + /// Configuration for the RL control plane. Inert unless `enabled`. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(default)] @@ -12,6 +14,9 @@ pub struct RlConfig { pub control_timeout_secs: u64, /// Maximum concurrent engine calls during a fan-out. pub fanout_concurrency: usize, + /// The default version policy, overridable per-request by the + /// `x-smg-version-policy` header. + pub version_policy: VersionPolicy, } impl Default for RlConfig { @@ -20,6 +25,7 @@ impl Default for RlConfig { enabled: false, control_timeout_secs: 600, fanout_concurrency: 32, + version_policy: VersionPolicy::Any, } } } @@ -28,6 +34,9 @@ impl RlConfig { /// Validate the values that matter when the control plane is enabled. pub fn validate(&self) -> Result<(), String> { if !self.enabled { + if self.version_policy != VersionPolicy::Any { + return Err("rl.version_policy requires rl.enabled".to_string()); + } return Ok(()); } if self.control_timeout_secs == 0 { @@ -50,6 +59,7 @@ mod tests { assert!(!cfg.enabled); assert_eq!(cfg.control_timeout_secs, 600); assert_eq!(cfg.fanout_concurrency, 32); + assert_eq!(cfg.version_policy, VersionPolicy::Any); assert_eq!(cfg.validate(), Ok(())); } @@ -68,6 +78,7 @@ mod tests { enabled: true, control_timeout_secs: 0, fanout_concurrency: 32, + ..RlConfig::default() }; assert!(cfg.validate().unwrap_err().contains("control_timeout_secs")); cfg.control_timeout_secs = 600; @@ -76,4 +87,16 @@ mod tests { cfg.enabled = false; assert_eq!(cfg.validate(), Ok(())); } + + #[test] + fn version_policy_needs_the_flag() { + let mut cfg = RlConfig { + enabled: false, + version_policy: VersionPolicy::LatestOnly, + ..RlConfig::default() + }; + assert!(cfg.validate().unwrap_err().contains("version_policy")); + cfg.enabled = true; + assert_eq!(cfg.validate(), Ok(())); + } } diff --git a/crates/rl/src/fanout.rs b/crates/rl/src/fanout.rs index 0ba22557f5..c62cd3bc6d 100644 --- a/crates/rl/src/fanout.rs +++ b/crates/rl/src/fanout.rs @@ -173,6 +173,7 @@ mod tests { enabled: true, control_timeout_secs: 5, fanout_concurrency: concurrency, + ..RlConfig::default() }; Arc::new(RlState::new(Arc::new(FakeView(workers)), cfg)) } diff --git a/crates/rl/src/lib.rs b/crates/rl/src/lib.rs index acfb58c00b..f6083e9dbd 100644 --- a/crates/rl/src/lib.rs +++ b/crates/rl/src/lib.rs @@ -13,6 +13,7 @@ pub mod policy; pub mod proxy; pub mod selector; pub mod state; +pub mod table; #[cfg(test)] pub(crate) mod testing; pub mod version; @@ -26,6 +27,7 @@ pub use error::RlError; pub use metrics::init_rl_metrics; pub use policy::{VersionPolicy, VersionPolicyError, VERSION_POLICY_HEADER}; pub use state::RlState; +pub use table::{ControlState, RlTable, VersionEvictionSink, VersionSource}; pub use version::Version; pub use view::{RlWorkerInfo, RlWorkerView}; diff --git a/crates/rl/src/metrics.rs b/crates/rl/src/metrics.rs index 6874999973..12299883fb 100644 --- a/crates/rl/src/metrics.rs +++ b/crates/rl/src/metrics.rs @@ -4,12 +4,18 @@ use std::time::Duration; -use metrics::{counter, describe_counter, describe_histogram, histogram}; +use metrics::{counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram}; const CONTROL_CALLS: &str = "smg_rl_control_calls_total"; const CONTROL_CALL_DURATION: &str = "smg_rl_control_call_duration_seconds"; const FANOUT: &str = "smg_rl_fanout_total"; const FANOUT_DURATION: &str = "smg_rl_fanout_duration_seconds"; +const WORKER_WEIGHT_VERSION: &str = "smg_rl_worker_weight_version"; +const WORKER_CONTROL_STATE: &str = "smg_rl_worker_control_state"; +const CANDIDATES_FILTERED: &str = "smg_rl_candidates_filtered_total"; +const REQUESTS_UNROUTABLE: &str = "smg_rl_requests_unroutable_total"; +const MIXED_VERSION: &str = "smg_rl_mixed_version_total"; +const VERSION_CHANGE_WITH_INFLIGHT: &str = "smg_rl_version_change_with_inflight_total"; /// Known engine control operations. Anything else is reported as `other` /// so a caller cannot mint unbounded label values by inventing paths. @@ -48,6 +54,30 @@ pub fn init_rl_metrics() { ); describe_counter!(FANOUT, "RL fan-out requests, by result"); describe_histogram!(FANOUT_DURATION, "Wall time of one RL fan-out request"); + describe_gauge!( + WORKER_WEIGHT_VERSION, + "Weight version SMG believes each engine holds (numeric versions only)" + ); + describe_gauge!( + WORKER_CONTROL_STATE, + "Control state SMG believes each engine is in: 0 active, 1 paused, 2 asleep" + ); + describe_counter!( + CANDIDATES_FILTERED, + "Routing candidates dropped by the RL pre-filter, by reason" + ); + describe_counter!( + REQUESTS_UNROUTABLE, + "Worker selections whose RL-eligible set was empty" + ); + describe_counter!( + MIXED_VERSION, + "Buffered /generate responses that spanned more than one weight version" + ); + describe_counter!( + VERSION_CHANGE_WITH_INFLIGHT, + "Weight version changes observed while the engine had requests in flight" + ); } /// Bounded metric label for an engine path: the whole path when it is a @@ -70,6 +100,30 @@ pub fn record_fanout(result: &'static str, elapsed: Duration) { histogram!(FANOUT_DURATION).record(elapsed.as_secs_f64()); } +pub fn set_worker_weight_version(worker: &str, version: u64) { + gauge!(WORKER_WEIGHT_VERSION, "worker" => worker.to_string()).set(version as f64); +} + +pub fn set_worker_control_state(worker: &str, value: f64) { + gauge!(WORKER_CONTROL_STATE, "worker" => worker.to_string()).set(value); +} + +pub fn record_candidate_filtered(reason: &'static str) { + counter!(CANDIDATES_FILTERED, "reason" => reason).increment(1); +} + +pub fn record_request_unroutable() { + counter!(REQUESTS_UNROUTABLE).increment(1); +} + +pub fn record_mixed_version() { + counter!(MIXED_VERSION).increment(1); +} + +pub fn record_version_change_with_inflight() { + counter!(VERSION_CHANGE_WITH_INFLIGHT).increment(1); +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/rl/src/proxy.rs b/crates/rl/src/proxy.rs index 898e118237..ad267baa07 100644 --- a/crates/rl/src/proxy.rs +++ b/crates/rl/src/proxy.rs @@ -283,6 +283,7 @@ mod tests { enabled: true, control_timeout_secs: timeout_secs, fanout_concurrency: 4, + ..RlConfig::default() }; Arc::new(RlState::new(Arc::new(FakeView(workers)), cfg)) } diff --git a/crates/rl/src/state.rs b/crates/rl/src/state.rs index 6789e4db40..4ecdf6f14d 100644 --- a/crates/rl/src/state.rs +++ b/crates/rl/src/state.rs @@ -1,23 +1,103 @@ -//! Process-wide state for the RL control plane: the registry view and the -//! configuration. +//! Process-wide state for the RL control plane: the registry view, the +//! configuration, and the version/control-state table. -use std::sync::Arc; +use std::{borrow::Cow, sync::Arc}; -use crate::{config::RlConfig, view::RlWorkerView}; +use http::HeaderMap; +use tracing::info; + +use crate::{ + config::RlConfig, + metrics, + policy::VersionPolicy, + table::{ControlState, NoopEvictionSink, RlTable, VersionEvictionSink, VersionSource}, + version::Version, + view::{RlWorkerInfo, RlWorkerView}, +}; pub struct RlState { pub(crate) view: Arc, pub(crate) config: RlConfig, + pub(crate) table: RlTable, } impl RlState { - /// Build the state. The control plane owns no HTTP client: every call - /// goes through the client the gateway negotiated for the target worker - /// (see [`crate::RlWorkerInfo::http_client`]), with the control deadline - /// applied per request. Breakers and load counters live on the gateway's - /// worker objects, not on the client, so control calls leave them - /// untouched. + /// Build the state with a table that evicts nothing on version changes + /// (tests and crate-only use). The control plane owns no HTTP client: + /// every call goes through the client the gateway negotiated for the + /// target worker (see [`RlWorkerInfo::http_client`]), with the control + /// deadline applied per request. Breakers and load counters live on the + /// gateway's worker objects, not on the client, so control calls leave + /// them untouched. pub fn new(view: Arc, config: RlConfig) -> Self { - Self { view, config } + Self::with_sink(view, config, Arc::new(NoopEvictionSink)) + } + + /// Build the state with the gateway's eviction sink. + pub fn with_sink( + view: Arc, + config: RlConfig, + sink: Arc, + ) -> Self { + Self { + view, + config, + table: RlTable::new(sink), + } + } + + pub fn table(&self) -> &RlTable { + &self.table + } + + pub fn config(&self) -> &RlConfig { + &self.config + } + + /// The policy for one request: a valid header override, else the + /// configured policy. (The middleware already rejected invalid headers.) + pub fn policy_for(&self, headers: Option<&HeaderMap>) -> Cow<'_, VersionPolicy> { + match headers.and_then(|h| VersionPolicy::from_headers(h).ok().flatten()) { + Some(policy) => Cow::Owned(policy), + None => Cow::Borrowed(&self.config.version_policy), + } + } + + /// Record a version for `worker`'s engine, with the inflight accounting. + pub fn apply_version( + &self, + worker: &RlWorkerInfo, + version: Version, + source: VersionSource, + ) -> bool { + let changed = + self.table + .set_version(&worker.base_url, &worker.model_id, version.clone(), source); + if changed { + if self.view.inflight(&worker.base_url) > 0 { + metrics::record_version_change_with_inflight(); + } + info!( + target: "smg_rl", + base_url = %worker.base_url, version = %version, source = source.as_str(), + "rl.observe version" + ); + } + changed + } + + /// Record a control state for `worker`'s engine. + pub fn apply_control(&self, worker: &RlWorkerInfo, control: ControlState) -> bool { + let changed = self + .table + .set_control(&worker.base_url, &worker.model_id, control); + if changed { + info!( + target: "smg_rl", + base_url = %worker.base_url, control = control.as_str(), + "rl.observe control" + ); + } + changed } } diff --git a/crates/rl/src/table.rs b/crates/rl/src/table.rs new file mode 100644 index 0000000000..d01b96cb74 --- /dev/null +++ b/crates/rl/src/table.rs @@ -0,0 +1,436 @@ +//! The RL side table: per-engine weight version and control state, keyed by +//! engine base URL, plus the per-model fleet maximum the pre-filter compares +//! against. Labels are never written; this table is the mutable state. + +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; + +use dashmap::DashMap; + +use crate::{metrics, version::Version}; + +/// Whether an engine can serve requests, as last observed through SMG. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ControlState { + #[default] + Active, + Paused, + Asleep, +} + +impl ControlState { + pub fn as_str(self) -> &'static str { + match self { + Self::Active => "active", + Self::Paused => "paused", + Self::Asleep => "asleep", + } + } + + fn gauge_value(self) -> f64 { + match self { + Self::Active => 0.0, + Self::Paused => 1.0, + Self::Asleep => 2.0, + } + } +} + +/// How SMG learned an engine's version. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum VersionSource { + Registration, + Passthrough, + Api, +} + +impl VersionSource { + pub fn as_str(self) -> &'static str { + match self { + Self::Registration => "registration", + Self::Passthrough => "passthrough", + Self::Api => "api", + } + } +} + +#[derive(Clone, Debug)] +pub struct EngineState { + /// `Worker::model_id()` at seed time; the fleet maximum is per model. + pub model: Arc, + pub version: Option, + pub version_source: Option, + pub control: ControlState, +} + +/// Told when an engine's version changed so the gateway can forget what the +/// engine's cache held (coupling surface (a)). +pub trait VersionEvictionSink: Send + Sync { + fn on_version_changed(&self, model_id: &str, base_url: &str); +} + +/// A sink that does nothing; for crate-only construction and tests. +pub struct NoopEvictionSink; + +impl VersionEvictionSink for NoopEvictionSink { + fn on_version_changed(&self, _model_id: &str, _base_url: &str) {} +} + +/// Strip the `@` suffix a DP-aware worker URL carries, giving the +/// engine base URL the table is keyed by. +pub fn base_url_of(worker_url: &str) -> &str { + match worker_url.rsplit_once('@') { + Some((base, rank)) if !rank.is_empty() && rank.bytes().all(|b| b.is_ascii_digit()) => base, + _ => worker_url, + } +} + +pub struct RlTable { + entries: DashMap, EngineState>, + fleet_max: DashMap, Version>, + /// Entries whose control state is not `Active`; lets `any` skip the map. + inactive: AtomicUsize, + sink: Arc, +} + +impl RlTable { + pub fn new(sink: Arc) -> Self { + Self { + entries: DashMap::new(), + fleet_max: DashMap::new(), + inactive: AtomicUsize::new(0), + sink, + } + } + + /// Insert an entry from a registration label unless one exists. Returns + /// whether an entry was inserted. + pub fn seed(&self, base_url: &str, model: &str, label: Option<&str>) -> bool { + if self.entries.contains_key(base_url) { + return false; + } + let version = Version::from_label(label); + self.entries.insert( + Arc::from(base_url), + EngineState { + model: Arc::from(model), + version_source: version.as_ref().map(|_| VersionSource::Registration), + version, + control: ControlState::Active, + }, + ); + self.recompute_fleet_max(model); + true + } + + /// Overwrite the version from a (changed) registration label, keeping + /// the control state. + pub fn reseed(&self, base_url: &str, model: &str, label: Option<&str>) { + let version = Version::from_label(label); + let changed = { + let mut entry = + self.entries + .entry(Arc::from(base_url)) + .or_insert_with(|| EngineState { + model: Arc::from(model), + version: None, + version_source: None, + control: ControlState::Active, + }); + let changed = entry.version != version; + entry.version_source = version.as_ref().map(|_| VersionSource::Registration); + entry.version.clone_from(&version); + changed + }; + self.recompute_fleet_max(model); + if changed { + self.after_version_change(model, base_url, version.as_ref()); + } + } + + pub fn remove(&self, base_url: &str) { + let Some((_, state)) = self.entries.remove(base_url) else { + return; + }; + self.recompute_fleet_max(&state.model); + self.recompute_inactive(); + } + + /// Drop every entry whose base URL fails `keep` (resync after a lagged + /// event stream). + pub fn retain(&self, keep: impl Fn(&str) -> bool) { + let dropped: Vec> = self + .entries + .iter() + .filter(|e| !keep(e.key())) + .map(|e| Arc::clone(e.key())) + .collect(); + for base_url in dropped { + self.remove(&base_url); + } + } + + /// Record a version. Returns whether it differs from what was held. + pub fn set_version( + &self, + base_url: &str, + model: &str, + version: Version, + source: VersionSource, + ) -> bool { + let (changed, model) = { + let mut entry = + self.entries + .entry(Arc::from(base_url)) + .or_insert_with(|| EngineState { + model: Arc::from(model), + version: None, + version_source: None, + control: ControlState::Active, + }); + let changed = entry.version.as_ref() != Some(&version); + entry.version = Some(version.clone()); + entry.version_source = Some(source); + (changed, Arc::clone(&entry.model)) + }; + self.recompute_fleet_max(&model); + if changed { + self.after_version_change(&model, base_url, Some(&version)); + } + changed + } + + /// Record a control state. Returns whether it changed. + pub fn set_control(&self, base_url: &str, model: &str, control: ControlState) -> bool { + let changed = { + let mut entry = + self.entries + .entry(Arc::from(base_url)) + .or_insert_with(|| EngineState { + model: Arc::from(model), + version: None, + version_source: None, + control: ControlState::Active, + }); + let changed = entry.control != control; + entry.control = control; + changed + }; + if changed { + self.recompute_inactive(); + metrics::set_worker_control_state(base_url, control.gauge_value()); + } + changed + } + + pub fn get(&self, base_url: &str) -> Option { + self.entries.get(base_url).map(|e| e.value().clone()) + } + + pub fn version_of(&self, base_url: &str) -> Option { + self.entries.get(base_url).and_then(|e| e.version.clone()) + } + + pub fn control_of(&self, base_url: &str) -> ControlState { + self.entries + .get(base_url) + .map_or(ControlState::Active, |e| e.control) + } + + pub fn fleet_max(&self, model: &str) -> Option { + self.fleet_max.get(model).map(|v| v.value().clone()) + } + + pub fn inactive_count(&self) -> usize { + self.inactive.load(Ordering::Relaxed) + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + fn after_version_change(&self, model: &str, base_url: &str, version: Option<&Version>) { + if let Some(numeric) = version.and_then(Version::numeric) { + metrics::set_worker_weight_version(base_url, numeric); + } + self.sink.on_version_changed(model, base_url); + } + + /// Scan the model's entries; fleets are small and this runs on writes only. + fn recompute_fleet_max(&self, model: &str) { + let max = self + .entries + .iter() + .filter(|e| &*e.model == model) + .filter_map(|e| e.version.clone()) + .max(); + match max { + Some(v) => { + self.fleet_max.insert(Arc::from(model), v); + } + None => { + self.fleet_max.remove(model); + } + } + } + + fn recompute_inactive(&self) { + let n = self + .entries + .iter() + .filter(|e| e.control != ControlState::Active) + .count(); + self.inactive.store(n, Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use super::*; + + struct RecordingSink(Mutex>); + + impl VersionEvictionSink for RecordingSink { + fn on_version_changed(&self, model_id: &str, base_url: &str) { + self.0 + .lock() + .unwrap() + .push((model_id.to_string(), base_url.to_string())); + } + } + + fn table() -> (RlTable, Arc) { + let sink = Arc::new(RecordingSink(Mutex::new(Vec::new()))); + (RlTable::new(sink.clone()), sink) + } + + #[test] + fn base_url_strips_only_a_numeric_rank_suffix() { + assert_eq!(base_url_of("http://a:1@3"), "http://a:1"); + assert_eq!(base_url_of("http://a:1"), "http://a:1"); + assert_eq!(base_url_of("http://u:p@h:1"), "http://u:p@h:1"); + assert_eq!(base_url_of("http://a:1@"), "http://a:1@"); + } + + #[test] + fn seed_inserts_once_and_treats_default_as_unversioned() { + let (t, sink) = table(); + assert!(t.seed("http://a:1", "m", Some("default"))); + assert!( + !t.seed("http://a:1", "m", Some("9")), + "second seed is a no-op" + ); + let s = t.get("http://a:1").unwrap(); + assert_eq!(s.version, None); + assert_eq!(s.version_source, None); + assert_eq!(s.control, ControlState::Active); + assert_eq!(t.fleet_max("m"), None); + assert!(t.seed("http://b:1", "m", Some("3"))); + assert_eq!( + t.get("http://b:1").unwrap().version_source, + Some(VersionSource::Registration) + ); + assert_eq!(t.fleet_max("m").unwrap().as_str(), "3"); + assert!(sink.0.lock().unwrap().is_empty(), "seeding never evicts"); + } + + #[test] + fn set_version_maintains_fleet_max_and_calls_the_sink_on_change() { + let (t, sink) = table(); + t.seed("http://a:1", "m", None); + t.seed("http://b:1", "m", None); + t.seed("http://c:1", "other", None); + assert!(t.set_version("http://a:1", "m", Version::parse("1"), VersionSource::Api)); + assert!(t.set_version( + "http://b:1", + "m", + Version::parse("2"), + VersionSource::Passthrough + )); + assert_eq!(t.fleet_max("m").unwrap().as_str(), "2"); + assert_eq!(t.fleet_max("other"), None); + assert!( + !t.set_version("http://b:1", "m", Version::parse("2"), VersionSource::Api), + "same version: unchanged" + ); + assert_eq!( + t.get("http://b:1").unwrap().version_source, + Some(VersionSource::Api) + ); + assert_eq!( + *sink.0.lock().unwrap(), + vec![ + ("m".to_string(), "http://a:1".to_string()), + ("m".to_string(), "http://b:1".to_string()) + ] + ); + t.remove("http://b:1"); + assert_eq!(t.fleet_max("m").unwrap().as_str(), "1"); + t.remove("http://a:1"); + assert_eq!(t.fleet_max("m"), None); + assert_eq!(t.len(), 1); + } + + #[test] + fn set_version_on_an_unknown_engine_creates_the_entry() { + let (t, _) = table(); + assert!(t.set_version("http://x:1", "m", Version::parse("5"), VersionSource::Api)); + assert_eq!(t.get("http://x:1").unwrap().model.as_ref(), "m"); + assert_eq!(t.version_of("http://x:1").unwrap().as_str(), "5"); + } + + #[test] + fn reseed_overwrites_from_the_label() { + let (t, sink) = table(); + t.set_version("http://a:1", "m", Version::parse("7"), VersionSource::Api); + t.set_control("http://a:1", "m", ControlState::Paused); + t.reseed("http://a:1", "m", Some("8")); + let s = t.get("http://a:1").unwrap(); + assert_eq!(s.version.unwrap().as_str(), "8"); + assert_eq!(s.version_source, Some(VersionSource::Registration)); + assert_eq!( + s.control, + ControlState::Paused, + "reseed keeps control state" + ); + assert_eq!(sink.0.lock().unwrap().len(), 2); + } + + #[test] + fn control_state_tracks_the_inactive_count() { + let (t, _) = table(); + t.seed("http://a:1", "m", None); + t.seed("http://b:1", "m", None); + assert_eq!(t.inactive_count(), 0); + assert!(t.set_control("http://a:1", "m", ControlState::Paused)); + assert!(!t.set_control("http://a:1", "m", ControlState::Paused)); + assert!(t.set_control("http://b:1", "m", ControlState::Asleep)); + assert_eq!(t.inactive_count(), 2); + assert_eq!(t.control_of("http://a:1"), ControlState::Paused); + assert_eq!(t.control_of("http://zz:1"), ControlState::Active); + t.set_control("http://a:1", "m", ControlState::Active); + assert_eq!(t.inactive_count(), 1); + t.remove("http://b:1"); + assert_eq!(t.inactive_count(), 0); + } + + #[test] + fn retain_drops_entries_outside_the_registry() { + let (t, _) = table(); + t.seed("http://a:1", "m", Some("1")); + t.seed("http://b:1", "m", Some("2")); + t.set_control("http://b:1", "m", ControlState::Asleep); + t.retain(|url| url == "http://a:1"); + assert_eq!(t.len(), 1); + assert_eq!(t.fleet_max("m").unwrap().as_str(), "1"); + assert_eq!(t.inactive_count(), 0); + } +} diff --git a/crates/rl/src/view.rs b/crates/rl/src/view.rs index 1b08fe3add..66dec99549 100644 --- a/crates/rl/src/view.rs +++ b/crates/rl/src/view.rs @@ -57,6 +57,13 @@ pub trait RlWorkerView: Send + Sync { fn list(&self) -> Vec; /// One worker by registry UUID. fn get(&self, id: &str) -> Option; + + /// Requests currently in flight on the engine at `base_url` (all DP + /// ranks). Used for the mixed-version accounting metric; the default + /// reports none. + fn inflight(&self, _base_url: &str) -> usize { + 0 + } } #[cfg(test)] From 7accd70d5686f13d78ca63d2e21a590d07d5158a Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 14:04:40 -0700 Subject: [PATCH 03/27] fix(rl): carry version_policy through the gateway and bindings config literals Signed-off-by: key4ng --- bindings/python/src/lib.rs | 1 + model_gateway/src/config/validation.rs | 1 + model_gateway/src/main.rs | 1 + 3 files changed, 3 insertions(+) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 0cae550345..f3f5260aa3 100755 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -944,6 +944,7 @@ impl Router { enabled: self.enable_rl, control_timeout_secs: self.rl_control_timeout_secs, fanout_concurrency: self.rl_fanout_concurrency, + version_policy: smg_rl::VersionPolicy::default(), }) .maybe_client_cert_and_key( self.client_cert_path.as_ref(), diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index 7d2354a9aa..cf4a2ed259 100644 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -2292,6 +2292,7 @@ mod tests { enabled: true, control_timeout_secs: 0, fanout_concurrency: 32, + version_policy: smg_rl::VersionPolicy::default(), }) .build_unchecked(); let err = ConfigValidator::validate(&bad).unwrap_err().to_string(); diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 9f135d63e9..55b9a5d826 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -1958,6 +1958,7 @@ impl CliArgs { enabled: self.enable_rl, control_timeout_secs: self.rl_control_timeout_secs, fanout_concurrency: self.rl_fanout_concurrency, + version_policy: smg_rl::VersionPolicy::default(), }) .dp_minimum_tokens_scheduler(self.dp_minimum_tokens_scheduler) .maybe_server_cert_and_key(self.tls_cert_path.as_ref(), self.tls_key_path.as_ref()); From 6d36b0810f315f7790ef63b0f64d2720db99d968 Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 14:14:30 -0700 Subject: [PATCH 04/27] fix(rl): derive reseed's model from the entry and serialize table writes Signed-off-by: key4ng --- crates/rl/src/table.rs | 140 ++++++++++++++++++++++++++++++----------- 1 file changed, 104 insertions(+), 36 deletions(-) diff --git a/crates/rl/src/table.rs b/crates/rl/src/table.rs index d01b96cb74..0ddea32743 100644 --- a/crates/rl/src/table.rs +++ b/crates/rl/src/table.rs @@ -4,7 +4,7 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, - Arc, + Arc, Mutex, MutexGuard, PoisonError, }; use dashmap::DashMap; @@ -87,12 +87,30 @@ pub fn base_url_of(worker_url: &str) -> &str { } } +/// A freshly observed engine: unversioned and active until a write says +/// otherwise. +fn fresh(model: &str) -> EngineState { + EngineState { + model: Arc::from(model), + version: None, + version_source: None, + control: ControlState::Active, + } +} + pub struct RlTable { entries: DashMap, EngineState>, fleet_max: DashMap, Version>, /// Entries whose control state is not `Active`; lets `any` skip the map. inactive: AtomicUsize, sink: Arc, + /// Serializes every write. `recompute_fleet_max` is a scan-then-store + /// across two `DashMap`s with no synchronization of its own, so two + /// concurrent writers for the same model could otherwise leave + /// `fleet_max` durably stale. Writes are control-path and rare (a + /// refit, a pause, a registration), so one uncontended lock is the + /// whole cost; reads never take it. + writes: Mutex<()>, } impl RlTable { @@ -102,12 +120,18 @@ impl RlTable { fleet_max: DashMap::new(), inactive: AtomicUsize::new(0), sink, + writes: Mutex::new(()), } } + fn lock_writes(&self) -> MutexGuard<'_, ()> { + self.writes.lock().unwrap_or_else(PoisonError::into_inner) + } + /// Insert an entry from a registration label unless one exists. Returns /// whether an entry was inserted. pub fn seed(&self, base_url: &str, model: &str, label: Option<&str>) -> bool { + let _guard = self.lock_writes(); if self.entries.contains_key(base_url) { return false; } @@ -115,10 +139,9 @@ impl RlTable { self.entries.insert( Arc::from(base_url), EngineState { - model: Arc::from(model), version_source: version.as_ref().map(|_| VersionSource::Registration), version, - control: ControlState::Active, + ..fresh(model) }, ); self.recompute_fleet_max(model); @@ -128,29 +151,32 @@ impl RlTable { /// Overwrite the version from a (changed) registration label, keeping /// the control state. pub fn reseed(&self, base_url: &str, model: &str, label: Option<&str>) { + let _guard = self.lock_writes(); let version = Version::from_label(label); - let changed = { - let mut entry = - self.entries - .entry(Arc::from(base_url)) - .or_insert_with(|| EngineState { - model: Arc::from(model), - version: None, - version_source: None, - control: ControlState::Active, - }); + let (changed, model) = { + let mut entry = self + .entries + .entry(Arc::from(base_url)) + .or_insert_with(|| fresh(model)); let changed = entry.version != version; entry.version_source = version.as_ref().map(|_| VersionSource::Registration); entry.version.clone_from(&version); - changed + (changed, Arc::clone(&entry.model)) }; - self.recompute_fleet_max(model); + self.recompute_fleet_max(&model); if changed { - self.after_version_change(model, base_url, version.as_ref()); + self.after_version_change(&model, base_url, version.as_ref()); } } pub fn remove(&self, base_url: &str) { + let _guard = self.lock_writes(); + self.remove_locked(base_url); + } + + /// `remove`'s body, for callers that already hold `writes` (namely + /// `retain`, which must not re-lock the mutex it is already holding). + fn remove_locked(&self, base_url: &str) { let Some((_, state)) = self.entries.remove(base_url) else { return; }; @@ -161,6 +187,7 @@ impl RlTable { /// Drop every entry whose base URL fails `keep` (resync after a lagged /// event stream). pub fn retain(&self, keep: impl Fn(&str) -> bool) { + let _guard = self.lock_writes(); let dropped: Vec> = self .entries .iter() @@ -168,7 +195,7 @@ impl RlTable { .map(|e| Arc::clone(e.key())) .collect(); for base_url in dropped { - self.remove(&base_url); + self.remove_locked(&base_url); } } @@ -180,16 +207,12 @@ impl RlTable { version: Version, source: VersionSource, ) -> bool { + let _guard = self.lock_writes(); let (changed, model) = { - let mut entry = - self.entries - .entry(Arc::from(base_url)) - .or_insert_with(|| EngineState { - model: Arc::from(model), - version: None, - version_source: None, - control: ControlState::Active, - }); + let mut entry = self + .entries + .entry(Arc::from(base_url)) + .or_insert_with(|| fresh(model)); let changed = entry.version.as_ref() != Some(&version); entry.version = Some(version.clone()); entry.version_source = Some(source); @@ -204,16 +227,12 @@ impl RlTable { /// Record a control state. Returns whether it changed. pub fn set_control(&self, base_url: &str, model: &str, control: ControlState) -> bool { + let _guard = self.lock_writes(); let changed = { - let mut entry = - self.entries - .entry(Arc::from(base_url)) - .or_insert_with(|| EngineState { - model: Arc::from(model), - version: None, - version_source: None, - control: ControlState::Active, - }); + let mut entry = self + .entries + .entry(Arc::from(base_url)) + .or_insert_with(|| fresh(model)); let changed = entry.control != control; entry.control = control; changed @@ -262,7 +281,9 @@ impl RlTable { self.sink.on_version_changed(model, base_url); } - /// Scan the model's entries; fleets are small and this runs on writes only. + /// Scan the model's entries; fleets are small and this runs on writes + /// only. Always called with `writes` held: the scan-then-store here is + /// not atomic on its own. fn recompute_fleet_max(&self, model: &str) { let max = self .entries @@ -433,4 +454,51 @@ mod tests { assert_eq!(t.fleet_max("m").unwrap().as_str(), "1"); assert_eq!(t.inactive_count(), 0); } + + #[test] + fn reseed_recomputes_using_the_entrys_stored_model_not_the_callers_argument() { + let (t, sink) = table(); + t.seed("http://a:1", "real-model", Some("5")); + assert_eq!(t.fleet_max("real-model").unwrap().as_str(), "5"); + + // A caller passes a model that doesn't match what the entry was + // seeded under; the stored model must win for both the fleet-max + // recompute and the eviction-sink call. + t.reseed("http://a:1", "wrong-model", Some("9")); + + assert_eq!(t.get("http://a:1").unwrap().model.as_ref(), "real-model"); + assert_eq!(t.fleet_max("real-model").unwrap().as_str(), "9"); + assert_eq!(t.fleet_max("wrong-model"), None); + assert_eq!( + sink.0.lock().unwrap().last(), + Some(&("real-model".to_string(), "http://a:1".to_string())) + ); + } + + #[test] + fn concurrent_writers_leave_fleet_max_consistent() { + let (t, _) = table(); + let t = Arc::new(t); + let handles: Vec<_> = (0..8) + .map(|i| { + let t = Arc::clone(&t); + std::thread::spawn(move || { + let base_url = format!("http://w{i}:1"); + for v in 1..=200 { + t.set_version( + &base_url, + "m", + Version::parse(&v.to_string()), + VersionSource::Api, + ); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + assert_eq!(t.fleet_max("m").unwrap().as_str(), "200"); + assert_eq!(t.len(), 8); + } } From 76951aeb4ea034064eebcf0ca086e0cd73c97e86 Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 14:29:14 -0700 Subject: [PATCH 05/27] feat(rl): version-policy eligibility over the table, with a bench Adds RlTable::eligible, the lock-free read path that filters routing candidates by ControlState and VersionPolicy against the per-model fleet max, plus FilterReason (the smg_rl_candidates_filtered_total reason label) and its judge() decision function for LatestOnly, MaxStaleness(k), and MinVersion(floor). Also fixes recompute_fleet_max: it picked the fleet max via Version's Ord, which falls back to a lexical compare across a numeric/text pair, so a stray checkpoint label like "ckpt-7" could outrank every numbered release just because 'c' > '1' byte-wise, corrupting the staleness baseline for the whole model. It now prefers the numeric maximum whenever the model has any numeric version. Bench (cargo bench -p smg-rl --bench filter), 8 candidates: eligible_8_candidates/any: 1.77 ns eligible_8_candidates/latest_only: 318.05 ns eligible_8_candidates/max_staleness_1: 189.91 ns Signed-off-by: key4ng --- Cargo.lock | 1 + crates/rl/Cargo.toml | 5 + crates/rl/benches/filter.rs | 35 +++++ crates/rl/src/lib.rs | 2 +- crates/rl/src/table.rs | 247 +++++++++++++++++++++++++++++++++++- 5 files changed, 282 insertions(+), 8 deletions(-) create mode 100644 crates/rl/benches/filter.rs diff --git a/Cargo.lock b/Cargo.lock index 02b8791a3f..70c8decd03 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7730,6 +7730,7 @@ name = "smg-rl" version = "0.1.0" dependencies = [ "axum", + "criterion", "dashmap", "futures", "http", diff --git a/crates/rl/Cargo.toml b/crates/rl/Cargo.toml index 26059294a7..e5b0b8814d 100644 --- a/crates/rl/Cargo.toml +++ b/crates/rl/Cargo.toml @@ -32,6 +32,11 @@ metrics = "0.24.6" tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net"] } tower = { version = "0.5", features = ["util"] } http-body-util = "0.1" +criterion = { version = "0.8", features = ["html_reports"] } + +[[bench]] +name = "filter" +harness = false [lints] workspace = true diff --git a/crates/rl/benches/filter.rs b/crates/rl/benches/filter.rs new file mode 100644 index 0000000000..ad2cbb3f1e --- /dev/null +++ b/crates/rl/benches/filter.rs @@ -0,0 +1,35 @@ +//! Pins the cost of the RL pre-filter at a typical rollout fleet size. + +use std::{hint::black_box, sync::Arc}; + +use criterion::{criterion_group, criterion_main, Criterion}; +use smg_rl::{ + table::{NoopEvictionSink, RlTable, VersionSource}, + version::Version, + VersionPolicy, +}; + +fn bench_eligible(c: &mut Criterion) { + let table = RlTable::new(Arc::new(NoopEvictionSink)); + let urls: Vec = (0..8).map(|i| format!("http://engine-{i}:30000")).collect(); + for (i, url) in urls.iter().enumerate() { + let version = if i % 2 == 0 { "42" } else { "41" }; + table.set_version(url, "m", Version::parse(version), VersionSource::Api); + } + let candidates: Vec<(&str, &str)> = urls.iter().map(|u| ("m", u.as_str())).collect(); + + let mut group = c.benchmark_group("eligible_8_candidates"); + for (name, policy) in [ + ("any", VersionPolicy::Any), + ("latest_only", VersionPolicy::LatestOnly), + ("max_staleness_1", VersionPolicy::MaxStaleness(1)), + ] { + group.bench_function(name, |b| { + b.iter(|| table.eligible(black_box(candidates.iter().copied()), black_box(&policy))); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_eligible); +criterion_main!(benches); diff --git a/crates/rl/src/lib.rs b/crates/rl/src/lib.rs index f6083e9dbd..d690931c30 100644 --- a/crates/rl/src/lib.rs +++ b/crates/rl/src/lib.rs @@ -27,7 +27,7 @@ pub use error::RlError; pub use metrics::init_rl_metrics; pub use policy::{VersionPolicy, VersionPolicyError, VERSION_POLICY_HEADER}; pub use state::RlState; -pub use table::{ControlState, RlTable, VersionEvictionSink, VersionSource}; +pub use table::{ControlState, FilterReason, RlTable, VersionEvictionSink, VersionSource}; pub use version::Version; pub use view::{RlWorkerInfo, RlWorkerView}; diff --git a/crates/rl/src/table.rs b/crates/rl/src/table.rs index 0ddea32743..ded1738994 100644 --- a/crates/rl/src/table.rs +++ b/crates/rl/src/table.rs @@ -9,7 +9,7 @@ use std::sync::{ use dashmap::DashMap; -use crate::{metrics, version::Version}; +use crate::{metrics, policy::VersionPolicy, version::Version}; /// Whether an engine can serve requests, as last observed through SMG. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] @@ -98,6 +98,27 @@ fn fresh(model: &str) -> EngineState { } } +/// Why a candidate was dropped; the `reason` label of +/// `smg_rl_candidates_filtered_total`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FilterReason { + Paused, + Asleep, + Stale, + Unversioned, +} + +impl FilterReason { + pub fn as_str(self) -> &'static str { + match self { + Self::Paused => "paused", + Self::Asleep => "asleep", + Self::Stale => "stale", + Self::Unversioned => "unversioned", + } + } +} + pub struct RlTable { entries: DashMap, EngineState>, fleet_max: DashMap, Version>, @@ -284,13 +305,26 @@ impl RlTable { /// Scan the model's entries; fleets are small and this runs on writes /// only. Always called with `writes` held: the scan-then-store here is /// not atomic on its own. + /// + /// Prefers the numeric maximum when the fleet has any numeric version: + /// `Version`'s `Ord` falls back to a lexical compare across a + /// numeric/text pair, so a stray checkpoint label like `ckpt-7` would + /// otherwise outrank every numbered release just because `'c' > '1'` + /// byte-wise, corrupting the staleness baseline for the whole model. + /// Only when every version for the model is text does that lexical + /// order stand, matching `Version`'s own documented text-vs-text + /// ordering. fn recompute_fleet_max(&self, model: &str) { - let max = self - .entries - .iter() - .filter(|e| &*e.model == model) - .filter_map(|e| e.version.clone()) - .max(); + let versions_of_model = || { + self.entries + .iter() + .filter(|e| &*e.model == model) + .filter_map(|e| e.version.clone()) + }; + let max = versions_of_model() + .filter(|v| v.numeric().is_some()) + .max() + .or_else(|| versions_of_model().max()); match max { Some(v) => { self.fleet_max.insert(Arc::from(model), v); @@ -309,6 +343,87 @@ impl RlTable { .count(); self.inactive.store(n, Ordering::Relaxed); } + + /// Indices of `candidates` (as `(model_id, base_url)` pairs) that stay + /// eligible under `policy`; `None` when every candidate is eligible, so + /// the common case allocates nothing. Paused and asleep engines are + /// dropped under every policy. + pub fn eligible<'a, I>(&self, candidates: I, policy: &VersionPolicy) -> Option> + where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, + { + if *policy == VersionPolicy::Any && self.inactive.load(Ordering::Relaxed) == 0 { + return None; + } + let iter = candidates.into_iter(); + let total = iter.len(); + let mut keep = Vec::with_capacity(total); + let mut cached_max: Option<(&'a str, Option)> = None; + for (idx, (model, base_url)) in iter.enumerate() { + let (control, version) = match self.entries.get(base_url) { + Some(e) => (e.control, e.version.clone()), + None => (ControlState::Active, None), + }; + let verdict = match control { + ControlState::Paused => Err(FilterReason::Paused), + ControlState::Asleep => Err(FilterReason::Asleep), + ControlState::Active if *policy == VersionPolicy::Any => Ok(()), + ControlState::Active => { + let fleet_max = match &cached_max { + Some((m, max)) if *m == model => max.clone(), + _ => { + let max = self.fleet_max(model); + cached_max = Some((model, max.clone())); + max + } + }; + Self::judge(policy, version.as_ref(), fleet_max.as_ref()) + } + }; + match verdict { + Ok(()) => keep.push(idx), + Err(reason) => metrics::record_candidate_filtered(reason.as_str()), + } + } + if keep.is_empty() { + metrics::record_request_unroutable(); + } + if keep.len() == total { + None + } else { + Some(keep) + } + } + + fn judge( + policy: &VersionPolicy, + version: Option<&Version>, + fleet_max: Option<&Version>, + ) -> Result<(), FilterReason> { + match policy { + VersionPolicy::Any => Ok(()), + VersionPolicy::LatestOnly => match (version, fleet_max) { + (_, None) => Ok(()), + (None, Some(_)) => Err(FilterReason::Unversioned), + (Some(v), Some(max)) if v == max => Ok(()), + (Some(_), Some(_)) => Err(FilterReason::Stale), + }, + VersionPolicy::MaxStaleness(k) => match (version, fleet_max) { + (_, None) => Ok(()), + (None, Some(_)) => Err(FilterReason::Unversioned), + (Some(v), Some(max)) => match (v.numeric(), max.numeric()) { + (Some(v), Some(max)) if max.saturating_sub(v) <= *k => Ok(()), + _ => Err(FilterReason::Stale), + }, + }, + VersionPolicy::MinVersion(floor) => match version { + None => Err(FilterReason::Unversioned), + Some(v) if v >= floor => Ok(()), + Some(_) => Err(FilterReason::Stale), + }, + } + } } #[cfg(test)] @@ -475,6 +590,124 @@ mod tests { ); } + fn candidates<'a>(urls: &'a [&'a str]) -> Vec<(&'a str, &'a str)> { + urls.iter().map(|u| ("m", *u)).collect() + } + + #[test] + fn any_with_every_engine_active_is_the_allocation_free_path() { + let (t, _) = table(); + t.seed("http://a:1", "m", Some("1")); + t.seed("http://b:1", "m", Some("2")); + assert_eq!( + t.eligible( + candidates(&["http://a:1", "http://b:1"]), + &VersionPolicy::Any + ), + None + ); + // Unknown engines are eligible too. + assert_eq!( + t.eligible(candidates(&["http://zz:1"]), &VersionPolicy::Any), + None + ); + } + + #[test] + fn paused_and_asleep_engines_are_dropped_under_every_policy() { + let (t, _) = table(); + t.seed("http://a:1", "m", Some("2")); + t.seed("http://b:1", "m", Some("2")); + t.seed("http://c:1", "m", Some("2")); + t.set_control("http://a:1", "m", ControlState::Paused); + t.set_control("http://c:1", "m", ControlState::Asleep); + let c = candidates(&["http://a:1", "http://b:1", "http://c:1"]); + assert_eq!(t.eligible(c.clone(), &VersionPolicy::Any), Some(vec![1])); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::LatestOnly), + Some(vec![1]) + ); + t.set_control("http://b:1", "m", ControlState::Paused); + assert_eq!(t.eligible(c, &VersionPolicy::Any), Some(vec![])); + } + + #[test] + fn latest_only_keeps_the_fleet_max_and_everyone_when_nobody_is_versioned() { + let (t, _) = table(); + t.seed("http://a:1", "m", None); + t.seed("http://b:1", "m", None); + let c = candidates(&["http://a:1", "http://b:1"]); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::LatestOnly), + None, + "nothing is stale yet" + ); + t.set_version("http://a:1", "m", Version::parse("1"), VersionSource::Api); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::LatestOnly), + Some(vec![0]), + "unversioned b is behind" + ); + t.set_version("http://b:1", "m", Version::parse("2"), VersionSource::Api); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::LatestOnly), + Some(vec![1]), + "a is stale" + ); + t.set_version("http://a:1", "m", Version::parse("2"), VersionSource::Api); + assert_eq!(t.eligible(c, &VersionPolicy::LatestOnly), None); + } + + #[test] + fn max_staleness_and_min_version_evaluate_numerically() { + let (t, _) = table(); + for (url, v) in [ + ("http://a:1", "10"), + ("http://b:1", "9"), + ("http://c:1", "7"), + ] { + t.set_version(url, "m", Version::parse(v), VersionSource::Api); + } + t.seed("http://d:1", "m", None); + let c = candidates(&["http://a:1", "http://b:1", "http://c:1", "http://d:1"]); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::MaxStaleness(1)), + Some(vec![0, 1]) + ); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::MaxStaleness(3)), + Some(vec![0, 1, 2]) + ); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::MinVersion(Version::parse("9"))), + Some(vec![0, 1]) + ); + assert_eq!( + t.eligible(c.clone(), &VersionPolicy::MinVersion(Version::parse("7"))), + Some(vec![0, 1, 2]) + ); + // A text version cannot satisfy a numeric staleness bound. + t.set_version( + "http://c:1", + "m", + Version::parse("ckpt-7"), + VersionSource::Api, + ); + assert_eq!( + t.eligible(c, &VersionPolicy::MaxStaleness(100)), + Some(vec![0, 1]) + ); + } + + #[test] + fn fleet_max_is_per_model() { + let (t, _) = table(); + t.set_version("http://a:1", "m", Version::parse("5"), VersionSource::Api); + t.set_version("http://b:1", "n", Version::parse("1"), VersionSource::Api); + let c = vec![("m", "http://a:1"), ("n", "http://b:1")]; + assert_eq!(t.eligible(c, &VersionPolicy::LatestOnly), None); + } + #[test] fn concurrent_writers_leave_fleet_max_consistent() { let (t, _) = table(); From f163f2d17e40345fe62acf24b06feeddfd4a807a Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 14:45:19 -0700 Subject: [PATCH 06/27] feat(rl): control-state and version-source wire types; discovery reads the table Signed-off-by: key4ng --- bindings/python/tests/test_rl_client.py | 2 + crates/protocols/src/rl.rs | 72 ++++++++++++++- crates/rl/src/discovery.rs | 112 +++++++++++++++++++----- crates/rl/src/fanout.rs | 11 ++- crates/rl/src/table.rs | 31 +++++++ 5 files changed, 203 insertions(+), 25 deletions(-) diff --git a/bindings/python/tests/test_rl_client.py b/bindings/python/tests/test_rl_client.py index e34aa9a77a..6c2f7216a0 100644 --- a/bindings/python/tests/test_rl_client.py +++ b/bindings/python/tests/test_rl_client.py @@ -26,6 +26,8 @@ "role": None, "health": "ready", "weight_version": "7", + "version_source": "api", + "control": "active", "labels": {"tp_size": "1"}, "capabilities": { "source": "static", diff --git a/crates/protocols/src/rl.rs b/crates/protocols/src/rl.rs index e5275fdfd9..bfbd417adc 100644 --- a/crates/protocols/src/rl.rs +++ b/crates/protocols/src/rl.rs @@ -1,5 +1,6 @@ //! Wire types of the RL control plane (`/v1/rl/*`): worker discovery rows, -//! per-call outcomes, and fan-out envelopes. +//! version and control-state writes, per-call outcomes, and fan-out +//! envelopes. //! //! The logic that produces them lives in the `smg-rl` crate; this module is //! the contract that clients and the OpenAPI generator depend on, in the same @@ -40,6 +41,38 @@ pub struct RlCapabilities { pub reports_weight_version: bool, } +/// Whether an engine can serve requests, as SMG last observed it. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum RlControlState { + #[default] + Active, + Paused, + Asleep, +} + +/// How SMG learned an engine's weight version. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum RlVersionSource { + Registration, + Passthrough, + Api, +} + +/// Body of `POST /v1/rl/workers/{id}/version` and `POST /v1/rl/version`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +pub struct RlSetVersionRequest { + /// Non-empty, at most 128 bytes; compared numerically when it parses. + pub weight_version: String, +} + +/// Body of `POST /v1/rl/workers/{id}/state` and `POST /v1/rl/state`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +pub struct RlSetControlRequest { + pub control: RlControlState, +} + /// One row of `GET /v1/rl/workers`. DP-aware ranks that share a `base_url` /// collapse into one row; `dp_ranks` counts them. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] @@ -66,6 +99,11 @@ pub struct RlWorkerEntry { /// Registry status (`ready`, `not_ready`, ...). pub health: String, pub weight_version: Option, + /// How SMG learned `weight_version`; `None` when unversioned. + pub version_source: Option, + /// Pause/sleep state as last observed through SMG; `active` when never observed. + #[serde(default)] + pub control: RlControlState, /// Registration labels, discovered metadata merged with caller labels. pub labels: HashMap, pub capabilities: RlCapabilities, @@ -157,6 +195,8 @@ mod tests { role: None, health: "ready".to_string(), weight_version: Some("default".to_string()), + version_source: None, + control: RlControlState::Active, labels: HashMap::new(), capabilities: RlCapabilities { source: RlCapabilitySource::Static, @@ -249,6 +289,34 @@ mod tests { assert_eq!(empty, RlFanoutResponse::default()); } + #[test] + fn control_state_and_source_serialize_lowercase_and_default() { + assert_eq!( + serde_json::to_value(RlControlState::Asleep).unwrap(), + json!("asleep") + ); + assert_eq!( + serde_json::to_value(RlVersionSource::Passthrough).unwrap(), + json!("passthrough") + ); + let mut v = serde_json::to_value(sample_entry()).unwrap(); + assert_eq!(v["control"], "active"); + assert_eq!(v["version_source"], Value::Null); + v.as_object_mut().unwrap().remove("control"); + let back: RlWorkerEntry = serde_json::from_value(v).unwrap(); + assert_eq!( + back.control, + RlControlState::Active, + "older gateways omit control" + ); + let req: RlSetControlRequest = + serde_json::from_value(json!({"control": "paused"})).unwrap(); + assert_eq!(req.control, RlControlState::Paused); + assert!( + serde_json::from_value::(json!({"control": "Paused"})).is_err() + ); + } + #[test] fn every_wire_type_has_a_json_schema() { for schema in [ @@ -256,6 +324,8 @@ mod tests { schema_for!(RlWorkerEntry), schema_for!(RlCallOutcome), schema_for!(RlFanoutResponse), + schema_for!(RlSetVersionRequest), + schema_for!(RlSetControlRequest), ] { assert!(schema.to_value().get("title").is_some()); } diff --git a/crates/rl/src/discovery.rs b/crates/rl/src/discovery.rs index 63fa66b611..88e58f3ae4 100644 --- a/crates/rl/src/discovery.rs +++ b/crates/rl/src/discovery.rs @@ -9,10 +9,18 @@ use axum::{ response::{IntoResponse, Response}, Json, }; -use openai_protocol::rl::{RlWorkerEntry, RlWorkersResponse, RL_PROTOCOL_VERSION}; +use openai_protocol::rl::{ + RlControlState, RlVersionSource, RlWorkerEntry, RlWorkersResponse, RL_PROTOCOL_VERSION, +}; use serde::Serialize; -use crate::{capability::capabilities_for, error::RlError, state::RlState, view::RlWorkerInfo}; +use crate::{ + capability::capabilities_for, + error::RlError, + state::RlState, + table::RlTable, + view::{RlWorkerInfo, RlWorkerView}, +}; /// Serialize an enum through serde to get its canonical wire string. pub(crate) fn enum_str(v: &T) -> String { @@ -74,8 +82,28 @@ pub fn collapse(workers: Vec) -> Vec<(RlWorkerInfo, usize)> { rows } +/// The version routing uses for `info`'s engine: the table's entry when there +/// is one (an unversioned entry is `None`), else the registration label. +fn table_version( + table: &RlTable, + info: &RlWorkerInfo, +) -> (Option, Option, RlControlState) { + match table.get(&info.base_url) { + Some(state) => ( + state.version.as_ref().map(|v| v.as_str().to_string()), + state.version_source.map(Into::into), + state.control.into(), + ), + None => ( + info.labels.get("weight_version").cloned(), + None, + RlControlState::Active, + ), + } +} + /// Labels plus synthetic keys, for selector matching. Synthetic keys shadow. -pub fn merged_labels(info: &RlWorkerInfo) -> HashMap { +pub fn merged_labels(table: &RlTable, info: &RlWorkerInfo) -> HashMap { let mut m = info.labels.clone(); m.insert("id".to_string(), info.id.clone()); m.insert("url".to_string(), info.url.clone()); @@ -88,17 +116,32 @@ pub fn merged_labels(info: &RlWorkerInfo) -> HashMap { enum_str(&info.connection_mode), ); m.insert("health".to_string(), enum_str(&info.status)); - if let Some(v) = info.labels.get("weight_version") { - m.insert("weight_version".to_string(), v.clone()); + let (weight_version, _, control) = table_version(table, info); + if let Some(v) = weight_version { + m.insert("weight_version".to_string(), v); } + m.insert("control".to_string(), enum_str(&control)); if let Some(v) = info.labels.get("role") { m.insert("role".to_string(), v.clone()); } m } +/// DP ranks that share `info`'s engine (1 for a non-DP worker). +pub fn dp_ranks(view: &dyn RlWorkerView, info: &RlWorkerInfo) -> usize { + if !info.is_dp_aware { + return 1; + } + view.list() + .iter() + .filter(|other| other.is_dp_aware && other.base_url == info.base_url) + .count() + .max(1) +} + /// One row of `GET /workers` for `info`, with `dp_ranks` collapsed into it. -pub fn entry(info: &RlWorkerInfo, dp_ranks: usize) -> RlWorkerEntry { +pub fn entry(table: &RlTable, info: &RlWorkerInfo, dp_ranks: usize) -> RlWorkerEntry { + let (weight_version, version_source, control) = table_version(table, info); RlWorkerEntry { id: info.id.clone(), url: info.url.clone(), @@ -114,7 +157,9 @@ pub fn entry(info: &RlWorkerInfo, dp_ranks: usize) -> RlWorkerEntry { dp_ranks, role: info.labels.get("role").cloned(), health: enum_str(&info.status), - weight_version: info.labels.get("weight_version").cloned(), + weight_version, + version_source, + control, labels: info.labels.clone(), capabilities: capabilities_for(info.runtime, &info.labels), } @@ -123,7 +168,7 @@ pub fn entry(info: &RlWorkerInfo, dp_ranks: usize) -> RlWorkerEntry { pub(crate) async fn list_workers(State(state): State>) -> Response { let workers: Vec = collapse(state.view.list()) .iter() - .map(|(w, n)| entry(w, *n)) + .map(|(w, n)| entry(&state.table, w, *n)) .collect(); let total = workers.len(); ( @@ -143,18 +188,8 @@ pub(crate) async fn get_worker( ) -> Response { match state.view.get(&id) { Some(w) => { - let ranks = if w.is_dp_aware { - state - .view - .list() - .iter() - .filter(|other| other.is_dp_aware && other.base_url == w.base_url) - .count() - .max(1) - } else { - 1 - }; - (StatusCode::OK, Json(entry(&w, ranks))).into_response() + let ranks = dp_ranks(state.view.as_ref(), &w); + (StatusCode::OK, Json(entry(&state.table, &w, ranks))).into_response() } None => RlError::WorkerNotFound(id).into_response(), } @@ -173,7 +208,9 @@ mod tests { use crate::{ config::RlConfig, state::RlState, + table::{ControlState, NoopEvictionSink, VersionSource}, testing::{worker, FakeView}, + version::Version, }; fn state(workers: Vec) -> Arc { @@ -240,17 +277,50 @@ mod tests { let mut w = worker("w1", "http://a:1", RuntimeType::Vllm); w.labels.insert("engine".to_string(), "spoofed".to_string()); w.labels.insert("role".to_string(), "reward".to_string()); - let m = merged_labels(&w); + let table = RlTable::new(Arc::new(NoopEvictionSink)); + let m = merged_labels(&table, &w); assert_eq!(m["engine"], "vllm"); assert_eq!(m["id"], "w1"); assert_eq!(m["url"], "http://a:1"); assert_eq!(m["model"], "mock-model"); assert_eq!(m["health"], "ready"); assert_eq!(m["weight_version"], "default"); + assert_eq!(m["control"], "active"); assert_eq!(m["role"], "reward"); assert_eq!(m["tp_size"], "1"); } + #[test] + fn table_version_and_control_win_over_labels() { + let table = RlTable::new(Arc::new(NoopEvictionSink)); + let mut w = worker("w1", "http://a:1", RuntimeType::Vllm); + table.set_version( + "http://a:1", + &w.model_id, + Version::parse("9"), + VersionSource::Api, + ); + table.set_control("http://a:1", &w.model_id, ControlState::Paused); + + let e = entry(&table, &w, 1); + assert_eq!(e.weight_version.as_deref(), Some("9")); + assert_eq!(e.version_source, Some(RlVersionSource::Api)); + assert_eq!(e.control, RlControlState::Paused); + + let m = merged_labels(&table, &w); + assert_eq!(m["weight_version"], "9"); + assert_eq!(m["control"], "paused"); + + // A seeded-but-unversioned engine reports weight_version as null, + // not the label, once it has a table entry. + w.id = "w2".to_string(); + w.url = "http://b:1".to_string(); + w.base_url = "http://b:1".to_string(); + table.seed("http://b:1", &w.model_id, None); + let e2 = entry(&table, &w, 1); + assert_eq!(e2.weight_version, None); + } + #[tokio::test] async fn list_reports_the_protocol_version() { let app = crate::router::<()>(state(vec![worker("w1", "http://a:1", RuntimeType::Sglang)])); diff --git a/crates/rl/src/fanout.rs b/crates/rl/src/fanout.rs index c62cd3bc6d..978b98781e 100644 --- a/crates/rl/src/fanout.rs +++ b/crates/rl/src/fanout.rs @@ -21,6 +21,7 @@ use crate::{ proxy::{call_worker, ProxyRequest}, selector::Selector, state::RlState, + table::RlTable, view::{RlWorkerInfo, RlWorkerView}, }; @@ -30,11 +31,15 @@ pub(crate) struct FanoutQuery { } /// DP-collapsed workers matching `selector`, sorted by base URL. -pub fn resolve_targets(view: &dyn RlWorkerView, selector: &Selector) -> Vec { +pub fn resolve_targets( + view: &dyn RlWorkerView, + table: &RlTable, + selector: &Selector, +) -> Vec { collapse(view.list()) .into_iter() .map(|(w, _)| w) - .filter(|w| selector.matches(&merged_labels(w))) + .filter(|w| selector.matches(&merged_labels(table, w))) .collect() } @@ -128,7 +133,7 @@ pub(crate) async fn fanout_handler( Ok(r) => r, Err(e) => return e.into_response(), }; - let targets = resolve_targets(state.view.as_ref(), &selector); + let targets = resolve_targets(state.view.as_ref(), &state.table, &selector); if targets.is_empty() { record_fanout("no_match", started.elapsed()); return RlError::NoWorkersMatch(selector.source().to_string()).into_response(); diff --git a/crates/rl/src/table.rs b/crates/rl/src/table.rs index ded1738994..65c746b004 100644 --- a/crates/rl/src/table.rs +++ b/crates/rl/src/table.rs @@ -8,6 +8,7 @@ use std::sync::{ }; use dashmap::DashMap; +use openai_protocol::rl::{RlControlState, RlVersionSource}; use crate::{metrics, policy::VersionPolicy, version::Version}; @@ -56,6 +57,36 @@ impl VersionSource { } } +impl From for RlControlState { + fn from(c: ControlState) -> Self { + match c { + ControlState::Active => Self::Active, + ControlState::Paused => Self::Paused, + ControlState::Asleep => Self::Asleep, + } + } +} + +impl From for ControlState { + fn from(c: RlControlState) -> Self { + match c { + RlControlState::Active => Self::Active, + RlControlState::Paused => Self::Paused, + RlControlState::Asleep => Self::Asleep, + } + } +} + +impl From for RlVersionSource { + fn from(s: VersionSource) -> Self { + match s { + VersionSource::Registration => Self::Registration, + VersionSource::Passthrough => Self::Passthrough, + VersionSource::Api => Self::Api, + } + } +} + #[derive(Clone, Debug)] pub struct EngineState { /// `Worker::model_id()` at seed time; the fleet maximum is per model. From 0f2f7b32e186ca33fff951d4572328af636ccb8f Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 14:55:05 -0700 Subject: [PATCH 07/27] feat(rl): version and state routes on /v1/rl, registered in the OpenAPI spec Signed-off-by: key4ng --- clients/openapi-gen/src/main.rs | 67 +++++++ crates/rl/src/control.rs | 327 ++++++++++++++++++++++++++++++++ crates/rl/src/error.rs | 14 +- crates/rl/src/fanout.rs | 19 +- crates/rl/src/lib.rs | 10 +- 5 files changed, 428 insertions(+), 9 deletions(-) create mode 100644 crates/rl/src/control.rs diff --git a/clients/openapi-gen/src/main.rs b/clients/openapi-gen/src/main.rs index 445cd2969f..db40cfb266 100644 --- a/clients/openapi-gen/src/main.rs +++ b/clients/openapi-gen/src/main.rs @@ -639,6 +639,8 @@ fn main() -> anyhow::Result<()> { let rl_entry_name = collect_schema(schema_for!(RlWorkerEntry), &mut schemas)?; let rl_outcome_name = collect_schema(schema_for!(RlCallOutcome), &mut schemas)?; let rl_fanout_name = collect_schema(schema_for!(RlFanoutResponse), &mut schemas)?; + let rl_set_version_name = collect_schema(schema_for!(RlSetVersionRequest), &mut schemas)?; + let rl_set_control_name = collect_schema(schema_for!(RlSetControlRequest), &mut schemas)?; let json_content = |name: &str| -> BTreeMap { let mut content = BTreeMap::new(); content.insert( @@ -748,6 +750,71 @@ fn main() -> anyhow::Result<()> { }, ); + let json_body = |name: &str| RequestBody { + required: true, + content: json_content(name), + }; + let selector_param = || Parameter { + name: "selector".to_string(), + location: "query".to_string(), + required: true, + schema: ParameterSchema { + schema_type: "string".to_string(), + }, + }; + paths.insert( + "/v1/rl/workers/{worker_id}/version".to_string(), + PathItem { + post: Some(operation( + "setRlWorkerVersion", + "Record the weight version one engine holds (no engine call)", + Some(vec![path_param("worker_id")]), + Some(json_body(&rl_set_version_name)), + json_response(&rl_entry_name, "The updated worker"), + )), + ..PathItem::default() + }, + ); + paths.insert( + "/v1/rl/workers/{worker_id}/state".to_string(), + PathItem { + post: Some(operation( + "setRlWorkerState", + "Record whether one engine is active, paused, or asleep (no engine call)", + Some(vec![path_param("worker_id")]), + Some(json_body(&rl_set_control_name)), + json_response(&rl_entry_name, "The updated worker"), + )), + ..PathItem::default() + }, + ); + paths.insert( + "/v1/rl/version".to_string(), + PathItem { + post: Some(operation( + "setRlFleetVersion", + "Record the weight version every worker matching `selector` holds", + Some(vec![selector_param()]), + Some(json_body(&rl_set_version_name)), + json_response(&rl_workers_name, "The updated workers"), + )), + ..PathItem::default() + }, + ); + paths.insert( + "/v1/rl/state".to_string(), + PathItem { + post: Some(operation( + "setRlFleetState", + "Record the control state of every worker matching `selector`", + Some(vec![selector_param()]), + Some(json_body(&rl_set_control_name)), + json_response(&rl_workers_name, "The updated workers"), + )), + ..PathItem::default() + }, + ); + // ---- Generate (SGLang native) ---- use openai_protocol::generate::*; let req_name = collect_schema(schema_for!(GenerateRequest), &mut schemas)?; diff --git a/crates/rl/src/control.rs b/crates/rl/src/control.rs new file mode 100644 index 0000000000..c283b929c5 --- /dev/null +++ b/crates/rl/src/control.rs @@ -0,0 +1,327 @@ +//! Explicit version and control-state writes: local table updates, no engine call. + +use std::sync::Arc; + +use axum::{ + extract::{rejection::JsonRejection, Path, Query, State}, + http::StatusCode, + response::{IntoResponse, Response}, + Json, +}; +use openai_protocol::rl::{ + RlSetControlRequest, RlSetVersionRequest, RlWorkerEntry, RlWorkersResponse, RL_PROTOCOL_VERSION, +}; + +use crate::{ + discovery::{dp_ranks, entry}, + error::RlError, + fanout::{parse_selector, resolve_targets, FanoutQuery}, + state::RlState, + table::{ControlState, VersionSource}, + version::Version, + view::RlWorkerInfo, +}; + +/// Cap on an API-supplied version string. +const MAX_VERSION_BYTES: usize = 128; + +fn parse_version(raw: &str) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err(RlError::InvalidVersion( + "weight_version is empty".to_string(), + )); + } + if trimmed.len() > MAX_VERSION_BYTES { + return Err(RlError::InvalidVersion(format!( + "weight_version exceeds {MAX_VERSION_BYTES} bytes" + ))); + } + Ok(Version::parse(trimmed)) +} + +fn body(body: Result, JsonRejection>) -> Result { + body.map(|Json(v)| v) + .map_err(|e| RlError::InvalidBody(e.body_text())) +} + +fn row(state: &RlState, worker: &RlWorkerInfo) -> RlWorkerEntry { + entry(&state.table, worker, dp_ranks(state.view.as_ref(), worker)) +} + +fn rows(state: &RlState, workers: &[RlWorkerInfo]) -> Response { + let workers: Vec = workers.iter().map(|w| row(state, w)).collect(); + let total = workers.len(); + ( + StatusCode::OK, + Json(RlWorkersResponse { + protocol_version: RL_PROTOCOL_VERSION, + workers, + total, + }), + ) + .into_response() +} + +fn targets(state: &RlState, selector: Option) -> Result, RlError> { + let selector = parse_selector(selector)?; + let targets = resolve_targets(state.view.as_ref(), &state.table, &selector); + if targets.is_empty() { + return Err(RlError::NoWorkersMatch(selector.source().to_string())); + } + Ok(targets) +} + +pub(crate) async fn set_worker_version( + State(state): State>, + Path(id): Path, + payload: Result, JsonRejection>, +) -> Response { + let version = match body(payload).and_then(|b| parse_version(&b.weight_version)) { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let Some(worker) = state.view.get(&id) else { + return RlError::WorkerNotFound(id).into_response(); + }; + state.apply_version(&worker, version, VersionSource::Api); + (StatusCode::OK, Json(row(&state, &worker))).into_response() +} + +pub(crate) async fn set_worker_state( + State(state): State>, + Path(id): Path, + payload: Result, JsonRejection>, +) -> Response { + let control: ControlState = match body(payload) { + Ok(b) => b.control.into(), + Err(e) => return e.into_response(), + }; + let Some(worker) = state.view.get(&id) else { + return RlError::WorkerNotFound(id).into_response(); + }; + state.apply_control(&worker, control); + (StatusCode::OK, Json(row(&state, &worker))).into_response() +} + +pub(crate) async fn set_fleet_version( + State(state): State>, + Query(FanoutQuery { selector }): Query, + payload: Result, JsonRejection>, +) -> Response { + let version = match body(payload).and_then(|b| parse_version(&b.weight_version)) { + Ok(v) => v, + Err(e) => return e.into_response(), + }; + let workers = match targets(&state, selector) { + Ok(w) => w, + Err(e) => return e.into_response(), + }; + for worker in &workers { + state.apply_version(worker, version.clone(), VersionSource::Api); + } + rows(&state, &workers) +} + +pub(crate) async fn set_fleet_state( + State(state): State>, + Query(FanoutQuery { selector }): Query, + payload: Result, JsonRejection>, +) -> Response { + let control: ControlState = match body(payload) { + Ok(b) => b.control.into(), + Err(e) => return e.into_response(), + }; + let workers = match targets(&state, selector) { + Ok(w) => w, + Err(e) => return e.into_response(), + }; + for worker in &workers { + state.apply_control(worker, control); + } + rows(&state, &workers) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use axum::{ + body::Body, + http::{Request, StatusCode}, + response::Response, + }; + use http_body_util::BodyExt; + use openai_protocol::worker::RuntimeType; + use serde_json::Value; + use tower::ServiceExt; + + use crate::{ + config::RlConfig, + state::RlState, + testing::{worker, FakeView}, + view::RlWorkerInfo, + }; + + fn state(workers: Vec) -> Arc { + Arc::new(RlState::new( + Arc::new(FakeView(workers)), + RlConfig::default(), + )) + } + + async fn json_body(resp: Response) -> Value { + serde_json::from_slice(&resp.into_body().collect().await.unwrap().to_bytes()).unwrap() + } + + #[tokio::test] + async fn set_worker_version_updates_the_row() { + let app = crate::router::<()>(state(vec![worker("w1", "http://a:1", RuntimeType::Sglang)])); + let resp = app + .clone() + .oneshot( + Request::post("/workers/w1/version") + .header("content-type", "application/json") + .body(Body::from(r#"{"weight_version": "42"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let row = json_body(resp).await; + assert_eq!(row["id"], "w1"); + assert_eq!(row["weight_version"], "42"); + assert_eq!(row["version_source"], "api"); + let listed = json_body( + app.oneshot(Request::get("/workers").body(Body::empty()).unwrap()) + .await + .unwrap(), + ) + .await; + assert_eq!(listed["workers"][0]["weight_version"], "42"); + } + + #[tokio::test] + async fn set_worker_state_and_errors() { + let app = crate::router::<()>(state(vec![worker("w1", "http://a:1", RuntimeType::Sglang)])); + let resp = app + .clone() + .oneshot( + Request::post("/workers/w1/state") + .header("content-type", "application/json") + .body(Body::from(r#"{"control": "asleep"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(json_body(resp).await["control"], "asleep"); + + let resp = app + .clone() + .oneshot( + Request::post("/workers/nope/state") + .header("content-type", "application/json") + .body(Body::from(r#"{"control": "paused"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + let resp = app + .clone() + .oneshot( + Request::post("/workers/w1/state") + .header("content-type", "application/json") + .body(Body::from(r#"{"control": "napping"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + assert_eq!(json_body(resp).await["error"], "invalid_body"); + + let resp = app + .clone() + .oneshot( + Request::post("/workers/w1/version") + .header("content-type", "application/json") + .body(Body::from(r#"{"weight_version": " "}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + assert_eq!(json_body(resp).await["error"], "invalid_version"); + + let long = "x".repeat(129); + let resp = app + .oneshot( + Request::post("/workers/w1/version") + .header("content-type", "application/json") + .body(Body::from(format!(r#"{{"weight_version": "{long}"}}"#))) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + } + + #[tokio::test] + async fn fleet_forms_write_one_entry_per_engine_and_need_a_selector() { + let app = crate::router::<()>(state(vec![ + worker("r0", "http://b:1@0", RuntimeType::Sglang), + worker("r1", "http://b:1@1", RuntimeType::Sglang), + worker("v1", "http://c:1", RuntimeType::Vllm), + ])); + let resp = app + .clone() + .oneshot( + Request::post("/version?selector=engine%3Dsglang") + .header("content-type", "application/json") + .body(Body::from(r#"{"weight_version": "3"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let body = json_body(resp).await; + assert_eq!(body["total"], 1, "DP ranks collapse to one engine row"); + assert_eq!(body["workers"][0]["weight_version"], "3"); + assert_eq!(body["workers"][0]["dp_ranks"], 2); + let listed = json_body( + app.clone() + .oneshot(Request::get("/workers/v1").body(Body::empty()).unwrap()) + .await + .unwrap(), + ) + .await; + assert_eq!(listed["weight_version"], "default", "vllm worker untouched"); + + let resp = app + .clone() + .oneshot( + Request::post("/state") + .header("content-type", "application/json") + .body(Body::from(r#"{"control": "paused"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + assert_eq!(json_body(resp).await["error"], "selector_required"); + + let resp = app + .oneshot( + Request::post("/state?selector=engine%3Dtrtllm") + .header("content-type", "application/json") + .body(Body::from(r#"{"control": "paused"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + assert_eq!(json_body(resp).await["error"], "no_workers_match"); + } +} diff --git a/crates/rl/src/error.rs b/crates/rl/src/error.rs index 5b07ed45f4..a57714a554 100644 --- a/crates/rl/src/error.rs +++ b/crates/rl/src/error.rs @@ -21,6 +21,12 @@ pub enum RlError { NoWorkersMatch(String), #[error("worker `{0}` not found")] WorkerNotFound(String), + #[error("{0}")] + InvalidVersion(String), + #[error("{0}")] + InvalidBody(String), + #[error("{0}")] + InvalidVersionPolicy(String), #[error("worker `{worker_id}` uses connection mode `{mode}`, which cannot be proxied")] UnsupportedConnectionMode { worker_id: String, @@ -50,6 +56,9 @@ impl RlError { Self::InvalidSelector { .. } => "invalid_selector", Self::NoWorkersMatch(_) => "no_workers_match", Self::WorkerNotFound(_) => "worker_not_found", + Self::InvalidVersion(_) => "invalid_version", + Self::InvalidBody(_) => "invalid_body", + Self::InvalidVersionPolicy(_) => "invalid_version_policy", Self::UnsupportedConnectionMode { .. } => "unsupported_connection_mode", Self::UpstreamUnreachable { .. } => "upstream_unreachable", Self::UpstreamTimeout { .. } => "upstream_timeout", @@ -61,7 +70,10 @@ impl RlError { Self::InvalidEnginePath(_) | Self::SelectorRequired | Self::InvalidSelector { .. } - | Self::NoWorkersMatch(_) => StatusCode::BAD_REQUEST, + | Self::NoWorkersMatch(_) + | Self::InvalidVersion(_) + | Self::InvalidBody(_) + | Self::InvalidVersionPolicy(_) => StatusCode::BAD_REQUEST, Self::WorkerNotFound(_) => StatusCode::NOT_FOUND, Self::UnsupportedConnectionMode { .. } => StatusCode::UNPROCESSABLE_ENTITY, Self::UpstreamUnreachable { .. } => StatusCode::BAD_GATEWAY, diff --git a/crates/rl/src/fanout.rs b/crates/rl/src/fanout.rs index 978b98781e..1ee32d40c7 100644 --- a/crates/rl/src/fanout.rs +++ b/crates/rl/src/fanout.rs @@ -27,7 +27,15 @@ use crate::{ #[derive(Deserialize)] pub(crate) struct FanoutQuery { - selector: Option, + pub(crate) selector: Option, +} + +/// The `selector` query parameter, required and non-empty. +pub(crate) fn parse_selector(selector: Option) -> Result { + match selector.as_deref().map(str::trim) { + Some(s) if !s.is_empty() => Selector::parse(s), + _ => Err(RlError::SelectorRequired), + } } /// DP-collapsed workers matching `selector`, sorted by base URL. @@ -121,12 +129,9 @@ pub(crate) async fn fanout_handler( body: Bytes, ) -> Response { let started = Instant::now(); - let selector = match selector.as_deref().map(str::trim) { - Some(s) if !s.is_empty() => match Selector::parse(s) { - Ok(sel) => sel, - Err(e) => return e.into_response(), - }, - _ => return RlError::SelectorRequired.into_response(), + let selector = match parse_selector(selector) { + Ok(sel) => sel, + Err(e) => return e.into_response(), }; let req = match ProxyRequest::from_parts(method, &raw_path, raw_query.as_deref(), &headers, body) { diff --git a/crates/rl/src/lib.rs b/crates/rl/src/lib.rs index d690931c30..cd783ae67a 100644 --- a/crates/rl/src/lib.rs +++ b/crates/rl/src/lib.rs @@ -4,6 +4,7 @@ pub mod capability; pub mod config; +pub mod control; pub mod discovery; pub mod error; pub mod fanout; @@ -21,7 +22,10 @@ pub mod view; use std::sync::Arc; -use axum::{routing::get, Router}; +use axum::{ + routing::{get, post}, + Router, +}; pub use config::RlConfig; pub use error::RlError; pub use metrics::init_rl_metrics; @@ -45,5 +49,9 @@ pub fn router(state: Arc) -> Router Date: Thu, 10 Sep 2026 15:02:57 -0700 Subject: [PATCH 08/27] feat(rl): observe versions and pause/sleep transitions on proxied control calls Signed-off-by: key4ng --- crates/rl/src/fanout.rs | 22 +++++++ crates/rl/src/lib.rs | 1 + crates/rl/src/observe.rs | 135 +++++++++++++++++++++++++++++++++++++++ crates/rl/src/proxy.rs | 84 ++++++++++++++++++++++++ crates/rl/src/state.rs | 18 ++++++ 5 files changed, 260 insertions(+) create mode 100644 crates/rl/src/observe.rs diff --git a/crates/rl/src/fanout.rs b/crates/rl/src/fanout.rs index 1ee32d40c7..7df799ce7c 100644 --- a/crates/rl/src/fanout.rs +++ b/crates/rl/src/fanout.rs @@ -230,6 +230,28 @@ mod tests { ); } + #[tokio::test] + async fn a_fanned_out_refit_records_the_version_on_every_target() { + let e1 = FakeEngine::start(StatusCode::OK, json!({}), 0).await; + let e2 = FakeEngine::start(StatusCode::OK, json!({}), 0).await; + let state = state( + vec![ + worker("w1", &e1.url, RuntimeType::Sglang), + worker("w2", &e2.url, RuntimeType::Sglang), + ], + 8, + ); + let app = crate::router::<()>(Arc::clone(&state)); + let req = Request::post("/engine/update_weight_version?selector=engine%3Dsglang") + .header("content-type", "application/json") + .body(Body::from(r#"{"new_version": "9"}"#)) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(state.table().version_of(&e1.url).unwrap().as_str(), "9"); + assert_eq!(state.table().version_of(&e2.url).unwrap().as_str(), "9"); + } + #[tokio::test] async fn one_failure_is_207_and_names_the_worker() { let good = FakeEngine::start(StatusCode::OK, json!({}), 0).await; diff --git a/crates/rl/src/lib.rs b/crates/rl/src/lib.rs index cd783ae67a..6771cae204 100644 --- a/crates/rl/src/lib.rs +++ b/crates/rl/src/lib.rs @@ -9,6 +9,7 @@ pub mod discovery; pub mod error; pub mod fanout; pub mod metrics; +pub mod observe; pub mod path; pub mod policy; pub mod proxy; diff --git a/crates/rl/src/observe.rs b/crates/rl/src/observe.rs new file mode 100644 index 0000000000..1e313882ea --- /dev/null +++ b/crates/rl/src/observe.rs @@ -0,0 +1,135 @@ +//! Read-only observation of engine control calls that pass through the +//! proxy: what a 2xx answer to a route tells SMG about the engine. The +//! forwarded bytes are never modified. + +use serde::Deserialize; +use serde_json::Value; + +use crate::{table::ControlState, version::Version}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Observation { + Version(Version), + Control(ControlState), +} + +/// Only the version fields are read; everything else in a refit body is +/// skipped without allocation. +#[derive(Deserialize)] +struct VersionProbe { + #[serde(default)] + weight_version: Option, + #[serde(default)] + new_version: Option, +} + +fn version_field(value: Option) -> Option { + match value? { + Value::String(s) if !s.trim().is_empty() => Some(Version::parse(&s)), + Value::Number(n) => Some(Version::parse(&n.to_string())), + _ => None, + } +} + +/// What a successful call to `path` (the proxied engine route, no leading +/// slash) with request `body` says about the engine, if anything. +pub fn observe(path: &str, body: &[u8]) -> Option { + match path { + "update_weights_from_disk" + | "update_weights_from_tensor" + | "update_weights_from_distributed" => { + let probe: VersionProbe = serde_json::from_slice(body).ok()?; + version_field(probe.weight_version).map(Observation::Version) + } + "update_weight_version" => { + let probe: VersionProbe = serde_json::from_slice(body).ok()?; + version_field(probe.new_version).map(Observation::Version) + } + "pause_generation" | "pause" => Some(Observation::Control(ControlState::Paused)), + "continue_generation" | "resume" | "resume_memory_occupation" | "wake_up" => { + Some(Observation::Control(ControlState::Active)) + } + "release_memory_occupation" | "sleep" => Some(Observation::Control(ControlState::Asleep)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn refit_routes_report_the_version_field_when_present() { + let v = observe( + "update_weights_from_disk", + br#"{"model_path": "/ckpt", "weight_version": "42"}"#, + ); + assert_eq!(v, Some(Observation::Version(Version::parse("42")))); + let v = observe( + "update_weights_from_tensor", + br#"{"serialized_named_tensors": ["..."], "weight_version": 7}"#, + ); + assert_eq!(v, Some(Observation::Version(Version::parse("7")))); + let v = observe( + "update_weights_from_distributed", + br#"{"names": [], "weight_version": "s3"}"#, + ); + assert_eq!(v, Some(Observation::Version(Version::parse("s3")))); + assert_eq!( + observe( + "update_weight_version", + br#"{"new_version": "43", "abort_all_requests": false}"# + ), + Some(Observation::Version(Version::parse("43"))) + ); + } + + #[test] + fn refit_routes_without_a_version_observe_nothing() { + assert_eq!( + observe("update_weights_from_disk", br#"{"model_path": "/ckpt"}"#), + None + ); + assert_eq!( + observe("update_weights_from_disk", br#"{"weight_version": ""}"#), + None + ); + assert_eq!( + observe("update_weights_from_disk", br#"{"weight_version": null}"#), + None + ); + assert_eq!(observe("update_weights_from_disk", b"not json"), None); + assert_eq!( + observe("update_weight_version", br#"{"weight_version": "1"}"#), + None, + "wrong field for this route" + ); + } + + #[test] + fn control_routes_map_to_states_regardless_of_body() { + for (path, state) in [ + ("pause_generation", ControlState::Paused), + ("pause", ControlState::Paused), + ("continue_generation", ControlState::Active), + ("resume", ControlState::Active), + ("release_memory_occupation", ControlState::Asleep), + ("sleep", ControlState::Asleep), + ("resume_memory_occupation", ControlState::Active), + ("wake_up", ControlState::Active), + ] { + assert_eq!( + observe(path, b"{}"), + Some(Observation::Control(state)), + "{path}" + ); + assert_eq!( + observe(path, b""), + Some(Observation::Control(state)), + "{path}" + ); + } + assert_eq!(observe("flush_cache", b"{}"), None); + assert_eq!(observe("server_info", b""), None); + } +} diff --git a/crates/rl/src/proxy.rs b/crates/rl/src/proxy.rs index ad267baa07..e5539bb8f7 100644 --- a/crates/rl/src/proxy.rs +++ b/crates/rl/src/proxy.rs @@ -23,8 +23,10 @@ use crate::{ discovery::enum_str, error::RlError, metrics::{op_label, record_control_call}, + observe::observe, path::{passthrough_query, validate_engine_path}, state::RlState, + table::VersionSource, view::RlWorkerInfo, }; @@ -220,6 +222,11 @@ pub async fn call_worker( body, body_truncated, }; + if outcome.is_success() { + if let Some(observation) = observe(&req.path, &req.body) { + state.apply_observation(worker, observation, VersionSource::Passthrough); + } + } record_control_call( op, if outcome.is_success() { @@ -522,6 +529,83 @@ mod tests { assert_eq!(v.as_str().unwrap().len(), BODY_CAP); } + #[tokio::test] + async fn a_successful_refit_records_the_version_and_a_failed_one_does_not() { + let ok = FakeEngine::start(StatusCode::OK, json!({"success": true}), 0).await; + let bad = FakeEngine::start( + StatusCode::INTERNAL_SERVER_ERROR, + json!({"success": false}), + 0, + ) + .await; + let state = state( + vec![ + worker("w1", &ok.url, RuntimeType::Sglang), + worker("w2", &bad.url, RuntimeType::Sglang), + ], + 5, + ); + let app = crate::router::<()>(Arc::clone(&state)); + for id in ["w1", "w2"] { + let _ = app + .clone() + .oneshot( + Request::post(format!("/workers/{id}/engine/update_weights_from_disk")) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model_path": "/ckpt", "weight_version": "5"}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + } + assert_eq!(state.table().version_of(&ok.url).unwrap().as_str(), "5"); + assert_eq!(state.table().version_of(&bad.url), None); + assert_eq!( + state.table().get(&ok.url).unwrap().version_source, + Some(VersionSource::Passthrough) + ); + } + + #[tokio::test] + async fn pause_and_resume_flip_the_control_state() { + let engine = FakeEngine::start(StatusCode::OK, json!({}), 0).await; + let state = state(vec![worker("w1", &engine.url, RuntimeType::Sglang)], 5); + let app = crate::router::<()>(Arc::clone(&state)); + let post = |path: &str| { + Request::post(format!("/workers/w1/engine/{path}")) + .body(Body::from("{}")) + .unwrap() + }; + app.clone().oneshot(post("pause_generation")).await.unwrap(); + assert_eq!( + state.table().control_of(&engine.url), + crate::table::ControlState::Paused + ); + app.clone() + .oneshot(post("continue_generation")) + .await + .unwrap(); + assert_eq!( + state.table().control_of(&engine.url), + crate::table::ControlState::Active + ); + app.clone() + .oneshot(post("release_memory_occupation")) + .await + .unwrap(); + assert_eq!( + state.table().control_of(&engine.url), + crate::table::ControlState::Asleep + ); + app.oneshot(post("resume_memory_occupation")).await.unwrap(); + assert_eq!( + state.table().control_of(&engine.url), + crate::table::ControlState::Active + ); + } + #[tokio::test] async fn large_json_bodies_are_truncated_and_flagged() { let engine = FakeEngine::start( diff --git a/crates/rl/src/state.rs b/crates/rl/src/state.rs index 4ecdf6f14d..3b051e3ec3 100644 --- a/crates/rl/src/state.rs +++ b/crates/rl/src/state.rs @@ -9,6 +9,7 @@ use tracing::info; use crate::{ config::RlConfig, metrics, + observe::Observation, policy::VersionPolicy, table::{ControlState, NoopEvictionSink, RlTable, VersionEvictionSink, VersionSource}, version::Version, @@ -86,6 +87,23 @@ impl RlState { changed } + /// Apply what a successful proxied call revealed about `worker`'s engine. + pub fn apply_observation( + &self, + worker: &RlWorkerInfo, + observation: Observation, + source: VersionSource, + ) { + match observation { + Observation::Version(version) => { + self.apply_version(worker, version, source); + } + Observation::Control(control) => { + self.apply_control(worker, control); + } + } + } + /// Record a control state for `worker`'s engine. pub fn apply_control(&self, worker: &RlWorkerInfo, control: ControlState) -> bool { let changed = self From fead9cae11e73d879b818d3d60c3dc26c4d6d7b7 Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 15:08:12 -0700 Subject: [PATCH 09/27] feat(rl): mixed-version probe for buffered generate bodies Signed-off-by: key4ng --- crates/rl/src/lib.rs | 1 + crates/rl/src/stamp.rs | 126 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 crates/rl/src/stamp.rs diff --git a/crates/rl/src/lib.rs b/crates/rl/src/lib.rs index 6771cae204..64dd2b6c25 100644 --- a/crates/rl/src/lib.rs +++ b/crates/rl/src/lib.rs @@ -14,6 +14,7 @@ pub mod path; pub mod policy; pub mod proxy; pub mod selector; +pub mod stamp; pub mod state; pub mod table; #[cfg(test)] diff --git a/crates/rl/src/stamp.rs b/crates/rl/src/stamp.rs new file mode 100644 index 0000000000..db0a9c1695 --- /dev/null +++ b/crates/rl/src/stamp.rs @@ -0,0 +1,126 @@ +//! Response stamping: the header names the gateway sets, and the probe that +//! decides whether a buffered `/generate` body spanned two weight versions. + +use serde::{de::IgnoredAny, Deserialize}; +use serde_json::Value; + +use crate::version::Version; + +/// The version SMG believed the routed engine held when it emitted the +/// response head. +pub const WEIGHT_VERSION_HEADER: &str = "x-smg-weight-version"; +/// Set to `true` on a buffered `/generate` response whose tokens came from +/// more than one version, or whose engine-reported version differs from the +/// stamped one. +pub const MIXED_VERSION_HEADER: &str = "x-smg-mixed-version"; + +#[derive(Deserialize, Default)] +struct MetaProbe { + #[serde(default)] + weight_version: Option, + #[serde(default)] + weight_versions: Option>, +} + +#[derive(Deserialize)] +struct GenerateProbe { + #[serde(default)] + meta_info: Option, +} + +fn reported_version(value: &Value) -> Option { + match value { + Value::String(s) if !s.trim().is_empty() => Some(Version::parse(s)), + Value::Number(n) => Some(Version::parse(&n.to_string())), + _ => None, + } +} + +fn is_mixed(meta: &MetaProbe, expected: Option<&Version>) -> bool { + if meta + .weight_versions + .as_ref() + .is_some_and(|spans| spans.len() > 1) + { + return true; + } + match ( + meta.weight_version.as_ref().and_then(reported_version), + expected, + ) { + (Some(reported), Some(expected)) => reported != *expected, + _ => false, + } +} + +/// Whether a buffered `/generate` body (one object, or an array for `n > 1`) +/// spans more than one weight version or reports a version other than +/// `expected`. Only `meta_info.weight_version` and the span count are read; +/// an unparseable body is not mixed. +pub fn generate_is_mixed(body: &[u8], expected: Option<&Version>) -> bool { + let first = body.iter().find(|b| !b.is_ascii_whitespace()); + match first { + Some(b'[') => serde_json::from_slice::>(body) + .map(|items| { + items + .iter() + .filter_map(|p| p.meta_info.as_ref()) + .any(|m| is_mixed(m, expected)) + }) + .unwrap_or(false), + Some(b'{') => serde_json::from_slice::(body) + .ok() + .and_then(|p| p.meta_info) + .is_some_and(|m| is_mixed(&m, expected)), + _ => false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const TWO_SPANS: &[u8] = br#"{"text":"x","meta_info":{"weight_version":"42","weight_versions":[{"version":"41","start":0,"end":57},{"version":"42","start":57,"end":128}]}}"#; + const ONE_SPAN: &[u8] = br#"{"text":"x","meta_info":{"weight_version":"42","weight_versions":[{"version":"42","start":0,"end":128}]}}"#; + + #[test] + fn two_spans_are_mixed_whatever_the_expectation() { + assert!(generate_is_mixed(TWO_SPANS, None)); + assert!(generate_is_mixed(TWO_SPANS, Some(&Version::parse("42")))); + } + + #[test] + fn one_span_is_mixed_only_when_the_reported_version_differs() { + assert!(!generate_is_mixed(ONE_SPAN, None)); + assert!(!generate_is_mixed(ONE_SPAN, Some(&Version::parse("42")))); + assert!(!generate_is_mixed(ONE_SPAN, Some(&Version::parse("042")))); + assert!(generate_is_mixed(ONE_SPAN, Some(&Version::parse("41")))); + let no_spans = br#"{"text":"x","meta_info":{"weight_version":"default"}}"#; + assert!(!generate_is_mixed(no_spans, None)); + assert!(generate_is_mixed(no_spans, Some(&Version::parse("3")))); + } + + #[test] + fn batches_and_garbage() { + let batch = format!( + "[{}, {}]", + std::str::from_utf8(ONE_SPAN).unwrap(), + std::str::from_utf8(TWO_SPANS).unwrap() + ); + assert!(generate_is_mixed( + batch.as_bytes(), + Some(&Version::parse("42")) + )); + let clean = format!(" [{}]", std::str::from_utf8(ONE_SPAN).unwrap()); + assert!(!generate_is_mixed( + clean.as_bytes(), + Some(&Version::parse("42")) + )); + assert!(!generate_is_mixed(b"", None)); + assert!(!generate_is_mixed(b"not json", Some(&Version::parse("1")))); + assert!(!generate_is_mixed( + br#"{"text":"no meta"}"#, + Some(&Version::parse("1")) + )); + } +} From a06ca3ee572527c8f5d6eb1c346b2ee249a393de Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 15:45:25 -0700 Subject: [PATCH 10/27] feat(rl): --rl-version-policy flag with Python parity Signed-off-by: key4ng --- bindings/python/src/lib.rs | 17 ++++++++++++++++- bindings/python/src/smg/router_args.py | 11 +++++++++++ bindings/python/tests/test_arg_parser.py | 6 ++++++ model_gateway/src/config/validation.rs | 16 ++++++++++++++-- model_gateway/src/main.rs | 19 ++++++++++++++++++- 5 files changed, 65 insertions(+), 4 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index f3f5260aa3..baadc0fbce 100755 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -534,6 +534,7 @@ struct Router { enable_rl: bool, rl_control_timeout_secs: u64, rl_fanout_concurrency: usize, + rl_version_policy: String, } impl Router { @@ -832,6 +833,17 @@ impl Router { None }; + let rl_version_policy: smg_rl::VersionPolicy = + self.rl_version_policy + .parse() + .map_err( + |e: smg_rl::VersionPolicyError| config::ConfigError::InvalidValue { + field: "rl-version-policy".to_string(), + value: self.rl_version_policy.clone(), + reason: e.reason.to_string(), + }, + )?; + config::RouterConfig::builder() .mode(mode) .policy(policy) @@ -944,7 +956,7 @@ impl Router { enabled: self.enable_rl, control_timeout_secs: self.rl_control_timeout_secs, fanout_concurrency: self.rl_fanout_concurrency, - version_policy: smg_rl::VersionPolicy::default(), + version_policy: rl_version_policy, }) .maybe_client_cert_and_key( self.client_cert_path.as_ref(), @@ -1118,6 +1130,7 @@ impl Router { enable_rl = false, rl_control_timeout_secs = 600, rl_fanout_concurrency = 32, + rl_version_policy = String::from("any"), ))] #[expect(clippy::too_many_arguments)] #[expect( @@ -1277,6 +1290,7 @@ impl Router { enable_rl: bool, rl_control_timeout_secs: u64, rl_fanout_concurrency: usize, + rl_version_policy: String, ) -> PyResult { let mut all_urls = worker_urls.clone(); @@ -1448,6 +1462,7 @@ impl Router { enable_rl, rl_control_timeout_secs, rl_fanout_concurrency, + rl_version_policy, }) } diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index e58ae44088..f0bf9dd9ef 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -257,6 +257,7 @@ class RouterArgs: enable_rl: bool = False # Mount the RL control plane under /v1/rl rl_control_timeout_secs: int = 600 # Timeout for one proxied engine control call rl_fanout_concurrency: int = 32 # Max concurrent engine calls in one fan-out + rl_version_policy: str = "any" # any | latest-only | min-version: | max-staleness: @staticmethod def add_cli_args( @@ -813,6 +814,16 @@ def add_cli_args( default=RouterArgs.rl_fanout_concurrency, help="Maximum concurrent engine calls in one fan-out (default: 32)", ) + rl_group.add_argument( + f"--{prefix}rl-version-policy", + type=str, + default=RouterArgs.rl_version_policy, + help=( + "Which engines a request may be routed to, judged against the model's newest " + "known weight version: any | latest-only | min-version: | max-staleness: " + "(default: any). Per-request override: the x-smg-version-policy header." + ), + ) # PD/EPD-specific arguments pd_group.add_argument( diff --git a/bindings/python/tests/test_arg_parser.py b/bindings/python/tests/test_arg_parser.py index f0f08c3296..1b994e9b49 100644 --- a/bindings/python/tests/test_arg_parser.py +++ b/bindings/python/tests/test_arg_parser.py @@ -50,6 +50,7 @@ def test_default_values(self): assert args.enable_rl is False assert args.rl_control_timeout_secs == 600 assert args.rl_fanout_concurrency == 32 + assert args.rl_version_policy == "any" def test_rl_flags_parse(self): """RL control-plane flags land on the dataclass.""" @@ -63,12 +64,15 @@ def test_rl_flags_parse(self): "1200", "--rl-fanout-concurrency", "8", + "--rl-version-policy", + "max-staleness:2", ] ) ) assert args.enable_rl is True assert args.rl_control_timeout_secs == 1200 assert args.rl_fanout_concurrency == 8 + assert args.rl_version_policy == "max-staleness:2" def test_parse_selector_valid(self): """Test parsing valid selector arguments.""" @@ -1468,6 +1472,7 @@ class TestRouterArgsFieldOrder: "enable_rl", "rl_control_timeout_secs", "rl_fanout_concurrency", + "rl_version_policy", ] def test_complete_field_sequence_is_frozen(self): @@ -1506,6 +1511,7 @@ def test_new_fields_appended_after_positional_reserve(self): "enable_rl", "rl_control_timeout_secs", "rl_fanout_concurrency", + "rl_version_policy", ): assert names.index(appended) > marker, ( f"{appended} must be appended after worker_startup_delay to " diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index cf4a2ed259..4a1159ee40 100644 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -844,8 +844,10 @@ impl ConfigValidator { .map_err(|reason| ConfigError::InvalidValue { field: "rl".to_string(), value: format!( - "control_timeout_secs={} fanout_concurrency={}", - config.rl.control_timeout_secs, config.rl.fanout_concurrency + "control_timeout_secs={} fanout_concurrency={} version_policy={}", + config.rl.control_timeout_secs, + config.rl.fanout_concurrency, + config.rl.version_policy ), reason, }) @@ -2308,4 +2310,14 @@ mod tests { .build_unchecked(); assert!(ConfigValidator::validate(&ok).is_ok()); } + + #[test] + fn rl_version_policy_requires_the_flag() { + let mut config = RouterConfig::default(); + config.rl.version_policy = smg_rl::VersionPolicy::LatestOnly; + let err = ConfigValidator::validate(&config).unwrap_err().to_string(); + assert!(err.contains("version_policy"), "{err}"); + config.rl.enabled = true; + assert!(ConfigValidator::validate(&config).is_ok()); + } } diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 55b9a5d826..a766a64051 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -485,6 +485,14 @@ struct CliArgs { #[arg(long, default_value_t = 32, help_heading = "RL Control Plane")] rl_fanout_concurrency: usize, + /// Which engines a request may be routed to, judged against the model's + /// newest known weight version: any | latest-only | min-version: | + /// max-staleness:. Per-request override: the x-smg-version-policy + /// header. An empty eligible set is a retryable 503; raise + /// --retry-max-retries / --retry-max-backoff-ms to wait out a refit. + #[arg(long, default_value = "any", help_heading = "RL Control Plane")] + rl_version_policy: String, + // ==================== PD Disaggregation ==================== /// Enable PD (Prefill-Decode) disaggregated mode #[arg(long, default_value_t = false, help_heading = "PD Disaggregation")] @@ -1832,6 +1840,15 @@ impl CliArgs { } } + let rl_version_policy: smg_rl::VersionPolicy = + self.rl_version_policy + .parse() + .map_err(|e: smg_rl::VersionPolicyError| ConfigError::InvalidValue { + field: "rl-version-policy".to_string(), + value: self.rl_version_policy.clone(), + reason: e.reason.to_string(), + })?; + let builder = RouterConfig::builder() .mode(mode) .policy(policy) @@ -1958,7 +1975,7 @@ impl CliArgs { enabled: self.enable_rl, control_timeout_secs: self.rl_control_timeout_secs, fanout_concurrency: self.rl_fanout_concurrency, - version_policy: smg_rl::VersionPolicy::default(), + version_policy: rl_version_policy, }) .dp_minimum_tokens_scheduler(self.dp_minimum_tokens_scheduler) .maybe_server_cert_and_key(self.tls_cert_path.as_ref(), self.tls_key_path.as_ref()); From 25effb3a8bff1bf77c1a51b60b99c63a629974a5 Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 16:04:47 -0700 Subject: [PATCH 11/27] feat(policies): candidate filter slot ahead of the sticky override and the policy Signed-off-by: key4ng --- model_gateway/src/policies/mod.rs | 15 +++ model_gateway/src/policies/registry.rs | 134 ++++++++++++++++++++++++- 2 files changed, 145 insertions(+), 4 deletions(-) diff --git a/model_gateway/src/policies/mod.rs b/model_gateway/src/policies/mod.rs index 9bf79d16c4..5f839a8f2a 100644 --- a/model_gateway/src/policies/mod.rs +++ b/model_gateway/src/policies/mod.rs @@ -110,6 +110,21 @@ pub trait LoadBalancingPolicy: Send + Sync + Debug { fn as_any(&self) -> &dyn std::any::Any; } +/// Narrows the candidate set before the sticky override and the policy run. +/// Installed once on the [`PolicyRegistry`] (the RL control plane uses it to +/// drop paused, asleep, or too-stale engines); every policy-driven selection +/// path goes through it. +pub trait CandidateFilter: Send + Sync { + /// Indices into `workers` that stay eligible; `None` when every candidate + /// is eligible (the allocation-free path). An empty vector makes the + /// selection fail, which callers already treat as "no available worker". + fn eligible( + &self, + workers: &[Arc], + headers: Option<&http::HeaderMap>, + ) -> Option>; +} + /// Whether a built-in policy applies the complete [`Worker::is_available`] /// predicate before selecting. Other policies keep the caller-side filter. pub(crate) fn policy_filters_unavailable_workers(policy: &dyn LoadBalancingPolicy) -> bool { diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 9fd45004a7..f681793c6d 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -17,8 +17,8 @@ use tracing::{debug, info, warn}; use super::{ get_healthy_worker_indices, manual::{ExecutionBranch, PinState}, - BucketPolicy, CacheAwarePolicy, DPRankLoadPolicy, LoadBalancingPolicy, ManualConfig, - ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, + BucketPolicy, CacheAwarePolicy, CandidateFilter, DPRankLoadPolicy, LoadBalancingPolicy, + ManualConfig, ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, }; use crate::{ config::types::{ManualAssignmentMode, PdPairingMode, PolicyConfig, RoutingKeyOverrideConfig}, @@ -75,6 +75,10 @@ pub struct PolicyRegistry { // DP-rank policy: Supports the selection of dp-rank outside the engine. dp_rank_policy: Arc>>, + /// Optional pre-policy candidate filter (see [`CandidateFilter`]), set + /// once at startup like the PD leg policies. + candidate_filter: Arc>>, + /// Shared sticky selector for the routing-key override. `Some` when the /// override is enabled; consulted (instead of the configured policy) for keyed /// requests via [`PolicyRegistry::select_worker`]. @@ -161,6 +165,7 @@ impl PolicyRegistry { load_rx: Arc::new(RwLock::new(None)), mesh_tree_sync: Arc::new(RwLock::new(None)), dp_rank_policy: Arc::new(OnceLock::new()), + candidate_filter: Arc::new(OnceLock::new()), routing_key_sticky, routing_key_headers: Arc::new(routing_key_headers), pd_pairing_mode: PdPairingMode::default(), @@ -251,12 +256,39 @@ impl PolicyRegistry { /// enabled, the request carries a key from the configured source, and the /// configured policy does not already honor the key (`manual` / /// `consistent_hashing`). Otherwise delegates to `policy`. `policy.name()` - /// stays the real policy (for metrics). + /// stays the real policy (for metrics). A candidate filter, when + /// installed, narrows `workers` first; the returned index always refers + /// to the caller's slice. pub fn select_worker( &self, policy: &Arc, workers: &[Arc], info: &SelectWorkerInfo, + ) -> Option { + if let Some(filter) = self.candidate_filter.get() { + if let Some(keep) = filter.eligible(workers, info.headers) { + if keep.is_empty() { + return None; + } + if keep.len() < workers.len() { + let subset: Vec> = + keep.iter().map(|&i| Arc::clone(&workers[i])).collect(); + return self + .select_unfiltered(policy, &subset, info) + .and_then(|i| keep.get(i).copied()); + } + } + } + self.select_unfiltered(policy, workers, info) + } + + /// Today's selection: the sticky routing-key override when it applies, + /// else the policy. + fn select_unfiltered( + &self, + policy: &Arc, + workers: &[Arc], + info: &SelectWorkerInfo, ) -> Option { if let Some(sticky) = self.routing_key_sticky.as_ref() { if Self::routing_key_override_applies(policy.name()) { @@ -693,6 +725,16 @@ impl PolicyRegistry { self.dp_rank_policy.get().map(Arc::clone) } + /// Install the candidate filter. Returns `false` (and keeps the first) + /// when one was already set. + pub fn set_candidate_filter(&self, filter: Arc) -> bool { + self.candidate_filter.set(filter).is_ok() + } + + pub fn has_candidate_filter(&self) -> bool { + self.candidate_filter.get().is_some() + } + /// Set the decode policy for PD mode (lock-free, set once at startup) pub fn set_decode_policy(&self, policy: Arc) { // OnceLock::set returns Err if already set, which we ignore since @@ -979,7 +1021,7 @@ impl std::fmt::Debug for PolicyRegistry { #[cfg(test)] mod tests { - use openai_protocol::worker::HealthCheckConfig; + use openai_protocol::worker::{HealthCheckConfig, WorkerStatus}; use tracing_test::traced_test; use super::*; @@ -2059,4 +2101,88 @@ mod tests { registry.remove_worker_from_pd_cache_aware("http://prefill-1:8000"); registry.remove_worker_from_pd_cache_aware("http://decode-1:8000"); } + + struct KeepIndices(Option>); + + impl CandidateFilter for KeepIndices { + fn eligible( + &self, + _workers: &[Arc], + _headers: Option<&HeaderMap>, + ) -> Option> { + self.0.clone() + } + } + + fn three_workers() -> Vec> { + ["http://w1:1", "http://w2:1", "http://w3:1"] + .iter() + .map(|url| { + let w: Arc = Arc::new( + BasicWorkerBuilder::new(*url) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(), + ); + w.set_status(WorkerStatus::Ready); + w + }) + .collect() + } + + #[test] + fn candidate_filter_remaps_indices_into_the_callers_slice() { + let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); + assert!(registry.set_candidate_filter(Arc::new(KeepIndices(Some(vec![2]))))); + assert!( + !registry.set_candidate_filter(Arc::new(KeepIndices(None))), + "set once" + ); + let workers = three_workers(); + let policy = registry.get_policy_or_default("m"); + for _ in 0..4 { + assert_eq!( + registry.select_worker(&policy, &workers, &SelectWorkerInfo::default()), + Some(2) + ); + } + } + + #[test] + fn empty_eligible_set_selects_nothing() { + let registry = PolicyRegistry::new(PolicyConfig::Random); + registry.set_candidate_filter(Arc::new(KeepIndices(Some(vec![])))); + let workers = three_workers(); + let policy = registry.get_policy_or_default("m"); + assert_eq!( + registry.select_worker(&policy, &workers, &SelectWorkerInfo::default()), + None + ); + } + + #[test] + fn a_filter_that_keeps_everything_changes_nothing() { + for config in [ + PolicyConfig::RoundRobin, + PolicyConfig::Random, + PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }, + ] { + let plain = PolicyRegistry::new(config.clone()); + let filtered = PolicyRegistry::new(config); + filtered.set_candidate_filter(Arc::new(KeepIndices(None))); + let workers = three_workers(); + let info = SelectWorkerInfo::default(); + // Round robin is deterministic; random/p2 are checked for validity only. + let a = plain.select_worker(&plain.get_policy_or_default("m"), &workers, &info); + let b = filtered.select_worker(&filtered.get_policy_or_default("m"), &workers, &info); + assert!(a.is_some() && b.is_some()); + if plain.get_policy_or_default("m").name() == "round_robin" { + assert_eq!(a, b); + } + } + } } From 9c1752ff09f9f55fcd359ed26bf537b99b048c41 Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 16:25:57 -0700 Subject: [PATCH 12/27] feat(policies): reset a worker's cache-aware entries without dropping its load state Signed-off-by: key4ng --- model_gateway/src/policies/cache_aware.rs | 158 ++++++++++++++++++++++ model_gateway/src/policies/registry.rs | 84 ++++++++++++ 2 files changed, 242 insertions(+) diff --git a/model_gateway/src/policies/cache_aware.rs b/model_gateway/src/policies/cache_aware.rs index 38c097e3d1..d4b9a46063 100644 --- a/model_gateway/src/policies/cache_aware.rs +++ b/model_gateway/src/policies/cache_aware.rs @@ -680,6 +680,65 @@ impl CacheAwarePolicy { } } + /// The worker's engine flushed its cache (a weight refit): forget what the + /// trees and hash placements hold for it, keep it as a tenant by re-seeding + /// the roots the way [`Self::add_worker`] does, and leave the load scorer + /// alone. Unlike [`Self::remove_worker_by_url`] the worker is not going + /// away, so its backend report and since-poll dispatch credit stay valid. + pub fn reset_worker_cache(&self, worker: &dyn Worker) { + let url = worker.url(); + let tenant: TenantId = Arc::from(url); + let tree_key = normalize_model_key(worker.model_id()); + if let Some(tree) = self.string_trees.get(tree_key) { + tree.remove_tenant_all(&tenant); + tree.insert_text("", url); + } + if let Some(tree) = self.token_trees.get(tree_key) { + tree.remove_tenant_all(&tenant); + tree.insert_tokens(&[], url); + } + // Hash mode (`cache_index = hash`) keeps no trees, so dropping the + // worker's placement holders is the whole reset there. + let placement_maps: Vec> = self + .placement_index + .iter() + .map(|model| Arc::clone(model.value())) + .collect(); + for placements in placement_maps { + placements.retain(|_, holders| { + holders.retain(|h| h.worker_url != url); + !holders.is_empty() + }); + } + } + + /// Test-only insert into one model's string tree, creating the tree when + /// the model has none yet (as the worker-add paths do). + #[cfg(test)] + pub(crate) fn insert_text_for_test(&self, model_id: &str, text: &str, tenant: &str) { + let tree_key = normalize_model_key(model_id).to_string(); + let tree = self + .string_trees + .entry(tree_key) + .or_insert_with(|| Arc::new(Tree::new())); + tree.insert_text(text, tenant); + } + + /// Test-only view of how much of `text` one model's string tree still + /// matches for `tenant`; empty when the model has no string tree. + #[cfg(test)] + pub(crate) fn string_prefix_for_tenant( + &self, + model_id: &str, + text: &str, + tenant: &str, + ) -> String { + self.string_trees + .get(normalize_model_key(model_id)) + .map(|tree| tree.prefix_match_tenant(text, tenant)) + .unwrap_or_default() + } + /// Remove several workers from one model. Tree cleanup remains precise per /// tenant, while hash-placement state is scanned only once for the batch. pub(crate) fn remove_workers_from_model(&self, model_id: &str, worker_urls: &HashSet) { @@ -3460,6 +3519,77 @@ mod tests { assert!(token_tree.get_tenant_token_counts().is_empty()); } + #[test] + fn reset_worker_cache_forgets_entries_but_keeps_the_tenant_and_its_load() { + // A weight refit flushes the engine's KV cache: the trees must forget + // everything the worker held, keep it as a tenant so it can repopulate, + // and leave the load scorer's view of it untouched. + let policy = CacheAwarePolicy::with_config(test_config()); + let workers = make_workers(&["http://w1:8000", "http://w2:8000"]); + policy.init_workers(&workers); + let model = normalize_model_key(workers[0].model_id()).to_string(); + let refit = workers[0].url(); + let neighbor = workers[1].url(); + + let long = "the quick brown fox jumps over the lazy dog ".repeat(4); + policy.insert_text_for_test(&model, &long, refit); + policy.insert_text_for_test(&model, "keep me around", neighbor); + let token_tree = Arc::clone(policy.token_trees.get(&model).unwrap().value()); + let tokens: Vec = (0..32).collect(); + token_tree.insert_tokens(&tokens, refit); + // Give w1 live scorer state: a backend report plus the since-poll + // dispatch credit from one request. w2 carries a heavy queue, so the + // pick is w1 whether or not its cached prefix is consulted. + update_expected_wait_loads(&policy, &workers, &[0, 10_000]); + assert_eq!( + policy.select_worker( + &workers, + &SelectWorkerInfo { + request_text: Some(&long), + ..Default::default() + }, + ), + Some(0) + ); + + assert_eq!(policy.string_prefix_for_tenant(&model, &long, refit), long); + let load_before = policy.load_scorer.load_state_for_test(refit); + assert!( + load_before.0 && load_before.2 > 0, + "precondition: the scorer holds w1's report and its dispatch credit" + ); + + policy.reset_worker_cache(workers[0].as_ref()); + + assert_eq!( + policy.string_prefix_for_tenant(&model, &long, refit), + "", + "a flushed worker must match nothing it held before" + ); + assert!( + !token_tree.get_tenant_token_counts().contains_key(refit), + "the token tree must forget the flushed worker's sequences" + ); + // The string root is re-seeded, so the worker stays a tenant and can + // repopulate. (The token re-seed mirrors `add_worker`, whose empty + // insert is a no-op below one page, so no token tenant entry either + // way.) + let string_tree = Arc::clone(policy.string_trees.get(&model).unwrap().value()); + assert!(string_tree.get_tenant_char_count().contains_key(refit)); + // Only the flushed worker is affected. + assert_eq!( + policy.string_prefix_for_tenant(&model, "keep me around", neighbor), + "keep me around" + ); + // Unlike removal, the scorer's report and dispatch credit survive. + assert_eq!(policy.load_scorer.load_state_for_test(refit), load_before); + policy.remove_worker_by_url(refit); + assert!( + !policy.load_scorer.load_state_for_test(refit).0, + "contrast: removing the worker does drop its load state" + ); + } + #[test] fn test_remove_workers_from_model_does_not_mutate_other_models() { let policy = CacheAwarePolicy::with_config(CacheAwareConfig { @@ -4921,6 +5051,34 @@ mod tests { assert_only_final_worker_credited(&policy, &workers, 0, 16); } + #[test] + fn reset_worker_cache_drops_hash_placements_only_for_that_worker() { + // Hash mode keeps no trees, so dropping the flushed worker's + // placements is the whole reset there. + let policy = CacheAwarePolicy::with_config(hash_config(&[16])); + let workers = make_workers(&["http://w1:8000", "http://w2:8000"]); + let model = normalize_model_key(workers[0].model_id()); + let refit: Vec = (0..16).collect(); + let kept: Vec = (100..116).collect(); + policy.record_placement(model, None, &refit, &[16], workers[0].url(), Instant::now()); + policy.record_placement(model, None, &kept, &[16], workers[1].url(), Instant::now()); + + policy.reset_worker_cache(workers[0].as_ref()); + + let placements = Arc::clone(policy.placement_index.get(model).unwrap().value()); + let holders: Vec = placements + .iter() + .flat_map(|entry| { + entry + .value() + .iter() + .map(|holder| holder.worker_url.clone()) + .collect::>() + }) + .collect(); + assert_eq!(holders, vec![workers[1].url().to_string()]); + } + #[test] fn hash_miss_uses_expected_wait_and_records_final_holder() { let policy = CacheAwarePolicy::with_config(hash_config(&[16])); diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index f681793c6d..1de6663634 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -938,6 +938,26 @@ impl PolicyRegistry { } } + /// The worker's engine flushed its cache (a weight refit): reset its + /// entries in every cache-aware policy that could route the model. + /// + /// A model can route through its explicit, default, or PD/EPD leg policy; + /// `policies_for_model` covers all of them and deduplicates shared + /// instances, so each tree is reset at most once. The worker keeps its + /// load state — it is still serving, just with a cold cache. + pub fn reset_worker_cache(&self, worker: &dyn Worker) { + for policy in self.policies_for_model(worker.model_id()) { + if let Some(cache_aware) = policy.as_any().downcast_ref::() { + cache_aware.reset_worker_cache(worker); + debug!( + worker = worker.url(), + policy = policy.name(), + "Reset cache-aware entries after a weight version change" + ); + } + } + } + /// Drop a removed worker's cached load report from all load-aware /// policies. /// @@ -1137,6 +1157,70 @@ mod tests { } } + /// Downcast a registry policy to the cache-aware implementation under test. + fn as_cache_aware(policy: &Arc) -> &CacheAwarePolicy { + policy.as_any().downcast_ref::().unwrap() + } + + #[test] + fn reset_worker_cache_clears_the_model_and_leg_policies() { + let reg = PolicyRegistry::new(cache_aware_config()); + let leg = cache_aware_policy(); + reg.set_prefill_policy(Arc::clone(&leg)); + let workers = vec![worker("http://w1:8000", WorkerType::Regular)]; + let model = workers[0].model_id().to_string(); + reg.on_worker_added(&model, None); + reg.init_cache_aware_policy(&model, &workers); + let model_policy = reg.get_policy(&model).unwrap(); + + let text = "a long shared instruction block this worker has served before"; + for policy in [&model_policy, &leg] { + as_cache_aware(policy).insert_text_for_test(&model, text, workers[0].url()); + assert_eq!( + as_cache_aware(policy).string_prefix_for_tenant(&model, text, workers[0].url()), + text + ); + } + + reg.reset_worker_cache(workers[0].as_ref()); + + for policy in [&model_policy, &leg] { + assert_eq!( + as_cache_aware(policy).string_prefix_for_tenant(&model, text, workers[0].url()), + "", + "every cache-aware policy that may route the model must forget the flush" + ); + } + } + + #[test] + fn reset_worker_cache_reaches_a_model_served_by_the_default_policy() { + // A model with no explicit entry routes through the default policy, so + // reaching only `model_policies` would leave a flushed worker's tree + // stale (the bug `policies_for_model` exists to prevent). + let reg = PolicyRegistry::new(cache_aware_config()); + let workers = [worker("http://w1:8000", WorkerType::Regular)]; + let model = workers[0].model_id().to_string(); + let default_policy = reg.get_default_policy(); + assert!( + reg.get_policy(&model).is_none(), + "precondition: no per-model policy" + ); + let text = "default-policy deployment prefix"; + as_cache_aware(&default_policy).insert_text_for_test(&model, text, workers[0].url()); + + reg.reset_worker_cache(workers[0].as_ref()); + + assert_eq!( + as_cache_aware(&default_policy).string_prefix_for_tenant( + &model, + text, + workers[0].url() + ), + "" + ); + } + fn headers_with_tokens(value: &str) -> HeaderMap { let mut h = HeaderMap::new(); h.insert("x-smg-routing-tokens", value.parse().unwrap()); From 3f82457699732b7d502b78cde3ab9dc178c307dd Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 16:44:26 -0700 Subject: [PATCH 13/27] feat(router): carry the routed worker as a response extension alongside its header Signed-off-by: key4ng --- crates/external_router/src/header_utils.rs | 35 ++++++++++++++++++++- model_gateway/src/routers/http/pd_router.rs | 11 ++++--- model_gateway/src/routers/http/router.rs | 12 +++---- 3 files changed, 46 insertions(+), 12 deletions(-) diff --git a/crates/external_router/src/header_utils.rs b/crates/external_router/src/header_utils.rs index c59ef72bc1..4e75e84e88 100644 --- a/crates/external_router/src/header_utils.rs +++ b/crates/external_router/src/header_utils.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use axum::{ body::Body, @@ -144,6 +144,22 @@ pub fn insert_routed_worker_id(headers: &mut HeaderMap, worker_url: &str) { } } +/// Response extension naming the worker that served the request. Set only by +/// the gateway, never derived from an upstream header, so a layer keyed on it +/// cannot be fooled through [`preserve_response_headers`] (the same reason +/// the error code travels as an extension). +#[derive(Clone, Debug)] +pub struct RoutedWorker(pub Arc); + +/// Stamp `x-smg-routed-worker-id` and the [`RoutedWorker`] extension on a +/// response. Call after the headers are in place so the gateway's value wins. +pub fn stamp_routed_worker(response: &mut http::Response, worker_url: &str) { + insert_routed_worker_id(response.headers_mut(), worker_url); + response + .extensions_mut() + .insert(RoutedWorker(Arc::from(worker_url))); +} + /// Whether `name` is SMG's own error-code response header. /// /// Backend responses must not mint this marker: downstream gateways use it to @@ -409,6 +425,23 @@ pub fn should_forward_request_header(name: &str) -> bool { mod tests { use super::*; + #[test] + fn stamp_routed_worker_sets_header_and_extension() { + let mut response = http::Response::new(()); + response + .headers_mut() + .insert("x-smg-routed-worker-id", "http://forged".parse().unwrap()); + stamp_routed_worker(&mut response, "http://w1:1@2"); + assert_eq!( + response.headers()["x-smg-routed-worker-id"], + "http://w1:1@2" + ); + assert_eq!( + &*response.extensions().get::().unwrap().0, + "http://w1:1@2" + ); + } + #[test] fn preserve_response_headers_drops_smg_owned_error_code() { let mut upstream = HeaderMap::new(); diff --git a/model_gateway/src/routers/http/pd_router.rs b/model_gateway/src/routers/http/pd_router.rs index b47c11a3eb..53a6c3b160 100644 --- a/model_gateway/src/routers/http/pd_router.rs +++ b/model_gateway/src/routers/http/pd_router.rs @@ -1248,11 +1248,10 @@ impl PDRouter { None }; - let mut response_headers = + let response_headers = header_utils::preserve_response_headers(decode_response.headers()); - header_utils::insert_routed_worker_id(&mut response_headers, decode.url()); - self.create_streaming_response( + let mut response = self.create_streaming_response( decode_response.bytes_stream(), status, prefill_logprobs, @@ -1260,7 +1259,9 @@ impl PDRouter { None, Some(response_headers), load_guards, - ) + ); + header_utils::stamp_routed_worker(&mut response, decode.url()); + response } else { // Non-streaming response let mut response = if context.return_logprob { @@ -1292,7 +1293,7 @@ impl PDRouter { // The decode worker is the one that produced the body the client // sees, on both the merged-logprob and passthrough paths. - header_utils::insert_routed_worker_id(response.headers_mut(), decode.url()); + header_utils::stamp_routed_worker(&mut response, decode.url()); response } } diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 2099b5fea9..c4588c40c1 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -935,8 +935,7 @@ impl Router { // JSON body (success or 4xx error). Don't relabel non-SSE // responses as SSE; leave that judgment to whatever the worker // set. - let mut response_headers = header_utils::preserve_response_headers(res.headers()); - header_utils::insert_routed_worker_id(&mut response_headers, worker.url()); + let response_headers = header_utils::preserve_response_headers(res.headers()); let stream = res.bytes_stream(); // Bounded channel applies backpressure: if the downstream client // is slow, the upstream relay awaits on `send` rather than piling @@ -1045,16 +1044,17 @@ impl Router { let mut response = Response::new(body); *response.status_mut() = status; *response.headers_mut() = response_headers; + header_utils::stamp_routed_worker(&mut response, worker.url()); response = AttachedBody::wrap_response(response, load_guard); response } else { - let mut response_headers = header_utils::preserve_response_headers(res.headers()); - header_utils::insert_routed_worker_id(&mut response_headers, worker.url()); + let response_headers = header_utils::preserve_response_headers(res.headers()); match res.bytes().await { Ok(body) => { let mut response = Response::new(Body::from(body)); *response.status_mut() = status; *response.headers_mut() = response_headers; + header_utils::stamp_routed_worker(&mut response, worker.url()); response } Err(e) => error::internal_error( @@ -1195,7 +1195,6 @@ impl Router { if is_stream { // Preserve headers for streaming response let mut response_headers = header_utils::preserve_response_headers(res.headers()); - header_utils::insert_routed_worker_id(&mut response_headers, worker_url); if status.is_success() && !response_headers.contains_key(CONTENT_TYPE) { response_headers .insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); @@ -1245,6 +1244,7 @@ impl Router { let mut response = Response::new(body); *response.status_mut() = status; *response.headers_mut() = response_headers; + header_utils::stamp_routed_worker(&mut response, worker_url); // Attach load guard to response body for proper RAII lifecycle // Guard is dropped when response body is consumed or client disconnects @@ -1270,7 +1270,7 @@ impl Router { } Err(error_response) => error_response, }; - header_utils::insert_routed_worker_id(response.headers_mut(), worker_url); + header_utils::stamp_routed_worker(&mut response, worker_url); // load_guard dropped here automatically after response body is read response From 214f1af988e7a27ac03aa96780b0a80cd5e0b2ae Mon Sep 17 00:00:00 2001 From: key4ng Date: Thu, 10 Sep 2026 16:49:01 -0700 Subject: [PATCH 14/27] feat(router): stamp the routed worker on Responses API streams Signed-off-by: key4ng --- crates/external_router/src/openai/responses/streaming.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/external_router/src/openai/responses/streaming.rs b/crates/external_router/src/openai/responses/streaming.rs index cb634ef3f8..211eba7776 100644 --- a/crates/external_router/src/openai/responses/streaming.rs +++ b/crates/external_router/src/openai/responses/streaming.rs @@ -45,7 +45,7 @@ const SSE_DONE: &str = "data: [DONE]\n\n"; use crate::{ error, header_utils::{ - extract_forwardable_request_headers, insert_routed_worker_id, preserve_response_headers, + extract_forwardable_request_headers, preserve_response_headers, stamp_routed_worker, ApiProvider, }, mcp_utils::DEFAULT_MAX_ITERATIONS, @@ -567,8 +567,7 @@ pub(super) async fn handle_simple_streaming_passthrough( return (status_code, error_body).into_response(); } - let mut preserved_headers = preserve_response_headers(response.headers()); - insert_routed_worker_id(&mut preserved_headers, worker.url()); + let preserved_headers = preserve_response_headers(response.headers()); let mut upstream_stream = response.bytes_stream(); let (tx, rx) = sse_channel(); @@ -678,6 +677,8 @@ pub(super) async fn handle_simple_streaming_passthrough( headers_mut.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); } + stamp_routed_worker(&mut response, worker.url()); + response } From a6bdb55c1fd2c0fdf51c066894571657bb280f97 Mon Sep 17 00:00:00 2001 From: key4ng Date: Fri, 11 Sep 2026 21:33:58 -0700 Subject: [PATCH 15/27] feat(rl): route around paused, asleep, and stale engines through the candidate filter Signed-off-by: key4ng --- model_gateway/src/app_context.rs | 10 +- model_gateway/src/rl_adapter.rs | 351 +++++++++++++++++- model_gateway/tests/common/mock_worker.rs | 2 + .../tests/rl_version_routing_test.rs | 312 ++++++++++++++++ 4 files changed, 664 insertions(+), 11 deletions(-) create mode 100644 model_gateway/tests/rl_version_routing_test.rs diff --git a/model_gateway/src/app_context.rs b/model_gateway/src/app_context.rs index 2ef11d77ed..fdd70e0a31 100644 --- a/model_gateway/src/app_context.rs +++ b/model_gateway/src/app_context.rs @@ -373,7 +373,11 @@ impl AppContextBuilder { &router_config.tenant_api_keys, ); - let rl = crate::rl_adapter::build_rl_state(&worker_registry, &router_config); + let policy_registry = self + .policy_registry + .ok_or(AppContextBuildError::MissingField("policy_registry"))?; + let rl = + crate::rl_adapter::build_rl_state(&worker_registry, &policy_registry, &router_config); Ok(AppContext { gateway_auth, @@ -390,9 +394,7 @@ impl AppContextBuilder { reasoning_parser_factory: self.reasoning_parser_factory, tool_parser_factory: self.tool_parser_factory, worker_registry, - policy_registry: self - .policy_registry - .ok_or(AppContextBuildError::MissingField("policy_registry"))?, + policy_registry, gateway: self.gateway, response_storage: self .response_storage diff --git a/model_gateway/src/rl_adapter.rs b/model_gateway/src/rl_adapter.rs index 66a61ae55d..ea89fa748c 100644 --- a/model_gateway/src/rl_adapter.rs +++ b/model_gateway/src/rl_adapter.rs @@ -1,14 +1,18 @@ -//! Glue between the RL control plane crate and the gateway registry. This -//! file is the whole of coupling touchpoint (a); see `crates/rl/COUPLING.md`. +//! Glue between the RL control plane crate and the gateway. This file is the +//! whole of coupling surfaces (a) and (b); see `crates/rl/COUPLING.md`. use std::sync::Arc; +use http::HeaderMap; use openai_protocol::worker::ConnectionMode; -use smg_rl::{RlState, RlWorkerInfo, RlWorkerView}; +use smg_rl::{RlState, RlWorkerInfo, RlWorkerView, VersionEvictionSink}; +use tokio::sync::broadcast::error::RecvError; +use tracing::{debug, warn}; use crate::{ config::RouterConfig, - worker::{registry::WorkerId, Worker, WorkerRegistry}, + policies::{CandidateFilter, PolicyRegistry}, + worker::{event::WorkerEvent, registry::WorkerId, Worker, WorkerRegistry}, }; /// Read-only registry view for the RL crate. @@ -62,19 +66,163 @@ impl RlWorkerView for RegistryRlView { let worker = self.registry.get(&WorkerId::from_string(id.to_string()))?; self.info(&worker) } + + /// One engine is one table entry but can be many registry entries (a + /// DP-aware group registers one worker per rank), so the engine's + /// inflight count is the sum over every worker sharing its base URL. + fn inflight(&self, base_url: &str) -> usize { + self.registry + .get_all() + .iter() + .filter(|w| w.base_url() == base_url) + .map(|w| w.load()) + .sum() + } +} + +/// Coupling (a): a version change resets every rank's cache-aware entries. +/// +/// Called with the RL table's write mutex held, so it must never write back +/// into the table — resetting policy caches is the whole of its work. +struct RegistryEvictionSink { + registry: Arc, + policy_registry: Arc, +} + +impl VersionEvictionSink for RegistryEvictionSink { + fn on_version_changed(&self, _model_id: &str, base_url: &str) { + for worker in self + .registry + .get_all() + .iter() + .filter(|w| w.base_url() == base_url) + { + self.policy_registry.reset_worker_cache(worker.as_ref()); + } + } +} + +/// Coupling (b): the RL pre-filter installed on the policy registry. +pub struct RlCandidateFilter { + rl: Arc, +} + +impl CandidateFilter for RlCandidateFilter { + fn eligible( + &self, + workers: &[Arc], + headers: Option<&HeaderMap>, + ) -> Option> { + let policy = self.rl.policy_for(headers); + self.rl.table().eligible( + workers.iter().map(|w| (w.model_id(), w.base_url())), + &policy, + ) + } +} + +/// The weight version a worker's discovered labels report, if any. +fn version_label(worker: &Arc) -> Option<&str> { + worker + .metadata() + .spec + .labels + .get("weight_version") + .map(String::as_str) } -/// Build the RL state when `config.rl.enabled`; `None` otherwise, so the +fn seed(rl: &RlState, worker: &Arc) { + rl.table() + .seed(worker.base_url(), worker.model_id(), version_label(worker)); +} + +/// Re-derive the table from the registry after a lagged event stream. +fn resync(rl: &RlState, registry: &WorkerRegistry) { + let workers = registry.get_all(); + for worker in &workers { + seed(rl, worker); + } + rl.table() + .retain(|base_url| workers.iter().any(|w| w.base_url() == base_url)); +} + +/// Keep the table in step with registrations: seed on `Registered`, reseed +/// on `Replaced` only when the discovered label changed (a property patch +/// must not reset an observed version), drop on the last `Removed` for an +/// engine. Runs only when a Tokio runtime is current. +fn spawn_table_maintainer(rl: Arc, registry: Arc) { + if tokio::runtime::Handle::try_current().is_err() { + warn!("RL table maintainer not started: no async runtime; versions seed lazily"); + return; + } + // Subscribe before the initial pass so a registration racing this call is + // either already in `get_all()` or still queued on the receiver. + let mut rx = registry.subscribe_events(); + resync(&rl, ®istry); + let registry = Arc::downgrade(®istry); + #[expect( + clippy::disallowed_methods, + reason = "runs for the lifetime of the gateway and exits when the registry is dropped" + )] + tokio::spawn(async move { + loop { + match rx.recv().await { + Ok(WorkerEvent::Registered { worker, .. }) => seed(&rl, &worker), + Ok(WorkerEvent::Replaced { old, new, .. }) => { + if version_label(&old) == version_label(&new) { + seed(&rl, &new); + } else { + rl.table() + .reseed(new.base_url(), new.model_id(), version_label(&new)); + } + } + Ok(WorkerEvent::Removed { worker, .. }) => { + let Some(registry) = registry.upgrade() else { + break; + }; + let base_url = worker.base_url(); + if !registry.get_all().iter().any(|w| w.base_url() == base_url) { + rl.table().remove(base_url); + } + } + Ok(WorkerEvent::StatusChanged { .. }) => {} + Err(RecvError::Lagged(n)) => { + debug!(missed = n, "RL table maintainer lagged; resyncing"); + let Some(registry) = registry.upgrade() else { + break; + }; + resync(&rl, ®istry); + } + Err(RecvError::Closed) => break, + } + } + }); +} + +/// Build the RL state when `config.rl.enabled`, install the candidate +/// filter, and start the table maintainer; `None` otherwise, so the /// disabled path constructs nothing. pub fn build_rl_state( registry: &Arc, + policy_registry: &Arc, config: &RouterConfig, ) -> Option> { if !config.rl.enabled { return None; } let view = Arc::new(RegistryRlView::new(Arc::clone(registry))); - Some(Arc::new(RlState::new(view, config.rl.clone()))) + let sink = Arc::new(RegistryEvictionSink { + registry: Arc::clone(registry), + policy_registry: Arc::clone(policy_registry), + }); + let rl = Arc::new(RlState::with_sink(view, config.rl.clone(), sink)); + if !policy_registry.set_candidate_filter(Arc::new(RlCandidateFilter { + rl: Arc::clone(&rl), + })) { + warn!("a candidate filter was already installed; RL routing filter not active"); + } + spawn_table_maintainer(Arc::clone(&rl), Arc::clone(registry)); + Some(rl) } #[cfg(test)] @@ -82,9 +230,14 @@ mod tests { use std::sync::Arc; use openai_protocol::worker::ConnectionMode; + use smg_rl::{RlConfig, Version, VersionSource}; use super::*; - use crate::worker::BasicWorkerBuilder; + use crate::{ + config::types::PolicyConfig, + policies::CacheAwarePolicy, + worker::{BasicWorkerBuilder, WorkerType}, + }; fn registry_with(worker: Arc) -> Arc { let registry = Arc::new(WorkerRegistry::new()); @@ -92,6 +245,20 @@ mod tests { registry } + fn config_with_rl(rl: RlConfig) -> RouterConfig { + RouterConfig { + rl, + ..Default::default() + } + } + + fn rl_enabled_config() -> RouterConfig { + config_with_rl(RlConfig { + enabled: true, + ..RlConfig::default() + }) + } + /// Control calls must use the client the gateway negotiated for the /// worker (HTTP version, TLS identity, pool tuning), not a second one. #[test] @@ -123,4 +290,174 @@ mod tests { let info = view.list().pop().expect("one worker"); assert!(info.http_client.is_none()); } + + /// The mixed-version accounting is per engine, so a DP group's ranks — + /// separate registry entries sharing one base URL — must add up, and a + /// second engine's load must not leak in. + #[test] + fn inflight_sums_ranks_sharing_a_base_url() { + let registry = Arc::new(WorkerRegistry::new()); + let rank0: Arc = Arc::new( + BasicWorkerBuilder::new("http://engine:30000") + .dp_config(0, 2) + .build(), + ); + let rank1: Arc = Arc::new( + BasicWorkerBuilder::new("http://engine:30000") + .dp_config(1, 2) + .build(), + ); + let other: Arc = + Arc::new(BasicWorkerBuilder::new("http://other:30000").build()); + for worker in [&rank0, &rank1, &other] { + registry.register(Arc::clone(worker)); + } + rank0.increment_load(); + rank1.increment_load(); + other.increment_load(); + + let view = RegistryRlView::new(Arc::clone(®istry)); + assert_eq!(view.inflight("http://engine:30000"), 2); + assert_eq!(view.inflight("http://other:30000"), 1); + assert_eq!(view.inflight("http://absent:30000"), 0); + } + + /// Registration labels seed the table and the last removal for an engine + /// drops the entry, so the filter never judges a worker the registry no + /// longer has. + #[tokio::test] + async fn maintainer_seeds_and_drops_entries() { + let registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + let rl = build_rl_state(®istry, &policy_registry, &rl_enabled_config()) + .expect("rl is enabled"); + assert!(policy_registry.has_candidate_filter()); + + let versioned: Arc = Arc::new( + BasicWorkerBuilder::new("http://engine-a:30000") + .label("weight_version", "5") + .build(), + ); + let unversioned: Arc = + Arc::new(BasicWorkerBuilder::new("http://engine-b:30000").build()); + registry.register(Arc::clone(&versioned)); + registry.register(Arc::clone(&unversioned)); + settle().await; + + let entry = rl + .table() + .get("http://engine-a:30000") + .expect("the label seeded an entry"); + assert_eq!(entry.version, Some(Version::parse("5"))); + assert_eq!(entry.version_source, Some(VersionSource::Registration)); + let entry = rl + .table() + .get("http://engine-b:30000") + .expect("an unlabeled worker still gets an entry"); + assert_eq!(entry.version, None); + assert_eq!(entry.version_source, None); + + registry.remove_by_url("http://engine-a:30000"); + settle().await; + assert!(rl.table().get("http://engine-a:30000").is_none()); + assert!(rl.table().get("http://engine-b:30000").is_some()); + } + + /// Let the maintainer task drain the event stream. + async fn settle() { + for _ in 0..16 { + tokio::task::yield_now().await; + } + } + + /// Coupling (a): a refit invalidates the engine's KV cache, so the + /// cache-aware trees must forget what it held or the next request with + /// that prefix is routed to a worker that no longer has it. + #[tokio::test] + async fn version_change_resets_cache_aware_entries() { + let registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::CacheAware { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 0, + max_tree_size: 4096, + block_size: 16, + balance_token_usage_threshold: 1.0, + overload_token_usage_threshold: 1.0, + overlap_decay: 0.0, + selection_temperature: 0.0, + cache_index: Default::default(), + cache_ttl_secs: 180, + cache_boundaries: Vec::new(), + })); + let rl = build_rl_state(®istry, &policy_registry, &rl_enabled_config()) + .expect("rl is enabled"); + + let worker: Arc = Arc::new( + BasicWorkerBuilder::new("http://engine:30000") + .worker_type(WorkerType::Regular) + .build(), + ); + let id = registry + .register(Arc::clone(&worker)) + .expect("registration returns an id"); + let model = worker.model_id().to_string(); + policy_registry.on_worker_added(&model, None); + policy_registry.init_cache_aware_policy(&model, &[Arc::clone(&worker)]); + let policy = policy_registry + .get_policy(&model) + .expect("the model has a policy"); + let cache_aware = policy + .as_any() + .downcast_ref::() + .expect("cache-aware policy"); + + let text = "a long shared instruction block this engine has served before"; + cache_aware.insert_text_for_test(&model, text, worker.url()); + assert_eq!( + cache_aware.string_prefix_for_tenant(&model, text, worker.url()), + text, + "precondition: the tree holds the prefix" + ); + + let view = RegistryRlView::new(Arc::clone(®istry)); + let info = view.get(id.as_str()).expect("the worker is registered"); + assert!(rl.apply_version(&info, Version::parse("2"), VersionSource::Api)); + + assert_eq!( + cache_aware.string_prefix_for_tenant(&model, text, worker.url()), + "", + "a refit must clear what the engine's cache no longer holds" + ); + } + + /// The disabled path constructs nothing: no state, and no filter on the + /// policy registry to slow the hot path down. + #[test] + fn disabled_builds_nothing_and_installs_no_filter() { + let registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + let config = RouterConfig::default(); + assert!(!config.rl.enabled, "precondition: the flag defaults off"); + + assert!(build_rl_state(®istry, &policy_registry, &config).is_none()); + assert!(!policy_registry.has_candidate_filter()); + } + + /// The state carries the gateway's own RL settings, not the crate + /// defaults, so the control deadline a deployment configured applies. + #[test] + fn build_uses_the_configured_rl_settings() { + let registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + let config = config_with_rl(RlConfig { + enabled: true, + control_timeout_secs: 7, + ..RlConfig::default() + }); + + let rl = build_rl_state(®istry, &policy_registry, &config).expect("rl is enabled"); + assert_eq!(rl.config().control_timeout_secs, 7); + } } diff --git a/model_gateway/tests/common/mock_worker.rs b/model_gateway/tests/common/mock_worker.rs index 4fb3cf958a..c5a106df40 100755 --- a/model_gateway/tests/common/mock_worker.rs +++ b/model_gateway/tests/common/mock_worker.rs @@ -360,6 +360,8 @@ impl MockWorker { .route("/continue_generation", post(rl_control_handler)) .route("/update_weights_from_disk", post(rl_control_handler)) .route("/update_weight_version", post(rl_control_handler)) + .route("/release_memory_occupation", post(rl_control_handler)) + .route("/resume_memory_occupation", post(rl_control_handler)) .route("/pause", post(rl_control_handler)) .route("/resume", post(rl_control_handler)) .route("/v1/loads", get(loads_handler)) diff --git a/model_gateway/tests/rl_version_routing_test.rs b/model_gateway/tests/rl_version_routing_test.rs new file mode 100644 index 0000000000..5e9827f119 --- /dev/null +++ b/model_gateway/tests/rl_version_routing_test.rs @@ -0,0 +1,312 @@ +//! Gateway-level tests for RL M2: state-aware and version-aware routing. + +mod common; + +use std::collections::BTreeSet; + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use common::{ + mock_worker::{HealthStatus, MockWorkerConfig, WorkerType}, + AppTestContext, TestRouterConfig, +}; +use http_body_util::BodyExt; +use serde_json::{json, Value}; +use tower::ServiceExt; + +fn mock(port: u16) -> MockWorkerConfig { + MockWorkerConfig { + port, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + } +} + +#[expect( + clippy::expect_used, + reason = "test helper - panicking on failure is intentional" +)] +async fn ctx(policy: &str, workers: Vec) -> AppTestContext { + let mut config = TestRouterConfig::round_robin(0); + config.rl.enabled = true; + config.rl.control_timeout_secs = 5; + config.rl.version_policy = policy.parse().expect("policy"); + AppTestContext::new_with_config(config, workers).await +} + +#[expect( + clippy::unwrap_used, + reason = "test helper - panicking on failure is intentional" +)] +async fn json_of(resp: axum::response::Response) -> Value { + serde_json::from_slice(&resp.into_body().collect().await.unwrap().to_bytes()) + .unwrap_or(Value::Null) +} + +/// `(id, url)` of every registered worker, in `/v1/rl/workers` order. +#[expect( + clippy::unwrap_used, + reason = "test helper - panicking on failure is intentional" +)] +async fn workers(app: &axum::Router) -> Vec<(String, String)> { + let resp = app + .clone() + .oneshot(Request::get("/v1/rl/workers").body(Body::empty()).unwrap()) + .await + .unwrap(); + json_of(resp).await["workers"] + .as_array() + .unwrap() + .iter() + .map(|w| { + ( + w["id"].as_str().unwrap().to_string(), + w["url"].as_str().unwrap().to_string(), + ) + }) + .collect() +} + +#[expect( + clippy::unwrap_used, + reason = "test helper - panicking on failure is intentional" +)] +async fn post_json(app: &axum::Router, uri: &str, body: Value) -> axum::response::Response { + app.clone() + .oneshot( + Request::post(uri) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap() +} + +/// Routed-worker URLs seen over `n` buffered `/generate` requests. +#[expect( + clippy::unwrap_used, + reason = "test helper - panicking on failure is intentional" +)] +async fn routed_over( + app: &axum::Router, + n: usize, + extra_header: Option<(&str, &str)>, +) -> BTreeSet { + let mut seen = BTreeSet::new(); + for _ in 0..n { + let mut req = Request::post("/generate").header("content-type", "application/json"); + if let Some((k, v)) = extra_header { + req = req.header(k, v); + } + let body = json!({"text": "hi", "stream": false}).to_string(); + let resp = app + .clone() + .oneshot(req.body(Body::from(body)).unwrap()) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + seen.insert( + resp.headers()["x-smg-routed-worker-id"] + .to_str() + .unwrap() + .to_string(), + ); + } + seen +} + +#[tokio::test] +async fn asleep_via_api_and_via_intercepted_release_get_no_dispatches() { + let ctx = ctx("any", vec![mock(18921), mock(18922), mock(18923)]).await; + let app = ctx.create_app(); + let ws = workers(&app).await; + assert_eq!(ws.len(), 3); + assert_eq!( + routed_over(&app, 12, None).await.len(), + 3, + "round robin over all three first" + ); + + // w0 asleep through the API, w1 through an intercepted release. + let resp = post_json( + &app, + &format!("/v1/rl/workers/{}/state", ws[0].0), + json!({"control": "asleep"}), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + let resp = post_json( + &app, + &format!( + "/v1/rl/workers/{}/engine/release_memory_occupation", + ws[1].0 + ), + json!({"tags": ["weights"]}), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + + let routed = routed_over(&app, 12, None).await; + assert_eq!(routed, BTreeSet::from([ws[2].1.clone()])); + let listed = json_of( + app.clone() + .oneshot(Request::get("/workers").body(Body::empty()).unwrap()) + .await + .unwrap(), + ) + .await; + assert!( + listed.to_string().matches("\"ready\"").count() >= 3, + "still ready to the registry: {listed}" + ); + + // Wake them: API for w0, intercepted resume for w1. + post_json( + &app, + &format!("/v1/rl/workers/{}/state", ws[0].0), + json!({"control": "active"}), + ) + .await; + post_json( + &app, + &format!("/v1/rl/workers/{}/engine/resume_memory_occupation", ws[1].0), + json!({"tags": ["weights"]}), + ) + .await; + assert_eq!(routed_over(&app, 12, None).await.len(), 3); + ctx.shutdown().await; +} + +#[tokio::test] +async fn rolling_update_under_latest_only_never_dispatches_to_a_stale_worker() { + let ctx = ctx( + "latest-only", + vec![mock(18924), mock(18925), mock(18926), mock(18927)], + ) + .await; + let app = ctx.create_app(); + let ws = workers(&app).await; + assert_eq!( + routed_over(&app, 8, None).await.len(), + 4, + "nobody versioned: nobody stale" + ); + + // Flip one at a time: API, then a proxied update_weight_version, then API. + let resp = post_json( + &app, + &format!("/v1/rl/workers/{}/version", ws[0].0), + json!({"weight_version": "2"}), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!( + routed_over(&app, 8, None).await, + BTreeSet::from([ws[0].1.clone()]) + ); + + let resp = post_json( + &app, + &format!("/v1/rl/workers/{}/engine/update_weight_version", ws[1].0), + json!({"new_version": "2"}), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + let routed = routed_over(&app, 8, None).await; + assert!( + routed.is_subset(&BTreeSet::from([ws[0].1.clone(), ws[1].1.clone()])), + "{routed:?}" + ); + + let resp = post_json( + &app, + "/v1/rl/version?selector=engine%3Dsglang", + json!({"weight_version": "2"}), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(json_of(resp).await["total"], 4); + assert_eq!(routed_over(&app, 8, None).await.len(), 4); + + // A newer version on one worker makes the other three stale again. + post_json( + &app, + &format!("/v1/rl/workers/{}/version", ws[3].0), + json!({"weight_version": "3"}), + ) + .await; + assert_eq!( + routed_over(&app, 8, None).await, + BTreeSet::from([ws[3].1.clone()]) + ); + let listed = json_of( + app.clone() + .oneshot(Request::get("/v1/rl/workers").body(Body::empty()).unwrap()) + .await + .unwrap(), + ) + .await; + let sources: BTreeSet<&str> = listed["workers"] + .as_array() + .unwrap() + .iter() + .map(|w| w["version_source"].as_str().unwrap()) + .collect(); + assert_eq!(sources, BTreeSet::from(["api"])); + ctx.shutdown().await; +} + +#[tokio::test] +async fn all_workers_ineligible_is_a_503() { + let ctx = ctx("any", vec![mock(18930)]).await; + let app = ctx.create_app(); + let ws = workers(&app).await; + post_json( + &app, + &format!("/v1/rl/workers/{}/state", ws[0].0), + json!({"control": "paused"}), + ) + .await; + let resp = app + .clone() + .oneshot( + Request::post("/generate") + .header("content-type", "application/json") + .body(Body::from(json!({"text": "hi"}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); + ctx.shutdown().await; +} + +#[tokio::test] +async fn flag_off_ignores_the_policy_header_and_stamps_nothing() { + let mut config = TestRouterConfig::round_robin(0); + config.rl.enabled = false; + let ctx = AppTestContext::new_with_config(config, vec![mock(18931)]).await; + let app = ctx.create_app(); + let resp = app + .clone() + .oneshot( + Request::post("/generate") + .header("content-type", "application/json") + .header("x-smg-version-policy", "garbage") + .body(Body::from(json!({"text": "hi"}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert!(resp.headers().get("x-smg-weight-version").is_none()); + assert!(resp.headers().get("x-smg-mixed-version").is_none()); + assert!(resp.headers().get("x-smg-routed-worker-id").is_some()); + assert!(!ctx.app_context.policy_registry.has_candidate_filter()); + ctx.shutdown().await; +} From bad78960a43b8a078eaff827528cb7d04aa3fa04 Mon Sep 17 00:00:00 2001 From: key4ng Date: Fri, 11 Sep 2026 21:42:58 -0700 Subject: [PATCH 16/27] fix(rl): hold the registries weakly from the eviction sink and table maintainer Signed-off-by: key4ng --- model_gateway/src/rl_adapter.rs | 93 ++++++++++++++++++++++++++------- 1 file changed, 75 insertions(+), 18 deletions(-) diff --git a/model_gateway/src/rl_adapter.rs b/model_gateway/src/rl_adapter.rs index ea89fa748c..b003040e4f 100644 --- a/model_gateway/src/rl_adapter.rs +++ b/model_gateway/src/rl_adapter.rs @@ -1,7 +1,7 @@ //! Glue between the RL control plane crate and the gateway. This file is the //! whole of coupling surfaces (a) and (b); see `crates/rl/COUPLING.md`. -use std::sync::Arc; +use std::sync::{Arc, Weak}; use http::HeaderMap; use openai_protocol::worker::ConnectionMode; @@ -84,20 +84,31 @@ impl RlWorkerView for RegistryRlView { /// /// Called with the RL table's write mutex held, so it must never write back /// into the table — resetting policy caches is the whole of its work. +/// +/// Both handles are weak. The policy registry owns the candidate filter, +/// which owns the [`RlState`] this sink lives in, so a strong handle here +/// would close an ownership cycle and leak both registries (and the +/// maintainer task with them) for the life of the process. struct RegistryEvictionSink { - registry: Arc, - policy_registry: Arc, + registry: Weak, + policy_registry: Weak, } impl VersionEvictionSink for RegistryEvictionSink { fn on_version_changed(&self, _model_id: &str, base_url: &str) { - for worker in self - .registry + // Either one gone means the gateway is tearing down: there is no + // cache left to reset. + let (Some(registry), Some(policy_registry)) = + (self.registry.upgrade(), self.policy_registry.upgrade()) + else { + return; + }; + for worker in registry .get_all() .iter() .filter(|w| w.base_url() == base_url) { - self.policy_registry.reset_worker_cache(worker.as_ref()); + policy_registry.reset_worker_cache(worker.as_ref()); } } } @@ -150,7 +161,7 @@ fn resync(rl: &RlState, registry: &WorkerRegistry) { /// on `Replaced` only when the discovered label changed (a property patch /// must not reset an observed version), drop on the last `Removed` for an /// engine. Runs only when a Tokio runtime is current. -fn spawn_table_maintainer(rl: Arc, registry: Arc) { +fn spawn_table_maintainer(rl: Weak, registry: Arc) { if tokio::runtime::Handle::try_current().is_err() { warn!("RL table maintainer not started: no async runtime; versions seed lazily"); return; @@ -158,21 +169,33 @@ fn spawn_table_maintainer(rl: Arc, registry: Arc) { // Subscribe before the initial pass so a registration racing this call is // either already in `get_all()` or still queued on the receiver. let mut rx = registry.subscribe_events(); - resync(&rl, ®istry); + // The caller still holds the state, so the initial pass runs against a + // strong handle. It happens before the loop, so workers registered before + // this call are in the table by the time it returns. + if let Some(state) = rl.upgrade() { + resync(&state, ®istry); + } + // Every handle the task keeps is weak, so it never keeps the gateway + // alive: it exits when the registry drops (`Closed`) or the state does. let registry = Arc::downgrade(®istry); #[expect( clippy::disallowed_methods, - reason = "runs for the lifetime of the gateway and exits when the registry is dropped" + reason = "runs for the lifetime of the gateway and exits when the registry or the RL state is dropped" )] tokio::spawn(async move { loop { - match rx.recv().await { - Ok(WorkerEvent::Registered { worker, .. }) => seed(&rl, &worker), + let event = rx.recv().await; + let Some(state) = rl.upgrade() else { + break; + }; + match event { + Ok(WorkerEvent::Registered { worker, .. }) => seed(&state, &worker), Ok(WorkerEvent::Replaced { old, new, .. }) => { if version_label(&old) == version_label(&new) { - seed(&rl, &new); + seed(&state, &new); } else { - rl.table() + state + .table() .reseed(new.base_url(), new.model_id(), version_label(&new)); } } @@ -182,7 +205,7 @@ fn spawn_table_maintainer(rl: Arc, registry: Arc) { }; let base_url = worker.base_url(); if !registry.get_all().iter().any(|w| w.base_url() == base_url) { - rl.table().remove(base_url); + state.table().remove(base_url); } } Ok(WorkerEvent::StatusChanged { .. }) => {} @@ -191,7 +214,7 @@ fn spawn_table_maintainer(rl: Arc, registry: Arc) { let Some(registry) = registry.upgrade() else { break; }; - resync(&rl, ®istry); + resync(&state, ®istry); } Err(RecvError::Closed) => break, } @@ -212,8 +235,8 @@ pub fn build_rl_state( } let view = Arc::new(RegistryRlView::new(Arc::clone(registry))); let sink = Arc::new(RegistryEvictionSink { - registry: Arc::clone(registry), - policy_registry: Arc::clone(policy_registry), + registry: Arc::downgrade(registry), + policy_registry: Arc::downgrade(policy_registry), }); let rl = Arc::new(RlState::with_sink(view, config.rl.clone(), sink)); if !policy_registry.set_candidate_filter(Arc::new(RlCandidateFilter { @@ -221,7 +244,7 @@ pub fn build_rl_state( })) { warn!("a candidate filter was already installed; RL routing filter not active"); } - spawn_table_maintainer(Arc::clone(&rl), Arc::clone(registry)); + spawn_table_maintainer(Arc::downgrade(&rl), Arc::clone(registry)); Some(rl) } @@ -432,6 +455,40 @@ mod tests { ); } + /// The gateway must not outlive itself: the policy registry owns the + /// candidate filter, which owns the state, which owns the sink — so the + /// sink's handles back to both registries have to be weak or nothing is + /// ever freed and the maintainer task never exits. + #[tokio::test] + async fn dropping_the_context_frees_the_registries() { + let registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + let rl = build_rl_state(®istry, &policy_registry, &rl_enabled_config()) + .expect("rl is enabled"); + + let weak_rl = Arc::downgrade(&rl); + let weak_registry = Arc::downgrade(®istry); + let weak_policies = Arc::downgrade(&policy_registry); + + drop(rl); + drop(registry); + drop(policy_registry); + settle().await; + + assert!( + weak_policies.upgrade().is_none(), + "the policy registry outlived the context" + ); + assert!( + weak_rl.upgrade().is_none(), + "the RL state outlived the context" + ); + assert!( + weak_registry.upgrade().is_none(), + "the worker registry outlived the context" + ); + } + /// The disabled path constructs nothing: no state, and no filter on the /// policy registry to slow the hot path down. #[test] From 6a9baeecf13fac943edbdcd8a90f2f01fa69cc87 Mon Sep 17 00:00:00 2001 From: key4ng Date: Fri, 11 Sep 2026 21:58:13 -0700 Subject: [PATCH 17/27] test(rl): tighten the rolling-update and readiness assertions; cover DP-rank removal Signed-off-by: key4ng --- model_gateway/src/rl_adapter.rs | 46 +++++++++++++++++++ .../tests/rl_version_routing_test.rs | 22 +++++---- 2 files changed, 60 insertions(+), 8 deletions(-) diff --git a/model_gateway/src/rl_adapter.rs b/model_gateway/src/rl_adapter.rs index b003040e4f..9557bf6116 100644 --- a/model_gateway/src/rl_adapter.rs +++ b/model_gateway/src/rl_adapter.rs @@ -386,6 +386,52 @@ mod tests { assert!(rl.table().get("http://engine-b:30000").is_some()); } + /// One engine, many registry entries: a DP group's ranks share a base + /// URL and one table entry, so a single rank leaving must not take the + /// engine's version with it — only the last one does. + #[tokio::test] + async fn removing_one_dp_rank_keeps_the_engines_entry() { + let registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + let rl = build_rl_state(®istry, &policy_registry, &rl_enabled_config()) + .expect("rl is enabled"); + + for rank in 0..2 { + let worker: Arc = Arc::new( + BasicWorkerBuilder::new("http://engine:30000") + .dp_config(rank, 2) + .label("weight_version", "5") + .build(), + ); + registry.register(worker); + } + settle().await; + let entry = rl + .table() + .get("http://engine:30000") + .expect("both ranks seed the one engine entry"); + assert_eq!(entry.version, Some(Version::parse("5"))); + + registry.remove_by_url("http://engine:30000@0"); + settle().await; + let entry = rl + .table() + .get("http://engine:30000") + .expect("rank 1 still serves this engine"); + assert_eq!( + entry.version, + Some(Version::parse("5")), + "a departing rank must not reset the engine's version" + ); + + registry.remove_by_url("http://engine:30000@1"); + settle().await; + assert!( + rl.table().get("http://engine:30000").is_none(), + "the last rank leaving drops the engine" + ); + } + /// Let the maintainer task drain the event stream. async fn settle() { for _ in 0..16 { diff --git a/model_gateway/tests/rl_version_routing_test.rs b/model_gateway/tests/rl_version_routing_test.rs index 5e9827f119..8f2785841c 100644 --- a/model_gateway/tests/rl_version_routing_test.rs +++ b/model_gateway/tests/rl_version_routing_test.rs @@ -160,10 +160,14 @@ async fn asleep_via_api_and_via_intercepted_release_get_no_dispatches() { .unwrap(), ) .await; - assert!( - listed.to_string().matches("\"ready\"").count() >= 3, - "still ready to the registry: {listed}" - ); + // Control state is RL's alone: the health machine still calls all three + // ready, so nothing the gateway does to a sleeping engine leaks into the + // registry's own view of it. + let listed = listed["workers"].as_array().expect("workers array"); + assert_eq!(listed.len(), 3); + for worker in listed { + assert_eq!(worker["status"], "ready", "{worker}"); + } // Wake them: API for w0, intercepted resume for w1. post_json( @@ -217,10 +221,12 @@ async fn rolling_update_under_latest_only_never_dispatches_to_a_stale_worker() { ) .await; assert_eq!(resp.status(), StatusCode::OK); - let routed = routed_over(&app, 8, None).await; - assert!( - routed.is_subset(&BTreeSet::from([ws[0].1.clone(), ws[1].1.clone()])), - "{routed:?}" + // Both refit workers, and only those two: an equality here is what proves + // the interceptor recorded the proxied call — a subset check would also + // pass if `ws[1]` had never been observed at version 2. + assert_eq!( + routed_over(&app, 8, None).await, + BTreeSet::from([ws[0].1.clone(), ws[1].1.clone()]) ); let resp = post_json( From 7edc867882a27de9740cb246d566edc4c02fce79 Mon Sep 17 00:00:00 2001 From: key4ng Date: Fri, 11 Sep 2026 22:11:52 -0700 Subject: [PATCH 18/27] feat(rl): stamp x-smg-weight-version and flag mixed-version generate responses Signed-off-by: key4ng --- model_gateway/src/rl_adapter.rs | 85 ++++++++++- model_gateway/src/server.rs | 20 ++- model_gateway/tests/common/mock_worker.rs | 60 +++++--- .../tests/rl_version_routing_test.rs | 142 ++++++++++++++++++ 4 files changed, 282 insertions(+), 25 deletions(-) diff --git a/model_gateway/src/rl_adapter.rs b/model_gateway/src/rl_adapter.rs index 9557bf6116..4bed2a071b 100644 --- a/model_gateway/src/rl_adapter.rs +++ b/model_gateway/src/rl_adapter.rs @@ -1,17 +1,29 @@ //! Glue between the RL control plane crate and the gateway. This file is the -//! whole of coupling surfaces (a) and (b); see `crates/rl/COUPLING.md`. +//! whole of coupling surfaces (a), (b), and the stamping half of (c); see +//! `crates/rl/COUPLING.md`. use std::sync::{Arc, Weak}; -use http::HeaderMap; +use axum::{ + body::{to_bytes, Body}, + extract::{Request, State}, + middleware::Next, + response::{IntoResponse, Response}, +}; +use http::{header::CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue}; +use http_body::Body as _; use openai_protocol::worker::ConnectionMode; -use smg_rl::{RlState, RlWorkerInfo, RlWorkerView, VersionEvictionSink}; +use smg_rl::{ + stamp::generate_is_mixed, table::base_url_of, RlError, RlState, RlWorkerInfo, RlWorkerView, + VersionEvictionSink, VersionPolicy, +}; use tokio::sync::broadcast::error::RecvError; use tracing::{debug, warn}; use crate::{ config::RouterConfig, policies::{CandidateFilter, PolicyRegistry}, + routers::{common::header_utils::RoutedWorker, error}, worker::{event::WorkerEvent, registry::WorkerId, Worker, WorkerRegistry}, }; @@ -248,6 +260,73 @@ pub fn build_rl_state( Some(rl) } +/// The version SMG believed the routed engine held when it answered. +static WEIGHT_VERSION: HeaderName = HeaderName::from_static(smg_rl::stamp::WEIGHT_VERSION_HEADER); +/// Set when a buffered `/generate` body spanned more than one version. +static MIXED_VERSION: HeaderName = HeaderName::from_static(smg_rl::stamp::MIXED_VERSION_HEADER); + +fn is_json(headers: &HeaderMap) -> bool { + headers + .get(CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .is_some_and(|ct| ct.contains("json")) +} + +/// Coupling (c): validate the per-request version policy, stamp every served +/// response with the engine's version, and flag buffered `/generate` bodies +/// that spanned two versions. Mounted only when `--enable-rl` is on. +/// +/// A streamed body is never read: `size_hint().exact()` is `None` for it, so +/// the mixed-version probe is skipped rather than buffering a generation. +pub async fn rl_middleware( + State(rl): State>, + request: Request, + next: Next, +) -> Response { + if let Err(e) = VersionPolicy::from_headers(request.headers()) { + return RlError::InvalidVersionPolicy(e.to_string()).into_response(); + } + // The request is consumed by `next`, so decide on the path now. + let inspect_body = request.uri().path() == "/generate"; + let mut response = next.run(request).await; + // Set by the gateway alone, so it names the engine that actually served + // this response — no upstream header can forge it. + let Some(routed) = response.extensions().get::().cloned() else { + return response; + }; + let version = rl.table().version_of(base_url_of(&routed.0)); + if let Some(value) = version + .as_ref() + .and_then(|v| HeaderValue::from_str(v.as_str()).ok()) + { + response.headers_mut().insert(WEIGHT_VERSION.clone(), value); + } + let buffered = response.body().size_hint().exact().is_some(); + if !(inspect_body && response.status().is_success() && is_json(response.headers()) && buffered) + { + return response; + } + let (parts, body) = response.into_parts(); + let limit = body.size_hint().exact().map_or(0, |n| n as usize); + match to_bytes(body, limit).await { + Ok(bytes) => { + let mixed = generate_is_mixed(&bytes, version.as_ref()); + let mut response = Response::from_parts(parts, Body::from(bytes)); + if mixed { + response + .headers_mut() + .insert(MIXED_VERSION.clone(), HeaderValue::from_static("true")); + smg_rl::metrics::record_mixed_version(); + } + response + } + Err(e) => { + warn!(error = %e, "RL middleware could not re-read a buffered generate body"); + error::internal_error("read_response_body_failed", "Failed to read response body") + } + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/model_gateway/src/server.rs b/model_gateway/src/server.rs index c9f4dc1935..bffe28ed5d 100644 --- a/model_gateway/src/server.rs +++ b/model_gateway/src/server.rs @@ -1024,6 +1024,16 @@ pub fn build_app( app = app.merge(rl_routes); } + // RL M2: validate x-smg-version-policy and stamp served responses with + // the engine's weight version. Innermost layer, so it sees the router's + // response before logging and metrics do. + if let Some(rl) = app_state.context.rl.as_ref() { + app = app.layer(axum::middleware::from_fn_with_state( + Arc::clone(rl), + crate::rl_adapter::rl_middleware, + )); + } + Ok(app .layer(axum::extract::DefaultBodyLimit::max(max_payload_size)) .layer(tower_http::limit::RequestBodyLimitLayer::new( @@ -1614,7 +1624,15 @@ fn create_cors_layer(allowed_origins: Vec) -> tower_http::cors::CorsLaye http::Method::OPTIONS, ]) .allow_headers([http::header::CONTENT_TYPE, http::header::AUTHORIZATION]) - .expose_headers([http::header::HeaderName::from_static("x-request-id")]) + // A browser client reads none of these unless they are exposed: + // the request id it correlates on, and the routing/version stamps + // an RL trainer needs to tell which engine answered. + .expose_headers([ + http::header::HeaderName::from_static("x-request-id"), + http::header::HeaderName::from_static("x-smg-routed-worker-id"), + http::header::HeaderName::from_static("x-smg-weight-version"), + http::header::HeaderName::from_static("x-smg-mixed-version"), + ]) }; cors.max_age(Duration::from_secs(3600)) diff --git a/model_gateway/tests/common/mock_worker.rs b/model_gateway/tests/common/mock_worker.rs index c5a106df40..f37dbbfe71 100755 --- a/model_gateway/tests/common/mock_worker.rs +++ b/model_gateway/tests/common/mock_worker.rs @@ -578,6 +578,24 @@ async fn model_info_handler(State(config): State>>) .into_response() } +/// Test hook: merge a request's `mock_meta_info` object into a response's +/// `meta_info`, so a test can make the mock report engine-side fields (weight +/// versions, for one) without a per-field knob on `MockWorkerConfig`. Existing +/// keys stay unless the request overrides them by name. +fn merge_mock_meta_info(payload: &serde_json::Value, response: &mut serde_json::Value) { + let Some(extra) = payload.get("mock_meta_info").and_then(|v| v.as_object()) else { + return; + }; + if let Some(meta) = response + .get_mut("meta_info") + .and_then(|m| m.as_object_mut()) + { + for (key, value) in extra { + meta.insert(key.clone(), value.clone()); + } + } +} + #[expect( clippy::unwrap_used, reason = "test helper - panicking on failure is intentional" @@ -643,7 +661,7 @@ async fn generate_handler( .unwrap() .as_secs_f64(); - let data = json!({ + let mut data = json!({ "text": format!("Mock response {}", i + 1), "meta_info": { "prompt_tokens": 10, @@ -663,6 +681,7 @@ async fn generate_handler( }, "stage": "mid" }); + merge_mock_meta_info(&payload, &mut data); events.push(Ok::<_, Infallible>(Event::default().data(data.to_string()))); } @@ -678,27 +697,26 @@ async fn generate_handler( ) .into_response() } else { - ( - [("x-worker-id", worker_id)], - Json(json!({ - "text": "This is a mock response.", - "meta_info": { - "prompt_tokens": 10, - "completion_tokens": 5, - "completion_tokens_wo_jump_forward": 5, - "input_token_logprobs": null, - "output_token_logprobs": null, - "first_token_latency": config.response_delay_ms as f64 / 1000.0, - "time_to_first_token": config.response_delay_ms as f64 / 1000.0, - "time_per_output_token": 0.01, - "finish_reason": { - "type": "stop", - "reason": "length" - } + let mut body = json!({ + "text": "This is a mock response.", + "meta_info": { + "prompt_tokens": 10, + "completion_tokens": 5, + "completion_tokens_wo_jump_forward": 5, + "input_token_logprobs": null, + "output_token_logprobs": null, + "first_token_latency": config.response_delay_ms as f64 / 1000.0, + "time_to_first_token": config.response_delay_ms as f64 / 1000.0, + "time_per_output_token": 0.01, + "finish_reason": { + "type": "stop", + "reason": "length" } - })), - ) - .into_response() + } + }); + merge_mock_meta_info(&payload, &mut body); + + ([("x-worker-id", worker_id)], Json(body)).into_response() } } diff --git a/model_gateway/tests/rl_version_routing_test.rs b/model_gateway/tests/rl_version_routing_test.rs index 8f2785841c..04065b554a 100644 --- a/model_gateway/tests/rl_version_routing_test.rs +++ b/model_gateway/tests/rl_version_routing_test.rs @@ -316,3 +316,145 @@ async fn flag_off_ignores_the_policy_header_and_stamps_nothing() { assert!(!ctx.app_context.policy_registry.has_candidate_filter()); ctx.shutdown().await; } + +#[tokio::test] +async fn header_override_and_invalid_header() { + let ctx = ctx("any", vec![mock(18928), mock(18929)]).await; + let app = ctx.create_app(); + let ws = workers(&app).await; + post_json( + &app, + &format!("/v1/rl/workers/{}/version", ws[0].0), + json!({"weight_version": "2"}), + ) + .await; + post_json( + &app, + &format!("/v1/rl/workers/{}/version", ws[1].0), + json!({"weight_version": "1"}), + ) + .await; + assert_eq!( + routed_over(&app, 8, None).await.len(), + 2, + "configured policy is any" + ); + assert_eq!( + routed_over(&app, 8, Some(("x-smg-version-policy", "latest-only"))).await, + BTreeSet::from([ws[0].1.clone()]) + ); + assert_eq!( + routed_over(&app, 8, Some(("x-smg-version-policy", "max-staleness:1"))) + .await + .len(), + 2 + ); + let resp = app + .clone() + .oneshot( + Request::post("/generate") + .header("content-type", "application/json") + .header("x-smg-version-policy", "newest") + .body(Body::from(json!({"text": "hi"}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + assert_eq!(json_of(resp).await["error"], "invalid_version_policy"); + ctx.shutdown().await; +} + +#[tokio::test] +async fn responses_carry_the_weight_version_once_known() { + let ctx = ctx("any", vec![mock(18932)]).await; + let app = ctx.create_app(); + let ws = workers(&app).await; + let gen = |stream: bool| json!({"text": "hi", "stream": stream}); + + let resp = post_json(&app, "/generate", gen(false)).await; + assert_eq!(resp.status(), StatusCode::OK); + assert!( + resp.headers().get("x-smg-weight-version").is_none(), + "unversioned engine: no header" + ); + + post_json( + &app, + &format!("/v1/rl/workers/{}/version", ws[0].0), + json!({"weight_version": "9"}), + ) + .await; + for (uri, body) in [ + ("/generate", gen(false)), + ("/generate", gen(true)), + ( + "/v1/chat/completions", + json!({"model": "mock-model", "messages": [{"role": "user", "content": "hi"}]}), + ), + ( + "/v1/chat/completions", + json!({"model": "mock-model", "stream": true, + "messages": [{"role": "user", "content": "hi"}]}), + ), + ] { + let resp = post_json(&app, uri, body).await; + assert_eq!(resp.status(), StatusCode::OK, "{uri}"); + assert_eq!(resp.headers()["x-smg-weight-version"], "9", "{uri}"); + assert_eq!(resp.headers()["x-smg-routed-worker-id"], ws[0].1.as_str()); + } + ctx.shutdown().await; +} + +#[tokio::test] +async fn buffered_generate_flags_mixed_versions() { + let ctx = ctx("any", vec![mock(18933)]).await; + let app = ctx.create_app(); + let ws = workers(&app).await; + post_json( + &app, + &format!("/v1/rl/workers/{}/version", ws[0].0), + json!({"weight_version": "9"}), + ) + .await; + + let mixed = json!({"text": "hi", "mock_meta_info": {"weight_version": "9", "weight_versions": [ + {"version": "8", "start": 0, "end": 3}, {"version": "9", "start": 3, "end": 5}]}}); + let resp = post_json(&app, "/generate", mixed).await; + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(resp.headers()["x-smg-mixed-version"], "true"); + assert_eq!(resp.headers()["x-smg-weight-version"], "9"); + let body = json_of(resp).await; + assert_eq!( + body["meta_info"]["weight_versions"] + .as_array() + .unwrap() + .len(), + 2, + "body untouched" + ); + + let clean = json!({"text": "hi", "mock_meta_info": {"weight_version": "9"}}); + let resp = post_json(&app, "/generate", clean).await; + assert!(resp.headers().get("x-smg-mixed-version").is_none()); + + let drifted = json!({"text": "hi", "mock_meta_info": {"weight_version": "7"}}); + let resp = post_json(&app, "/generate", drifted).await; + assert_eq!(resp.headers()["x-smg-mixed-version"], "true"); + + let streamed = + json!({"text": "hi", "stream": true, "mock_meta_info": {"weight_versions": [1, 2]}}); + let resp = post_json(&app, "/generate", streamed).await; + assert!( + resp.headers().get("x-smg-mixed-version").is_none(), + "streams are never inspected" + ); + // The engine really did report two spans, so the header's absence is the + // middleware declining to read a stream, not the mock declining the hook. + let events = resp.into_body().collect().await.unwrap().to_bytes(); + assert!( + String::from_utf8_lossy(&events).contains("weight_versions"), + "the streamed events carry the mixed-version metadata" + ); + ctx.shutdown().await; +} From 02ccdb205c2140ad1566612bc5adfc741c6f11be Mon Sep 17 00:00:00 2001 From: key4ng Date: Fri, 11 Sep 2026 22:40:58 -0700 Subject: [PATCH 19/27] feat(grpc): stamp the routed worker and report the live weight version Every JSON and SSE response the gRPC pipeline produces now names the worker that served it, via the same routed-worker header and extension the HTTP routers set. The RL middleware keys on that extension, so a stamped gRPC response picks up x-smg-weight-version for free. Dispatch metadata now reads the RL table first for the routed engine's weight version and falls back to the registration label, then to the "default" sentinel. The table is the live view of what an engine holds; the label is only what it reported when it joined. With the control plane off the pipeline passes None and the behavior is unchanged. Signed-off-by: key4ng --- .../grpc/common/stages/dispatch_metadata.rs | 113 +++++++- model_gateway/src/routers/grpc/pipeline.rs | 56 +++- model_gateway/src/routers/grpc/router.rs | 2 + model_gateway/tests/rl_grpc_stamp_test.rs | 251 ++++++++++++++++++ 4 files changed, 407 insertions(+), 15 deletions(-) create mode 100644 model_gateway/tests/rl_grpc_stamp_test.rs diff --git a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs index d0e1d80013..ee7ad2cd06 100644 --- a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs +++ b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs @@ -3,22 +3,34 @@ use std::time::{SystemTime, UNIX_EPOCH}; +use smg_rl::RlState; + use crate::routers::grpc::context::{DispatchMetadata, ExecutionPlan, WorkerSelection}; /// Metadata for one dispatch attempt. `dispatch_model` was captured from the /// request at the build boundary (already canonical); the weight version /// comes from the attempt's selected worker. +/// +/// `rl` is `None` when the control plane is off, which leaves the label-only +/// behavior the gateway had before it. pub(crate) fn prepare_dispatch_metadata( plan: &ExecutionPlan, dispatch_model: &str, workers: Option<&WorkerSelection>, + rl: Option<&RlState>, ) -> DispatchMetadata { - let weight_version = workers - .map(|w| match w { - WorkerSelection::Single { worker } => worker, - WorkerSelection::Disaggregated { decode, .. } => decode, + let worker = workers.map(|w| match w { + WorkerSelection::Single { worker } => worker, + WorkerSelection::Disaggregated { decode, .. } => decode, + }); + // The RL table is the live view of an engine's version; the registration + // label is what it reported when it joined. + let weight_version = worker + .and_then(|w| { + rl.and_then(|rl| rl.table().version_of(w.base_url())) + .map(|v| v.as_str().to_string()) + .or_else(|| w.metadata().spec.labels.get("weight_version").cloned()) }) - .and_then(|w| w.metadata().spec.labels.get("weight_version").cloned()) .unwrap_or_else(|| "default".to_string()); let created = SystemTime::now() @@ -33,3 +45,94 @@ pub(crate) fn prepare_dispatch_metadata( weight_version: Some(weight_version), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use openai_protocol::{model_card::ModelCard, worker::HealthCheckConfig}; + use smg_rl::{RlConfig, Version, VersionSource}; + + use super::*; + use crate::{ + rl_adapter::RegistryRlView, + routers::grpc::context::ExecutionPlanKind, + worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, Worker, WorkerRegistry}, + }; + + const MODEL: &str = "dispatch-metadata-model"; + const URL: &str = "grpc://127.0.0.1:39001"; + + /// A plan whose only role here is to carry the request id. + fn plan() -> ExecutionPlan { + ExecutionPlan::Batch { + kind: ExecutionPlanKind::Single, + shared_request_id: "req-1".to_string(), + requests: vec![], + } + } + + /// One regular gRPC worker, optionally carrying a `weight_version` label. + fn worker(label: Option<&str>) -> WorkerSelection { + let mut builder = BasicWorkerBuilder::new(URL) + .connection_mode(ConnectionMode::Grpc) + .runtime_type(RuntimeType::TokenSpeed) + .model(ModelCard::new(MODEL)) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }); + if let Some(label) = label { + builder = builder.label("weight_version", label); + } + WorkerSelection::Single { + worker: Arc::new(builder.build()) as Arc, + } + } + + fn rl_state() -> Arc { + let registry = Arc::new(WorkerRegistry::new()); + Arc::new(RlState::new( + Arc::new(RegistryRlView::new(registry)), + RlConfig::default(), + )) + } + + #[test] + fn no_worker_reports_the_default_sentinel() { + let metadata = prepare_dispatch_metadata(&plan(), MODEL, None, None); + assert_eq!(metadata.weight_version.as_deref(), Some("default")); + assert_eq!(metadata.request_id, "req-1"); + assert_eq!(metadata.model, MODEL); + } + + #[test] + fn without_the_control_plane_the_registration_label_is_the_version() { + let metadata = prepare_dispatch_metadata(&plan(), MODEL, Some(&worker(Some("7"))), None); + assert_eq!(metadata.weight_version.as_deref(), Some("7")); + } + + #[test] + fn an_unlabeled_worker_reports_the_default_sentinel() { + let metadata = prepare_dispatch_metadata(&plan(), MODEL, Some(&worker(None)), None); + assert_eq!(metadata.weight_version.as_deref(), Some("default")); + } + + #[test] + fn the_table_version_wins_over_the_registration_label() { + let rl = rl_state(); + rl.table() + .set_version(URL, MODEL, Version::parse("11"), VersionSource::Api); + let metadata = + prepare_dispatch_metadata(&plan(), MODEL, Some(&worker(Some("7"))), Some(&rl)); + assert_eq!(metadata.weight_version.as_deref(), Some("11")); + } + + #[test] + fn an_engine_the_table_has_no_version_for_falls_back_to_the_label() { + let rl = rl_state(); + let metadata = + prepare_dispatch_metadata(&plan(), MODEL, Some(&worker(Some("7"))), Some(&rl)); + assert_eq!(metadata.weight_version.as_deref(), Some("7")); + } +} diff --git a/model_gateway/src/routers/grpc/pipeline.rs b/model_gateway/src/routers/grpc/pipeline.rs index 1edb5e9182..b19bf983cf 100644 --- a/model_gateway/src/routers/grpc/pipeline.rs +++ b/model_gateway/src/routers/grpc/pipeline.rs @@ -22,6 +22,7 @@ use openai_protocol::{ messages::CreateMessageRequest, }; use reasoning_parser::ParserFactory as ReasoningParserFactory; +use smg_rl::RlState; use tool_parser::ParserFactory as ToolParserFactory; use tracing::error; @@ -64,7 +65,10 @@ use crate::{ policies::PolicyRegistry, rate_limit::{RateLimitManager, UsageSettlement}, routers::{ - common::retry::{is_retryable_response, BackoffCalculator}, + common::{ + header_utils, + retry::{is_retryable_response, BackoffCalculator}, + }, error, }, worker::WorkerRegistry, @@ -100,11 +104,15 @@ pub(crate) struct PipelineDeps { /// `None` when tenant rate limiting is disabled; read only by the /// endpoints that insert `RateLimitReserveStage` (chat/messages/completion/harmony). rate_limit_manager: Option>, + /// RL control plane state; `None` unless `--enable-rl`. Read per attempt + /// for the routed engine's live weight version. + rl: Option>, } impl PipelineDeps { /// Full deps for the chat/messages/harmony endpoints, which consume the /// configured parser factories/overrides. + #[expect(clippy::too_many_arguments)] pub(crate) fn new( worker_registry: Arc, policy_registry: Arc, @@ -113,6 +121,7 @@ impl PipelineDeps { configured_tool_parser: Option, configured_reasoning_parser: Option, rate_limit_manager: Option>, + rl: Option>, ) -> Self { Self { worker_registry, @@ -122,6 +131,7 @@ impl PipelineDeps { configured_tool_parser, configured_reasoning_parser, rate_limit_manager, + rl, } } @@ -131,6 +141,7 @@ impl PipelineDeps { worker_registry: Arc, policy_registry: Arc, rate_limit_manager: Option>, + rl: Option>, ) -> Self { Self { worker_registry, @@ -140,6 +151,7 @@ impl PipelineDeps { configured_tool_parser: None, configured_reasoning_parser: None, rate_limit_manager, + rl, } } @@ -204,6 +216,7 @@ impl PipelineDeps { configured_tool_parser: None, configured_reasoning_parser: None, rate_limit_manager: None, + rl: None, } } } @@ -231,6 +244,8 @@ pub(crate) struct RequestPipeline { backend_type: &'static str, /// Disaggregation mode, for per-leg retry metric labels. mode: Mode, + /// RL control plane state; `None` unless `--enable-rl`. + rl: Option>, } /// Outcome of one full pipeline run. @@ -393,6 +408,7 @@ impl RequestPipeline { stages: Arc::new(stages), backend_type: backend, mode, + rl: deps.rl.clone(), }) } @@ -493,6 +509,7 @@ impl RequestPipeline { &attempt_plan, &dctx.dispatch_model, dctx.workers.as_ref(), + self.rl.as_deref(), )); execute_plan(dctx, attempt_plan).await?; @@ -588,7 +605,7 @@ impl RequestPipeline { attempt_start.elapsed(), ); } - return Ok(RunOutcome::Early(response)); + return Ok(RunOutcome::Early(Self::stamp_routed(&dctx, response))); } Ok(None) => return Ok(RunOutcome::Final(dctx, attempt_start)), Err(response) => response, @@ -630,6 +647,21 @@ impl RequestPipeline { } } + /// Stamp the worker that served this attempt (the single worker, or the + /// decode leg of a disaggregated pair -- the same choice + /// `prepare_dispatch_metadata` makes). The routed-worker extension is what + /// the RL middleware keys on to add `x-smg-weight-version`. + fn stamp_routed(dctx: &DispatchContext, mut response: Response) -> Response { + if let Some(workers) = dctx.workers.as_ref() { + let worker = match workers { + WorkerSelection::Single { worker } => worker, + WorkerSelection::Disaggregated { decode, .. } => decode, + }; + header_utils::stamp_routed_worker(&mut response, worker.url()); + } + response + } + fn record_error(&self, endpoint: Option<&'static str>, model: &str, response: &Response) { if let Some(endpoint) = endpoint { Metrics::record_router_error( @@ -774,7 +806,7 @@ impl RequestPipeline { ) .await; self.record_duration(ENDPOINT, &dctx.model_id, start); - axum::Json(response).into_response() + Self::stamp_routed(&dctx, axum::Json(response).into_response()) } Some(other) => self.wrong_response_type( "execute_chat", @@ -821,7 +853,7 @@ impl RequestPipeline { ) .await; self.record_duration(ENDPOINT, &dctx.model_id, start); - axum::Json(response).into_response() + Self::stamp_routed(&dctx, axum::Json(response).into_response()) } Some(other) => self.wrong_response_type( "execute_generate", @@ -865,7 +897,7 @@ impl RequestPipeline { ) .await; self.record_duration(ENDPOINT, &dctx.model_id, start); - axum::Json(response).into_response() + Self::stamp_routed(&dctx, axum::Json(response).into_response()) } Some(other) => self.wrong_response_type( "execute_completion", @@ -908,7 +940,7 @@ impl RequestPipeline { ) .await; self.record_duration(ENDPOINT, &dctx.model_id, start); - axum::Json(response).into_response() + Self::stamp_routed(&dctx, axum::Json(response).into_response()) } Some(other) => self.wrong_response_type( "execute_messages", @@ -941,7 +973,7 @@ impl RequestPipeline { Ok(RunOutcome::Final(mut dctx, start)) => match dctx.response.final_response.take() { Some(FinalResponse::Embedding(response)) => { self.record_duration(ENDPOINT, &dctx.model_id, start); - axum::Json(response).into_response() + Self::stamp_routed(&dctx, axum::Json(response).into_response()) } Some(other) => self.wrong_response_type( "execute_embeddings", @@ -978,7 +1010,10 @@ impl RequestPipeline { Ok(RunOutcome::Final(mut dctx, start)) => match dctx.response.final_response.take() { Some(FinalResponse::Transcription { text, format }) => { self.record_duration(ENDPOINT, &dctx.model_id, start); - super::regular::stages::transcription::render(format, text) + Self::stamp_routed( + &dctx, + super::regular::stages::transcription::render(format, text), + ) } Some(other) => self.wrong_response_type( "execute_transcription", @@ -1013,7 +1048,7 @@ impl RequestPipeline { Ok(RunOutcome::Final(mut dctx, start)) => match dctx.response.final_response.take() { Some(FinalResponse::Classify(response)) => { self.record_duration(ENDPOINT, &dctx.model_id, start); - axum::Json(response).into_response() + Self::stamp_routed(&dctx, axum::Json(response).into_response()) } Some(other) => self.wrong_response_type( "execute_classify", @@ -1354,7 +1389,7 @@ mod alias_pipeline_tests { .unwrap(); let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); - let deps = PipelineDeps::pair(worker_registry.clone(), policy_registry, None); + let deps = PipelineDeps::pair(worker_registry.clone(), policy_registry, None, None); let pipeline = RequestPipeline::build(Endpoint::Chat, Mode::PrefillDecode, &deps).unwrap(); let components = Arc::new(SharedComponents { tokenizer_registry, @@ -1738,6 +1773,7 @@ mod request_release_tests { worker_registry.clone(), Arc::new(PolicyRegistry::new(PolicyConfig::Random)), None, + None, ); RequestPipeline::build(Endpoint::Completion, mode, &deps).expect("completion pipeline") } diff --git a/model_gateway/src/routers/grpc/router.rs b/model_gateway/src/routers/grpc/router.rs index 7ff2259662..e31428097f 100644 --- a/model_gateway/src/routers/grpc/router.rs +++ b/model_gateway/src/routers/grpc/router.rs @@ -130,6 +130,7 @@ impl GrpcRouter { ctx.configured_tool_parser.clone(), ctx.configured_reasoning_parser.clone(), ctx.rate_limit_manager.clone(), + ctx.rl.clone(), ); // Deps for the parser-free endpoints (completion/embeddings/classify). // Only completion's stage list actually reads `rate_limit_manager`; @@ -138,6 +139,7 @@ impl GrpcRouter { worker_registry.clone(), policy_registry.clone(), ctx.rate_limit_manager.clone(), + ctx.rl.clone(), ); // Present in every mode: chat/generate, messages, completion. diff --git a/model_gateway/tests/rl_grpc_stamp_test.rs b/model_gateway/tests/rl_grpc_stamp_test.rs new file mode 100644 index 0000000000..0cfe314ed1 --- /dev/null +++ b/model_gateway/tests/rl_grpc_stamp_test.rs @@ -0,0 +1,251 @@ +//! RL M2: the gRPC pipeline names the worker that served each response, and +//! reports the version the RL table holds for that engine. +//! +//! The routers are driven directly (no axum app), so the RL middleware never +//! runs here and `x-smg-weight-version` is not expected: what these tests pin +//! is the pipeline's own half of the contract -- the routed-worker header and +//! extension that the middleware keys on, and the live table version reaching +//! the body as `system_fingerprint`. + +#[path = "common/mod.rs"] +mod common; + +use std::{sync::Arc, time::Duration}; + +use llm_tokenizer::{traits::Tokenizer, MockTokenizer, TokenizerRegistry}; +use openai_protocol::{ + chat::ChatCompletionRequest, model_card::ModelCard, worker::HealthCheckConfig, +}; +use smg::{ + app_context::AppContext, + config::RouterConfig, + middleware::TenantRequestMeta, + routers::{common::header_utils::RoutedWorker, RouterFactory, RouterTrait}, + tenant::TenantKey, + worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, WorkerType}, +}; +use tokio::net::TcpListener; + +const MODEL: &str = "rl-grpc-stamp-test-model"; +/// Tokens the canned mock emits per request. +const OUTPUT_TOKENS: u32 = 3; +/// The version the RL table is told to hold for the registered engine. The +/// worker registers with no `weight_version` label, so the pipeline can only +/// report this by reading the table. +const LIVE_VERSION: &str = "11"; + +/// Spawn a canned mock gRPC worker in-process. +#[expect( + clippy::expect_used, + clippy::disallowed_methods, + reason = "test helper - panicking on failure is intentional; the spawned \ + mock server task is fire-and-forget for the test process's lifetime" +)] +async fn start_mock_grpc_worker() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind mock gRPC worker"); + let port = listener + .local_addr() + .expect("mock gRPC worker address") + .port(); + let cfg = Arc::new(mock_worker::config::Config { + host: "127.0.0.1".to_string(), + http_base_port: 0, + http_count: 0, + grpc_base_port: port, + grpc_count: 1, + zmq_handshake: None, + zmq_count: 0, + zmq_start_index: 0, + model_id: MODEL.to_string(), + tokenizer_path: MODEL.to_string(), + gen_delay: Duration::ZERO, + output_tokens: OUTPUT_TOKENS, + realistic: false, + engine: mock_worker::engine::EngineParams::default(), + }); + tokio::spawn(mock_worker::grpc::serve_with_listener(cfg, listener)); + port +} + +/// A regular-mode gRPC router over one mock worker, with the RL control plane +/// enabled so the app context carries an `RlState`. +#[expect( + clippy::unwrap_used, + clippy::expect_used, + reason = "test helper - panicking on failure is intentional" +)] +async fn build_regular_router_with_rl(port: u16) -> (Box, Arc) { + let mut config = RouterConfig::builder() + .grpc_connection() + .regular_mode(vec![]) + .random_policy() + .host("127.0.0.1") + .port(0) + .max_payload_size(1024 * 1024) + .build_unchecked(); + config.health_check.disable_health_check = true; + config.rl.enabled = true; + + let tokenizer_registry = Arc::new(TokenizerRegistry::new()); + let tokenizer = Arc::new(MockTokenizer::new()) as Arc; + tokenizer_registry + .load( + "tokenizer-id", + MODEL, + "test", + || async move { Ok(tokenizer) }, + ) + .await + .unwrap(); + + let app_context = + common::create_test_context_with_tokenizer_registry(config, tokenizer_registry).await; + let worker = BasicWorkerBuilder::new(worker_url(port)) + .worker_type(WorkerType::Regular) + .connection_mode(ConnectionMode::Grpc) + .runtime_type(RuntimeType::TokenSpeed) + .model(ModelCard::new(MODEL)) + .health_config(HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(); + app_context + .worker_registry + .register(Arc::new(worker)) + .unwrap(); + + let router = RouterFactory::create_router(&app_context) + .await + .expect("regular gRPC router should build"); + (router, app_context) +} + +fn worker_url(port: u16) -> String { + format!("grpc://127.0.0.1:{port}") +} + +#[expect( + clippy::unwrap_used, + reason = "test helper - panicking on failure is intentional" +)] +fn chat_request(stream: bool) -> ChatCompletionRequest { + serde_json::from_value(serde_json::json!({ + "model": MODEL, + "messages": [{"role": "user", "content": "Hello world"}], + "max_tokens": 8, + "stream": stream, + })) + .unwrap() +} + +fn tenant() -> TenantRequestMeta { + TenantRequestMeta::new(TenantKey::new("rl-grpc-stamp-test-tenant")) +} + +#[expect( + clippy::expect_used, + reason = "test helper - panicking on failure is intentional" +)] +async fn read_body(response: axum::response::Response) -> String { + let bytes = axum::body::to_bytes(response.into_body(), 1 << 20) + .await + .expect("response body"); + String::from_utf8(bytes.to_vec()).expect("utf-8 body") +} + +/// A buffered gRPC chat response carries the routed-worker header and +/// extension, and reports the table's version rather than the registration +/// label (there is none) or the `"default"` fallback. +#[tokio::test] +async fn a_buffered_grpc_response_names_its_worker_and_reports_the_live_version() { + let port = start_mock_grpc_worker().await; + let (router, ctx) = build_regular_router_with_rl(port).await; + let url = worker_url(port); + let rl = ctx.rl.as_ref().expect("rl enabled"); + rl.table().set_version( + &url, + MODEL, + smg_rl::Version::parse(LIVE_VERSION), + smg_rl::VersionSource::Api, + ); + + let response = router + .route_chat(None, &tenant(), chat_request(false), MODEL) + .await; + + assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!( + response.headers()["x-smg-routed-worker-id"], + url.as_str(), + "the buffered response must name the worker that served it" + ); + assert!( + response.extensions().get::().is_some(), + "the RL middleware keys on the RoutedWorker extension" + ); + + let body = read_body(response).await; + let json: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(json["system_fingerprint"], LIVE_VERSION, "{body}"); +} + +/// The streamed variant of the same request: the SSE response is stamped too, +/// and its chunks carry the live version. +#[tokio::test] +async fn a_streamed_grpc_response_names_its_worker_and_reports_the_live_version() { + let port = start_mock_grpc_worker().await; + let (router, ctx) = build_regular_router_with_rl(port).await; + let url = worker_url(port); + let rl = ctx.rl.as_ref().expect("rl enabled"); + rl.table().set_version( + &url, + MODEL, + smg_rl::Version::parse(LIVE_VERSION), + smg_rl::VersionSource::Api, + ); + + let response = router + .route_chat(None, &tenant(), chat_request(true), MODEL) + .await; + + assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!( + response.headers()["x-smg-routed-worker-id"], + url.as_str(), + "the SSE response must name the worker that served it" + ); + assert!( + response.extensions().get::().is_some(), + "the RL middleware keys on the RoutedWorker extension" + ); + + let body = read_body(response).await; + assert!( + body.contains(&format!("\"system_fingerprint\":\"{LIVE_VERSION}\"")), + "streamed chunks must report the live version: {body}" + ); +} + +/// With no version in the table the pipeline falls back to what the worker +/// reported at registration -- here nothing, so the `"default"` sentinel -- +/// and the stamp is still applied. +#[tokio::test] +async fn an_unversioned_engine_still_stamps_and_falls_back_to_default() { + let port = start_mock_grpc_worker().await; + let (router, _ctx) = build_regular_router_with_rl(port).await; + let url = worker_url(port); + + let response = router + .route_chat(None, &tenant(), chat_request(false), MODEL) + .await; + + assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!(response.headers()["x-smg-routed-worker-id"], url.as_str()); + + let body = read_body(response).await; + let json: serde_json::Value = serde_json::from_str(&body).unwrap(); + assert_eq!(json["system_fingerprint"], "default", "{body}"); +} From 3eba7a5e692134b8787e9b599d2a034272eb1fe8 Mon Sep 17 00:00:00 2001 From: key4ng Date: Fri, 11 Sep 2026 22:57:06 -0700 Subject: [PATCH 20/27] docs(rl): version and state client calls, coupling contract for M2, launch profile Add smg.rl.RL.set_version/set_fleet_version/set_state/set_fleet_state and the control/version_source fields on Worker, with stub-server tests. Extend the refit_from_disk.py example to assert x-smg-weight-version matches the refit. Document the four new /v1/rl routes, the version/control-state table, and the retry-budget launch profile; update crates/rl/COUPLING.md for the M2 surfaces and crates/rl/NOTES.md for the SGLang/vLLM version-source drift. Also fix the two codespell hits introduced by M2: unparseable -> unparsable in stamp.rs, and the abc/abd test strings in version.rs. Signed-off-by: key4ng --- bindings/python/src/smg/rl.py | 40 +++++++++++++++++++++++++ bindings/python/tests/test_rl_client.py | 26 ++++++++++++++++ crates/rl/COUPLING.md | 15 ++++++++-- crates/rl/NOTES.md | 2 ++ crates/rl/README.md | 30 ++++++++++++++++++- crates/rl/src/stamp.rs | 2 +- crates/rl/src/version.rs | 2 +- examples/rl/README.md | 20 +++++++++++-- examples/rl/refit_from_disk.py | 14 ++++++--- 9 files changed, 140 insertions(+), 11 deletions(-) diff --git a/bindings/python/src/smg/rl.py b/bindings/python/src/smg/rl.py index 82fede79eb..70c46a154b 100644 --- a/bindings/python/src/smg/rl.py +++ b/bindings/python/src/smg/rl.py @@ -11,6 +11,10 @@ Only HTTP workers can be proxied. A gRPC or ZMQ worker matched by a selector is reported in `failed[]` as `unsupported_connection_mode`, which makes `fanout` raise `FanoutError` unless `allow_partial=True`. + +`set_fleet_version` (like `set_version`) records a version SMG did not observe +itself; it is a local table write, never an engine call. Every response SMG +serves for a versioned engine carries `x-smg-weight-version`. """ from __future__ import annotations @@ -54,12 +58,15 @@ class Worker: weight_version: str | None labels: dict[str, str] = field(default_factory=dict) capabilities: dict[str, Any] = field(default_factory=dict) + control: str = "active" + version_source: str | None = None @classmethod def from_json(cls, d: dict[str, Any]) -> Worker: kwargs = {k: d.get(k) for k in cls.__dataclass_fields__} kwargs["labels"] = d.get("labels") or {} kwargs["capabilities"] = d.get("capabilities") or {} + kwargs["control"] = d.get("control") or "active" return cls(**kwargs) # type: ignore[arg-type] @@ -177,6 +184,39 @@ def worker(self, worker_id: str) -> Worker: _raise_for(status, payload) return Worker.from_json(payload) + def set_version(self, worker_id: str, weight_version: str) -> Worker: + """Record the weight version one engine holds (SMG-local; no engine call).""" + status, payload = self._request( + "POST", f"/v1/rl/workers/{worker_id}/version", {"weight_version": str(weight_version)} + ) + _raise_for(status, payload) + return Worker.from_json(payload) + + def set_fleet_version(self, weight_version: str, *, selector: str) -> list[Worker]: + status, payload = self._request( + "POST", + "/v1/rl/version", + {"weight_version": str(weight_version)}, + {"selector": selector}, + ) + _raise_for(status, payload) + return [Worker.from_json(w) for w in payload.get("workers", [])] + + def set_state(self, worker_id: str, control: str) -> Worker: + """Mark one engine active, paused, or asleep for routing (SMG-local).""" + status, payload = self._request( + "POST", f"/v1/rl/workers/{worker_id}/state", {"control": control} + ) + _raise_for(status, payload) + return Worker.from_json(payload) + + def set_fleet_state(self, control: str, *, selector: str) -> list[Worker]: + status, payload = self._request( + "POST", "/v1/rl/state", {"control": control}, {"selector": selector} + ) + _raise_for(status, payload) + return [Worker.from_json(w) for w in payload.get("workers", [])] + def call( self, worker_id: str, diff --git a/bindings/python/tests/test_rl_client.py b/bindings/python/tests/test_rl_client.py index 6c2f7216a0..5420efa8c8 100644 --- a/bindings/python/tests/test_rl_client.py +++ b/bindings/python/tests/test_rl_client.py @@ -95,6 +95,32 @@ def test_workers_and_worker(stub): assert _Stub.seen[0]["auth"] == "Bearer k" +def test_set_version_and_state_post_bodies(stub): + _Stub.stub_responses["/v1/rl/workers/w1/version"] = (200, WORKER) + rl = RL(stub) + w = rl.set_version("w1", "42") + assert _Stub.seen[-1]["body"] == {"weight_version": "42"} + assert w.control == "active" and w.version_source == "api" + + +def test_fleet_forms_send_the_selector(stub): + _Stub.stub_responses["/v1/rl/state"] = ( + 200, + {"protocol_version": 1, "workers": [WORKER], "total": 1}, + ) + rl = RL(stub) + workers = rl.set_fleet_state("paused", selector="engine=sglang") + assert _Stub.seen[-1]["query"] == {"selector": ["engine=sglang"]} + assert len(workers) == 1 and workers[0].id == "w1" + + +def test_worker_from_json_defaults_control(): + d = dict(WORKER) + del d["control"] + w = Worker.from_json(d) + assert w.control == "active" + + def test_call_sends_method_params_body(stub): _Stub.stub_responses["/v1/rl/workers/w1/engine/pause"] = ( 200, diff --git a/crates/rl/COUPLING.md b/crates/rl/COUPLING.md index 9466c3705b..c16f98ff13 100644 --- a/crates/rl/COUPLING.md +++ b/crates/rl/COUPLING.md @@ -1,4 +1,4 @@ -# Coupling contract with model_gateway (M1) +# Coupling contract with model_gateway (M1, M2) The RL crate may touch the gateway only through the surfaces below. Any PR that adds a surface must update this file. @@ -11,6 +11,15 @@ that adds a surface must update this file. | (d) | metrics HELP registration | `src/observability/metrics.rs` | `smg_rl::init_rl_metrics()` | | (d) | config + flags | `src/config/{types,builder,validation}.rs`, `src/main.rs`, `bindings/python/src/smg/router_args.py`, `bindings/python/src/lib.rs` | `RouterConfig.rl`, three CLI flags | | (d) | manifests | `model_gateway/Cargo.toml`, `bindings/python/Cargo.toml` | `smg-rl` pulled in via `workspace = true` (router crate and the pyo3 bindings crate) | +| (a) | `RlWorkerView::inflight(base_url)` and `VersionEvictionSink` | `src/rl_adapter.rs` | inflight sums `Worker::load()` over the ranks; `VersionEvictionSink::on_version_changed` runs under the RL table's write mutex, so the gateway's `RegistryEvictionSink` only resets policy caches and never writes back into the table; it holds `Weak`/`Weak` handles (a strong pair would close an ownership cycle through the candidate filter and leak both registries) and calls `PolicyRegistry::reset_worker_cache`, which covers the model policy, the PD/EPD leg policies, and the default-policy fallback | +| (a) | registry event subscription | `src/rl_adapter.rs` | seeds on `Registered`, reseeds on `Replaced` only when the discovered label changed, drops on the last `Removed` for a base URL; a `Replaced` that changes a worker's `model_id` without changing its label calls `seed`, which is insert-if-absent, so the table entry keeps the *old* `model_id` until the worker is removed and re-registered | +| (a) | `CacheAwarePolicy::reset_worker_cache`, `PolicyRegistry::reset_worker_cache` | `src/policies/cache_aware.rs`, `src/policies/registry.rs` | tree purge + root re-seed, no load-scorer effect | +| (b) | `CandidateFilter` trait and `PolicyRegistry::{candidate_filter, set_candidate_filter}` | `src/policies/mod.rs`, `src/policies/registry.rs`, `src/app_context.rs` | set in `AppContextBuilder::build()` when `rl` is `Some`; EPD per-item encode assignment and external-provider routing call their own selection paths and bypass the candidate filter by design | +| (c) | `RoutedWorker` + `stamp_routed_worker` | `crates/external_router/src/header_utils.rs`, seven HTTP sites, `routers/grpc/pipeline.rs` | header and extension in one call; the seven HTTP sites include the Responses-API streaming path in `crates/external_router/src/openai/responses/streaming.rs`; gRPC stamps the eight `pipeline.rs` dispatch sites, while the gRPC Responses-API helpers (typed, non-streaming returns) stay unstamped in M2 | +| (c) | RL middleware | `src/rl_adapter.rs`, `src/server.rs` `build_app` | header validation, version stamp, mixed flag; added inside `build_app`, so `tests/common/test_app.rs` inherits it automatically | +| (c) | `dispatch_metadata.rs` reads the table first | `src/routers/grpc/common/stages/dispatch_metadata.rs` (+ the pipeline holding `Option>`) | label fallback unchanged | +| (c) | CORS expose list | `src/server.rs` `create_cors_layer` | three headers | +| (d) | `--rl-version-policy` | `src/config/{types,validation}.rs`, `src/main.rs`, Python bindings | | Test-only files that the mount also touches, none of them a new surface: `model_gateway/tests/rl_control_plane_test.rs` (gateway-level `/v1/rl` tests), @@ -48,4 +57,6 @@ so the duplicated loop costs nothing at runtime. The shared mechanism (the header allow-list and one bounded reader for the four readers now in the tree) is tracked in #2489. -Not touched: policies, routers, worker trait, response pipeline (M2). +Not touched: the worker trait, the health machine, `WorkerStatus::Draining`; +EPD per-item encode assignment and external-provider routing bypass the +candidate filter by design. diff --git a/crates/rl/NOTES.md b/crates/rl/NOTES.md index 62caeb4c02..465eb073a2 100644 --- a/crates/rl/NOTES.md +++ b/crates/rl/NOTES.md @@ -7,3 +7,5 @@ planning docs, with date, engine version, and what was done. |---|---|---|---|---| | 2026-09-03 | SGLang 0.5.15.post1 | `/get_weight_version` | 404; version is in `/model_info` | discovery reads the `weight_version` registration label | | 2026-09-03 | SGLang 0.5.15.post1 | `POST /pause_generation` with no body succeeds | 400 Bad Request; FastAPI requires a JSON body, `{}` is enough | callers must send `{}` on bodyless control routes; `examples/rl/refit_from_disk.py` needs `rl.fanout("pause_generation", {}, ...)` (the e2e test already passes `json={}`) | +| 2026-09-10 | SGLang 0.5.15.post1 | discovery reports a real version before the first refit | registration reports `weight_version: default` (the placeholder) until a refit lands | seeded as unversioned (`Version::from_label` treats `default` as `None`) so a first refit is not out-ordered by the placeholder | +| 2026-09-10 | vLLM (version unpinned) | refit routes carry a version field the passthrough observer can read | documented weight-transfer bodies carry no version field | passthrough observes control state only for vLLM; versions must come through the API (`set_version`/`set_fleet_version`); verify against a live vLLM when one is available | diff --git a/crates/rl/README.md b/crates/rl/README.md index 9a7641275a..bea7444798 100644 --- a/crates/rl/README.md +++ b/crates/rl/README.md @@ -8,6 +8,10 @@ RL control plane for the Shepherd Model Gateway. Enabled with `--enable-rl`. | `GET /v1/rl/workers/{id}` | one worker | | `GET\|POST /v1/rl/workers/{id}/engine/{path}` | proxy one engine-native route to one worker | | `GET\|POST /v1/rl/engine/{path}?selector=...` | the same call fanned out to every matching worker | +| `POST /v1/rl/workers/{id}/version` | record the weight version one worker holds (SMG-local; no engine call) | +| `POST /v1/rl/workers/{id}/state` | mark one worker `active`, `paused`, or `asleep` for routing (SMG-local) | +| `POST /v1/rl/version?selector=...` | the version write fanned out to every matching worker | +| `POST /v1/rl/state?selector=...` | the state write fanned out to every matching worker | Response bodies are the `openai_protocol::rl` types (`RlWorkersResponse`, `RlWorkerEntry`, `RlCallOutcome`, `RlFanoutResponse`), registered in @@ -19,9 +23,33 @@ reported in `failed[]` as `unsupported_connection_mode` (HTTP 422 on the per-wor route), so a fan-out over a mixed fleet answers 207 and `smg.rl.RL.fanout` raises `FanoutError` unless `allow_partial=True`. -Flags: `--enable-rl`, `--rl-control-timeout-secs` (600), `--rl-fanout-concurrency` (32). +Flags: `--enable-rl`, `--rl-control-timeout-secs` (600), `--rl-fanout-concurrency` (32), +`--rl-version-policy` (any). Recommended RL launch profile: `--enable-rl --disable-health-check --disable-circuit-breaker --request-timeout-secs 14400`. +## Versions and control state + +SMG keeps a side table of `(weight_version, control_state)` per engine, keyed +by base URL (DP ranks share an entry). Three sources feed it, last write wins: +the refit and pause/sleep calls proxied through `/v1/rl` (`version_source: +passthrough`), the explicit API above (`version_source: api`), and the +`weight_version` registration label at discovery time (`version_source: +registration`; SGLang's `default` placeholder seeds as unversioned). `GET +/v1/rl/workers` reports the table's current view: `weight_version` (`null` +when unversioned), `version_source` (`null` when unversioned), and `control` +(`active`, `paused`, or `asleep`). + +With `--enable-rl`, a worker the table shows `paused` or `asleep` is never +selected for inference; `--rl-version-policy` (`any`, `latest-only`, +`min-version:`, `max-staleness:`, or the per-request +`x-smg-version-policy` header) further narrows the candidates by version +against the model's known maximum. + +Two response headers carry the table's view of what served a request: +`x-smg-weight-version` on every response from a versioned engine, and +`x-smg-mixed-version: true` on a buffered `/generate` response whose tokens +spanned two versions. + ## Python client ```python diff --git a/crates/rl/src/stamp.rs b/crates/rl/src/stamp.rs index db0a9c1695..8d1a1b68e8 100644 --- a/crates/rl/src/stamp.rs +++ b/crates/rl/src/stamp.rs @@ -56,7 +56,7 @@ fn is_mixed(meta: &MetaProbe, expected: Option<&Version>) -> bool { /// Whether a buffered `/generate` body (one object, or an array for `n > 1`) /// spans more than one weight version or reports a version other than /// `expected`. Only `meta_info.weight_version` and the span count are read; -/// an unparseable body is not mixed. +/// an unparsable body is not mixed. pub fn generate_is_mixed(body: &[u8], expected: Option<&Version>) -> bool { let first = body.iter().find(|b| !b.is_ascii_whitespace()); match first { diff --git a/crates/rl/src/version.rs b/crates/rl/src/version.rs index 7ff562c562..660e85ffef 100644 --- a/crates/rl/src/version.rs +++ b/crates/rl/src/version.rs @@ -109,7 +109,7 @@ mod tests { fn text_versions_compare_lexically() { assert!(Version::parse("step-9") > Version::parse("step-10")); assert_eq!(Version::parse("abc").numeric(), None); - assert_ne!(Version::parse("abc"), Version::parse("abd")); + assert_ne!(Version::parse("abc"), Version::parse("acd")); } #[test] diff --git a/examples/rl/README.md b/examples/rl/README.md index 235e93bb9a..cd57c5cae6 100644 --- a/examples/rl/README.md +++ b/examples/rl/README.md @@ -4,7 +4,9 @@ ```bash smg launch --worker-urls http://rollout:30000 http://rollout:30001 --policy cache_aware \ - --enable-rl --disable-health-check --disable-circuit-breaker --request-timeout-secs 14400 + --enable-rl --disable-health-check --disable-circuit-breaker --request-timeout-secs 14400 \ + --rl-version-policy latest-only --retry-max-retries 30 --retry-initial-backoff-ms 100 \ + --retry-max-backoff-ms 1000 ``` `--disable-health-check` and `--disable-circuit-breaker` match what slime and vime @@ -12,6 +14,19 @@ set on their own routers (a transient engine error must not open a breaker for a whole training step); `--request-timeout-secs 14400` covers multi-hour agentic rollouts. Raise `--rl-control-timeout-secs` (default 600) if a disk refit takes longer. +With `--enable-rl`, SMG never routes to an engine it has observed paused or +asleep, whether that came from a call proxied through `/v1/rl` or from +`smg.rl.RL.set_state`. `--rl-version-policy` (or the per-request +`x-smg-version-policy` header) then narrows the remaining candidates by +version: `any` (default), `latest-only`, `min-version:`, or +`max-staleness:`. An empty eligible set answers a retryable 503, so the +retry knobs above are the wait budget across a refit — at the values shown, +about 27 s. Every response SMG serves for a versioned engine carries +`x-smg-weight-version`; a buffered `/generate` response whose tokens spanned +two versions also carries `x-smg-mixed-version: true`. SGLang's `default` +label counts as unversioned. An engine restarted behind the same URL without +re-registering keeps its table entry until the next refit or `set_version`. + ## refit_from_disk.py Pauses, refits, and resumes every worker matching `--selector`, then confirms the @@ -28,4 +43,5 @@ Exit code 0 on success, 1 on a version mismatch; a failed fan-out raises ## Selector cheatsheet `engine=sglang` · `engine in (sglang,vllm)` · `role!=reward` · `url=http://rollout:30000` -(single worker) · `model=Qwen/Qwen3-8B, tp_size=1` +(single worker) · `model=Qwen/Qwen3-8B, tp_size=1` · `control=active` · +`weight_version!=42` diff --git a/examples/rl/refit_from_disk.py b/examples/rl/refit_from_disk.py index 07c4f9eec6..6c14ef93a9 100755 --- a/examples/rl/refit_from_disk.py +++ b/examples/rl/refit_from_disk.py @@ -6,10 +6,11 @@ Sequence: pause_generation -> update_weights_from_disk -> continue_generation, each as one fan-out, then one /generate through SMG to confirm the engine -reports the new meta_info.weight_version. The pause/resume pair is -`smg.rl.paused`, so a failure at any stage still resumes the engines that did -pause. Only HTTP workers can be proxied: a gRPC or ZMQ worker matched by -`--selector` fails the fan-out with `unsupported_connection_mode`. +reports the new meta_info.weight_version, and that SMG stamps +`x-smg-weight-version` with it. The pause/resume pair is `smg.rl.paused`, so +a failure at any stage still resumes the engines that did pause. Only HTTP +workers can be proxied: a gRPC or ZMQ worker matched by `--selector` fails +the fan-out with `unsupported_connection_mode`. """ from __future__ import annotations @@ -64,11 +65,16 @@ def main() -> int: ) with urllib.request.urlopen(req, timeout=120) as resp: out = json.loads(resp.read()) + stamped = resp.headers.get("x-smg-weight-version") got = str(out.get("meta_info", {}).get("weight_version")) print(f"/generate reported weight_version={got}") if got != str(args.weight_version): print("MISMATCH: engine did not report the new version", file=sys.stderr) return 1 + if stamped != str(args.weight_version): + print(f"MISMATCH: gateway stamped x-smg-weight-version={stamped!r}", file=sys.stderr) + return 1 + print(f"x-smg-weight-version={stamped}") print("OK") return 0 From 9415c5edeb7045b9980e7c0d00bf69ff7a96c31b Mon Sep 17 00:00:00 2001 From: key4ng Date: Fri, 11 Sep 2026 23:10:50 -0700 Subject: [PATCH 21/27] test(rl): e2e coverage for version stamping and paused-engine routing Signed-off-by: key4ng --- e2e_test/router/test_rl_control_plane.py | 84 ++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/e2e_test/router/test_rl_control_plane.py b/e2e_test/router/test_rl_control_plane.py index 9a32228be1..9a24dd38ff 100644 --- a/e2e_test/router/test_rl_control_plane.py +++ b/e2e_test/router/test_rl_control_plane.py @@ -85,3 +85,87 @@ def test_v1_rl_is_404_without_flag(self, setup_backend): resp = httpx.get(f"{gateway.base_url}/v1/rl/workers", timeout=TIMEOUT) assert resp.status_code == 404 assert resp.content == b"" + + +@pytest.mark.engine("sglang") +@pytest.mark.gpu(1) +@pytest.mark.e2e +@pytest.mark.gateway( + policy="round_robin", + extra_args=[ + "--enable-rl", + "--rl-version-policy", + "latest-only", + "--retry-max-retries", + "3", + "--retry-initial-backoff-ms", + "50", + "--retry-max-backoff-ms", + "200", + ], +) +@pytest.mark.parametrize("setup_backend", ["http"], indirect=True) +class TestRlVersionRouting: + def _chat(self, gateway, model, **kwargs): + return httpx.post( + f"{gateway.base_url}/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": "Say hi"}], + "max_tokens": 4, + }, + timeout=TIMEOUT, + **kwargs, + ) + + def test_responses_are_stamped_and_follow_a_refit(self, setup_backend): + _backend, model, _client, gateway = setup_backend + # Primary path: fan out update_weight_version. If the engine in the + # harness is an older SGLang that 404s on this RPC, fall back to + # update_weights_from_disk with model_path=model and + # weight_version="7" (mirrors examples/rl/refit_from_disk.py). + resp = httpx.post( + f"{gateway.base_url}/v1/rl/engine/update_weight_version", + params={"selector": "engine=sglang"}, + json={"new_version": "7", "abort_all_requests": False}, + timeout=TIMEOUT, + ) + assert resp.status_code == 200, resp.text + resp = self._chat(gateway, model) + assert resp.status_code == 200, resp.text + assert resp.headers.get("x-smg-weight-version") == "7" + assert resp.headers.get("x-smg-routed-worker-id") + rows = httpx.get(f"{gateway.base_url}/v1/rl/workers", timeout=TIMEOUT).json()["workers"] + assert {w["weight_version"] for w in rows} == {"7"} + assert {w["version_source"] for w in rows} == {"passthrough"} + assert {w["control"] for w in rows} == {"active"} + + def test_a_paused_engine_is_not_routed(self, setup_backend): + _backend, model, _client, gateway = setup_backend + pause = httpx.post( + f"{gateway.base_url}/v1/rl/engine/pause_generation", + params={"selector": "engine=sglang"}, + json={}, + timeout=TIMEOUT, + ) + assert pause.status_code == 200, pause.text + try: + resp = self._chat(gateway, model) + assert resp.status_code == 503, resp.text + assert resp.headers.get("x-smg-routed-worker-id") is None + finally: + resume = httpx.post( + f"{gateway.base_url}/v1/rl/engine/continue_generation", + params={"selector": "engine=sglang"}, + json={}, + timeout=TIMEOUT, + ) + assert resume.status_code == 200, resume.text + resp = self._chat(gateway, model) + assert resp.status_code == 200, resp.text + + def test_invalid_policy_header_is_rejected(self, setup_backend): + _backend, model, _client, gateway = setup_backend + resp = self._chat(gateway, model, headers={"x-smg-version-policy": "freshest"}) + assert resp.status_code == 400 + assert resp.json()["error"] == "invalid_version_policy" From 619c9dfa3f8421cae1df828bc8524716659d3c9d Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 14 Sep 2026 13:50:50 -0700 Subject: [PATCH 22/27] fix(router): answer x-smg-target-worker against the caller's slice under a filter When the RL candidate filter narrows the worker list to a strict subset, `PolicyRegistry::select_worker` hands the subset to the policy and remaps the result. For `ConsistentHashingPolicy` -- the only policy that honors `x-smg-target-worker` -- that silently re-bases the header's index: target `2` against three workers filtered to `[0, 2]` would have selected worker 2 as "index 2 of the subset", which does not exist, and target `1` would have resolved to worker 2. The index the caller sent named a worker in the caller's list, and which worker it lands on would change request to request as the filter's verdict moves. The registry now decides the target-worker request itself before building the subset, and never re-bases the number: a named worker that survived the filter and is available is returned at the caller's own index; one the filter dropped (paused, asleep, or too stale for the version policy), one out of range, and a value that is not an index are all refused, each with a `debug!` naming the reason. Refusal is the existing retryable 503, which is the right answer: the caller asked for a specific engine and that engine is not currently routable. Policies that ignore the header are unaffected and keep selecting from the subset, and with the filter keeping every candidate there is no subset, so the policy reads the header exactly as before. Also corrects the CORS row in COUPLING.md: the expose list is three RL headers plus `x-request-id`, not three headers. Signed-off-by: key4ng --- crates/rl/COUPLING.md | 4 +- model_gateway/src/policies/registry.rs | 166 ++++++++++++++++++++++++- 2 files changed, 165 insertions(+), 5 deletions(-) diff --git a/crates/rl/COUPLING.md b/crates/rl/COUPLING.md index c16f98ff13..0262ee7b1e 100644 --- a/crates/rl/COUPLING.md +++ b/crates/rl/COUPLING.md @@ -14,11 +14,11 @@ that adds a surface must update this file. | (a) | `RlWorkerView::inflight(base_url)` and `VersionEvictionSink` | `src/rl_adapter.rs` | inflight sums `Worker::load()` over the ranks; `VersionEvictionSink::on_version_changed` runs under the RL table's write mutex, so the gateway's `RegistryEvictionSink` only resets policy caches and never writes back into the table; it holds `Weak`/`Weak` handles (a strong pair would close an ownership cycle through the candidate filter and leak both registries) and calls `PolicyRegistry::reset_worker_cache`, which covers the model policy, the PD/EPD leg policies, and the default-policy fallback | | (a) | registry event subscription | `src/rl_adapter.rs` | seeds on `Registered`, reseeds on `Replaced` only when the discovered label changed, drops on the last `Removed` for a base URL; a `Replaced` that changes a worker's `model_id` without changing its label calls `seed`, which is insert-if-absent, so the table entry keeps the *old* `model_id` until the worker is removed and re-registered | | (a) | `CacheAwarePolicy::reset_worker_cache`, `PolicyRegistry::reset_worker_cache` | `src/policies/cache_aware.rs`, `src/policies/registry.rs` | tree purge + root re-seed, no load-scorer effect | -| (b) | `CandidateFilter` trait and `PolicyRegistry::{candidate_filter, set_candidate_filter}` | `src/policies/mod.rs`, `src/policies/registry.rs`, `src/app_context.rs` | set in `AppContextBuilder::build()` when `rl` is `Some`; EPD per-item encode assignment and external-provider routing call their own selection paths and bypass the candidate filter by design | +| (b) | `CandidateFilter` trait and `PolicyRegistry::{candidate_filter, set_candidate_filter}` | `src/policies/mod.rs`, `src/policies/registry.rs`, `src/app_context.rs` | set in `AppContextBuilder::build()` when `rl` is `Some`; EPD per-item encode assignment and external-provider routing call their own selection paths and bypass the candidate filter by design. An `x-smg-target-worker` index names a worker in the caller's slice, never in the narrowed subset, so when the RL filter dropped the named worker the registry refuses the request (the retryable 503) instead of letting the policy re-read the index against the survivors | | (c) | `RoutedWorker` + `stamp_routed_worker` | `crates/external_router/src/header_utils.rs`, seven HTTP sites, `routers/grpc/pipeline.rs` | header and extension in one call; the seven HTTP sites include the Responses-API streaming path in `crates/external_router/src/openai/responses/streaming.rs`; gRPC stamps the eight `pipeline.rs` dispatch sites, while the gRPC Responses-API helpers (typed, non-streaming returns) stay unstamped in M2 | | (c) | RL middleware | `src/rl_adapter.rs`, `src/server.rs` `build_app` | header validation, version stamp, mixed flag; added inside `build_app`, so `tests/common/test_app.rs` inherits it automatically | | (c) | `dispatch_metadata.rs` reads the table first | `src/routers/grpc/common/stages/dispatch_metadata.rs` (+ the pipeline holding `Option>`) | label fallback unchanged | -| (c) | CORS expose list | `src/server.rs` `create_cors_layer` | three headers | +| (c) | CORS expose list | `src/server.rs` `create_cors_layer` | three RL headers plus `x-request-id` | | (d) | `--rl-version-policy` | `src/config/{types,validation}.rs`, `src/main.rs`, Python bindings | | Test-only files that the mount also touches, none of them a new surface: diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 1de6663634..157eef3f8f 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -17,8 +17,8 @@ use tracing::{debug, info, warn}; use super::{ get_healthy_worker_indices, manual::{ExecutionBranch, PinState}, - BucketPolicy, CacheAwarePolicy, CandidateFilter, DPRankLoadPolicy, LoadBalancingPolicy, - ManualConfig, ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, + BucketPolicy, CacheAwarePolicy, CandidateFilter, ConsistentHashingPolicy, DPRankLoadPolicy, + LoadBalancingPolicy, ManualConfig, ManualPolicy, PolicyFactory, SelectWorkerInfo, WorkerLeg, }; use crate::{ config::types::{ManualAssignmentMode, PdPairingMode, PolicyConfig, RoutingKeyOverrideConfig}, @@ -26,7 +26,8 @@ use crate::{ observability::metrics::Metrics, policies::cache_aware::LoadReceiver, routers::common::header_utils::{ - extract_routing_key_hint_named, parse_routing_tokens_hint, ROUTING_KEY_HINT_MAX_BYTES, + extract_routing_key_hint_named, extract_target_worker, parse_routing_tokens_hint, + ROUTING_KEY_HINT_MAX_BYTES, }, worker::{KvEventMonitor, Worker}, }; @@ -271,6 +272,15 @@ impl PolicyRegistry { return None; } if keep.len() < workers.len() { + // An explicit target index is answered against the + // caller's slice before the subset is built; handing it + // to the policy would silently rewrite which worker it + // names. See `target_worker_under_filter`. + if let Some(decision) = + Self::target_worker_under_filter(policy, workers, info, &keep) + { + return decision; + } let subset: Vec> = keep.iter().map(|&i| Arc::clone(&workers[i])).collect(); return self @@ -282,6 +292,71 @@ impl PolicyRegistry { self.select_unfiltered(policy, workers, info) } + /// Decide an `x-smg-target-worker` request against the *caller's* slice + /// when a candidate filter narrowed it to a strict subset. + /// + /// The header carries an index, and an index only means something + /// relative to the list it indexes. Passing the narrowed subset to a + /// policy that honors the header would make index `1` name a different + /// worker than the caller asked for — silently, and differently on every + /// request as the filter's verdict moves. So the registry answers it + /// here instead, and never re-bases the number. + /// + /// `Some(Some(idx))` when the named worker survived the filter and is + /// available; `Some(None)` when it did not, which the caller already + /// turns into the retryable 503. `None` means no target-worker decision + /// applies (this policy does not honor the header, or none was sent), + /// and the subset path continues. + /// + /// [`ConsistentHashingPolicy`] is the only policy that honors the header + /// today, and this mirrors its own rules: a value that is not an index + /// into the slice, or names an unavailable worker, refuses rather than + /// falling through to a different worker. + fn target_worker_under_filter( + policy: &Arc, + workers: &[Arc], + info: &SelectWorkerInfo, + keep: &[usize], + ) -> Option> { + if !policy.as_any().is::() { + return None; + } + let raw = extract_target_worker(info.headers)?; + let Ok(idx) = raw.parse::() else { + debug!( + target_worker = raw, + "x-smg-target-worker is not a worker index; refusing the request" + ); + return Some(None); + }; + if idx >= workers.len() { + debug!( + target_worker = idx, + workers = workers.len(), + "x-smg-target-worker is out of range; refusing the request" + ); + return Some(None); + } + if !keep.contains(&idx) { + debug!( + target_worker = idx, + worker = workers[idx].url(), + "x-smg-target-worker names a worker the candidate filter dropped \ + (paused, asleep, or too stale for the version policy); refusing the request" + ); + return Some(None); + } + if !workers[idx].is_healthy_and_eligible() { + debug!( + target_worker = idx, + worker = workers[idx].url(), + "x-smg-target-worker names an unavailable worker; refusing the request" + ); + return Some(None); + } + Some(Some(idx)) + } + /// Today's selection: the sticky routing-key override when it applies, /// else the policy. fn select_unfiltered( @@ -2246,6 +2321,91 @@ mod tests { ); } + fn headers_with_target(idx: &str) -> HeaderMap { + let mut h = HeaderMap::new(); + h.insert("x-smg-target-worker", idx.parse().unwrap()); + h + } + + /// `x-smg-target-worker` indexes the caller's slice. Under a narrowed + /// candidate set the registry answers it itself, so the number always + /// names the worker the caller meant: a survivor is routed to at its own + /// index, and a dropped one refuses rather than resolving to whichever + /// worker happens to sit at that offset in the subset. + #[test] + fn target_worker_index_stays_in_the_callers_slice_under_a_filter() { + let registry = PolicyRegistry::new(PolicyConfig::ConsistentHashing); + registry.set_candidate_filter(Arc::new(KeepIndices(Some(vec![0, 2])))); + let workers = three_workers(); + let policy = registry.get_policy_or_default("m"); + assert_eq!(policy.name(), "consistent_hashing"); + + for (target, expected, note) in [ + ( + "2", + Some(2), + "a survivor keeps the caller's index, not the subset's", + ), + ("0", Some(0), "the first survivor is unaffected"), + ( + "1", + None, + "a filtered-out worker refuses instead of sliding to another", + ), + ("3", None, "out of range refuses"), + ("nope", None, "a non-index value refuses"), + ] { + let headers = headers_with_target(target); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + assert_eq!( + registry.select_worker(&policy, &workers, &info), + expected, + "{note}" + ); + } + } + + /// With every candidate eligible there is no subset, so the policy reads + /// the header against the caller's slice exactly as it always has. + #[test] + fn target_worker_index_is_unchanged_when_the_filter_keeps_everything() { + let registry = PolicyRegistry::new(PolicyConfig::ConsistentHashing); + registry.set_candidate_filter(Arc::new(KeepIndices(None))); + let workers = three_workers(); + let policy = registry.get_policy_or_default("m"); + let headers = headers_with_target("1"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + assert_eq!(registry.select_worker(&policy, &workers, &info), Some(1)); + } + + /// Only a policy that honors the header gets this treatment. Round robin + /// ignores `x-smg-target-worker`, so it keeps selecting from the subset. + #[test] + fn a_policy_that_ignores_the_target_header_still_routes_within_the_subset() { + let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); + registry.set_candidate_filter(Arc::new(KeepIndices(Some(vec![0, 2])))); + let workers = three_workers(); + let policy = registry.get_policy_or_default("m"); + let headers = headers_with_target("1"); + let info = SelectWorkerInfo { + headers: Some(&headers), + ..Default::default() + }; + for _ in 0..4 { + let picked = registry.select_worker(&policy, &workers, &info); + assert!( + matches!(picked, Some(0) | Some(2)), + "round robin ignores the header and stays in the subset: {picked:?}" + ); + } + } + #[test] fn a_filter_that_keeps_everything_changes_nothing() { for config in [ From 145f41aa6308be738d50d4849f221607e51cfb67 Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 14 Sep 2026 13:51:00 -0700 Subject: [PATCH 23/27] refactor(rl): one validator for every weight version SMG accepts The explicit API write trimmed, rejected empty, and capped at 128 bytes; the passthrough observer only checked that a string was non-blank. So a version the control plane would answer `invalid_version` for -- 200 bytes, an interior space, a control character -- still reached the table through a proxied refit body, and then either could not be stamped into `x-smg-weight-version` (the adapter drops an unrepresentable value) or was stamped as something the caller never wrote. `Version::validated` is now the single gate: trim, then reject empty, over 128 bytes, or any byte outside visible ASCII (0x21..=0x7E). It returns a `VersionError` carrying both a message for the `invalid_version` body and a static `reason()` tag for logs. `control.rs` uses it for the per-worker and fleet writes; `observe.rs` uses it for both refit field shapes (`weight_version` and `new_version`, with a JSON number stringified first) and observes nothing when it refuses, logging `warn!(target: "smg_rl", route, reason)` so a silently ignored engine report is visible. A JSON type that is neither string nor number is logged by type name, never by value. Tests: the validator's accept/reject matrix and the stable reason tags in `version.rs`; whitespace-only, boolean, array, object, 129-byte, interior-space and control-character versions each observing nothing in `observe.rs`; and the over-cap API write now asserts the `invalid_version` code and message, with a control-character write added alongside. Signed-off-by: key4ng --- crates/rl/src/control.rs | 37 ++++++++------- crates/rl/src/lib.rs | 2 +- crates/rl/src/observe.rs | 86 ++++++++++++++++++++++++++++++++--- crates/rl/src/version.rs | 97 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 199 insertions(+), 23 deletions(-) diff --git a/crates/rl/src/control.rs b/crates/rl/src/control.rs index c283b929c5..ef77764fa8 100644 --- a/crates/rl/src/control.rs +++ b/crates/rl/src/control.rs @@ -22,22 +22,10 @@ use crate::{ view::RlWorkerInfo, }; -/// Cap on an API-supplied version string. -const MAX_VERSION_BYTES: usize = 128; - +/// The API write shares one validator with the passthrough observer, so a +/// version SMG would refuse here cannot slip in through a proxied refit body. fn parse_version(raw: &str) -> Result { - let trimmed = raw.trim(); - if trimmed.is_empty() { - return Err(RlError::InvalidVersion( - "weight_version is empty".to_string(), - )); - } - if trimmed.len() > MAX_VERSION_BYTES { - return Err(RlError::InvalidVersion(format!( - "weight_version exceeds {MAX_VERSION_BYTES} bytes" - ))); - } - Ok(Version::parse(trimmed)) + Version::validated(raw).map_err(|e| RlError::InvalidVersion(e.to_string())) } fn body(body: Result, JsonRejection>) -> Result { @@ -257,6 +245,7 @@ mod tests { let long = "x".repeat(129); let resp = app + .clone() .oneshot( Request::post("/workers/w1/version") .header("content-type", "application/json") @@ -266,6 +255,24 @@ mod tests { .await .unwrap(); assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + let body = json_body(resp).await; + assert_eq!(body["error"], "invalid_version"); + assert_eq!(body["message"], "weight_version exceeds 128 bytes"); + + // The shared validator also refuses anything outside printable ASCII, + // so a version that could never be stamped into a header is refused at + // the write rather than silently dropped at stamp time. + let resp = app + .oneshot( + Request::post("/workers/w1/version") + .header("content-type", "application/json") + .body(Body::from(r#"{"weight_version": "v1\u0007"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + assert_eq!(json_body(resp).await["error"], "invalid_version"); } #[tokio::test] diff --git a/crates/rl/src/lib.rs b/crates/rl/src/lib.rs index 64dd2b6c25..e928f17529 100644 --- a/crates/rl/src/lib.rs +++ b/crates/rl/src/lib.rs @@ -34,7 +34,7 @@ pub use metrics::init_rl_metrics; pub use policy::{VersionPolicy, VersionPolicyError, VERSION_POLICY_HEADER}; pub use state::RlState; pub use table::{ControlState, FilterReason, RlTable, VersionEvictionSink, VersionSource}; -pub use version::Version; +pub use version::{Version, VersionError}; pub use view::{RlWorkerInfo, RlWorkerView}; /// Build the `/v1/rl` router. `with_state` returns `Router` for any `S`, diff --git a/crates/rl/src/observe.rs b/crates/rl/src/observe.rs index 1e313882ea..0597c63b13 100644 --- a/crates/rl/src/observe.rs +++ b/crates/rl/src/observe.rs @@ -4,6 +4,7 @@ use serde::Deserialize; use serde_json::Value; +use tracing::warn; use crate::{table::ControlState, version::Version}; @@ -23,11 +24,46 @@ struct VersionProbe { new_version: Option, } -fn version_field(value: Option) -> Option { - match value? { - Value::String(s) if !s.trim().is_empty() => Some(Version::parse(&s)), - Value::Number(n) => Some(Version::parse(&n.to_string())), - _ => None, +/// A passthrough version gets exactly the validation the explicit API write +/// gets (see [`Version::validated`]), so a refit body cannot put a value into +/// the table that `POST /v1/rl/workers//version` would have refused. A +/// rejected value observes nothing and says why, named by the route that +/// carried it. An absent field stays silent: it says nothing about the engine. +fn version_field(route: &str, value: Option) -> Option { + let raw = match value? { + Value::String(s) => s, + Value::Number(n) => n.to_string(), + other => { + warn!( + target: "smg_rl", + route, reason = "not_a_string_or_number", kind = kind_of(&other), + "ignoring a weight version an engine reported in an unusable JSON type" + ); + return None; + } + }; + match Version::validated(&raw) { + Ok(version) => Some(version), + Err(e) => { + warn!( + target: "smg_rl", + route, reason = e.reason(), + "ignoring an invalid weight version an engine reported" + ); + None + } + } +} + +/// The JSON type name for the rejection log. The value itself is never +/// logged: a refit body is caller-supplied and may be large. +fn kind_of(value: &Value) -> &'static str { + match value { + Value::Null => "null", + Value::Bool(_) => "boolean", + Value::Array(_) => "array", + Value::Object(_) => "object", + Value::String(_) | Value::Number(_) => "scalar", } } @@ -39,11 +75,11 @@ pub fn observe(path: &str, body: &[u8]) -> Option { | "update_weights_from_tensor" | "update_weights_from_distributed" => { let probe: VersionProbe = serde_json::from_slice(body).ok()?; - version_field(probe.weight_version).map(Observation::Version) + version_field(path, probe.weight_version).map(Observation::Version) } "update_weight_version" => { let probe: VersionProbe = serde_json::from_slice(body).ok()?; - version_field(probe.new_version).map(Observation::Version) + version_field(path, probe.new_version).map(Observation::Version) } "pause_generation" | "pause" => Some(Observation::Control(ControlState::Paused)), "continue_generation" | "resume" | "resume_memory_occupation" | "wake_up" => { @@ -106,6 +142,42 @@ mod tests { ); } + /// Passthrough versions clear the same bar as the API write: a value the + /// control plane would answer `invalid_version` for observes nothing + /// rather than landing in the table through the proxy. + #[test] + fn a_rejected_passthrough_version_observes_nothing() { + let long = "x".repeat(129); + for body in [ + r#"{"weight_version": " "}"#.to_string(), + r#"{"weight_version": true}"#.to_string(), + r#"{"weight_version": ["42"]}"#.to_string(), + r#"{"weight_version": {"v": "42"}}"#.to_string(), + format!(r#"{{"weight_version": "{long}"}}"#), + r#"{"weight_version": "v1\u0007"}"#.to_string(), + r#"{"weight_version": "step 7"}"#.to_string(), + ] { + assert_eq!( + observe("update_weights_from_disk", body.as_bytes()), + None, + "{body}" + ); + } + // The same bar on the other field shape. + for body in [ + r#"{"new_version": " "}"#, + r#"{"new_version": false}"#, + r#"{"new_version": []}"#, + r#"{"new_version": "v1\u0007"}"#, + ] { + assert_eq!( + observe("update_weight_version", body.as_bytes()), + None, + "{body}" + ); + } + } + #[test] fn control_routes_map_to_states_regardless_of_body() { for (path, state) in [ diff --git a/crates/rl/src/version.rs b/crates/rl/src/version.rs index 660e85ffef..2ce5ee8124 100644 --- a/crates/rl/src/version.rs +++ b/crates/rl/src/version.rs @@ -10,11 +10,39 @@ use std::{ }, }; +use thiserror::Error; use tracing::warn; /// The label an engine reports before any refit; seeds as "no version". pub const UNVERSIONED_LABEL: &str = "default"; +/// Cap on a version string SMG will accept from any source. +pub const MAX_VERSION_BYTES: usize = 128; + +/// Why a candidate version string was refused. Every variant carries a +/// static [`VersionError::reason`] tag for structured logs, and a message +/// for the `invalid_version` error body. +#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)] +pub enum VersionError { + #[error("weight_version is empty")] + Empty, + #[error("weight_version exceeds {MAX_VERSION_BYTES} bytes")] + TooLong, + #[error("weight_version contains a byte outside printable ASCII")] + NotPrintableAscii, +} + +impl VersionError { + /// Stable tag for `reason = ...` log fields. + pub fn reason(self) -> &'static str { + match self { + Self::Empty => "empty", + Self::TooLong => "too_long", + Self::NotPrintableAscii => "not_printable_ascii", + } + } +} + /// A weight version: numeric when the text parses as `u64`, opaque text /// otherwise. Numeric versions compare by value, text versions lexically, /// and a numeric/text pair lexically with one warning per process. @@ -27,6 +55,26 @@ pub struct Version { static MIXED_ORDERING_WARNED: AtomicBool = AtomicBool::new(false); impl Version { + /// The one gate every externally supplied version passes through: the + /// API write, and every version observed in a proxied refit body. Trims, + /// then refuses an empty string, one over [`MAX_VERSION_BYTES`], and any + /// byte outside visible ASCII (`0x21..=0x7E`) -- which also rules out + /// interior spaces, control characters, and anything that would not + /// survive being stamped into `x-smg-weight-version`. + pub fn validated(raw: &str) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err(VersionError::Empty); + } + if trimmed.len() > MAX_VERSION_BYTES { + return Err(VersionError::TooLong); + } + if !trimmed.bytes().all(|b| (0x21..=0x7E).contains(&b)) { + return Err(VersionError::NotPrintableAscii); + } + Ok(Self::parse(trimmed)) + } + pub fn parse(raw: &str) -> Self { let raw = raw.trim(); Self { @@ -119,6 +167,55 @@ mod tests { assert!(Version::parse("1") < Version::parse("1a")); } + #[test] + fn the_validator_accepts_printable_ascii_within_the_cap() { + assert_eq!(Version::validated("42").map(|v| v.numeric()), Ok(Some(42))); + assert_eq!( + Version::validated(" step-7 ").map(|v| v.as_str().to_string()), + Ok("step-7".to_string()) + ); + assert_eq!( + Version::validated("!~").map(|v| v.as_str().to_string()), + Ok("!~".to_string()) + ); + assert!(Version::validated(&"x".repeat(MAX_VERSION_BYTES)).is_ok()); + } + + #[test] + fn the_validator_refuses_empty_over_cap_and_non_printable() { + assert_eq!(Version::validated(""), Err(VersionError::Empty)); + assert_eq!(Version::validated(" "), Err(VersionError::Empty)); + assert_eq!(Version::validated("\t\n"), Err(VersionError::Empty)); + assert_eq!( + Version::validated(&"x".repeat(MAX_VERSION_BYTES + 1)), + Err(VersionError::TooLong) + ); + // Trimming happens first, so padding cannot push a legal value over. + assert!(Version::validated(&format!(" {} ", "x".repeat(MAX_VERSION_BYTES))).is_ok()); + for raw in ["a b", "v\u{1}", "v\u{7f}", "v\u{e9}", "v\n1"] { + assert_eq!( + Version::validated(raw), + Err(VersionError::NotPrintableAscii), + "{raw:?}" + ); + } + } + + #[test] + fn rejection_reasons_are_stable_tags() { + assert_eq!(VersionError::Empty.reason(), "empty"); + assert_eq!(VersionError::TooLong.reason(), "too_long"); + assert_eq!( + VersionError::NotPrintableAscii.reason(), + "not_printable_ascii" + ); + assert_eq!(VersionError::Empty.to_string(), "weight_version is empty"); + assert_eq!( + VersionError::TooLong.to_string(), + "weight_version exceeds 128 bytes" + ); + } + #[test] fn labels_seed_versions_except_the_placeholder() { assert_eq!(Version::from_label(None), None); From f3eed10a1ebfc312a7f423238edf48d9a8ed4e9d Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 14 Sep 2026 13:51:06 -0700 Subject: [PATCH 24/27] test(rl): drive the gRPC pipeline through the app so the middleware runs Every test in this file called `RouterTrait` directly, so the RL middleware -- the only thing that writes `x-smg-weight-version` -- never ran. The routed-worker half of the contract was covered, the version-header half was not, on the gRPC path specifically. Adds a test that builds the same mock-gRPC-backed router, wraps it with `build_app` through the shared test-app helper, sets the worker's version the way an RL trainer does (`POST /v1/rl/workers/{id}/version` with `{"weight_version": "11"}`, addressed by the registry id the control plane lists), then posts `/v1/chat/completions` in both buffered and `stream: true` form and asserts `x-smg-weight-version` and `x-smg-routed-worker-id` on each response. Signed-off-by: key4ng --- model_gateway/tests/rl_grpc_stamp_test.rs | 92 +++++++++++++++++++++-- 1 file changed, 84 insertions(+), 8 deletions(-) diff --git a/model_gateway/tests/rl_grpc_stamp_test.rs b/model_gateway/tests/rl_grpc_stamp_test.rs index 0cfe314ed1..eee36cc6a8 100644 --- a/model_gateway/tests/rl_grpc_stamp_test.rs +++ b/model_gateway/tests/rl_grpc_stamp_test.rs @@ -1,17 +1,22 @@ //! RL M2: the gRPC pipeline names the worker that served each response, and //! reports the version the RL table holds for that engine. //! -//! The routers are driven directly (no axum app), so the RL middleware never -//! runs here and `x-smg-weight-version` is not expected: what these tests pin -//! is the pipeline's own half of the contract -- the routed-worker header and -//! extension that the middleware keys on, and the live table version reaching -//! the body as `system_fingerprint`. +//! Most tests here drive the routers directly (no axum app), so the RL +//! middleware never runs and `x-smg-weight-version` is not expected: what +//! they pin is the pipeline's own half of the contract -- the routed-worker +//! header and extension that the middleware keys on, and the live table +//! version reaching the body as `system_fingerprint`. The last test closes +//! the loop, driving the same pipeline through `build_app` so the middleware +//! does run and the version header is asserted end to end. #[path = "common/mod.rs"] mod common; use std::{sync::Arc, time::Duration}; +use axum::body::Body; +use http::{Request, StatusCode}; +use http_body_util::BodyExt as _; use llm_tokenizer::{traits::Tokenizer, MockTokenizer, TokenizerRegistry}; use openai_protocol::{ chat::ChatCompletionRequest, model_card::ModelCard, worker::HealthCheckConfig, @@ -25,6 +30,7 @@ use smg::{ worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, WorkerType}, }; use tokio::net::TcpListener; +use tower::ServiceExt as _; const MODEL: &str = "rl-grpc-stamp-test-model"; /// Tokens the canned mock emits per request. @@ -176,7 +182,7 @@ async fn a_buffered_grpc_response_names_its_worker_and_reports_the_live_version( .route_chat(None, &tenant(), chat_request(false), MODEL) .await; - assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!(response.status(), StatusCode::OK); assert_eq!( response.headers()["x-smg-routed-worker-id"], url.as_str(), @@ -211,7 +217,7 @@ async fn a_streamed_grpc_response_names_its_worker_and_reports_the_live_version( .route_chat(None, &tenant(), chat_request(true), MODEL) .await; - assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!(response.status(), StatusCode::OK); assert_eq!( response.headers()["x-smg-routed-worker-id"], url.as_str(), @@ -242,10 +248,80 @@ async fn an_unversioned_engine_still_stamps_and_falls_back_to_default() { .route_chat(None, &tenant(), chat_request(false), MODEL) .await; - assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!(response.status(), StatusCode::OK); assert_eq!(response.headers()["x-smg-routed-worker-id"], url.as_str()); let body = read_body(response).await; let json: serde_json::Value = serde_json::from_str(&body).unwrap(); assert_eq!(json["system_fingerprint"], "default", "{body}"); } + +/// The gRPC pipeline behind the real app, with every layer `build_app` +/// installs -- including the RL middleware. The version is set the way an RL +/// trainer sets it, through `POST /v1/rl/workers/{id}/version`, and comes +/// back on the response as `x-smg-weight-version` next to the routed-worker +/// stamp. The direct-router tests above cannot see this: they never build +/// the app, so the middleware that reads the table and writes the header +/// never runs. +#[tokio::test] +async fn the_app_stamps_the_control_plane_version_on_grpc_responses() { + let port = start_mock_grpc_worker().await; + let (router, ctx) = build_regular_router_with_rl(port).await; + let app = common::test_app::create_test_app_with_context(Arc::from(router), Arc::clone(&ctx)); + let url = worker_url(port); + + // The control plane addresses workers by registry id, not URL. + let listed = app + .clone() + .oneshot(Request::get("/v1/rl/workers").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(listed.status(), StatusCode::OK); + let listed: serde_json::Value = + serde_json::from_slice(&listed.into_body().collect().await.unwrap().to_bytes()).unwrap(); + let id = listed["workers"][0]["id"] + .as_str() + .expect("the gRPC worker is listed by the RL control plane") + .to_string(); + + let set = app + .clone() + .oneshot( + Request::post(format!("/v1/rl/workers/{id}/version")) + .header("content-type", "application/json") + .body(Body::from(format!( + r#"{{"weight_version": "{LIVE_VERSION}"}}"# + ))) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(set.status(), StatusCode::OK); + + for stream in [false, true] { + let response = app + .clone() + .oneshot( + Request::post("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + serde_json::to_vec(&chat_request(stream)).unwrap(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK, "stream={stream}"); + assert_eq!( + response.headers()["x-smg-weight-version"], + LIVE_VERSION, + "the RL middleware must stamp the table's version; stream={stream}" + ); + assert_eq!( + response.headers()["x-smg-routed-worker-id"], + url.as_str(), + "the response must still name the worker that served it; stream={stream}" + ); + } +} From f896077217357b5ab756f4c320702b5ac51c9c2a Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 14 Sep 2026 13:51:14 -0700 Subject: [PATCH 25/27] chore(rl): unroutable counter, lagged-resync log level, bench feature set Three small corrections found in final review: - `RlTable::eligible` counted a request unroutable whenever it kept nothing, including when it was handed no candidates at all. A caller with an empty worker list already had nothing to route; counting it blamed the filter for a decision it never made. Now gated on `total > 0`. - The RL table maintainer's lagged-resync message was `debug!`. A dropped registry event stream means the table was briefly out of step with the registry, which an operator should see; it is now `warn!`. - Dropped the `html_reports` feature from the `criterion` dev-dependency: the filter bench is run for its numbers, not its HTML, and the feature pulls extra dependencies into every build of the crate's dev tree. Also documents the max-staleness rule that non-numeric versions on either side of the comparison have no distance to measure and count as stale, dropped under `FilterReason::Stale` with no separate log or metric. Signed-off-by: key4ng --- crates/rl/Cargo.toml | 2 +- crates/rl/src/table.rs | 10 +++++++++- model_gateway/src/rl_adapter.rs | 4 ++-- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/crates/rl/Cargo.toml b/crates/rl/Cargo.toml index e5b0b8814d..3039947c97 100644 --- a/crates/rl/Cargo.toml +++ b/crates/rl/Cargo.toml @@ -32,7 +32,7 @@ metrics = "0.24.6" tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net"] } tower = { version = "0.5", features = ["util"] } http-body-util = "0.1" -criterion = { version = "0.8", features = ["html_reports"] } +criterion = { version = "0.8" } [[bench]] name = "filter" diff --git a/crates/rl/src/table.rs b/crates/rl/src/table.rs index 65c746b004..f27870af93 100644 --- a/crates/rl/src/table.rs +++ b/crates/rl/src/table.rs @@ -417,7 +417,10 @@ impl RlTable { Err(reason) => metrics::record_candidate_filtered(reason.as_str()), } } - if keep.is_empty() { + // An empty candidate list is only "unroutable" when there was + // something to drop: a caller with no candidates at all already had + // nothing to route, and counting it here would blame the filter. + if keep.is_empty() && total > 0 { metrics::record_request_unroutable(); } if keep.len() == total { @@ -440,6 +443,11 @@ impl RlTable { (Some(v), Some(max)) if v == max => Ok(()), (Some(_), Some(_)) => Err(FilterReason::Stale), }, + // Staleness is a numeric distance, so a version either side of + // the comparison that is not a `u64` has no distance to measure + // and counts as stale. That is the whole rule: the candidate is + // dropped under `FilterReason::Stale` like any other, with no + // separate log or metric for the non-numeric case. VersionPolicy::MaxStaleness(k) => match (version, fleet_max) { (_, None) => Ok(()), (None, Some(_)) => Err(FilterReason::Unversioned), diff --git a/model_gateway/src/rl_adapter.rs b/model_gateway/src/rl_adapter.rs index 4bed2a071b..02ea97067e 100644 --- a/model_gateway/src/rl_adapter.rs +++ b/model_gateway/src/rl_adapter.rs @@ -18,7 +18,7 @@ use smg_rl::{ VersionEvictionSink, VersionPolicy, }; use tokio::sync::broadcast::error::RecvError; -use tracing::{debug, warn}; +use tracing::warn; use crate::{ config::RouterConfig, @@ -222,7 +222,7 @@ fn spawn_table_maintainer(rl: Weak, registry: Arc) { } Ok(WorkerEvent::StatusChanged { .. }) => {} Err(RecvError::Lagged(n)) => { - debug!(missed = n, "RL table maintainer lagged; resyncing"); + warn!(missed = n, "RL table maintainer lagged; resyncing"); let Some(registry) = registry.upgrade() else { break; }; From e446b3baf32d0e33e7b758290ec9eb585b901f3c Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 14 Sep 2026 13:51:19 -0700 Subject: [PATCH 26/27] docs(rl): error codes, the routed-worker header/stamp pair, the refit comment - `insert_routed_worker_id` now says it is the header half of `stamp_routed_worker`, and that a caller with a whole response should use the latter: layers keyed on the `RoutedWorker` extension, the RL version stamp among them, see nothing when only the header is set. - The RL README lists the error codes the control plane answers with, including `invalid_body` (400), which had no mention anywhere. - The e2e refit test's comment claimed a fallback that is not implemented. It now reads as the instruction it is: if the engine lacks `update_weight_version`, switch this call to `update_weights_from_disk` with a `weight_version` field. Signed-off-by: key4ng --- crates/external_router/src/header_utils.rs | 5 +++++ crates/rl/README.md | 8 ++++++++ e2e_test/router/test_rl_control_plane.py | 8 ++++---- 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/crates/external_router/src/header_utils.rs b/crates/external_router/src/header_utils.rs index 4e75e84e88..302d2db95e 100644 --- a/crates/external_router/src/header_utils.rs +++ b/crates/external_router/src/header_utils.rs @@ -138,6 +138,11 @@ pub fn preserve_response_headers(reqwest_headers: &HeaderMap) -> HeaderMap { /// [`preserve_response_headers`] so the gateway's value wins over anything an /// upstream set under the same name. Skipped if the URL isn't representable as /// a header value. +/// +/// This is the header half of [`stamp_routed_worker`], which is what callers +/// with a whole response in hand should use: it sets this header and the +/// [`RoutedWorker`] extension together, and layers keyed on the extension +/// (the RL version stamp among them) see nothing when only the header is set. pub fn insert_routed_worker_id(headers: &mut HeaderMap, worker_url: &str) { if let Ok(value) = HeaderValue::from_str(worker_url) { headers.insert(HEADER_ROUTED_WORKER_ID.clone(), value); diff --git a/crates/rl/README.md b/crates/rl/README.md index bea7444798..3ce33df0f6 100644 --- a/crates/rl/README.md +++ b/crates/rl/README.md @@ -23,6 +23,14 @@ reported in `failed[]` as `unsupported_connection_mode` (HTTP 422 on the per-wor route), so a fan-out over a mixed fleet answers 207 and `smg.rl.RL.fanout` raises `FanoutError` unless `allow_partial=True`. +Every error answers `{"error": , "message": ...}` with a stable code: +`invalid_body` (400, a malformed or absent JSON body), `invalid_version` (400, +a `weight_version` that is empty, over 128 bytes, or not printable ASCII), +`invalid_version_policy` (400), `invalid_engine_path` (400), +`selector_required` / `invalid_selector` / `no_workers_match` (400), +`worker_not_found` (404), `unsupported_connection_mode` (422), +`upstream_unreachable` (502), and `upstream_timeout` (504). + Flags: `--enable-rl`, `--rl-control-timeout-secs` (600), `--rl-fanout-concurrency` (32), `--rl-version-policy` (any). Recommended RL launch profile: `--enable-rl --disable-health-check --disable-circuit-breaker --request-timeout-secs 14400`. diff --git a/e2e_test/router/test_rl_control_plane.py b/e2e_test/router/test_rl_control_plane.py index 9a24dd38ff..0a61dc3683 100644 --- a/e2e_test/router/test_rl_control_plane.py +++ b/e2e_test/router/test_rl_control_plane.py @@ -120,10 +120,10 @@ def _chat(self, gateway, model, **kwargs): def test_responses_are_stamped_and_follow_a_refit(self, setup_backend): _backend, model, _client, gateway = setup_backend - # Primary path: fan out update_weight_version. If the engine in the - # harness is an older SGLang that 404s on this RPC, fall back to - # update_weights_from_disk with model_path=model and - # weight_version="7" (mirrors examples/rl/refit_from_disk.py). + # Fans out update_weight_version. There is no fallback here: if the + # engine in the harness is an older SGLang that 404s on this RPC, + # switch this call to update_weights_from_disk with model_path=model + # and a weight_version field (mirrors examples/rl/refit_from_disk.py). resp = httpx.post( f"{gateway.base_url}/v1/rl/engine/update_weight_version", params={"selector": "engine=sglang"}, From 8d7ff36b7f59a8094ebf3b509792f35e24e4f047 Mon Sep 17 00:00:00 2001 From: key4ng Date: Mon, 14 Sep 2026 15:26:33 -0700 Subject: [PATCH 27/27] docs(rl): record SGLang 0.5.19 drift seen in the M2 acceptance run Signed-off-by: key4ng --- crates/rl/NOTES.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/rl/NOTES.md b/crates/rl/NOTES.md index 465eb073a2..3ab891e29f 100644 --- a/crates/rl/NOTES.md +++ b/crates/rl/NOTES.md @@ -9,3 +9,5 @@ planning docs, with date, engine version, and what was done. | 2026-09-03 | SGLang 0.5.15.post1 | `POST /pause_generation` with no body succeeds | 400 Bad Request; FastAPI requires a JSON body, `{}` is enough | callers must send `{}` on bodyless control routes; `examples/rl/refit_from_disk.py` needs `rl.fanout("pause_generation", {}, ...)` (the e2e test already passes `json={}`) | | 2026-09-10 | SGLang 0.5.15.post1 | discovery reports a real version before the first refit | registration reports `weight_version: default` (the placeholder) until a refit lands | seeded as unversioned (`Version::from_label` treats `default` as `None`) so a first refit is not out-ordered by the placeholder | | 2026-09-10 | vLLM (version unpinned) | refit routes carry a version field the passthrough observer can read | documented weight-transfer bodies carry no version field | passthrough observes control state only for vLLM; versions must come through the API (`set_version`/`set_fleet_version`); verify against a live vLLM when one is available | +| 2026-09-14 | SGLang 0.5.19 | `POST /release_memory_occupation` on a busy engine pauses it | a scheduler assertion aborts the whole engine process; idle, the route answers 200 in 3 ms | the proxy reports 502 `upstream_unreachable` and observes nothing (the table keeps `control: active`); pause the fleet before releasing memory, as the refit sequence already does | +| 2026-09-14 | SGLang 0.5.19 | `POST /flush_cache` answers 200 under traffic (0.5.15 did) | 400 on a busy server; a fleet flush under load is a 207 with every worker in `failed[]` | `flush_cache: true` inside a refit does not reliably flush under load; the version change still resets the gateway's cache-aware entries, so routing is unaffected |