From cc4af9de2779da10d9899d35b70c14a8639717e0 Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:15:18 -0700 Subject: [PATCH 1/8] feat(scheduler): derive partition capacity from workers Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- .../src/middleware/scheduler/config.rs | 10 +- .../src/middleware/scheduler/state.rs | 117 ++++++++++++++---- 2 files changed, 98 insertions(+), 29 deletions(-) diff --git a/model_gateway/src/middleware/scheduler/config.rs b/model_gateway/src/middleware/scheduler/config.rs index c0dbc8c3b..0e6820477 100644 --- a/model_gateway/src/middleware/scheduler/config.rs +++ b/model_gateway/src/middleware/scheduler/config.rs @@ -124,8 +124,9 @@ pub struct TenantPolicyConfig { /// A deployment chooses one capacity mode for every configured partition: /// /// - static: set `max_concurrent_requests` (the original behavior), or -/// - replica-aware: set `capacity_from_healthy_replicas: true`, configure -/// `max_concurrent_requests_per_healthy_replica`, and omit the static maximum. +/// - replica-aware: set `capacity_from_healthy_replicas: true`, omit the static +/// maximum, and either derive each worker's ceiling from its reported +/// `max_running_requests` or explicitly configure a per-replica override. /// /// Replica-aware partitions divide the live global scheduler capacity in /// proportion to the number of healthy workers assigned to each partition. @@ -141,8 +142,9 @@ pub struct AdmissionPartitionConfig { /// Derive this partition's share from its healthy replica count. #[serde(default)] pub capacity_from_healthy_replicas: bool, - /// Admission slots contributed by each healthy replica in replica-aware - /// mode. Omitted in static mode. + /// Optional admission slots contributed by each healthy replica in + /// replica-aware mode. When omitted, use worker-reported + /// `max_running_requests`. Omitted in static mode. #[serde(default, skip_serializing_if = "Option::is_none")] pub max_concurrent_requests_per_healthy_replica: Option, /// Work-conserving queue budget shared by the priority classes inside diff --git a/model_gateway/src/middleware/scheduler/state.rs b/model_gateway/src/middleware/scheduler/state.rs index 57c691185..89c257daf 100644 --- a/model_gateway/src/middleware/scheduler/state.rs +++ b/model_gateway/src/middleware/scheduler/state.rs @@ -406,10 +406,10 @@ fn validate_partitions( config.max_concurrent_requests, config.max_concurrent_requests_per_healthy_replica, ) { - (true, None, Some(per_replica)) if per_replica > 0 => {} + (true, None, None) | (true, None, Some(1..)) => {} (true, _, _) => { return Err(format!( - "partition {name}: replica-aware capacity must omit max_concurrent_requests and set max_concurrent_requests_per_healthy_replica > 0" + "partition {name}: replica-aware capacity must omit max_concurrent_requests; max_concurrent_requests_per_healthy_replica, when set, must be > 0" )); } (false, Some(limit), None) if limit > 0 => { @@ -478,17 +478,21 @@ fn allocate_weighted_capacities(weights: &[(String, u64)], target: u16) -> HashM .collect() } -/// Count each healthy worker once. An explicit control-plane label moves the -/// worker to that partition; otherwise its primary model id is used. Unknown -/// labels/models land in the configured default partition. -fn healthy_replica_counts( +#[derive(Debug, Clone, Copy, Default)] +struct PartitionWorkerCapacity { + replicas: u16, + reported_replicas: u16, + reported_capacity: u64, +} + +fn healthy_partition_worker_capacities( partition_names: &HashSet, default_partition: &str, registry: &WorkerRegistry, -) -> HashMap { - let mut counts: HashMap = partition_names +) -> HashMap { + let mut capacities: HashMap = partition_names .iter() - .map(|name| (name.clone(), 0)) + .map(|name| (name.clone(), PartitionWorkerCapacity::default())) .collect(); for worker in registry.get_all().into_iter().filter(|w| w.is_healthy()) { let requested = worker @@ -504,12 +508,16 @@ fn healthy_replica_counts( } else { default_partition }; - counts - .entry(partition.to_string()) - .and_modify(|count| *count = count.saturating_add(1)) - .or_insert(1); + let aggregate = capacities.entry(partition.to_string()).or_default(); + aggregate.replicas = aggregate.replicas.saturating_add(1); + if let Some(reported) = worker.max_running_requests() { + aggregate.reported_replicas = aggregate.reported_replicas.saturating_add(1); + aggregate.reported_capacity = aggregate + .reported_capacity + .saturating_add(u64::from(reported)); + } } - counts + capacities } /// Static mode preserves the original configured-cap behavior. Replica-aware @@ -546,23 +554,41 @@ fn allocate_partition_capacities( } let names: HashSet<_> = configs.iter().map(|(name, _)| name.clone()).collect(); - let counts = healthy_replica_counts(&names, default_partition, registry); + let worker_capacities = + healthy_partition_worker_capacities(&names, default_partition, registry); + let counts: HashMap = worker_capacities + .iter() + .map(|(name, capacity)| (name.clone(), capacity.replicas)) + .collect(); + let total_replicas: u64 = worker_capacities + .values() + .map(|capacity| u64::from(capacity.replicas)) + .sum(); let Some(weights): Option> = configs .iter() .map(|(name, config)| { - let replicas = counts.get(name).copied().unwrap_or(0); - config - .max_concurrent_requests_per_healthy_replica - .map(|per_replica| { - ( - name.clone(), - u64::from(replicas).saturating_mul(u64::from(per_replica)), - ) - }) + let capacity = worker_capacities.get(name).copied().unwrap_or_default(); + let weight = + if let Some(per_replica) = config.max_concurrent_requests_per_healthy_replica { + u64::from(capacity.replicas).saturating_mul(u64::from(per_replica)) + } else if capacity.reported_replicas > 0 { + // Scale the observed per-replica mean across any healthy + // workers whose metadata discovery is still catching up. + capacity + .reported_capacity + .saturating_mul(u64::from(capacity.replicas)) + / u64::from(capacity.reported_replicas) + } else if total_replicas > 0 { + u64::from(global_capacity).saturating_mul(u64::from(capacity.replicas)) + / total_replicas + } else { + 0 + }; + Some((name.clone(), weight)) }) .collect() else { - error!("validated replica-aware admission partition is missing its per-replica capacity"); + error!("validated replica-aware admission partition is missing its capacity source"); return (HashMap::new(), counts); }; let desired: u64 = weights.iter().map(|(_, weight)| *weight).sum(); @@ -836,6 +862,47 @@ mod tests { assert_eq!(allocation.values().copied().sum::(), 500); } + #[test] + fn replica_aware_allocation_uses_worker_reported_limits_without_override() { + let registry = WorkerRegistry::new(); + let ready = openai_protocol::worker::WorkerStatus::Ready; + for (url, model, limit) in [ + ("http://k3-a:8000", "kimi-k3", "64"), + ("http://k3-b:8000", "kimi-k3", "64"), + ("http://dsv4:8000", "dsv4", "256"), + ] { + let mut labels = HashMap::new(); + labels.insert("max_running_requests".to_string(), limit.to_string()); + let worker = Arc::new( + BasicWorkerBuilder::new(url) + .model(openai_protocol::model_card::ModelCard::new(model)) + .labels(labels) + .status(ready) + .build(), + ); + registry.register(worker).unwrap(); + } + + let dynamic = |queue_size| super::super::AdmissionPartitionConfig { + max_concurrent_requests: None, + capacity_from_healthy_replicas: true, + max_concurrent_requests_per_healthy_replica: None, + queue_size, + }; + let configs = vec![ + ("default".to_string(), dynamic(1)), + ("dsv4".to_string(), dynamic(1)), + ("kimi-k3".to_string(), dynamic(1)), + ]; + let (allocation, counts) = + allocate_partition_capacities(&configs, "default", 384, ®istry); + assert_eq!(counts["kimi-k3"], 2); + assert_eq!(counts["dsv4"], 1); + assert_eq!(allocation["kimi-k3"], 128); + assert_eq!(allocation["dsv4"], 256); + assert_eq!(allocation["default"], 0); + } + #[test] fn replica_aware_allocation_sends_unconfigured_models_to_default() { let registry = WorkerRegistry::new(); From c283a9e4f71a8135c27bf9816b915d980a0d8808 Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:15:31 -0700 Subject: [PATCH 2/8] feat(admission): learn capacity from engine feedback Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- model_gateway/src/app_context.rs | 5 +- model_gateway/src/config/types.rs | 68 +++ model_gateway/src/config/validation.rs | 20 + model_gateway/src/main.rs | 98 +++- model_gateway/src/observability/metrics.rs | 6 +- .../src/routers/grpc/adaptive_admission.rs | 474 +++++++++++++++--- 6 files changed, 602 insertions(+), 69 deletions(-) diff --git a/model_gateway/src/app_context.rs b/model_gateway/src/app_context.rs index 963a26df6..fbd18733e 100644 --- a/model_gateway/src/app_context.rs +++ b/model_gateway/src/app_context.rs @@ -22,8 +22,9 @@ use crate::{ rate_limit::RateLimitManager, routers::{ common::{openai_bridge::FormatRegistry, realtime::RealtimeRegistry}, - grpc::adaptive_admission::AdaptiveAdmissionController, - grpc::multimodal::MultimodalConfigRegistry, + grpc::{ + adaptive_admission::AdaptiveAdmissionController, multimodal::MultimodalConfigRegistry, + }, router_manager::RouterManager, }, wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager}, diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index c79aab00c..5791d2fb1 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -26,6 +26,34 @@ pub enum AdaptiveAdmissionMode { Enforce, } +/// Signal used to make adaptive admission decisions. +/// +/// `predicted_work` preserves the original output-token predictor. The +/// `engine_feedback` strategy instead learns the running-concurrency knee from +/// live engine throughput and backs off when the engines report queue or KV +/// pressure. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AdaptiveAdmissionStrategy { + #[default] + PredictedWork, + EngineFeedback, +} + +impl std::str::FromStr for AdaptiveAdmissionStrategy { + type Err = String; + + fn from_str(value: &str) -> Result { + match value.trim().to_ascii_lowercase().as_str() { + "predicted_work" => Ok(Self::PredictedWork), + "engine_feedback" => Ok(Self::EngineFeedback), + _ => Err(format!( + "adaptive admission strategy must be one of predicted_work, engine_feedback; got {value:?}" + )), + } + } +} + impl std::str::FromStr for AdaptiveAdmissionMode { type Err = String; @@ -65,6 +93,22 @@ fn default_adaptive_cold_start_output_tokens() -> u32 { 4096 } +fn default_feedback_probe_requests_per_healthy_replica() -> u32 { + 2 +} + +fn default_feedback_max_waiting_requests_per_healthy_replica() -> u32 { + 2 +} + +fn default_feedback_max_token_usage() -> f64 { + 0.9 +} + +fn default_feedback_throughput_improvement_ratio() -> f64 { + 0.02 +} + /// Predictive token-work admission settings. /// /// The work horizon is an operator-facing latency objective rather than a @@ -77,6 +121,8 @@ fn default_adaptive_cold_start_output_tokens() -> u32 { pub struct AdaptiveAdmissionConfig { #[serde(default)] pub mode: AdaptiveAdmissionMode, + #[serde(default)] + pub strategy: AdaptiveAdmissionStrategy, #[serde(default = "default_adaptive_work_horizon_secs")] pub work_horizon_secs: f64, #[serde(default = "default_adaptive_estimator_half_life_secs")] @@ -89,18 +135,40 @@ pub struct AdaptiveAdmissionConfig { pub min_load_coverage: f64, #[serde(default = "default_adaptive_cold_start_output_tokens")] pub cold_start_output_tokens: u32, + /// Additional per-replica requests allowed above the learned + /// running-concurrency knee so the controller can discover more capacity. + #[serde(default = "default_feedback_probe_requests_per_healthy_replica")] + pub feedback_probe_requests_per_healthy_replica: u32, + /// Engine waiting-queue threshold that closes admission until pressure + /// falls. This is evaluated against workers with fresh load telemetry. + #[serde(default = "default_feedback_max_waiting_requests_per_healthy_replica")] + pub feedback_max_waiting_requests_per_healthy_replica: u32, + /// Maximum engine token/KV usage before feedback admission closes. + #[serde(default = "default_feedback_max_token_usage")] + pub feedback_max_token_usage: f64, + /// Minimum relative throughput gain required to move the learned + /// concurrency knee upward. Near-equal throughput may move it downward. + #[serde(default = "default_feedback_throughput_improvement_ratio")] + pub feedback_throughput_improvement_ratio: f64, } impl Default for AdaptiveAdmissionConfig { fn default() -> Self { Self { mode: AdaptiveAdmissionMode::Off, + strategy: AdaptiveAdmissionStrategy::PredictedWork, work_horizon_secs: default_adaptive_work_horizon_secs(), estimator_half_life_secs: default_adaptive_estimator_half_life_secs(), prior_observations: default_adaptive_prior_observations(), max_segments: default_adaptive_max_segments(), min_load_coverage: default_adaptive_min_load_coverage(), cold_start_output_tokens: default_adaptive_cold_start_output_tokens(), + feedback_probe_requests_per_healthy_replica: + default_feedback_probe_requests_per_healthy_replica(), + feedback_max_waiting_requests_per_healthy_replica: + default_feedback_max_waiting_requests_per_healthy_replica(), + feedback_max_token_usage: default_feedback_max_token_usage(), + feedback_throughput_improvement_ratio: default_feedback_throughput_improvement_ratio(), } } } diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index b5a3071e3..793f8cac1 100755 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -772,6 +772,26 @@ impl ConfigValidator { reason: "Must be > 0".to_string(), }); } + if !adaptive.feedback_max_token_usage.is_finite() + || adaptive.feedback_max_token_usage <= 0.0 + || adaptive.feedback_max_token_usage > 1.0 + { + return Err(ConfigError::InvalidValue { + field: "adaptive_admission.feedback_max_token_usage".to_string(), + value: adaptive.feedback_max_token_usage.to_string(), + reason: "Must be finite and in (0, 1]".to_string(), + }); + } + if !adaptive.feedback_throughput_improvement_ratio.is_finite() + || adaptive.feedback_throughput_improvement_ratio < 0.0 + || adaptive.feedback_throughput_improvement_ratio > 1.0 + { + return Err(ConfigError::InvalidValue { + field: "adaptive_admission.feedback_throughput_improvement_ratio".to_string(), + value: adaptive.feedback_throughput_improvement_ratio.to_string(), + reason: "Must be finite and in [0, 1]".to_string(), + }); + } Ok(()) } diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 01b9507e6..b91243b5d 100755 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -6,10 +6,11 @@ use rand::{distr::Alphanumeric, RngExt}; use smg::{ config::{ validate_mesh_server_name, AdaptiveAdmissionConfig, AdaptiveAdmissionMode, - CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig, - HistoryBackend, ManualAssignmentMode, MetricsConfig, OracleConfig, PolicyConfig, - PostgresConfig, RedisConfig, RetryConfig, RouterConfig, RoutingKeyOverrideConfig, - RoutingMode, SchemaConfig, TenantApiKeyEntry, TokenizerCacheConfig, TraceConfig, + AdaptiveAdmissionStrategy, CircuitBreakerConfig, ConfigError, ConfigResult, + DiscoveryConfig, HealthCheckConfig, HistoryBackend, ManualAssignmentMode, MetricsConfig, + OracleConfig, PolicyConfig, PostgresConfig, RedisConfig, RetryConfig, RouterConfig, + RoutingKeyOverrideConfig, RoutingMode, SchemaConfig, TenantApiKeyEntry, + TokenizerCacheConfig, TraceConfig, }, observability::{ metrics::PrometheusConfig, @@ -125,6 +126,31 @@ impl From for AdaptiveAdmissionMode { } } +#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, ValueEnum)] +enum AdaptiveAdmissionCliStrategy { + #[default] + PredictedWork, + EngineFeedback, +} + +impl std::fmt::Display for AdaptiveAdmissionCliStrategy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::PredictedWork => "predicted_work", + Self::EngineFeedback => "engine_feedback", + }) + } +} + +impl From for AdaptiveAdmissionStrategy { + fn from(value: AdaptiveAdmissionCliStrategy) -> Self { + match value { + AdaptiveAdmissionCliStrategy::PredictedWork => Self::PredictedWork, + AdaptiveAdmissionCliStrategy::EngineFeedback => Self::EngineFeedback, + } + } +} + #[derive(Parser, Debug)] #[command(name = "shepherd-model-gateway", alias = "smg", alias = "amg")] #[command(about = "Shepherd Model Gateway - High-performance inference gateway")] @@ -527,6 +553,16 @@ struct CliArgs { )] adaptive_admission_mode: AdaptiveAdmissionCliMode, + /// Signal used for admission: predicted output work or direct engine + /// throughput, waiting-queue, and token-pressure feedback. + #[arg( + long, + value_enum, + default_value_t = AdaptiveAdmissionCliStrategy::PredictedWork, + help_heading = "Adaptive Admission" + )] + adaptive_admission_strategy: AdaptiveAdmissionCliStrategy, + /// Maximum predicted outstanding decode-work horizon in seconds. #[arg(long, default_value_t = 30.0, help_heading = "Adaptive Admission")] adaptive_admission_work_horizon_secs: f64, @@ -552,6 +588,22 @@ struct CliArgs { #[arg(long, default_value_t = 4096, help_heading = "Adaptive Admission")] adaptive_admission_cold_start_output_tokens: u32, + /// Per-replica exploration margin above the learned throughput knee. + #[arg(long, default_value_t = 2, help_heading = "Adaptive Admission")] + adaptive_admission_feedback_probe_requests_per_healthy_replica: u32, + + /// Per-replica engine waiting queue that closes feedback admission. + #[arg(long, default_value_t = 2, help_heading = "Adaptive Admission")] + adaptive_admission_feedback_max_waiting_requests_per_healthy_replica: u32, + + /// Engine token/KV usage ratio that closes feedback admission. + #[arg(long, default_value_t = 0.9, help_heading = "Adaptive Admission")] + adaptive_admission_feedback_max_token_usage: f64, + + /// Relative generation-throughput gain required to raise the learned knee. + #[arg(long, default_value_t = 0.02, help_heading = "Adaptive Admission")] + adaptive_admission_feedback_throughput_improvement_ratio: f64, + // ==================== Tenant Rate Limit ==================== /// Enable per-tenant LLM token/request rate limiting. When unset /// (default), no rate limiter is constructed. @@ -1556,12 +1608,20 @@ impl CliArgs { .priority_scheduler_tenant_metric_top_n(self.priority_scheduler_tenant_metric_top_n) .adaptive_admission(AdaptiveAdmissionConfig { mode: self.adaptive_admission_mode.into(), + strategy: self.adaptive_admission_strategy.into(), work_horizon_secs: self.adaptive_admission_work_horizon_secs, estimator_half_life_secs: self.adaptive_admission_estimator_half_life_secs, prior_observations: self.adaptive_admission_prior_observations, max_segments: self.adaptive_admission_max_segments, min_load_coverage: self.adaptive_admission_min_load_coverage, cold_start_output_tokens: self.adaptive_admission_cold_start_output_tokens, + feedback_probe_requests_per_healthy_replica: self + .adaptive_admission_feedback_probe_requests_per_healthy_replica, + feedback_max_waiting_requests_per_healthy_replica: self + .adaptive_admission_feedback_max_waiting_requests_per_healthy_replica, + feedback_max_token_usage: self.adaptive_admission_feedback_max_token_usage, + feedback_throughput_improvement_ratio: self + .adaptive_admission_feedback_throughput_improvement_ratio, }) .tenant_rate_limit_enabled(self.tenant_rate_limit_enabled) .tenant_rate_limit_config(self.tenant_rate_limit_config.clone()) @@ -1930,6 +1990,36 @@ mod tests { ); } + #[test] + fn engine_feedback_admission_options_flow_into_router_config() { + let cli = cli_args_from(&[ + "--adaptive-admission-mode", + "shadow", + "--adaptive-admission-strategy", + "engine-feedback", + "--adaptive-admission-feedback-probe-requests-per-healthy-replica", + "3", + "--adaptive-admission-feedback-max-waiting-requests-per-healthy-replica", + "4", + "--adaptive-admission-feedback-max-token-usage", + "0.85", + "--adaptive-admission-feedback-throughput-improvement-ratio", + "0.03", + ]); + + let router_config = cli.to_router_config(vec![], vec![]).unwrap(); + let adaptive = router_config.adaptive_admission; + assert_eq!(adaptive.mode, AdaptiveAdmissionMode::Shadow); + assert_eq!(adaptive.strategy, AdaptiveAdmissionStrategy::EngineFeedback); + assert_eq!(adaptive.feedback_probe_requests_per_healthy_replica, 3); + assert_eq!( + adaptive.feedback_max_waiting_requests_per_healthy_replica, + 4 + ); + assert_eq!(adaptive.feedback_max_token_usage, 0.85); + assert_eq!(adaptive.feedback_throughput_improvement_ratio, 0.03); + } + /// The multimodal transport flags must reach both `RouterConfig` and the /// wrapped `ServerConfig.router_config`. Two-path config-plumbing guard. #[test] diff --git a/model_gateway/src/observability/metrics.rs b/model_gateway/src/observability/metrics.rs index e78ae1716..eba71993e 100644 --- a/model_gateway/src/observability/metrics.rs +++ b/model_gateway/src/observability/metrics.rs @@ -479,8 +479,10 @@ pub(crate) fn init_metrics() { // Priority scheduler metrics (no-op at scrape time unless the scheduler // is enabled and recording). - use crate::middleware::scheduler::metrics as scheduler_metrics; - use crate::routers::grpc::adaptive_admission as adaptive_admission_metrics; + use crate::{ + middleware::scheduler::metrics as scheduler_metrics, + routers::grpc::adaptive_admission as adaptive_admission_metrics, + }; scheduler_metrics::describe(); adaptive_admission_metrics::describe_metrics(); } diff --git a/model_gateway/src/routers/grpc/adaptive_admission.rs b/model_gateway/src/routers/grpc/adaptive_admission.rs index 73f7726bf..067dfd39d 100644 --- a/model_gateway/src/routers/grpc/adaptive_admission.rs +++ b/model_gateway/src/routers/grpc/adaptive_admission.rs @@ -1,10 +1,11 @@ -//! Adaptive, token-work admission shared by the gRPC and HTTP serving paths. +//! Adaptive admission shared by the gRPC and HTTP serving paths. //! //! The existing priority scheduler remains the infrastructure safety layer. -//! This controller estimates request work after tokenization, learns output -//! length online, and compares predicted outstanding decode work with fresh -//! aggregate engine throughput. Shadow mode exercises the complete state -//! machine without delaying or rejecting traffic. +//! The original strategy predicts output-token work. The engine-feedback +//! strategy instead learns each partition's useful running-concurrency knee +//! from live throughput, probes just above it, and backs off on engine queue or +//! KV pressure. Shadow mode exercises either state machine without delaying or +//! rejecting traffic. use std::{ collections::HashMap, @@ -21,7 +22,7 @@ use parking_lot::Mutex; use tokio::sync::watch; use crate::{ - config::{AdaptiveAdmissionConfig, AdaptiveAdmissionMode}, + config::{AdaptiveAdmissionConfig, AdaptiveAdmissionMode, AdaptiveAdmissionStrategy}, observability::metrics::intern_string, worker::WorkerRegistry, }; @@ -45,6 +46,11 @@ const ENGINE_RUNNING: &str = "smg_adaptive_admission_engine_running_requests"; const ENGINE_WAITING: &str = "smg_adaptive_admission_engine_waiting_requests"; const ENGINE_WAITING_TOKENS: &str = "smg_adaptive_admission_engine_waiting_uncached_tokens"; const ENGINE_TOKEN_USAGE: &str = "smg_adaptive_admission_engine_max_token_usage"; +const ENGINE_MEAN_TOKEN_USAGE: &str = "smg_adaptive_admission_engine_mean_token_usage"; +const ENGINE_MAX_RUNNING: &str = "smg_adaptive_admission_engine_max_running_requests"; +const ROUTER_OUTSTANDING_REQUESTS: &str = "smg_adaptive_admission_router_outstanding_requests"; +const FEEDBACK_RUNNING_LIMIT: &str = "smg_adaptive_admission_feedback_running_limit"; +const FEEDBACK_KNEE_PER_REPLICA: &str = "smg_adaptive_admission_feedback_knee_requests_per_replica"; const SEGMENTS: &str = "smg_adaptive_admission_estimator_segments"; pub(crate) const FLAG_MULTIPLE_COMPLETIONS: u16 = 1 << 0; @@ -122,6 +128,26 @@ pub(crate) fn describe_metrics() { ENGINE_TOKEN_USAGE, "Maximum engine-reported token usage in an admission partition" ); + describe_gauge!( + ENGINE_MEAN_TOKEN_USAGE, + "Mean engine-reported token usage in an admission partition" + ); + describe_gauge!( + ENGINE_MAX_RUNNING, + "Sum of engine-reported maximum running requests in an admission partition" + ); + describe_gauge!( + ROUTER_OUTSTANDING_REQUESTS, + "Requests currently tracked by this router process" + ); + describe_gauge!( + FEEDBACK_RUNNING_LIMIT, + "Dynamic request limit selected by engine-feedback admission" + ); + describe_gauge!( + FEEDBACK_KNEE_PER_REPLICA, + "Learned running requests per replica at the throughput knee" + ); describe_gauge!( SEGMENTS, "Current bounded in-memory output estimator segment count" @@ -393,6 +419,8 @@ struct PartitionLoad { waiting_requests: i64, waiting_uncached_tokens: i64, max_token_usage: f64, + token_usage_sum: f64, + max_running_requests: i64, } #[derive(Debug, Clone)] @@ -415,6 +443,44 @@ impl CapacityEstimate { } } +#[derive(Debug, Clone)] +struct FeedbackEstimate { + peak_tokens_per_second_per_replica: f64, + running_requests_per_replica_at_peak: f64, + last_update: Instant, +} + +impl FeedbackEstimate { + fn effective_peak(&self, now: Instant, half_life_secs: f64) -> f64 { + let elapsed = now + .saturating_duration_since(self.last_update) + .as_secs_f64(); + self.peak_tokens_per_second_per_replica * 2.0_f64.powf(-elapsed / half_life_secs) + } + + fn observe( + &mut self, + throughput_per_replica: f64, + running_per_replica: f64, + now: Instant, + half_life_secs: f64, + improvement_ratio: f64, + ) { + let effective_peak = self.effective_peak(now, half_life_secs); + let raises_peak = + throughput_per_replica > effective_peak * (1.0 + improvement_ratio.max(0.0)); + let same_plateau = + throughput_per_replica >= effective_peak * (1.0 - improvement_ratio.clamp(0.0, 1.0)); + if raises_peak || same_plateau { + self.peak_tokens_per_second_per_replica = effective_peak.max(throughput_per_replica); + if raises_peak || running_per_replica < self.running_requests_per_replica_at_peak { + self.running_requests_per_replica_at_peak = running_per_replica; + } + self.last_update = now; + } + } +} + impl PartitionLoad { fn coverage(&self) -> f64 { if self.healthy_replicas == 0 { @@ -423,13 +489,23 @@ impl PartitionLoad { f64::from(self.observed_replicas) / f64::from(self.healthy_replicas) } } + + fn mean_token_usage(&self) -> f64 { + if self.observed_replicas == 0 { + 0.0 + } else { + self.token_usage_sum / f64::from(self.observed_replicas) + } + } } #[derive(Debug, Default)] struct WorkState { outstanding_tokens: HashMap, + outstanding_requests: HashMap, loads: HashMap, capacities: HashMap, + feedback_estimates: HashMap, } #[derive(Debug)] @@ -552,15 +628,27 @@ impl AdaptiveAdmissionController { .map(|rank| i64::from(rank.num_waiting_reqs.max(0))) .sum::(); aggregate.waiting_uncached_tokens += load.total_waiting_uncached_tokens().max(0); - aggregate.max_token_usage = aggregate - .max_token_usage - .max(load.effective_token_usage().clamp(0.0, 1.0)); + let token_usage = load.effective_token_usage().clamp(0.0, 1.0); + aggregate.token_usage_sum += token_usage; + aggregate.max_token_usage = aggregate.max_token_usage.max(token_usage); + let reported_max_running = load + .loads + .iter() + .map(|rank| i64::from(rank.max_running_requests.max(0))) + .sum::(); + aggregate.max_running_requests += if reported_max_running > 0 { + reported_max_running + } else { + worker.max_running_requests().map_or(0, i64::from) + }; } let now = Instant::now(); let mut work = self.work.lock(); work.capacities .retain(|partition, _| partitions.contains_key(partition)); + work.feedback_estimates + .retain(|partition, _| partitions.contains_key(partition)); for (partition, load) in &mut partitions { if load.observed_replicas > 0 && load.generation_tokens_per_second > 0.0 { let per_replica = @@ -574,6 +662,26 @@ impl AdaptiveAdmissionController { per_replica_tokens_per_second: per_replica, last_update: now, }); + if load.running_requests > 0 { + let running_per_replica = + load.running_requests as f64 / f64::from(load.observed_replicas); + work.feedback_estimates + .entry(partition.clone()) + .and_modify(|estimate| { + estimate.observe( + per_replica, + running_per_replica, + now, + self.config.estimator_half_life_secs, + self.config.feedback_throughput_improvement_ratio, + ); + }) + .or_insert(FeedbackEstimate { + peak_tokens_per_second_per_replica: per_replica, + running_requests_per_replica_at_peak: running_per_replica, + last_update: now, + }); + } } if let Some(capacity) = work.capacities.get(partition) { load.learned_capacity_tokens_per_second = capacity @@ -595,7 +703,19 @@ impl AdaptiveAdmissionController { .set(load.waiting_requests as f64); gauge!(ENGINE_WAITING_TOKENS, "partition" => Arc::clone(&partition_label)) .set(load.waiting_uncached_tokens as f64); - gauge!(ENGINE_TOKEN_USAGE, "partition" => partition_label).set(load.max_token_usage); + gauge!(ENGINE_TOKEN_USAGE, "partition" => Arc::clone(&partition_label)) + .set(load.max_token_usage); + gauge!(ENGINE_MEAN_TOKEN_USAGE, "partition" => Arc::clone(&partition_label)) + .set(load.mean_token_usage()); + gauge!(ENGINE_MAX_RUNNING, "partition" => Arc::clone(&partition_label)) + .set(load.max_running_requests as f64); + gauge!(FEEDBACK_KNEE_PER_REPLICA, "partition" => partition_label).set( + work.feedback_estimates + .get(partition) + .map_or(0.0, |estimate| { + estimate.running_requests_per_replica_at_peak + }), + ); } } @@ -617,72 +737,171 @@ impl AdaptiveAdmissionController { } }; let now = Instant::now(); - let prediction = self.predictor.lock().predict_at(&features, now); - let model_label = intern_string(&features.model); - counter!( - PREDICTIONS_TOTAL, - "model" => Arc::clone(&model_label), - "source" => prediction.source.as_str() - ) - .increment(1); - histogram!(PREDICTED_OUTPUT_TOKENS, "model" => model_label) - .record(f64::from(prediction.output_tokens)); + let prediction = if self.config.strategy == AdaptiveAdmissionStrategy::PredictedWork { + let prediction = self.predictor.lock().predict_at(&features, now); + let model_label = intern_string(&features.model); + counter!( + PREDICTIONS_TOTAL, + "model" => Arc::clone(&model_label), + "source" => prediction.source.as_str() + ) + .increment(1); + histogram!(PREDICTED_OUTPUT_TOKENS, "model" => model_label) + .record(f64::from(prediction.output_tokens)); + prediction + } else { + // Engine feedback does not predict completion length. Retain a + // zero reservation only so both strategies share the same tracker + // lifecycle without adding predictor lock contention. + Prediction { + output_tokens: 0, + model_output_tokens: 0, + source: PredictionSource::ColdStart, + } + }; let decision = { let mut work = self.work.lock(); let load = work.loads.get(&partition).cloned().unwrap_or_default(); + let feedback_estimate = work.feedback_estimates.get(&partition).cloned(); let outstanding = work .outstanding_tokens .entry(partition.clone()) .or_default(); let prior_outstanding = *outstanding; *outstanding = outstanding.saturating_add(u64::from(prediction.output_tokens)); - let coverage = load.coverage(); - let telemetry_usable = coverage >= self.config.min_load_coverage - && load.learned_capacity_tokens_per_second.is_finite() - && load.learned_capacity_tokens_per_second > 0.0; - let budget = if telemetry_usable { - load.learned_capacity_tokens_per_second * self.config.work_horizon_secs - } else { - f64::INFINITY - }; let router_outstanding = *outstanding as f64; + let outstanding_requests = work + .outstanding_requests + .entry(partition.clone()) + .or_default(); + *outstanding_requests = outstanding_requests.saturating_add(1); + let router_outstanding_requests = *outstanding_requests as f64; + let coverage = load.coverage(); let engine_request_count = load .running_requests .saturating_add(load.waiting_requests) .max(0) as f64; - let engine_estimated = engine_request_count * f64::from(prediction.model_output_tokens); - // `max` avoids counting work both in the router reservation table - // and in a later engine poll. The incoming request is not yet in - // the engine snapshot, so include it in the engine-side bound. - let projected = - router_outstanding.max(engine_estimated + f64::from(prediction.output_tokens)); - let drain_seconds = if telemetry_usable { - projected / load.learned_capacity_tokens_per_second - } else { - 0.0 - }; - let would_admit = !telemetry_usable || projected <= budget || prior_outstanding == 0; - let partition_label = intern_string(&partition); - gauge!(OUTSTANDING_TOKENS, "partition" => Arc::clone(&partition_label)).set(projected); - gauge!(ROUTER_OUTSTANDING_TOKENS, "partition" => Arc::clone(&partition_label)) - .set(router_outstanding); - gauge!(ENGINE_ESTIMATED_TOKENS, "partition" => Arc::clone(&partition_label)) - .set(engine_estimated); - gauge!(WORK_BUDGET_TOKENS, "partition" => Arc::clone(&partition_label)) - .set(if budget.is_finite() { budget } else { 0.0 }); - gauge!(DRAIN_SECONDS, "partition" => partition_label).set(drain_seconds); - AdmissionDecision { - would_admit, - telemetry_usable, - retry_after_secs: if would_admit || !telemetry_usable { - 0 - } else { - ((projected - budget) / load.learned_capacity_tokens_per_second) + gauge!(ROUTER_OUTSTANDING_REQUESTS, "partition" => Arc::clone(&partition_label)) + .set(router_outstanding_requests); + + match self.config.strategy { + AdaptiveAdmissionStrategy::PredictedWork => { + let telemetry_usable = coverage >= self.config.min_load_coverage + && load.learned_capacity_tokens_per_second.is_finite() + && load.learned_capacity_tokens_per_second > 0.0; + let budget = if telemetry_usable { + load.learned_capacity_tokens_per_second * self.config.work_horizon_secs + } else { + f64::INFINITY + }; + let engine_estimated = + engine_request_count * f64::from(prediction.model_output_tokens); + // `max` avoids counting work both in router reservations + // and in a later engine poll. Include the incoming request + // because it is not in the engine snapshot yet. + let projected = router_outstanding + .max(engine_estimated + f64::from(prediction.output_tokens)); + let drain_seconds = if telemetry_usable { + projected / load.learned_capacity_tokens_per_second + } else { + 0.0 + }; + let would_admit = + !telemetry_usable || projected <= budget || prior_outstanding == 0; + gauge!(OUTSTANDING_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(projected); + gauge!(ROUTER_OUTSTANDING_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(router_outstanding); + gauge!(ENGINE_ESTIMATED_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(engine_estimated); + gauge!(WORK_BUDGET_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(if budget.is_finite() { budget } else { 0.0 }); + gauge!(DRAIN_SECONDS, "partition" => Arc::clone(&partition_label)) + .set(drain_seconds); + gauge!(FEEDBACK_RUNNING_LIMIT, "partition" => partition_label).set(0.0); + AdmissionDecision { + would_admit, + telemetry_usable, + reason: if !telemetry_usable { + "telemetry_fallback" + } else if would_admit { + "within_work_budget" + } else { + "work_budget" + }, + retry_after_secs: if would_admit || !telemetry_usable { + 0 + } else { + ((projected - budget) / load.learned_capacity_tokens_per_second) + .ceil() + .clamp(1.0, f64::from(u32::MAX)) as u32 + }, + } + } + AdaptiveAdmissionStrategy::EngineFeedback => { + let telemetry_usable = coverage >= self.config.min_load_coverage + && load.observed_replicas > 0 + && load.max_running_requests > 0; + let engine_limit = if telemetry_usable { + (load.max_running_requests as f64 * f64::from(load.healthy_replicas) + / f64::from(load.observed_replicas)) + .floor() + } else { + 0.0 + }; + let learned_limit = feedback_estimate.as_ref().map(|estimate| { + (estimate.running_requests_per_replica_at_peak + * f64::from(load.healthy_replicas)) .ceil() - .clamp(1.0, f64::from(u32::MAX)) as u32 - }, + + f64::from( + self.config + .feedback_probe_requests_per_healthy_replica + .saturating_mul(load.healthy_replicas), + ) + }); + // Before a busy sample exists, the engine's own hard + // running limit is the bounded cold-start ceiling. + let running_limit = learned_limit + .map_or(engine_limit, |learned| learned.max(1.0).min(engine_limit)); + let projected_requests = + router_outstanding_requests.max(engine_request_count + 1.0); + let waiting_limit = i64::from( + self.config + .feedback_max_waiting_requests_per_healthy_replica, + ) * i64::from(load.observed_replicas); + let reason = if !telemetry_usable { + "telemetry_fallback" + } else if load.mean_token_usage() >= self.config.feedback_max_token_usage { + "token_pressure" + } else if load.waiting_requests > waiting_limit { + "engine_waiting" + } else if projected_requests > running_limit { + "running_limit" + } else { + "within_feedback_limit" + }; + let would_admit = + matches!(reason, "telemetry_fallback" | "within_feedback_limit"); + gauge!(FEEDBACK_RUNNING_LIMIT, "partition" => Arc::clone(&partition_label)) + .set(running_limit); + gauge!(OUTSTANDING_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(router_outstanding); + gauge!(ROUTER_OUTSTANDING_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(router_outstanding); + gauge!(ENGINE_ESTIMATED_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(0.0); + gauge!(WORK_BUDGET_TOKENS, "partition" => Arc::clone(&partition_label)) + .set(0.0); + gauge!(DRAIN_SECONDS, "partition" => partition_label).set(0.0); + AdmissionDecision { + would_admit, + telemetry_usable, + reason, + retry_after_secs: u32::from(!would_admit), + } + } } }; @@ -701,6 +920,11 @@ impl AdaptiveAdmissionController { AdaptiveAdmissionMode::Shadow => "shadow", AdaptiveAdmissionMode::Enforce => "enforce", }, + "strategy" => match self.config.strategy { + AdaptiveAdmissionStrategy::PredictedWork => "predicted_work", + AdaptiveAdmissionStrategy::EngineFeedback => "engine_feedback", + }, + "reason" => decision.reason, "outcome" => outcome ) .increment(1); @@ -727,10 +951,20 @@ impl AdaptiveAdmissionController { *outstanding = outstanding.saturating_sub(u64::from(inner.prediction.output_tokens)); gauge!(OUTSTANDING_TOKENS, "partition" => intern_string(&inner.partition)) .set(*outstanding as f64); + let outstanding_requests = work + .outstanding_requests + .entry(inner.partition.clone()) + .or_default(); + *outstanding_requests = outstanding_requests.saturating_sub(1); + gauge!(ROUTER_OUTSTANDING_REQUESTS, "partition" => intern_string(&inner.partition)) + .set(*outstanding_requests as f64); } let Some(observed) = observed_output_tokens else { return; }; + if self.config.strategy == AdaptiveAdmissionStrategy::EngineFeedback { + return; + } self.predictor .lock() .observe_at(&inner.features, observed, Instant::now()); @@ -764,6 +998,7 @@ impl AdaptiveAdmissionController { struct AdmissionDecision { would_admit: bool, telemetry_usable: bool, + reason: &'static str, retry_after_secs: u32, } @@ -777,8 +1012,9 @@ struct TrackerInner { } /// Per-request adaptive-admission state. Dropping an unfinished tracker -/// releases its predicted work without teaching the estimator from a partial -/// or failed response. +/// releases its request reservation. Predicted-work mode also releases its +/// token reservation without teaching the estimator from a partial or failed +/// response. pub(crate) struct AdaptiveRequestTracker { inner: Option, } @@ -833,12 +1069,17 @@ mod tests { fn config() -> AdaptiveAdmissionConfig { AdaptiveAdmissionConfig { mode: AdaptiveAdmissionMode::Shadow, + strategy: AdaptiveAdmissionStrategy::PredictedWork, work_horizon_secs: 10.0, estimator_half_life_secs: 60.0, prior_observations: 2.0, max_segments: 20, min_load_coverage: 0.8, cold_start_output_tokens: 100, + feedback_probe_requests_per_healthy_replica: 2, + feedback_max_waiting_requests_per_healthy_replica: 2, + feedback_max_token_usage: 0.9, + feedback_throughput_improvement_ratio: 0.02, } } @@ -1038,4 +1279,115 @@ mod tests { let next = controller.begin("model".to_string(), features("b", 10, None)); assert!(next.should_reject()); } + + #[test] + fn engine_feedback_uses_learned_knee_plus_probe_margin() { + let mut settings = config(); + settings.mode = AdaptiveAdmissionMode::Enforce; + settings.strategy = AdaptiveAdmissionStrategy::EngineFeedback; + let controller = + AdaptiveAdmissionController::new(settings, Arc::new(WorkerRegistry::new())); + let now = Instant::now(); + let mut work = controller.work.lock(); + work.loads.insert( + "model".to_string(), + PartitionLoad { + healthy_replicas: 1, + observed_replicas: 1, + generation_tokens_per_second: 100.0, + running_requests: 10, + max_running_requests: 64, + max_token_usage: 0.5, + ..PartitionLoad::default() + }, + ); + work.feedback_estimates.insert( + "model".to_string(), + FeedbackEstimate { + peak_tokens_per_second_per_replica: 100.0, + running_requests_per_replica_at_peak: 10.0, + last_update: now, + }, + ); + drop(work); + + let trackers: Vec<_> = (0..12) + .map(|i| controller.begin("model".to_string(), features(&format!("u-{i}"), 10, None))) + .collect(); + assert!(trackers.iter().all(|tracker| !tracker.should_reject())); + let excess = controller.begin("model".to_string(), features("excess", 10, None)); + assert!(excess.should_reject()); + assert_eq!(excess.retry_after_secs(), 1); + } + + #[test] + fn engine_feedback_backs_off_on_waiting_or_token_pressure() { + for (waiting_requests, token_usage, expected_reason) in + [(3, 0.5, "engine_waiting"), (0, 0.91, "token_pressure")] + { + let mut settings = config(); + settings.mode = AdaptiveAdmissionMode::Enforce; + settings.strategy = AdaptiveAdmissionStrategy::EngineFeedback; + let controller = + AdaptiveAdmissionController::new(settings, Arc::new(WorkerRegistry::new())); + controller.work.lock().loads.insert( + "model".to_string(), + PartitionLoad { + healthy_replicas: 1, + observed_replicas: 1, + generation_tokens_per_second: 100.0, + running_requests: 1, + waiting_requests, + max_running_requests: 64, + max_token_usage: token_usage, + token_usage_sum: token_usage, + ..PartitionLoad::default() + }, + ); + + let tracker = controller.begin("model".to_string(), features("u", 10, None)); + assert!(tracker.should_reject()); + assert_eq!( + tracker.inner.as_ref().unwrap().decision.reason, + expected_reason + ); + } + } + + #[test] + fn engine_feedback_cold_start_is_bounded_by_engine_limit() { + let mut settings = config(); + settings.mode = AdaptiveAdmissionMode::Enforce; + settings.strategy = AdaptiveAdmissionStrategy::EngineFeedback; + let controller = + AdaptiveAdmissionController::new(settings, Arc::new(WorkerRegistry::new())); + controller.work.lock().loads.insert( + "model".to_string(), + PartitionLoad { + healthy_replicas: 1, + observed_replicas: 1, + max_running_requests: 2, + ..PartitionLoad::default() + }, + ); + + let first = controller.begin("model".to_string(), features("a", 10, None)); + let second = controller.begin("model".to_string(), features("b", 10, None)); + let third = controller.begin("model".to_string(), features("c", 10, None)); + assert!(!first.should_reject()); + assert!(!second.should_reject()); + assert!(third.should_reject()); + } + + #[test] + fn feedback_knee_moves_down_on_same_throughput_plateau() { + let start = Instant::now(); + let mut estimate = FeedbackEstimate { + peak_tokens_per_second_per_replica: 100.0, + running_requests_per_replica_at_peak: 20.0, + last_update: start, + }; + estimate.observe(99.0, 12.0, start + Duration::from_secs(1), 60.0, 0.02); + assert_eq!(estimate.running_requests_per_replica_at_peak, 12.0); + } } From d327a42be5910c66a03338694512d9651ec0897a Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:15:40 -0700 Subject: [PATCH 3/8] feat(python): expose engine-feedback admission Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- bindings/python/src/lib.rs | 36 +++++++++++++++++++ bindings/python/src/smg/router_args.py | 45 ++++++++++++++++++++---- bindings/python/tests/test_arg_parser.py | 25 ++++++++++--- 3 files changed, 96 insertions(+), 10 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index d1f0afefa..b5942a2c7 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -505,6 +505,11 @@ struct Router { adaptive_admission_max_segments: usize, adaptive_admission_min_load_coverage: f64, adaptive_admission_cold_start_output_tokens: u32, + adaptive_admission_strategy: String, + adaptive_admission_feedback_probe_requests_per_healthy_replica: u32, + adaptive_admission_feedback_max_waiting_requests_per_healthy_replica: u32, + adaptive_admission_feedback_max_token_usage: f64, + adaptive_admission_feedback_throughput_improvement_ratio: f64, } impl Router { @@ -714,6 +719,14 @@ impl Router { reason, } })?; + let adaptive_admission_strategy = + self.adaptive_admission_strategy.parse().map_err(|reason| { + config::ConfigError::InvalidValue { + field: "adaptive_admission_strategy".to_string(), + value: self.adaptive_admission_strategy.clone(), + reason, + } + })?; let history_backend = match self.history_backend { HistoryBackendType::Memory => config::HistoryBackend::Memory, @@ -793,12 +806,20 @@ impl Router { .priority_scheduler_tenant_metric_top_n(self.priority_scheduler_tenant_metric_top_n) .adaptive_admission(config::AdaptiveAdmissionConfig { mode: adaptive_admission_mode, + strategy: adaptive_admission_strategy, work_horizon_secs: self.adaptive_admission_work_horizon_secs, estimator_half_life_secs: self.adaptive_admission_estimator_half_life_secs, prior_observations: self.adaptive_admission_prior_observations, max_segments: self.adaptive_admission_max_segments, min_load_coverage: self.adaptive_admission_min_load_coverage, cold_start_output_tokens: self.adaptive_admission_cold_start_output_tokens, + feedback_probe_requests_per_healthy_replica: self + .adaptive_admission_feedback_probe_requests_per_healthy_replica, + feedback_max_waiting_requests_per_healthy_replica: self + .adaptive_admission_feedback_max_waiting_requests_per_healthy_replica, + feedback_max_token_usage: self.adaptive_admission_feedback_max_token_usage, + feedback_throughput_improvement_ratio: self + .adaptive_admission_feedback_throughput_improvement_ratio, }) .cors_allowed_origins(self.cors_allowed_origins.clone()) .retry_config(config::RetryConfig { @@ -1026,6 +1047,11 @@ impl Router { adaptive_admission_max_segments = 50000, adaptive_admission_min_load_coverage = 0.8, adaptive_admission_cold_start_output_tokens = 4096, + adaptive_admission_strategy = String::from("predicted_work"), + adaptive_admission_feedback_probe_requests_per_healthy_replica = 2, + adaptive_admission_feedback_max_waiting_requests_per_healthy_replica = 2, + adaptive_admission_feedback_max_token_usage = 0.9, + adaptive_admission_feedback_throughput_improvement_ratio = 0.02, ))] #[expect(clippy::too_many_arguments)] #[expect( @@ -1170,6 +1196,11 @@ impl Router { adaptive_admission_max_segments: usize, adaptive_admission_min_load_coverage: f64, adaptive_admission_cold_start_output_tokens: u32, + adaptive_admission_strategy: String, + adaptive_admission_feedback_probe_requests_per_healthy_replica: u32, + adaptive_admission_feedback_max_waiting_requests_per_healthy_replica: u32, + adaptive_admission_feedback_max_token_usage: f64, + adaptive_admission_feedback_throughput_improvement_ratio: f64, ) -> PyResult { let mut all_urls = worker_urls.clone(); @@ -1328,6 +1359,11 @@ impl Router { adaptive_admission_max_segments, adaptive_admission_min_load_coverage, adaptive_admission_cold_start_output_tokens, + adaptive_admission_strategy, + adaptive_admission_feedback_probe_requests_per_healthy_replica, + adaptive_admission_feedback_max_waiting_requests_per_healthy_replica, + adaptive_admission_feedback_max_token_usage, + adaptive_admission_feedback_throughput_improvement_ratio, }) } diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index de7805549..9c613ba6e 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -124,12 +124,17 @@ class RouterArgs: # Engine telemetry and predictive token-work admission. engine_metrics: bool = False adaptive_admission_mode: str = "off" + adaptive_admission_strategy: str = "predicted_work" adaptive_admission_work_horizon_secs: float = 30.0 adaptive_admission_estimator_half_life_secs: float = 900.0 adaptive_admission_prior_observations: float = 20.0 adaptive_admission_max_segments: int = 50_000 adaptive_admission_min_load_coverage: float = 0.8 adaptive_admission_cold_start_output_tokens: int = 4096 + adaptive_admission_feedback_probe_requests_per_healthy_replica: int = 2 + adaptive_admission_feedback_max_waiting_requests_per_healthy_replica: int = 2 + adaptive_admission_feedback_max_token_usage: float = 0.9 + adaptive_admission_feedback_throughput_improvement_ratio: float = 0.02 # Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests rate_limit_tokens_per_second: int | None = None # Cluster-wide requests-per-second ceiling. Requires mesh and the same value on every gateway. @@ -891,6 +896,12 @@ def add_cli_args( default=RouterArgs.adaptive_admission_mode, help="Predictive token-work admission mode", ) + adaptive_admission_group.add_argument( + f"--{prefix}adaptive-admission-strategy", + choices=["predicted_work", "engine_feedback"], + default=RouterArgs.adaptive_admission_strategy, + help="Admission signal: output-work prediction or direct engine feedback", + ) adaptive_admission_group.add_argument( f"--{prefix}adaptive-admission-work-horizon-secs", type=float, @@ -927,6 +938,32 @@ def add_cli_args( default=RouterArgs.adaptive_admission_cold_start_output_tokens, help="Cold-start output-token prediction before observations", ) + adaptive_admission_group.add_argument( + f"--{prefix}adaptive-admission-feedback-probe-requests-per-healthy-replica", + type=int, + default=RouterArgs.adaptive_admission_feedback_probe_requests_per_healthy_replica, + help="Per-replica exploration margin above the learned throughput knee", + ) + adaptive_admission_group.add_argument( + f"--{prefix}adaptive-admission-feedback-max-waiting-requests-per-healthy-replica", + type=int, + default=( + RouterArgs.adaptive_admission_feedback_max_waiting_requests_per_healthy_replica + ), + help="Per-replica engine waiting queue that closes feedback admission", + ) + adaptive_admission_group.add_argument( + f"--{prefix}adaptive-admission-feedback-max-token-usage", + type=float, + default=RouterArgs.adaptive_admission_feedback_max_token_usage, + help="Engine token/KV usage ratio that closes feedback admission", + ) + adaptive_admission_group.add_argument( + f"--{prefix}adaptive-admission-feedback-throughput-improvement-ratio", + type=float, + default=RouterArgs.adaptive_admission_feedback_throughput_improvement_ratio, + help="Relative throughput gain required to raise the learned concurrency knee", + ) # Retry configuration retry_group.add_argument( @@ -1517,13 +1554,9 @@ def _parse_model_policies(values: list[str] | None) -> dict[str, str]: def _validate_router_args(self): if self.global_rate_limit_requests_per_second is not None: if self.global_rate_limit_requests_per_second <= 0: - raise ValueError( - "global_rate_limit_requests_per_second must be greater than zero" - ) + raise ValueError("global_rate_limit_requests_per_second must be greater than zero") if not self.enable_mesh: - raise ValueError( - "global_rate_limit_requests_per_second requires enable_mesh=True" - ) + raise ValueError("global_rate_limit_requests_per_second requires enable_mesh=True") # Validate configuration based on mode if self.epd_disaggregation: diff --git a/bindings/python/tests/test_arg_parser.py b/bindings/python/tests/test_arg_parser.py index 20ceb6d5a..ab84aec2c 100644 --- a/bindings/python/tests/test_arg_parser.py +++ b/bindings/python/tests/test_arg_parser.py @@ -53,12 +53,17 @@ def test_default_values(self): assert args.priority_scheduler_tenant_metric_top_n == 32 assert args.engine_metrics is False assert args.adaptive_admission_mode == "off" + assert args.adaptive_admission_strategy == "predicted_work" assert args.adaptive_admission_work_horizon_secs == 30.0 assert args.adaptive_admission_estimator_half_life_secs == 900.0 assert args.adaptive_admission_prior_observations == 20.0 assert args.adaptive_admission_max_segments == 50_000 assert args.adaptive_admission_min_load_coverage == 0.8 assert args.adaptive_admission_cold_start_output_tokens == 4096 + assert args.adaptive_admission_feedback_probe_requests_per_healthy_replica == 2 + assert args.adaptive_admission_feedback_max_waiting_requests_per_healthy_replica == 2 + assert args.adaptive_admission_feedback_max_token_usage == 0.9 + assert args.adaptive_admission_feedback_throughput_improvement_ratio == 0.02 def test_parse_priority_scheduler_options(self): args = parse_router_args( @@ -84,6 +89,8 @@ def test_parse_adaptive_admission_options(self): "--engine-metrics", "--adaptive-admission-mode", "shadow", + "--adaptive-admission-strategy", + "engine_feedback", "--adaptive-admission-work-horizon-secs", "45", "--adaptive-admission-estimator-half-life-secs", @@ -96,17 +103,30 @@ def test_parse_adaptive_admission_options(self): "0.75", "--adaptive-admission-cold-start-output-tokens", "2048", + "--adaptive-admission-feedback-probe-requests-per-healthy-replica", + "3", + "--adaptive-admission-feedback-max-waiting-requests-per-healthy-replica", + "4", + "--adaptive-admission-feedback-max-token-usage", + "0.85", + "--adaptive-admission-feedback-throughput-improvement-ratio", + "0.03", ] ) assert args.engine_metrics is True assert args.adaptive_admission_mode == "shadow" + assert args.adaptive_admission_strategy == "engine_feedback" assert args.adaptive_admission_work_horizon_secs == 45.0 assert args.adaptive_admission_estimator_half_life_secs == 600.0 assert args.adaptive_admission_prior_observations == 12.0 assert args.adaptive_admission_max_segments == 12_345 assert args.adaptive_admission_min_load_coverage == 0.75 assert args.adaptive_admission_cold_start_output_tokens == 2048 + assert args.adaptive_admission_feedback_probe_requests_per_healthy_replica == 3 + assert args.adaptive_admission_feedback_max_waiting_requests_per_healthy_replica == 4 + assert args.adaptive_admission_feedback_max_token_usage == 0.85 + assert args.adaptive_admission_feedback_throughput_improvement_ratio == 0.03 def test_parse_selector_valid(self): """Test parsing valid selector arguments.""" @@ -540,10 +560,7 @@ def test_valid_policies(self): assert policy_from_str("round_robin") == PolicyType.RoundRobin assert policy_from_str("cache_aware") == PolicyType.CacheAware assert policy_from_str("power_of_two") == PolicyType.PowerOfTwo - assert ( - policy_from_str("size_aware_power_of_two") - == PolicyType.SizeAwarePowerOfTwo - ) + assert policy_from_str("size_aware_power_of_two") == PolicyType.SizeAwarePowerOfTwo assert policy_from_str("consistent_hashing") == PolicyType.ConsistentHashing assert policy_from_str("prefix_hash") == PolicyType.PrefixHash From 5632e41947f5a8631e5540836a83318733b983d6 Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:36:33 -0700 Subject: [PATCH 4/8] fix(admission): require representative engine ceilings Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- .../src/routers/grpc/adaptive_admission.rs | 95 +++++++++++++++++-- 1 file changed, 89 insertions(+), 6 deletions(-) diff --git a/model_gateway/src/routers/grpc/adaptive_admission.rs b/model_gateway/src/routers/grpc/adaptive_admission.rs index 067dfd39d..528176718 100644 --- a/model_gateway/src/routers/grpc/adaptive_admission.rs +++ b/model_gateway/src/routers/grpc/adaptive_admission.rs @@ -48,6 +48,8 @@ const ENGINE_WAITING_TOKENS: &str = "smg_adaptive_admission_engine_waiting_uncac const ENGINE_TOKEN_USAGE: &str = "smg_adaptive_admission_engine_max_token_usage"; const ENGINE_MEAN_TOKEN_USAGE: &str = "smg_adaptive_admission_engine_mean_token_usage"; const ENGINE_MAX_RUNNING: &str = "smg_adaptive_admission_engine_max_running_requests"; +const ENGINE_MAX_RUNNING_COVERAGE: &str = + "smg_adaptive_admission_engine_max_running_requests_coverage"; const ROUTER_OUTSTANDING_REQUESTS: &str = "smg_adaptive_admission_router_outstanding_requests"; const FEEDBACK_RUNNING_LIMIT: &str = "smg_adaptive_admission_feedback_running_limit"; const FEEDBACK_KNEE_PER_REPLICA: &str = "smg_adaptive_admission_feedback_knee_requests_per_replica"; @@ -136,6 +138,10 @@ pub(crate) fn describe_metrics() { ENGINE_MAX_RUNNING, "Sum of engine-reported maximum running requests in an admission partition" ); + describe_gauge!( + ENGINE_MAX_RUNNING_COVERAGE, + "Fraction of healthy replicas contributing a maximum-running-requests ceiling" + ); describe_gauge!( ROUTER_OUTSTANDING_REQUESTS, "Requests currently tracked by this router process" @@ -421,6 +427,7 @@ struct PartitionLoad { max_token_usage: f64, token_usage_sum: f64, max_running_requests: i64, + max_running_observed_replicas: u32, } #[derive(Debug, Clone)] @@ -497,6 +504,24 @@ impl PartitionLoad { self.token_usage_sum / f64::from(self.observed_replicas) } } + + fn max_running_coverage(&self) -> f64 { + if self.healthy_replicas == 0 { + 0.0 + } else { + f64::from(self.max_running_observed_replicas) / f64::from(self.healthy_replicas) + } + } + + fn scaled_max_running_requests(&self) -> f64 { + if self.max_running_observed_replicas == 0 { + 0.0 + } else { + (self.max_running_requests as f64 * f64::from(self.healthy_replicas) + / f64::from(self.max_running_observed_replicas)) + .floor() + } + } } #[derive(Debug, Default)] @@ -636,11 +661,16 @@ impl AdaptiveAdmissionController { .iter() .map(|rank| i64::from(rank.max_running_requests.max(0))) .sum::(); - aggregate.max_running_requests += if reported_max_running > 0 { + let max_running_requests = if reported_max_running > 0 { reported_max_running } else { worker.max_running_requests().map_or(0, i64::from) }; + if max_running_requests > 0 { + aggregate.max_running_requests += max_running_requests; + aggregate.max_running_observed_replicas = + aggregate.max_running_observed_replicas.saturating_add(1); + } } let now = Instant::now(); @@ -709,6 +739,8 @@ impl AdaptiveAdmissionController { .set(load.mean_token_usage()); gauge!(ENGINE_MAX_RUNNING, "partition" => Arc::clone(&partition_label)) .set(load.max_running_requests as f64); + gauge!(ENGINE_MAX_RUNNING_COVERAGE, "partition" => Arc::clone(&partition_label)) + .set(load.max_running_coverage()); gauge!(FEEDBACK_KNEE_PER_REPLICA, "partition" => partition_label).set( work.feedback_estimates .get(partition) @@ -842,12 +874,9 @@ impl AdaptiveAdmissionController { } AdaptiveAdmissionStrategy::EngineFeedback => { let telemetry_usable = coverage >= self.config.min_load_coverage - && load.observed_replicas > 0 - && load.max_running_requests > 0; + && load.max_running_coverage() >= self.config.min_load_coverage; let engine_limit = if telemetry_usable { - (load.max_running_requests as f64 * f64::from(load.healthy_replicas) - / f64::from(load.observed_replicas)) - .floor() + load.scaled_max_running_requests() } else { 0.0 }; @@ -1297,6 +1326,7 @@ mod tests { generation_tokens_per_second: 100.0, running_requests: 10, max_running_requests: 64, + max_running_observed_replicas: 1, max_token_usage: 0.5, ..PartitionLoad::default() }, @@ -1339,6 +1369,7 @@ mod tests { running_requests: 1, waiting_requests, max_running_requests: 64, + max_running_observed_replicas: 1, max_token_usage: token_usage, token_usage_sum: token_usage, ..PartitionLoad::default() @@ -1367,6 +1398,7 @@ mod tests { healthy_replicas: 1, observed_replicas: 1, max_running_requests: 2, + max_running_observed_replicas: 1, ..PartitionLoad::default() }, ); @@ -1379,6 +1411,57 @@ mod tests { assert!(third.should_reject()); } + #[test] + fn engine_feedback_scales_only_sufficient_max_running_coverage() { + let sufficiently_covered = PartitionLoad { + healthy_replicas: 5, + observed_replicas: 4, + max_running_requests: 256, + max_running_observed_replicas: 4, + ..PartitionLoad::default() + }; + assert_eq!(sufficiently_covered.coverage(), 0.8); + assert_eq!(sufficiently_covered.max_running_coverage(), 0.8); + assert_eq!(sufficiently_covered.scaled_max_running_requests(), 320.0); + + let insufficiently_covered = PartitionLoad { + healthy_replicas: 5, + observed_replicas: 5, + max_running_requests: 192, + max_running_observed_replicas: 3, + ..PartitionLoad::default() + }; + assert_eq!(insufficiently_covered.coverage(), 1.0); + assert_eq!(insufficiently_covered.max_running_coverage(), 0.6); + } + + #[test] + fn engine_feedback_fails_open_when_max_running_coverage_is_incomplete() { + let mut settings = config(); + settings.mode = AdaptiveAdmissionMode::Enforce; + settings.strategy = AdaptiveAdmissionStrategy::EngineFeedback; + let controller = + AdaptiveAdmissionController::new(settings, Arc::new(WorkerRegistry::new())); + controller.work.lock().loads.insert( + "model".to_string(), + PartitionLoad { + healthy_replicas: 5, + observed_replicas: 5, + max_running_requests: 192, + max_running_observed_replicas: 3, + ..PartitionLoad::default() + }, + ); + + let tracker = controller.begin("model".to_string(), features("u", 10, None)); + assert!(!tracker.should_reject()); + assert!(!tracker.inner.as_ref().unwrap().decision.telemetry_usable); + assert_eq!( + tracker.inner.as_ref().unwrap().decision.reason, + "telemetry_fallback" + ); + } + #[test] fn feedback_knee_moves_down_on_same_throughput_plateau() { let start = Instant::now(); From b5db3d07e25b824f920d05fc6212e5aa8414d5bf Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:06:11 -0700 Subject: [PATCH 5/8] feat(scheduler): add global tenant fair sharing Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- bindings/python/src/lib.rs | 5 + bindings/python/src/smg/router_args.py | 22 + bindings/python/tests/test_arg_parser.py | 17 + docs/reference/priority-scheduler.md | 82 +- model_gateway/src/config/builder.rs | 5 + model_gateway/src/config/types.rs | 2 + model_gateway/src/config/validation.rs | 21 + model_gateway/src/main.rs | 6 + .../src/middleware/scheduler/admission.rs | 87 +- .../src/middleware/scheduler/body.rs | 266 +++++- .../src/middleware/scheduler/config.rs | 110 ++- .../src/middleware/scheduler/engine.rs | 319 ++++++- .../src/middleware/scheduler/fair_share.rs | 894 ++++++++++++++++++ .../src/middleware/scheduler/metrics.rs | 71 ++ model_gateway/src/middleware/scheduler/mod.rs | 7 +- .../src/middleware/scheduler/output_tokens.rs | 97 ++ .../src/middleware/scheduler/queue.rs | 146 ++- .../src/middleware/scheduler/state.rs | 20 +- .../src/middleware/tenant_resolution.rs | 87 ++ 19 files changed, 2220 insertions(+), 44 deletions(-) create mode 100644 model_gateway/src/middleware/scheduler/fair_share.rs create mode 100644 model_gateway/src/middleware/scheduler/output_tokens.rs diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index b5942a2c7..0582c2980 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -415,6 +415,7 @@ struct Router { shutdown_grace_period_secs: u64, request_id_headers: Option>, trust_tenant_header: bool, + prefer_trusted_tenant_header: bool, tenant_header_name: String, storage_context_headers: HashMap, pd_disaggregation: bool, @@ -861,6 +862,7 @@ impl Router { .maybe_log_level(self.log_level.as_ref()) .maybe_request_id_headers(self.request_id_headers.clone()) .trust_tenant_header(self.trust_tenant_header) + .prefer_trusted_tenant_header(self.prefer_trusted_tenant_header) .tenant_header_name(&self.tenant_header_name) .maybe_storage_context_headers( (!self.storage_context_headers.is_empty()) @@ -956,6 +958,7 @@ impl Router { shutdown_grace_period_secs = 180, request_id_headers = None, trust_tenant_header = false, + prefer_trusted_tenant_header = false, tenant_header_name = String::from("x-smg-tenant-id"), storage_context_headers = HashMap::new(), pd_disaggregation = false, @@ -1106,6 +1109,7 @@ impl Router { shutdown_grace_period_secs: u64, request_id_headers: Option>, trust_tenant_header: bool, + prefer_trusted_tenant_header: bool, tenant_header_name: String, storage_context_headers: HashMap, pd_disaggregation: bool, @@ -1272,6 +1276,7 @@ impl Router { shutdown_grace_period_secs, request_id_headers, trust_tenant_header, + prefer_trusted_tenant_header, tenant_header_name, storage_context_headers, pd_disaggregation, diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index 9c613ba6e..47dbf17e5 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -104,6 +104,9 @@ class RouterArgs: prometheus_duration_buckets: list[float] | None = None # Request ID headers configuration request_id_headers: list[str] | None = None + trust_tenant_header: bool = False + prefer_trusted_tenant_header: bool = False + tenant_header_name: str = "x-smg-tenant-id" # HTTP header to storage hook context mapping storage_context_headers: dict[str, str] = dataclasses.field(default_factory=dict) # Request timeout in seconds @@ -787,6 +790,25 @@ def add_cli_args( " If not specified, uses common defaults." ), ) + request_group.add_argument( + f"--{prefix}trust-tenant-header", + action="store_true", + help="Trust the configured upstream tenant identity header", + ) + request_group.add_argument( + f"--{prefix}prefer-trusted-tenant-header", + action="store_true", + help=( + "Prefer the trusted tenant header over authenticated proxy identity; " + "requires --trust-tenant-header" + ), + ) + request_group.add_argument( + f"--{prefix}tenant-header-name", + type=str, + default=RouterArgs.tenant_header_name, + help="Trusted tenant identity header name", + ) request_group.add_argument( f"--{prefix}storage-context-headers", type=str, diff --git a/bindings/python/tests/test_arg_parser.py b/bindings/python/tests/test_arg_parser.py index ab84aec2c..0452c7a02 100644 --- a/bindings/python/tests/test_arg_parser.py +++ b/bindings/python/tests/test_arg_parser.py @@ -51,6 +51,9 @@ def test_default_values(self): assert args.priority_scheduler_default_max_class == "default" assert args.priority_scheduler_config is None assert args.priority_scheduler_tenant_metric_top_n == 32 + assert args.trust_tenant_header is False + assert args.prefer_trusted_tenant_header is False + assert args.tenant_header_name == "x-smg-tenant-id" assert args.engine_metrics is False assert args.adaptive_admission_mode == "off" assert args.adaptive_admission_strategy == "predicted_work" @@ -83,6 +86,20 @@ def test_parse_priority_scheduler_options(self): assert args.priority_scheduler_config == "/tmp/priority.yaml" assert args.priority_scheduler_tenant_metric_top_n == 16 + def test_parse_preferred_trusted_tenant_header_options(self): + args = parse_router_args( + [ + "--trust-tenant-header", + "--prefer-trusted-tenant-header", + "--tenant-header-name", + "x-comet-user", + ] + ) + + assert args.trust_tenant_header is True + assert args.prefer_trusted_tenant_header is True + assert args.tenant_header_name == "x-comet-user" + def test_parse_adaptive_admission_options(self): args = parse_router_args( [ diff --git a/docs/reference/priority-scheduler.md b/docs/reference/priority-scheduler.md index 55ea81294..b97c0776c 100644 --- a/docs/reference/priority-scheduler.md +++ b/docs/reference/priority-scheduler.md @@ -73,7 +73,7 @@ smg \ | `--priority-scheduler-enabled` | `false` | Master switch. When unset, the legacy concurrency-limit middleware stays wired and no scheduler is constructed. | | `--priority-scheduler-default-max-class` | `default` | Maximum class for tenants not listed in the YAML (`system` \| `interactive` \| `default` \| `bulk`). Parsed with the same rules as the header — an unknown value falls back to `default`. | | `--priority-scheduler-config` | unset | Path to the optional priority-scheduler YAML (per-class overrides + per-tenant policy). Absent → built-in defaults and an empty tenant policy map. | -| `--priority-scheduler-tenant-metric-top-n` | `32` | Intended cap on per-tenant metric label cardinality. **Not yet enforced** — the value is stored but no top-N bucketing is applied today; per-tenant counters currently intern the raw tenant. | +| `--priority-scheduler-tenant-metric-top-n` | `32` | Cap on configured tenants emitted as distinct fair-share metric labels. Remaining tenants use `tenant="other"`. Existing non-fair-share tenant counters still intern their raw tenant label. | !!! warning "Fail-safe startup" If the scheduler is enabled but cannot start — unparsable YAML, or class reservation floors + shares that sum to more than the live backend capacity — the gateway logs at `ERROR` and **falls back to legacy admission** instead of aborting. It does not take the data plane down. @@ -108,6 +108,15 @@ tenant_policies: max_class: interactive "auth:internal-cron": max_class: system + +# Optional weighted sharing among tenants that contend in the same model pool. +fair_share: + default_weight: 1 + default_output_tokens: 256 + trust_output_token_estimate_header: false + tenant_weights: + "header:alice": 10 + "header:bob": 5 ``` Class keys and `max_class` values are lowercase: `system`, `interactive`, `default`, `bulk`. An unknown class name in the YAML is a parse error (which triggers the fail-safe fallback above), unlike the lenient request header. @@ -150,6 +159,71 @@ Any validation failure triggers the [fail-safe fallback to legacy admission](#en --- +## Weighted fair sharing by output tokens + +The optional `fair_share` map replaces FIFO ordering within each priority-class +queue with weighted ordering by output tokens. Weights are relative and do not +need to sum to 100. For example, weights `10` and `5` target a 2:1 token split +while both tenants remain backlogged and eligible for the same model pool. +Priority class selection remains the outer policy. + +| Field | Default | Meaning | +|-------|---------|---------| +| `default_weight` | `1.0` | Weight assigned to a resolved tenant absent from `tenant_weights`. Must be finite and greater than zero. | +| `default_output_tokens` | `256` | Provisional output-token charge used when no trusted estimate is available. Must be greater than zero. | +| `trust_output_token_estimate_header` | `false` | Honor `x-smg-output-token-estimate`. Enable only behind a proxy that strips client copies and injects a validated estimate. | +| `tenant_weights` | `{}` | Relative weights keyed by canonical tenant key. Every value must be finite and greater than zero. | + +The scheduler reserves the estimate when a request is admitted. A trustworthy +terminal usage record replaces that estimate with actual output tokens. If a +client disconnects, the backend fails, or terminal usage is missing or +truncated, the provisional charge remains. This prevents cancellation from +evading fair-share accounting. + +Fairness credit accrues only during active contention. When a new or returning +tenant becomes active, its virtual finish is rebased to the current active-set +virtual time, so idle tenants do not bank unlimited catch-up credit. The queue +is work-conserving: a model partition with a free slot and an eligible local +request never idles for an underserved tenant that can use only another model. + +One `GlobalFairShare` instance aggregates actual-token metrics and canonical +tenant virtual finish across every partition built by one SMG process. Each +partition chooses among only the requests eligible for its non-fungible model +pool. Consequently: + +- Batch and interactive requests resolve to the same tenant ledger when they + enter the same SMG process with the same canonical tenant identity. Priority + classes still decide which class is considered first. +- Configured percentages converge when tenants are simultaneously backlogged + for the same constrained pool. Exact fleet-wide percentages are not + enforceable for tenants targeting disjoint pools without idling capacity. +- The ledger does not coordinate separate M1 and M2 gateways, or overlapping + old and new gateway processes during a rollout. Strict cross-gateway fairness + requires a distributed ledger or one authoritative admission front door. + +### Preferred trusted tenant identity + +By default, tenant resolution remains authenticated caller, then trusted +header, then client IP, then anonymous. A deployment with a shared authenticated +proxy identity can deliberately make the proxy-injected end-user header +canonical by setting all three flags: + +```bash +--trust-tenant-header \ +--prefer-trusted-tenant-header \ +--tenant-header-name x-comet-user +``` + +`--prefer-trusted-tenant-header` requires `--trust-tenant-header`. A missing, +empty, or invalid preferred header falls back to the existing authenticated +caller path. The option is disabled by default. Enabling it changes the +canonical `RouteRequestMeta` tenant for every tenant-aware subsystem, including +rate-limit policy lookup, tenant metrics, priority clamps, and fair sharing. +The upstream proxy must strip caller-supplied copies and reinject only its +authenticated user identity. + +--- + ## Tenant policy A tenant's priority ceiling is resolved per request: @@ -182,6 +256,12 @@ The scheduler exposes these Prometheus metrics (see the [Metrics Reference](metr | `smg_scheduler_queue_size_limit` | Gauge | `class` | Configured queue limit per class. | | `smg_scheduler_utilization` | Gauge | — | Total in-flight divided by backend capacity. | | `smg_scheduler_class_capacity_pressure` | Gauge | `class` | Normalized 0.0–1.0 pressure (worse of queue and slot pressure). | +| `smg_fair_share_charged_output_tokens_total` | Counter | `tenant` | Actual or conservative fallback output tokens charged across the process-local ledger. | +| `smg_fair_share_virtual_finish` | Gauge | `tenant` | Process-global active-set normalized virtual finish used for local eligible-candidate dispatch. | +| `smg_fair_share_reserved_output_tokens` | Gauge | `tenant` | Provisional output-token charges held by active requests. | +| `smg_fair_share_queue_wait_seconds` | Histogram | `tenant`, `class` | Fair-share queue wait by tenant and priority class. | +| `smg_fair_share_fallback_total` | Counter | `reason` | Settlement or estimate paths that used a configured fallback. | +| `smg_fair_share_unknown_tenant_total` | Counter | `tenant` | Requests whose canonical tenant has no explicit configured weight. | --- diff --git a/model_gateway/src/config/builder.rs b/model_gateway/src/config/builder.rs index 4641c30f7..681fda702 100644 --- a/model_gateway/src/config/builder.rs +++ b/model_gateway/src/config/builder.rs @@ -445,6 +445,11 @@ impl RouterConfigBuilder { self } + pub fn prefer_trusted_tenant_header(mut self, prefer: bool) -> Self { + self.config.tenant_resolution.prefer_trusted_tenant_header = prefer; + self + } + pub fn tenant_header_name>(mut self, header_name: S) -> Self { self.config.tenant_resolution.tenant_header_name = header_name.into(); self diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index 5791d2fb1..50e59868a 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -345,6 +345,7 @@ pub struct RouterConfig { #[serde(default)] pub struct TenantResolutionConfig { pub trust_tenant_header: bool, + pub prefer_trusted_tenant_header: bool, pub tenant_header_name: String, } @@ -369,6 +370,7 @@ impl Default for TenantResolutionConfig { fn default() -> Self { Self { trust_tenant_header: false, + prefer_trusted_tenant_header: false, tenant_header_name: DEFAULT_TENANT_HEADER_NAME.to_string(), } } diff --git a/model_gateway/src/config/validation.rs b/model_gateway/src/config/validation.rs index 793f8cac1..01d08a567 100755 --- a/model_gateway/src/config/validation.rs +++ b/model_gateway/src/config/validation.rs @@ -162,6 +162,15 @@ impl ConfigValidator { } fn validate_tenant_resolution(config: &RouterConfig) -> ConfigResult<()> { + if config.tenant_resolution.prefer_trusted_tenant_header + && !config.tenant_resolution.trust_tenant_header + { + return Err(ConfigError::ValidationFailed { + reason: + "tenant_resolution.prefer_trusted_tenant_header requires trust_tenant_header" + .to_string(), + }); + } let header_name = config.tenant_resolution.tenant_header_name.trim(); if header_name.is_empty() { return Err(ConfigError::ValidationFailed { @@ -1235,6 +1244,18 @@ mod tests { ) } + #[test] + fn preferred_tenant_header_requires_trust() { + let mut config = regular_mode_config(); + config.tenant_resolution.prefer_trusted_tenant_header = true; + + let err = ConfigValidator::validate(&config).unwrap_err(); + assert!(matches!(err, ConfigError::ValidationFailed { .. })); + assert!(err + .to_string() + .contains("prefer_trusted_tenant_header requires trust_tenant_header")); + } + #[test] fn test_validate_distinct_tenant_api_keys_accepted() { let mut config = regular_mode_config(); diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index b91243b5d..e5327ec4d 100755 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -486,6 +486,11 @@ struct CliArgs { #[arg(long, default_value_t = false, help_heading = "Request Handling")] trust_tenant_header: bool, + /// Prefer the trusted tenant header over an authenticated shared proxy + /// identity. Requires --trust-tenant-header. + #[arg(long, default_value_t = false, help_heading = "Request Handling")] + prefer_trusted_tenant_header: bool, + /// Header name to use when --trust-tenant-header is enabled. #[arg( long, @@ -1672,6 +1677,7 @@ impl CliArgs { .then(|| Self::parse_selector(&self.storage_context_headers)), ) .trust_tenant_header(self.trust_tenant_header) + .prefer_trusted_tenant_header(self.prefer_trusted_tenant_header) .tenant_header_name(&self.tenant_header_name) .maybe_rate_limit_tokens_per_second(self.rate_limit_tokens_per_second) .maybe_global_rate_limit_requests_per_second(self.global_rate_limit_requests_per_second) diff --git a/model_gateway/src/middleware/scheduler/admission.rs b/model_gateway/src/middleware/scheduler/admission.rs index 9e265bae6..fe63190e5 100644 --- a/model_gateway/src/middleware/scheduler/admission.rs +++ b/model_gateway/src/middleware/scheduler/admission.rs @@ -13,14 +13,21 @@ use std::sync::{ Arc, }; -use axum::{body::Body, extract::State, http::Request, middleware::Next, response::Response}; +use axum::{ + body::Body, + extract::State, + http::{HeaderMap, Request}, + middleware::Next, + response::Response, +}; use smg_auth::RequestId; use tokio_util::sync::CancellationToken; use tracing::trace; use super::{ - metrics as sched_metrics, state::SchedulerState, AdmitOutcome, Class, RejectionReason, - SchedulerError, SchedulerGuardBody, HEADER_X_SMG_PREEMPTED, PRIORITY_HEADER, + metrics as sched_metrics, state::SchedulerState, AdmitOutcome, Class, GlobalFairShare, + RejectionReason, SchedulerError, SchedulerGuardBody, HEADER_X_SMG_PREEMPTED, + OUTPUT_TOKEN_ESTIMATE_HEADER, PRIORITY_HEADER, }; use crate::{ middleware::{ @@ -89,6 +96,32 @@ fn rejection_outcome(reason: RejectionReason) -> &'static str { } } +fn output_token_estimate(headers: &HeaderMap, ledger: &GlobalFairShare) -> u32 { + let configured_default = ledger.default_output_tokens(); + let header = headers.get(OUTPUT_TOKEN_ESTIMATE_HEADER); + if !ledger.trusts_output_token_estimate_header() { + if header.is_some() { + sched_metrics::record_fair_share_fallback("untrusted_estimate"); + } + return configured_default; + } + match header { + Some(value) => value + .to_str() + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or_else(|| { + sched_metrics::record_fair_share_fallback("invalid_estimate"); + configured_default + }), + None => { + sched_metrics::record_fair_share_fallback("missing_estimate"); + configured_default + } + } +} + pub async fn priority_admission_middleware( State(state): State>, mut req: Request, @@ -134,7 +167,22 @@ pub async fn priority_admission_middleware( // once admitted, releasing the slot. let cancel = CancellationToken::new(); - match partition.scheduler.admit(class, request_id, cancel).await { + let estimated_output_tokens = match partition.scheduler.fair_share() { + Some(ledger) => output_token_estimate(req.headers(), ledger), + None => 1, + }; + + match partition + .scheduler + .admit_for_tenant( + class, + request_id, + cancel, + tenant.clone(), + estimated_output_tokens, + ) + .await + { AdmitOutcome::Admitted(permit) => { pending_guard.resolve(); Metrics::record_http_admission_admitted(); @@ -196,3 +244,34 @@ pub async fn priority_admission_middleware( } } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::middleware::scheduler::FairShareConfig; + + fn ledger(trust_header: bool) -> GlobalFairShare { + GlobalFairShare::from_config(&FairShareConfig { + default_weight: 1.0, + default_output_tokens: 256, + trust_output_token_estimate_header: trust_header, + tenant_weights: HashMap::new(), + }) + } + + #[test] + fn untrusted_client_estimate_cannot_reduce_reservation() { + let mut headers = HeaderMap::new(); + headers.insert(OUTPUT_TOKEN_ESTIMATE_HEADER, "1".parse().unwrap()); + assert_eq!(output_token_estimate(&headers, &ledger(false)), 256); + } + + #[test] + fn trusted_proxy_estimate_is_honored() { + let mut headers = HeaderMap::new(); + headers.insert(OUTPUT_TOKEN_ESTIMATE_HEADER, "128".parse().unwrap()); + assert_eq!(output_token_estimate(&headers, &ledger(true)), 128); + } +} diff --git a/model_gateway/src/middleware/scheduler/body.rs b/model_gateway/src/middleware/scheduler/body.rs index 8abf0feac..b68dda204 100644 --- a/model_gateway/src/middleware/scheduler/body.rs +++ b/model_gateway/src/middleware/scheduler/body.rs @@ -3,6 +3,7 @@ //! when the body finishes draining (or is dropped). use std::{ + collections::VecDeque, pin::Pin, task::{Context, Poll}, }; @@ -11,9 +12,11 @@ use axum::{body::Body, http::StatusCode}; use bytes::Bytes; use http_body::{Body as HttpBody, Frame}; -use super::engine::SchedulerPermit; +use super::{engine::SchedulerPermit, output_tokens::observed_output_tokens, SettlementKind}; use crate::middleware::admission_metrics::AdmissionActiveGuard; +const FAIR_SHARE_RESPONSE_TAIL_LIMIT: usize = 64 * 1024; + /// Wraps a response [`Body`], holding the request's [`SchedulerPermit`] /// for the lifetime of the stream. /// @@ -43,6 +46,8 @@ pub struct SchedulerGuardBody { /// `is_end_stream` consistent with `poll_frame`: the inner body still /// has data, but this wrapper has already signalled `None`. terminated: bool, + response_tail: Option>, + fair_share_settled: bool, } impl SchedulerGuardBody { @@ -53,6 +58,10 @@ impl SchedulerGuardBody { status_code: StatusCode, ) -> Self { let completed = inner.is_end_stream(); + let response_tail = permit + .has_fair_share_reservation() + .then(|| VecDeque::with_capacity(FAIR_SHARE_RESPONSE_TAIL_LIMIT)); + let fair_share_settled = response_tail.is_none(); Self { inner, permit, @@ -61,13 +70,56 @@ impl SchedulerGuardBody { completed, ttft_marked: false, terminated: false, + response_tail, + fair_share_settled, + } + } + + fn append_response_tail(&mut self, data: &[u8]) { + let Some(response_tail) = self.response_tail.as_mut() else { + return; + }; + if data.len() >= FAIR_SHARE_RESPONSE_TAIL_LIMIT { + response_tail.clear(); + response_tail.extend( + data[data.len().saturating_sub(FAIR_SHARE_RESPONSE_TAIL_LIMIT)..] + .iter() + .copied(), + ); + return; + } + let overflow = response_tail + .len() + .saturating_add(data.len()) + .saturating_sub(FAIR_SHARE_RESPONSE_TAIL_LIMIT); + for _ in 0..overflow { + response_tail.pop_front(); + } + response_tail.extend(data.iter().copied()); + } + + fn settle_completed_response(&mut self) { + if self.fair_share_settled { + return; } + let observed = self + .response_tail + .as_mut() + .and_then(|tail| observed_output_tokens(tail.make_contiguous())); + let kind = if observed.is_some() { + SettlementKind::Observed + } else { + SettlementKind::MissingUsage + }; + self.permit.settle_output_tokens(observed, kind); + self.fair_share_settled = true; } } impl Drop for SchedulerGuardBody { fn drop(&mut self) { if self.completed { + self.settle_completed_response(); self.active_guard.record_outcome(self.status_code.as_u16()); } else { self.active_guard.record_interrupted(); @@ -108,8 +160,16 @@ impl http_body::Body for SchedulerGuardBody { } } } + if let Poll::Ready(Some(Ok(frame))) = &polled { + if let Some(data) = frame.data_ref() { + this.append_response_tail(data); + } + } match &polled { - Poll::Ready(None) => this.completed = true, + Poll::Ready(None) => { + this.completed = true; + this.settle_completed_response(); + } Poll::Ready(Some(Err(_))) => this.completed = false, _ => {} } @@ -129,13 +189,22 @@ impl http_body::Body for SchedulerGuardBody { #[cfg(test)] mod tests { + use std::{collections::HashMap, sync::Arc}; + use http_body_util::BodyExt; use smg_auth::RequestId; + use tokio_util::sync::CancellationToken; use super::*; - use crate::middleware::scheduler::{Class, PriorityScheduler, SchedulerSettings}; + use crate::{ + middleware::scheduler::{ + AdmitOutcome, Class, ClassConfig, FairShareConfig, GlobalFairShare, PriorityScheduler, + PrioritySchedulerYaml, SchedulerSettings, + }, + tenant::TenantKey, + }; - fn scheduler() -> std::sync::Arc { + fn scheduler() -> Arc { let settings = SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, None).unwrap(); // Default reservations sum to 160 (Interactive 128 + System 32), @@ -143,7 +212,7 @@ mod tests { PriorityScheduler::new(&settings, 256).unwrap() } - fn permit(sched: &std::sync::Arc, id: &str) -> SchedulerPermit { + fn permit(sched: &Arc, id: &str) -> SchedulerPermit { sched .acquire_inflight(Class::Default, RequestId(id.to_string())) .expect("slot available") @@ -153,11 +222,65 @@ mod tests { SchedulerGuardBody::new(inner, permit, AdmissionActiveGuard::new(), StatusCode::OK) } + fn fair_scheduler( + default_output_tokens: u32, + ) -> (Arc, Arc, TenantKey, u64) { + let mut classes = HashMap::new(); + for class in Class::ALL { + let mut config = ClassConfig::default_for(class); + config.reserved_floor = 0; + config.reserved_per_slot = 0.0; + config.queue_size = 8; + classes.insert(class, config); + } + let yaml = PrioritySchedulerYaml { + classes, + fair_share: Some(FairShareConfig { + default_weight: 1.0, + default_output_tokens, + trust_output_token_estimate_header: false, + tenant_weights: HashMap::from([("header:alice".to_string(), 1.0)]), + }), + ..Default::default() + }; + let settings = + SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&yaml)).unwrap(); + let ledger = Arc::new(GlobalFairShare::from_settings(&settings).unwrap()); + let scheduler = + PriorityScheduler::new_with_fair_share(&settings, 1, Some(Arc::clone(&ledger))) + .unwrap(); + let scope_id = scheduler + .fair_share_scope_for_test() + .expect("fair-share scheduler has a scope"); + (scheduler, ledger, TenantKey::new("header:alice"), scope_id) + } + + async fn fair_permit( + scheduler: &Arc, + tenant: &TenantKey, + id: &str, + estimate: u32, + ) -> SchedulerPermit { + let outcome = scheduler + .admit_for_tenant( + Class::Default, + RequestId(id.to_string()), + CancellationToken::new(), + tenant.clone(), + estimate, + ) + .await; + let AdmitOutcome::Admitted(permit) = outcome else { + panic!("request should be admitted"); + }; + permit + } + #[tokio::test] async fn test_first_data_frame_marks_ttft() { let sched = scheduler(); let p = permit(&sched, "req-ttft"); - let handle = std::sync::Arc::clone(p.handle()); + let handle = Arc::clone(p.handle()); let mut guarded = guarded(Body::from("hello world"), p); assert!(handle.is_preemptible(), "pre-TTFT before first frame"); @@ -176,7 +299,7 @@ mod tests { async fn test_preempted_before_first_byte_ends_stream() { let sched = scheduler(); let p = permit(&sched, "req-preempt"); - let handle = std::sync::Arc::clone(p.handle()); + let handle = Arc::clone(p.handle()); // Scheduler wins the preempt CAS before any byte is polled. assert!(handle.try_mark_preempted()); @@ -203,7 +326,7 @@ mod tests { async fn test_trailer_only_response_never_marks_ttft() { let sched = scheduler(); let p = permit(&sched, "req-empty"); - let handle = std::sync::Arc::clone(p.handle()); + let handle = Arc::clone(p.handle()); // Empty body: no data frames at all. let guarded = guarded(Body::empty(), p); let _ = guarded.collect().await; // exhaust @@ -227,4 +350,131 @@ mod tests { "dropping the guarded body releases the slot" ); } + + #[tokio::test] + async fn disabled_mode_never_allocates_a_response_tail() { + let sched = scheduler(); + let p = permit(&sched, "disabled-tail"); + let mut guarded = guarded(Body::from("hello"), p); + assert!(guarded.response_tail.is_none()); + assert!(guarded.frame().await.unwrap().unwrap().is_data()); + assert!(guarded.response_tail.is_none()); + } + + #[tokio::test] + async fn terminal_usage_settles_actual_output_tokens() { + let (scheduler, ledger, tenant, scope_id) = fair_scheduler(100); + let permit = fair_permit(&scheduler, &tenant, "usage", 100).await; + let guarded = SchedulerGuardBody::new( + Body::from(r#"{"usage":{"completion_tokens":37}}"#), + permit, + AdmissionActiveGuard::new(), + StatusCode::OK, + ); + guarded.collect().await.unwrap(); + + assert_eq!(ledger.snapshot(scope_id, &tenant).0, 37); + } + + #[tokio::test] + async fn completed_sse_usage_settles_actual_output_tokens() { + let (scheduler, ledger, tenant, scope_id) = fair_scheduler(100); + let permit = fair_permit(&scheduler, &tenant, "sse-usage", 100).await; + let body = Body::from( + "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n\ + data: {\"choices\":[],\"usage\":{\"completion_tokens\":19}}\n\n\ + data: [DONE]\n\n", + ); + guarded(body, permit).collect().await.unwrap(); + assert_eq!(ledger.snapshot(scope_id, &tenant).0, 19); + } + + #[tokio::test] + async fn split_sse_frames_are_reassembled_for_terminal_usage() { + let (scheduler, ledger, tenant, scope_id) = fair_scheduler(100); + let permit = fair_permit(&scheduler, &tenant, "split-sse", 100).await; + let frames: Vec> = [ + "data: {\"choices\":[]", + ",\"usage\":{\"completion_", + "tokens\":23}}\n\n", + "data: [DONE]\n\n", + ] + .into_iter() + .map(|chunk| Ok(Bytes::from(chunk))) + .collect(); + let body = Body::from_stream(tokio_stream::iter(frames)); + guarded(body, permit).collect().await.unwrap(); + assert_eq!(ledger.snapshot(scope_id, &tenant).0, 23); + } + + #[tokio::test] + async fn client_drop_before_usage_keeps_the_provisional_charge() { + let (scheduler, ledger, tenant, scope_id) = fair_scheduler(100); + let permit = fair_permit(&scheduler, &tenant, "client-drop", 100).await; + let mut body = guarded(Body::from("data: {\"choices\":[]}\n\n"), permit); + assert!(body.frame().await.unwrap().unwrap().is_data()); + drop(body); + assert_eq!(ledger.snapshot(scope_id, &tenant).0, 100); + } + + #[tokio::test] + async fn backend_error_keeps_the_provisional_charge() { + let (scheduler, ledger, tenant, scope_id) = fair_scheduler(100); + let permit = fair_permit(&scheduler, &tenant, "backend-error", 100).await; + let frames = vec![ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: partial\n\n")), + Err(std::io::Error::other("backend failed")), + ]; + let mut body = guarded(Body::from_stream(tokio_stream::iter(frames)), permit); + assert!(body.frame().await.unwrap().unwrap().is_data()); + assert!(body.frame().await.unwrap().is_err()); + drop(body); + assert_eq!(ledger.snapshot(scope_id, &tenant).0, 100); + } + + #[tokio::test] + async fn permit_settlement_and_drop_are_idempotent() { + let (scheduler, ledger, tenant, scope_id) = fair_scheduler(100); + let mut permit = fair_permit(&scheduler, &tenant, "double-settle", 100).await; + permit.settle_output_tokens(Some(31), SettlementKind::Observed); + permit.settle_output_tokens(Some(999), SettlementKind::Observed); + drop(permit); + assert_eq!(ledger.snapshot(scope_id, &tenant).0, 31); + } + + #[tokio::test] + async fn many_small_frames_stay_bounded_and_parse_usage_after_ring_wrap() { + let (scheduler, ledger, tenant, scope_id) = fair_scheduler(100); + let permit = fair_permit(&scheduler, &tenant, "ring-wrap", 100).await; + let mut frames: Vec> = (0..9_000) + .map(|_| Ok(Bytes::from_static(b"data: filler\n\n"))) + .collect(); + for chunk in [ + "data: {\"choices\":[],\"usage\":{\"comple", + "tion_tokens\":47}}\n\n", + "data: [DONE]\n\n", + ] { + frames.push(Ok(Bytes::from(chunk))); + } + + let mut body = guarded(Body::from_stream(tokio_stream::iter(frames)), permit); + let mut frame_count = 0; + while let Some(frame) = body.frame().await { + frame.unwrap(); + frame_count += 1; + assert!( + body.response_tail + .as_ref() + .is_some_and(|tail| tail.len() <= FAIR_SHARE_RESPONSE_TAIL_LIMIT), + "ring buffer must stay within its fixed cap" + ); + } + assert_eq!(frame_count, 9_003); + assert_eq!( + body.response_tail.as_ref().map(VecDeque::len), + Some(FAIR_SHARE_RESPONSE_TAIL_LIMIT) + ); + drop(body); + assert_eq!(ledger.snapshot(scope_id, &tenant).0, 47); + } } diff --git a/model_gateway/src/middleware/scheduler/config.rs b/model_gateway/src/middleware/scheduler/config.rs index 0e6820477..661b6af68 100644 --- a/model_gateway/src/middleware/scheduler/config.rs +++ b/model_gateway/src/middleware/scheduler/config.rs @@ -119,6 +119,36 @@ pub struct TenantPolicyConfig { pub max_class: Class, } +fn default_fair_share_weight() -> f64 { + 1.0 +} + +fn default_fair_share_output_tokens() -> u32 { + 256 +} + +/// Process-wide weighted sharing with per-partition eligibility boundaries. +/// +/// Weights are relative and need not sum to 100. For example, weights 10 and +/// 5 give two continuously contending tenants a 2:1 output-token share. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct FairShareConfig { + /// Weight for a resolved tenant absent from `tenant_weights`. + #[serde(default = "default_fair_share_weight")] + pub default_weight: f64, + /// Provisional charge when the trusted request estimate is absent or + /// invalid. Terminal response usage replaces this estimate. + #[serde(default = "default_fair_share_output_tokens")] + pub default_output_tokens: u32, + /// Honor `x-smg-output-token-estimate`. Keep false unless a trusted proxy + /// strips client copies and injects a validated value. + #[serde(default)] + pub trust_output_token_estimate_header: bool, + /// Per-tenant relative weights keyed by canonical `TenantKey` string. + #[serde(default)] + pub tenant_weights: HashMap, +} + /// Admission budget for one trusted upstream partition selector. /// /// A deployment chooses one capacity mode for every configured partition: @@ -168,6 +198,10 @@ pub struct PrioritySchedulerYaml { pub classes: HashMap, #[serde(default)] pub tenant_policies: HashMap, + /// Optional global weighted output-token ledger. Absence preserves FIFO + /// queueing and all existing scheduler behavior. + #[serde(default)] + pub fair_share: Option, /// Optional admission partitions keyed by the exact value of the trusted /// `x-smg-admission-partition` header. An empty map preserves the original /// single global scheduler. @@ -184,6 +218,7 @@ impl Default for PrioritySchedulerYaml { Self { classes: HashMap::new(), tenant_policies: HashMap::new(), + fair_share: None, admission_partitions: HashMap::new(), default_admission_partition: default_admission_partition(), } @@ -203,6 +238,12 @@ pub enum SettingsValidationError { ZeroStarvationThreshold { class: Class }, #[error("class {class:?}: reserved_per_slot must be finite and >= 0")] InvalidReservedPerSlot { class: Class }, + #[error("fair_share.default_weight must be finite and > 0")] + InvalidFairShareDefaultWeight, + #[error("fair_share.default_output_tokens must be > 0")] + ZeroFairShareDefaultOutputTokens, + #[error("fair_share.tenant_weights[{tenant:?}] must be finite and > 0")] + InvalidFairShareTenantWeight { tenant: String }, } /// Runtime scheduler configuration assembled from CLI flags + the @@ -227,9 +268,10 @@ pub struct SchedulerSettings { /// Per-tenant clamp lookup. Keys come from /// [`crate::tenant::RouteRequestMeta::tenant_key`]. pub tenant_policies: HashMap, + fair_share: Option, /// Cap on the number of tenants emitted as labels for - /// `scheduler_tenant_*` gauges. Everything past the cap is - /// bucketed under `tenant="other"`. + /// tenant-labeled scheduler and fair-share metrics. Everything past the + /// cap is bucketed under `tenant="other"`. pub tenant_metric_top_n: u32, } @@ -239,6 +281,11 @@ impl SchedulerSettings { &self.classes[class as usize] } + #[must_use] + pub fn fair_share_config(&self) -> Option<&FairShareConfig> { + self.fair_share.as_ref() + } + /// Scale the built-in per-class queue weights to one exact global budget. /// This is used when no scheduler YAML overrides the class limits, so the /// legacy `--queue-size` flag remains the total queue contract after the @@ -329,6 +376,23 @@ impl SchedulerSettings { } } + let fair_share = yaml.and_then(|value| value.fair_share.clone()); + if let Some(config) = &fair_share { + if !config.default_weight.is_finite() || config.default_weight <= 0.0 { + return Err(SettingsValidationError::InvalidFairShareDefaultWeight); + } + if config.default_output_tokens == 0 { + return Err(SettingsValidationError::ZeroFairShareDefaultOutputTokens); + } + for (tenant, weight) in &config.tenant_weights { + if !weight.is_finite() || *weight <= 0.0 { + return Err(SettingsValidationError::InvalidFairShareTenantWeight { + tenant: tenant.clone(), + }); + } + } + } + let tenant_policies = yaml .map(|y| { y.tenant_policies @@ -343,6 +407,7 @@ impl SchedulerSettings { default_max_class, classes, tenant_policies, + fair_share, tenant_metric_top_n, }) } @@ -465,6 +530,47 @@ tenant_policies: ); } + #[test] + fn test_yaml_fair_share_weights_round_trip() { + let yaml = r#" +fair_share: + default_weight: 0.5 + default_output_tokens: 128 + trust_output_token_estimate_header: true + tenant_weights: + "header:alice": 10 + "header:bob": 5 +"#; + let parsed: PrioritySchedulerYaml = serde_yaml::from_str(yaml).unwrap(); + let fair_share = parsed.fair_share.as_ref().unwrap(); + assert_eq!(fair_share.default_weight, 0.5); + assert_eq!(fair_share.default_output_tokens, 128); + assert!(fair_share.trust_output_token_estimate_header); + assert_eq!(fair_share.tenant_weights["header:alice"], 10.0); + assert_eq!(fair_share.tenant_weights["header:bob"], 5.0); + + let settings = + SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&parsed)).unwrap(); + assert_eq!(settings.fair_share_config(), Some(fair_share)); + } + + #[test] + fn test_invalid_fair_share_weights_are_rejected() { + for weight in [0.0, -1.0, f64::INFINITY, f64::NAN] { + let mut yaml = PrioritySchedulerYaml::default(); + yaml.fair_share = Some(FairShareConfig { + default_weight: 1.0, + default_output_tokens: 128, + trust_output_token_estimate_header: false, + tenant_weights: HashMap::from([("header:alice".to_string(), weight)]), + }); + assert!(matches!( + SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&yaml)), + Err(SettingsValidationError::InvalidFairShareTenantWeight { .. }) + )); + } + } + #[test] fn test_yaml_admission_partitions_round_trip() { let yaml = r#" diff --git a/model_gateway/src/middleware/scheduler/engine.rs b/model_gateway/src/middleware/scheduler/engine.rs index ff3cbae1a..74ac06f40 100644 --- a/model_gateway/src/middleware/scheduler/engine.rs +++ b/model_gateway/src/middleware/scheduler/engine.rs @@ -23,14 +23,15 @@ use tokio_util::sync::CancellationToken; use tracing::{info, warn}; use super::{ + fair_share::{FairShareReservation, GlobalFairShare, SettlementKind}, inflight::InflightHandle, - queue::{ClassQueue, FifoClassQueue, QueueBudget, Waiter}, + queue::{ClassQueue, FairClassQueue, FifoClassQueue, QueueBudget, Waiter}, slots::SlotPool, Class, ClassRuntimeConfig, SchedulerSettings, }; use crate::{ middleware::admission_metrics::AdmissionQueuedGuard, observability::metrics::Metrics, - worker::WorkerCapacity, + tenant::TenantKey, worker::WorkerCapacity, }; /// Max time to wait, after firing a preemption cancel, for the victim's slot @@ -99,6 +100,9 @@ pub enum RejectionReason { pub struct PriorityScheduler { slot_pool: SlotPool, class_queues: [Arc; 4], + fair_share: Option>, + #[cfg(test)] + fair_share_scope: Option, /// One work-conserving occupancy ceiling shared by all class queues. queue_budget: Arc, inflight_registry: RwLock>>, @@ -150,13 +154,21 @@ impl PriorityScheduler { /// when they do not fit. The same rule is used for runtime capacity dips, /// so a small admission partition has identical startup and drain /// behavior. + pub fn new( + settings: &SchedulerSettings, + capacity: u16, + ) -> Result, SchedulerInitError> { + Self::new_with_fair_share(settings, capacity, None) + } + #[expect( clippy::unnecessary_wraps, - reason = "preserve the public constructor API while startup reservation overflow becomes a safe clamp" + reason = "preserve the approved constructor seam for compatibility with existing callers" )] - pub fn new( + pub fn new_with_fair_share( settings: &SchedulerSettings, capacity: u16, + fair_share: Option>, ) -> Result, SchedulerInitError> { let reserved_floor = Class::ALL.map(|c| settings.class_config(c).reserved_floor); let reserved_per_slot = Class::ALL.map(|c| settings.class_config(c).reserved_per_slot); @@ -178,14 +190,24 @@ impl PriorityScheduler { .map(|class| settings.class_config(*class).queue_size as usize) .fold(0_usize, usize::saturating_add); let queue_budget = Arc::new(QueueBudget::new(total_queue_capacity)); - let class_queues: [Arc; 4] = - Class::ALL.map(|c| queue_for(settings, c, Arc::clone(&queue_budget))); + let fair_share_scope = fair_share.as_ref().map(|ledger| ledger.new_scope()); + let class_queues: [Arc; 4] = Class::ALL.map(|c| { + queue_for( + settings, + c, + Arc::clone(&queue_budget), + fair_share.as_ref().zip(fair_share_scope), + ) + }); let class_config: [ClassRuntimeConfig; 4] = Class::ALL.map(|c| ClassRuntimeConfig::from_class_config(settings.class_config(c))); - Ok(Arc::new(Self { + let scheduler = Arc::new(Self { slot_pool: SlotPool::new(capacity, effective_reserved), class_queues, + fair_share, + #[cfg(test)] + fair_share_scope, queue_budget, inflight_registry: RwLock::new(HashMap::new()), release_notify: Arc::new(Notify::new()), @@ -198,7 +220,8 @@ impl PriorityScheduler { sampled_releases: 0, ewma_per_second: None, }), - })) + }); + Ok(scheduler) } /// Try to acquire a slot under `class` for the given request id. @@ -214,12 +237,17 @@ impl PriorityScheduler { if !self.slot_pool.try_acquire(class) { return None; } - Some(self.register_inflight(class, request_id)) + Some(self.register_inflight(class, request_id, None)) } /// Register a handle in the registry and wrap it in a permit. /// Caller already acquired a slot via the pool. - fn register_inflight(self: &Arc, class: Class, request_id: RequestId) -> SchedulerPermit { + fn register_inflight( + self: &Arc, + class: Class, + request_id: RequestId, + fair_share_reservation: Option, + ) -> SchedulerPermit { let handle = Arc::new(InflightHandle::new(class, request_id)); self.inflight_registry .write() @@ -227,9 +255,20 @@ impl PriorityScheduler { SchedulerPermit { scheduler: Arc::clone(self), handle, + fair_share_reservation, } } + #[must_use] + pub fn fair_share(&self) -> Option<&Arc> { + self.fair_share.as_ref() + } + + #[cfg(test)] + pub(crate) fn fair_share_scope_for_test(&self) -> Option { + self.fair_share_scope + } + /// Admit a request under `class`. Tries the fast path first; if the /// slot pool refuses, enqueues a waiter and awaits the dispatcher /// (or a queue timeout, or a client-side cancel). @@ -359,6 +398,83 @@ impl PriorityScheduler { outcome } + /// Admit one request using the shared output-token ledger when enabled. + /// + /// Fair-share requests enter the queue even when a slot is currently + /// free, so all locally eligible contenders are visible to the weighted + /// selector. The queue remains work-conserving: a free partition slot + /// always dispatches some local waiter. + pub async fn admit_for_tenant( + self: &Arc, + class: Class, + request_id: RequestId, + cancel: CancellationToken, + tenant: TenantKey, + estimated_output_tokens: u32, + ) -> AdmitOutcome { + if self.fair_share.is_none() { + return self.admit(class, request_id, cancel).await; + } + if cancel.is_cancelled() { + return AdmitOutcome::Rejected(RejectionReason::ClientCancelled); + } + + let (tx, rx) = oneshot::channel::(); + let waiter_cancel = cancel.child_token(); + let waiter = Waiter::new_fair( + class, + waiter_cancel.clone(), + request_id, + tx, + tenant, + estimated_output_tokens, + ); + if self.class_queues[class as usize] + .try_enqueue(waiter) + .is_err() + { + return AdmitOutcome::Rejected(RejectionReason::QueueFull); + } + let _queued_guard = AdmissionQueuedGuard::new(); + let enqueued_at = Instant::now(); + + // Drain synchronously once before yielding. This preserves the normal + // free-slot fast path while still making the request visible to the + // fair queue. If capacity is blocked, priority preemption only frees + // a slot; the fair queue still decides which eligible tenant receives + // it. + let made_progress = self.wake_next_waiter(); + if !made_progress && self.class_config[class as usize].can_preempt { + if let Some(victim) = self.find_preemption_victim(class) { + if victim.try_mark_preempted() { + info!( + victim_id = %victim.request_id().0, + victim_class = ?victim.class(), + preemptor_class = ?class, + "scheduler: preempting pre-TTFT request for fair-share queue" + ); + victim.cancel(); + super::metrics::record_preemption(victim.class(), class); + } + } + } + self.release_notify.notify_one(); + + let timeout = self.class_config[class as usize].queue_timeout; + let outcome = tokio::select! { + result = rx => match result { + Ok(permit) => AdmitOutcome::Admitted(permit), + Err(_) => AdmitOutcome::Rejected(RejectionReason::ClientCancelled), + }, + () = tokio::time::sleep(timeout) => AdmitOutcome::Rejected(RejectionReason::QueueTimeout), + () = cancel.cancelled() => AdmitOutcome::Rejected(RejectionReason::ClientCancelled), + }; + super::metrics::record_queue_wait(class, enqueued_at.elapsed()); + waiter_cancel.cancel(); + self.release_notify.notify_one(); + outcome + } + /// Remove a handle from the registry, release its slot, and notify /// the dispatcher. Called from [`SchedulerPermit`]'s `Drop`. fn release_inflight(&self, handle: &InflightHandle) { @@ -463,7 +579,7 @@ impl PriorityScheduler { // (this loop runs up to budget/poll-interval times per // preemption). if self.slot_pool.try_acquire(class) { - return Some(self.register_inflight(class, request_id)); + return Some(self.register_inflight(class, request_id, None)); } if Instant::now() >= deadline { return None; @@ -550,6 +666,7 @@ impl PriorityScheduler { let Some(Waiter { request_id, permit_tx, + fair_share_reservation, .. }) = self.class_queues[class as usize].pop_eligible() else { @@ -559,14 +676,19 @@ impl PriorityScheduler { // Receiver gone — skip the registry write and the // matched permit-drop release. Try the next waiter // using the same slot we already acquired. + if let Some(reservation) = fair_share_reservation { + reservation.cancel(); + } continue; } - let permit = self.register_inflight(class, request_id); + let permit = self.register_inflight(class, request_id, fair_share_reservation); // If the receiver was dropped between is_closed() above and - // send below (unlikely race window), the permit goes out of - // scope and its Drop releases the slot back; the caller's - // outer loop will try again. - let _ = permit_tx.send(permit); + // send below (unlikely race window), cancel the provisional + // charge because no backend work began, then let permit Drop + // release the slot. The caller's outer loop will try again. + if let Err(mut undelivered) = permit_tx.send(permit) { + undelivered.cancel_fair_share_reservation(); + } return true; } } @@ -888,11 +1010,19 @@ fn queue_for( settings: &SchedulerSettings, class: Class, queue_budget: Arc, + fair_share: Option<(&Arc, u64)>, ) -> Arc { - Arc::new(FifoClassQueue::with_shared_budget( - settings.class_config(class).queue_size as usize, - queue_budget, - )) + let soft_limit = settings.class_config(class).queue_size as usize; + match fair_share { + Some((ledger, scope_id)) => Arc::new(FairClassQueue::with_shared_budget( + class, + soft_limit, + queue_budget, + Arc::clone(ledger), + scope_id, + )), + None => Arc::new(FifoClassQueue::with_shared_budget(soft_limit, queue_budget)), + } } /// RAII handle on one admitted request. Holding a permit keeps the slot @@ -901,6 +1031,7 @@ fn queue_for( pub struct SchedulerPermit { scheduler: Arc, handle: Arc, + fair_share_reservation: Option, } impl SchedulerPermit { @@ -910,6 +1041,11 @@ impl SchedulerPermit { &self.handle } + #[must_use] + pub fn has_fair_share_reservation(&self) -> bool { + self.fair_share_reservation.is_some() + } + /// Mark the first response byte. Called by [`super::body::SchedulerGuardBody`] /// on the first data frame. Returns `false` if the scheduler already /// won the preemption CAS — the body wrapper treats that as @@ -927,6 +1063,24 @@ impl SchedulerPermit { pub fn cancel_token(&self) -> CancellationToken { self.handle.cancel_token() } + + /// Replace the provisional fair-share charge with terminal observed + /// output tokens. Calling this more than once is a no-op. + pub fn settle_output_tokens( + &mut self, + observed_output_tokens: Option, + kind: SettlementKind, + ) { + if let Some(reservation) = self.fair_share_reservation.take() { + reservation.settle(observed_output_tokens, kind); + } + } + + fn cancel_fair_share_reservation(&mut self) { + if let Some(reservation) = self.fair_share_reservation.take() { + reservation.cancel(); + } + } } impl std::fmt::Debug for SchedulerPermit { @@ -942,13 +1096,16 @@ impl std::fmt::Debug for SchedulerPermit { impl Drop for SchedulerPermit { fn drop(&mut self) { + if let Some(reservation) = self.fair_share_reservation.take() { + reservation.settle(None, SettlementKind::Interrupted); + } self.scheduler.release_inflight(&self.handle); } } #[cfg(test)] mod tests { - use std::{sync::Arc, time::Duration}; + use std::{collections::HashMap, sync::Arc, time::Duration}; use super::*; use crate::middleware::scheduler::{ClassConfig, PrioritySchedulerYaml}; @@ -1885,4 +2042,124 @@ mod tests { "can_preempt=false class must not cancel anyone" ); } + + fn fair_settings() -> SchedulerSettings { + let mut classes = HashMap::new(); + for class in Class::ALL { + let mut config = ClassConfig::default_for(class); + config.reserved_floor = 0; + config.reserved_per_slot = 0.0; + config.queue_size = 16; + classes.insert(class, config); + } + let yaml = PrioritySchedulerYaml { + classes, + fair_share: Some(crate::middleware::scheduler::FairShareConfig { + default_weight: 1.0, + default_output_tokens: 10, + trust_output_token_estimate_header: false, + tenant_weights: HashMap::from([ + ("header:a".to_string(), 1.0), + ("header:b".to_string(), 1.0), + ]), + }), + ..Default::default() + }; + SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&yaml)).unwrap() + } + + #[tokio::test] + async fn fair_share_free_slot_is_work_conserving_for_lone_user() { + let settings = fair_settings(); + let ledger = Arc::new(GlobalFairShare::from_settings(&settings).unwrap()); + let scheduler = PriorityScheduler::new_with_fair_share(&settings, 1, Some(ledger)).unwrap(); + + let outcome = scheduler + .admit_for_tenant( + Class::Default, + rid("a-1"), + CancellationToken::new(), + TenantKey::new("header:a"), + 10, + ) + .await; + let AdmitOutcome::Admitted(mut permit) = outcome else { + panic!("a lone eligible waiter must use the free partition slot"); + }; + permit.settle_output_tokens(Some(10), SettlementKind::Observed); + drop(permit); + assert_eq!(scheduler.inflight_for_test(Class::Default), 0); + } + + #[test] + fn fair_share_underserved_local_user_wins_contention() { + let settings = fair_settings(); + let ledger = Arc::new(GlobalFairShare::from_settings(&settings).unwrap()); + let scheduler = + PriorityScheduler::new_with_fair_share(&settings, 1, Some(Arc::clone(&ledger))) + .unwrap(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + let scope_id = scheduler + .fair_share_scope_for_test() + .expect("fair-share scheduler has a scope"); + let held = scheduler + .acquire_inflight(Class::Default, rid("held")) + .unwrap(); + let (a_tx, mut a_rx) = oneshot::channel(); + let (b_tx, mut b_rx) = oneshot::channel(); + scheduler.class_queues[Class::Default as usize] + .try_enqueue(Waiter::new_fair( + Class::Default, + CancellationToken::new(), + rid("b-queued"), + b_tx, + b, + 10, + )) + .unwrap(); + + let reservation = { + // b is already backlogged for this partition while a receives + // service, so a accumulates debt only under real contention. + ledger.register_waiter(scope_id, &a, Class::Default); + ledger + .reserve_local_candidate( + scope_id, + Class::Default, + &[ + crate::middleware::scheduler::fair_share::FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 50, + }, + ], + ) + .unwrap() + .reservation + }; + reservation.settle(Some(50), SettlementKind::Observed); + + scheduler.class_queues[Class::Default as usize] + .try_enqueue(Waiter::new_fair( + Class::Default, + CancellationToken::new(), + rid("a-queued"), + a_tx, + a, + 10, + )) + .unwrap(); + + drop(held); + assert!(scheduler.wake_next_waiter()); + let permit = b_rx.try_recv().expect("underserved local user admitted"); + assert!(matches!( + a_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + drop(permit); + assert!(scheduler.wake_next_waiter()); + drop(a_rx.try_recv().expect("remaining local waiter admitted")); + } } diff --git a/model_gateway/src/middleware/scheduler/fair_share.rs b/model_gateway/src/middleware/scheduler/fair_share.rs new file mode 100644 index 000000000..aaa450851 --- /dev/null +++ b/model_gateway/src/middleware/scheduler/fair_share.rs @@ -0,0 +1,894 @@ +//! Process-wide output-token accounting and weighted service with local queues. +//! +//! Every priority-scheduler partition receives the same [`Arc`]. +//! Lifetime charged-token accounting and canonical tenant virtual finish span +//! that process. Each partition stores only local queued/reservation membership. +//! Queued work is registered before a partition asks for its next local waiter, +//! and the ledger chooses the globally least-served tenant among candidates +//! eligible for that partition. This keeps every model pool work-conserving: +//! unrelated or non-fungible capacity is never idled to repay another tenant's +//! debt. +//! +//! The ledger is process-local. It spans all model/admission partitions and +//! all workload types that resolve to the same tenant key in one SMG process; +//! it does not coordinate separate gateways or overlapping blue/green +//! processes. Because model pools are non-fungible, it also cannot guarantee +//! exact aggregate percentages when users target disjoint pools. It enforces +//! the configured ratios whenever weighted tenants contend for substitutable +//! eligible capacity. Service and contention debt can follow a tenant across +//! pools, but cannot force a disjoint pool to idle. + +use std::{ + collections::{HashMap, HashSet}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }, + time::Duration, +}; + +use parking_lot::Mutex; + +use super::{Class, FairShareConfig, SchedulerSettings}; +use crate::tenant::TenantKey; + +/// Trusted estimate injected by the authenticated Comet proxy. +pub const OUTPUT_TOKEN_ESTIMATE_HEADER: &str = "x-smg-output-token-estimate"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SettlementKind { + Observed, + MissingUsage, + Interrupted, +} + +#[derive(Debug, Default)] +struct TenantAccounting { + charged_output_tokens: u64, + reserved_output_tokens: u64, +} + +#[derive(Debug, Default)] +struct TenantService { + active_reservations: u64, + queued: usize, + virtual_finish: f64, +} + +#[derive(Debug, Default)] +struct ScopeTenant { + active_reservations: u64, + queued: [usize; 4], +} + +#[derive(Debug, Default)] +struct ScopeLedger { + tenants: HashMap, +} + +#[derive(Debug, Default)] +struct LedgerState { + accounting: HashMap, + service: HashMap, + scopes: HashMap, + system_virtual_time: f64, +} + +/// One local candidate offered by a partition queue. +pub(crate) struct FairShareCandidate<'a> { + pub index: usize, + pub tenant: &'a TenantKey, + pub estimated_output_tokens: u32, +} + +/// The candidate selected by the shared ledger, plus its provisional charge. +pub(crate) struct FairShareSelection { + pub index: usize, + pub reservation: FairShareReservation, +} + +/// Shared weighted-service ledger for every scheduler partition in one router. +pub struct GlobalFairShare { + default_weight: f64, + default_output_tokens: u32, + trust_output_token_estimate_header: bool, + tenant_weights: HashMap, + metric_tenants: HashSet, + state: Mutex, + next_scope: AtomicU64, +} + +impl std::fmt::Debug for GlobalFairShare { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GlobalFairShare") + .field("default_weight", &self.default_weight) + .field("default_output_tokens", &self.default_output_tokens) + .field( + "trust_output_token_estimate_header", + &self.trust_output_token_estimate_header, + ) + .field("tenant_weights", &self.tenant_weights) + .finish_non_exhaustive() + } +} + +impl GlobalFairShare { + #[must_use] + pub fn from_settings(settings: &SchedulerSettings) -> Option { + settings.fair_share_config().map(|config| { + let mut ledger = Self::from_config(config); + let mut metric_tenants: Vec<_> = ledger.tenant_weights.keys().cloned().collect(); + metric_tenants.sort_by(|left, right| left.as_str().cmp(right.as_str())); + metric_tenants.truncate(settings.tenant_metric_top_n as usize); + ledger.metric_tenants = metric_tenants.into_iter().collect(); + ledger + }) + } + + #[must_use] + pub fn from_config(config: &FairShareConfig) -> Self { + let tenant_weights: HashMap<_, _> = config + .tenant_weights + .iter() + .map(|(tenant, weight)| (TenantKey::new(tenant), *weight)) + .collect(); + let metric_tenants = tenant_weights.keys().cloned().collect(); + Self { + default_weight: config.default_weight, + default_output_tokens: config.default_output_tokens, + trust_output_token_estimate_header: config.trust_output_token_estimate_header, + tenant_weights, + metric_tenants, + state: Mutex::new(LedgerState::default()), + next_scope: AtomicU64::new(0), + } + } + + #[must_use] + pub fn default_output_tokens(&self) -> u32 { + self.default_output_tokens + } + + #[must_use] + pub fn trusts_output_token_estimate_header(&self) -> bool { + self.trust_output_token_estimate_header + } + + pub(crate) fn new_scope(&self) -> u64 { + self.next_scope.fetch_add(1, Ordering::Relaxed) + } + + pub(crate) fn record_queue_wait(&self, tenant: &TenantKey, class: Class, wait: Duration) { + super::metrics::record_fair_share_queue_wait(self.metric_tenant(tenant), class, wait); + } + + pub(crate) fn register_waiter(&self, scope_id: u64, tenant: &TenantKey, class: Class) { + let unknown = !self.tenant_weights.contains_key(tenant); + let mut state = self.state.lock(); + let system_virtual_time = state.system_virtual_time; + let service = state.service.entry(tenant.clone()).or_default(); + if !Self::is_active(service) { + // Start-time fair queueing: idle tenants join the current epoch. + // They neither retain stale low credit nor inherit solo service as + // debt against users that were not backlogged at the time. + service.virtual_finish = service.virtual_finish.max(system_virtual_time); + } + service.queued = service.queued.saturating_add(1); + state + .scopes + .entry(scope_id) + .or_default() + .tenants + .entry(tenant.clone()) + .or_default() + .queued[class as usize] += 1; + Self::advance_system_virtual_time(&mut state); + drop(state); + if unknown { + super::metrics::record_fair_share_unknown_tenant(self.metric_tenant(tenant)); + } + } + + pub(crate) fn remove_waiter(&self, scope_id: u64, tenant: &TenantKey, class: Class) { + let mut state = self.state.lock(); + { + let Some(scope) = state.scopes.get_mut(&scope_id) else { + return; + }; + let Some(membership) = scope.tenants.get_mut(tenant) else { + return; + }; + let queued = &mut membership.queued[class as usize]; + debug_assert!(*queued > 0, "fair-share waiter removed below zero"); + if *queued == 0 { + return; + } + *queued -= 1; + } + if let Some(service) = state.service.get_mut(tenant) { + service.queued = service.queued.saturating_sub(1); + } + Self::advance_system_virtual_time(&mut state); + drop(state); + } + + /// Reserve the least-served candidate eligible for this local partition. + /// + /// Selection never consults tenants that are queued only in another + /// partition. That is the work-conserving boundary for non-fungible model + /// pools: a local slot is never idled for work that cannot use it. + pub(crate) fn reserve_local_candidate( + self: &Arc, + scope_id: u64, + class: Class, + candidates: &[FairShareCandidate<'_>], + ) -> Option { + let mut state = self.state.lock(); + let class_index = class as usize; + let selected = { + let scope = state.scopes.get(&scope_id)?; + candidates + .iter() + .filter_map(|candidate| { + let membership = scope.tenants.get(candidate.tenant)?; + if membership.queued[class_index] == 0 { + return None; + } + let service = state.service.get(candidate.tenant)?; + Some((candidate, service.virtual_finish)) + }) + .min_by(|(left, left_service), (right, right_service)| { + left_service + .total_cmp(right_service) + .then_with(|| left.tenant.as_str().cmp(right.tenant.as_str())) + .then_with(|| left.index.cmp(&right.index)) + })? + .0 + }; + + let tenant = selected.tenant.clone(); + let estimated_output_tokens = selected.estimated_output_tokens.max(1); + let membership = state.scopes.get_mut(&scope_id)?.tenants.get_mut(&tenant)?; + debug_assert!(membership.queued[class_index] > 0); + membership.queued[class_index] = membership.queued[class_index].saturating_sub(1); + membership.active_reservations = membership.active_reservations.saturating_add(1); + let service = state.service.get_mut(&tenant)?; + service.queued = service.queued.saturating_sub(1); + service.active_reservations = service.active_reservations.saturating_add(1); + service.virtual_finish += f64::from(estimated_output_tokens) / self.weight(&tenant); + let virtual_service = service.virtual_finish; + Self::advance_system_virtual_time(&mut state); + let accounting = state.accounting.entry(tenant.clone()).or_default(); + accounting.reserved_output_tokens = accounting + .reserved_output_tokens + .saturating_add(u64::from(estimated_output_tokens)); + let reserved_output_tokens = accounting.reserved_output_tokens; + drop(state); + + super::metrics::set_fair_share_virtual_finish(self.metric_tenant(&tenant), virtual_service); + super::metrics::set_fair_share_reserved_output_tokens( + self.metric_tenant(&tenant), + reserved_output_tokens, + ); + + Some(FairShareSelection { + index: selected.index, + reservation: FairShareReservation { + ledger: Arc::clone(self), + tenant, + scope_id, + estimated_output_tokens, + settled: false, + }, + }) + } + + fn weight(&self, tenant: &TenantKey) -> f64 { + self.tenant_weights + .get(tenant) + .copied() + .unwrap_or(self.default_weight) + } + + fn is_active(service: &TenantService) -> bool { + service.active_reservations > 0 || service.queued > 0 + } + + fn advance_system_virtual_time(state: &mut LedgerState) { + if let Some(active_minimum) = state + .service + .values() + .filter(|service| Self::is_active(service)) + .map(|service| service.virtual_finish) + .min_by(f64::total_cmp) + { + state.system_virtual_time = state.system_virtual_time.max(active_minimum); + } + } + + fn metric_tenant<'a>(&'a self, tenant: &'a TenantKey) -> &'a str { + if self.metric_tenants.contains(tenant) { + tenant.as_str() + } else { + "other" + } + } + + fn cancel_reservation(&self, scope_id: u64, tenant: &TenantKey, estimated_output_tokens: u32) { + let mut state = self.state.lock(); + { + let Some(membership) = state + .scopes + .get_mut(&scope_id) + .and_then(|scope| scope.tenants.get_mut(tenant)) + else { + return; + }; + if membership.active_reservations == 0 { + return; + } + membership.active_reservations -= 1; + } + let system_virtual_time = state.system_virtual_time; + let Some(service) = state.service.get_mut(tenant) else { + return; + }; + service.active_reservations = service.active_reservations.saturating_sub(1); + service.virtual_finish = (service.virtual_finish + - f64::from(estimated_output_tokens) / self.weight(tenant)) + .max(system_virtual_time); + let virtual_service = service.virtual_finish; + Self::advance_system_virtual_time(&mut state); + let accounting = state.accounting.entry(tenant.clone()).or_default(); + accounting.reserved_output_tokens = accounting + .reserved_output_tokens + .saturating_sub(u64::from(estimated_output_tokens)); + let reserved_output_tokens = accounting.reserved_output_tokens; + drop(state); + super::metrics::set_fair_share_virtual_finish(self.metric_tenant(tenant), virtual_service); + super::metrics::set_fair_share_reserved_output_tokens( + self.metric_tenant(tenant), + reserved_output_tokens, + ); + } + + fn settle_reservation( + &self, + scope_id: u64, + tenant: &TenantKey, + estimated_output_tokens: u32, + observed_output_tokens: Option, + kind: SettlementKind, + ) { + let charged = observed_output_tokens.unwrap_or(estimated_output_tokens); + let mut state = self.state.lock(); + { + let membership = state + .scopes + .entry(scope_id) + .or_default() + .tenants + .entry(tenant.clone()) + .or_default(); + membership.active_reservations = membership.active_reservations.saturating_sub(1); + } + let system_virtual_time = state.system_virtual_time; + let service = state.service.entry(tenant.clone()).or_default(); + service.active_reservations = service.active_reservations.saturating_sub(1); + let correction = + (f64::from(charged) - f64::from(estimated_output_tokens)) / self.weight(tenant); + service.virtual_finish = (service.virtual_finish + correction).max(system_virtual_time); + let virtual_service = service.virtual_finish; + Self::advance_system_virtual_time(&mut state); + let accounting = state.accounting.entry(tenant.clone()).or_default(); + accounting.reserved_output_tokens = accounting + .reserved_output_tokens + .saturating_sub(u64::from(estimated_output_tokens)); + accounting.charged_output_tokens = accounting + .charged_output_tokens + .saturating_add(u64::from(charged)); + let reserved_output_tokens = accounting.reserved_output_tokens; + drop(state); + + let metric_tenant = self.metric_tenant(tenant); + super::metrics::record_fair_share_charged_output_tokens(metric_tenant, charged); + super::metrics::set_fair_share_virtual_finish(metric_tenant, virtual_service); + super::metrics::set_fair_share_reserved_output_tokens( + metric_tenant, + reserved_output_tokens, + ); + if kind != SettlementKind::Observed { + super::metrics::record_fair_share_fallback(kind.as_str()); + } + } + + #[cfg(test)] + pub(crate) fn snapshot( + &self, + scope_id: u64, + tenant: &TenantKey, + ) -> (u64, u64, u64, [usize; 4]) { + let state = self.state.lock(); + let accounting = state + .accounting + .get(tenant) + .expect("tenant accounting exists"); + let membership = state + .scopes + .get(&scope_id) + .and_then(|scope| scope.tenants.get(tenant)) + .expect("tenant scope membership exists"); + ( + accounting.charged_output_tokens, + accounting.reserved_output_tokens, + membership.active_reservations, + membership.queued, + ) + } + + #[cfg(test)] + fn virtual_snapshot(&self, scope_id: u64, tenant: &TenantKey) -> (f64, f64) { + let state = self.state.lock(); + state + .scopes + .get(&scope_id) + .and_then(|scope| scope.tenants.get(tenant)) + .expect("tenant scope membership exists"); + ( + state + .service + .get(tenant) + .expect("tenant service exists") + .virtual_finish, + state.system_virtual_time, + ) + } +} + +impl SettlementKind { + fn as_str(self) -> &'static str { + match self { + Self::Observed => "observed", + Self::MissingUsage => "missing_usage", + Self::Interrupted => "interrupted", + } + } +} + +/// Provisional output-token charge attached to an admitted request. +/// +/// Explicit settlement replaces the estimate with observed terminal usage. +/// Dropping without terminal usage keeps the estimate charged, preventing a +/// disconnect from becoming a way to escape fair-share accounting. +pub struct FairShareReservation { + ledger: Arc, + tenant: TenantKey, + scope_id: u64, + estimated_output_tokens: u32, + settled: bool, +} + +impl std::fmt::Debug for FairShareReservation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FairShareReservation") + .field("tenant", &self.tenant) + .field("scope_id", &self.scope_id) + .field("estimated_output_tokens", &self.estimated_output_tokens) + .field("settled", &self.settled) + .finish() + } +} + +impl FairShareReservation { + pub fn settle(mut self, observed_output_tokens: Option, kind: SettlementKind) { + self.ledger.settle_reservation( + self.scope_id, + &self.tenant, + self.estimated_output_tokens, + observed_output_tokens, + kind, + ); + self.settled = true; + } + + pub fn cancel(mut self) { + self.ledger + .cancel_reservation(self.scope_id, &self.tenant, self.estimated_output_tokens); + self.settled = true; + } +} + +impl Drop for FairShareReservation { + fn drop(&mut self) { + if self.settled { + return; + } + self.ledger.settle_reservation( + self.scope_id, + &self.tenant, + self.estimated_output_tokens, + None, + SettlementKind::Interrupted, + ); + self.settled = true; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config(weights: &[(&str, f64)]) -> FairShareConfig { + FairShareConfig { + default_weight: 1.0, + default_output_tokens: 10, + trust_output_token_estimate_header: false, + tenant_weights: weights + .iter() + .map(|(tenant, weight)| ((*tenant).to_string(), *weight)) + .collect(), + } + } + + fn reserve_one( + ledger: &Arc, + scope_id: u64, + class: Class, + tenant: &TenantKey, + estimated_output_tokens: u32, + ) -> Option { + ledger.register_waiter(scope_id, tenant, class); + ledger + .reserve_local_candidate( + scope_id, + class, + &[FairShareCandidate { + index: 0, + tenant, + estimated_output_tokens, + }], + ) + .map(|selection| selection.reservation) + } + + #[test] + fn observed_tokens_replace_the_provisional_estimate_once() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[("header:a", 1.0)]))); + let scope_id = ledger.new_scope(); + let tenant = TenantKey::new("header:a"); + let reservation = reserve_one(&ledger, scope_id, Class::Default, &tenant, 100).unwrap(); + assert_eq!(ledger.snapshot(scope_id, &tenant), (0, 100, 1, [0; 4])); + + reservation.settle(Some(37), SettlementKind::Observed); + assert_eq!(ledger.snapshot(scope_id, &tenant), (37, 0, 0, [0; 4])); + } + + #[test] + fn cancelled_delivery_reverts_the_provisional_charge() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[("header:a", 1.0)]))); + let scope_id = ledger.new_scope(); + let tenant = TenantKey::new("header:a"); + let reservation = reserve_one(&ledger, scope_id, Class::Default, &tenant, 100).unwrap(); + reservation.cancel(); + assert_eq!(ledger.snapshot(scope_id, &tenant), (0, 0, 0, [0; 4])); + } + + #[test] + fn missing_terminal_usage_keeps_the_estimate_charged() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[("header:a", 1.0)]))); + let scope_id = ledger.new_scope(); + let tenant = TenantKey::new("header:a"); + let reservation = reserve_one(&ledger, scope_id, Class::Default, &tenant, 100).unwrap(); + drop(reservation); + assert_eq!(ledger.snapshot(scope_id, &tenant), (100, 0, 0, [0; 4])); + } + + #[test] + fn weighted_share_converges_when_users_contend_for_one_pool() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 2.0), + ("header:b", 1.0), + ]))); + let scope_id = ledger.new_scope(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + let mut admitted_a = 0; + let mut admitted_b = 0; + + // Keep both tenants continuously backlogged for every selection. + for _ in 0..30 { + ledger.register_waiter(scope_id, &a, Class::Default); + ledger.register_waiter(scope_id, &b, Class::Default); + } + for _ in 0..30 { + let candidates = [ + FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }, + FairShareCandidate { + index: 1, + tenant: &b, + estimated_output_tokens: 10, + }, + ]; + let selection = ledger + .reserve_local_candidate(scope_id, Class::Default, &candidates) + .expect("a local contender must be selected"); + if selection.index == 0 { + admitted_a += 1; + selection + .reservation + .settle(Some(10), SettlementKind::Observed); + } else { + admitted_b += 1; + selection + .reservation + .settle(Some(10), SettlementKind::Observed); + } + } + + assert_eq!((admitted_a, admitted_b), (20, 10)); + } + + #[test] + fn lone_backlogged_user_borrows_all_available_capacity() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 1.0), + ("header:b", 1.0), + ]))); + let scope_id = ledger.new_scope(); + let a = TenantKey::new("header:a"); + + for _ in 0..4 { + reserve_one(&ledger, scope_id, Class::Default, &a, 10) + .expect("a lone eligible user must never be idled") + .settle(Some(10), SettlementKind::Observed); + } + assert_eq!(ledger.snapshot(scope_id, &a).0, 40); + } + + #[test] + fn underserved_eligible_user_wins_under_simultaneous_contention() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 1.0), + ("header:b", 1.0), + ]))); + let scope_id = ledger.new_scope(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + + // b is already backlogged while a receives service, so this is real + // debt on a shared constrained resource rather than idle-time credit. + ledger.register_waiter(scope_id, &a, Class::Default); + ledger.register_waiter(scope_id, &b, Class::Default); + let candidates = [ + FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }, + FairShareCandidate { + index: 1, + tenant: &b, + estimated_output_tokens: 10, + }, + ]; + let first = ledger + .reserve_local_candidate(scope_id, Class::Default, &candidates) + .unwrap(); + assert_eq!(first.index, 0, "stable tie-break selects a first"); + first.reservation.settle(Some(50), SettlementKind::Observed); + ledger.register_waiter(scope_id, &a, Class::Default); + + let selected = ledger + .reserve_local_candidate(scope_id, Class::Default, &candidates) + .unwrap(); + assert_eq!(selected.index, 1, "the underserved eligible user wins"); + selected.reservation.cancel(); + } + + #[test] + fn solo_service_does_not_create_debt_against_an_idle_tenant() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 1.0), + ("header:b", 1.0), + ]))); + let scope_id = ledger.new_scope(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + + for _ in 0..100 { + reserve_one(&ledger, scope_id, Class::Default, &a, 10) + .unwrap() + .settle(Some(10), SettlementKind::Observed); + } + ledger.register_waiter(scope_id, &a, Class::Default); + ledger.register_waiter(scope_id, &b, Class::Default); + + let (a_finish, system_time) = ledger.virtual_snapshot(scope_id, &a); + let (b_finish, _) = ledger.virtual_snapshot(scope_id, &b); + assert_eq!(a_finish, system_time); + assert_eq!(b_finish, system_time); + + let candidates = [ + FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }, + FairShareCandidate { + index: 1, + tenant: &b, + estimated_output_tokens: 10, + }, + ]; + let selected = ledger + .reserve_local_candidate(scope_id, Class::Default, &candidates) + .unwrap(); + assert_eq!( + selected.index, 0, + "new b must not monopolize service to match a's lifetime total" + ); + selected.reservation.cancel(); + } + + #[test] + fn idle_returning_tenant_is_prompt_without_unbounded_catch_up() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 1.0), + ("header:b", 1.0), + ]))); + let scope_id = ledger.new_scope(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + + for _ in 0..20 { + reserve_one(&ledger, scope_id, Class::Default, &a, 10) + .unwrap() + .settle(Some(10), SettlementKind::Observed); + } + for _ in 0..4 { + ledger.register_waiter(scope_id, &a, Class::Default); + ledger.register_waiter(scope_id, &b, Class::Default); + } + let candidates = [ + FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }, + FairShareCandidate { + index: 1, + tenant: &b, + estimated_output_tokens: 10, + }, + ]; + let mut order = Vec::new(); + for _ in 0..4 { + let selected = ledger + .reserve_local_candidate(scope_id, Class::Default, &candidates) + .unwrap(); + order.push(selected.index); + selected + .reservation + .settle(Some(10), SettlementKind::Observed); + } + assert_eq!(order, vec![0, 1, 0, 1]); + } + + #[test] + fn unrelated_model_backlog_never_blocks_local_capacity() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 1.0), + ("header:b", 1.0), + ]))); + let local_scope = ledger.new_scope(); + let other_scope = ledger.new_scope(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + + // b is queued only for another model. It cannot consume this scope's + // slot or cause local idling. + ledger.register_waiter(other_scope, &b, Class::Default); + ledger.register_waiter(local_scope, &a, Class::Default); + let selected = ledger + .reserve_local_candidate( + local_scope, + Class::Default, + &[FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }], + ) + .expect("an eligible local request must keep the pool work-conserving"); + selected.reservation.cancel(); + } + + #[test] + fn contended_service_in_one_pool_affects_ordering_in_another_pool() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 1.0), + ("header:b", 1.0), + ]))); + let first_pool = ledger.new_scope(); + let second_pool = ledger.new_scope(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + + // b is backlogged while a receives service in the first pool, so a's + // canonical process-global virtual finish moves ahead of b's. + ledger.register_waiter(first_pool, &b, Class::Default); + reserve_one(&ledger, first_pool, Class::Default, &a, 50) + .unwrap() + .settle(Some(50), SettlementKind::Observed); + + ledger.register_waiter(second_pool, &a, Class::Default); + ledger.register_waiter(second_pool, &b, Class::Default); + let candidates = [ + FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }, + FairShareCandidate { + index: 1, + tenant: &b, + estimated_output_tokens: 10, + }, + ]; + let selected = ledger + .reserve_local_candidate(second_pool, Class::Default, &candidates) + .unwrap(); + assert_eq!( + selected.index, 1, + "contended service debt must follow a tenant across model pools" + ); + selected.reservation.cancel(); + assert_eq!(ledger.snapshot(first_pool, &a).0, 50); + } + + #[test] + fn batch_and_interactive_share_one_partition_ledger() { + let ledger = Arc::new(GlobalFairShare::from_config(&config(&[ + ("header:a", 1.0), + ("header:b", 1.0), + ]))); + let scope_id = ledger.new_scope(); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + + // b is backlogged in the partition while a receives bulk service. + // Class priority remains the outer policy, but the tenant's virtual + // finish is shared when a later interactive contention is resolved. + ledger.register_waiter(scope_id, &b, Class::Default); + let bulk = reserve_one(&ledger, scope_id, Class::Bulk, &a, 50).unwrap(); + bulk.settle(Some(50), SettlementKind::Observed); + + ledger.register_waiter(scope_id, &a, Class::Interactive); + ledger.register_waiter(scope_id, &b, Class::Interactive); + let candidates = [ + FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }, + FairShareCandidate { + index: 1, + tenant: &b, + estimated_output_tokens: 10, + }, + ]; + let selected = ledger + .reserve_local_candidate(scope_id, Class::Interactive, &candidates) + .unwrap(); + assert_eq!( + selected.index, 1, + "bulk output remains visible to later interactive contention" + ); + selected.reservation.cancel(); + assert_eq!(ledger.snapshot(scope_id, &a).0, 50); + } +} diff --git a/model_gateway/src/middleware/scheduler/metrics.rs b/model_gateway/src/middleware/scheduler/metrics.rs index dc18f44dc..1ed9d56a2 100644 --- a/model_gateway/src/middleware/scheduler/metrics.rs +++ b/model_gateway/src/middleware/scheduler/metrics.rs @@ -25,6 +25,12 @@ const CLAMP_TOTAL: &str = "smg_scheduler_clamp_total"; const UNKNOWN_PRIORITY_TOTAL: &str = "smg_scheduler_unknown_priority_value_total"; const STARVATION_PROMOTION_TOTAL: &str = "smg_scheduler_starvation_promotion_total"; const PARTITION_ADMIT_TOTAL: &str = "smg_scheduler_partition_admit_total"; +const FAIR_SHARE_CHARGED_OUTPUT_TOKENS_TOTAL: &str = "smg_fair_share_charged_output_tokens_total"; +const FAIR_SHARE_VIRTUAL_FINISH: &str = "smg_fair_share_virtual_finish"; +const FAIR_SHARE_RESERVED_OUTPUT_TOKENS: &str = "smg_fair_share_reserved_output_tokens"; +const FAIR_SHARE_QUEUE_WAIT_SECONDS: &str = "smg_fair_share_queue_wait_seconds"; +const FAIR_SHARE_FALLBACK_TOTAL: &str = "smg_fair_share_fallback_total"; +const FAIR_SHARE_UNKNOWN_TENANT_TOTAL: &str = "smg_fair_share_unknown_tenant_total"; // Capacity / autoscaling gauges, refreshed by the sampler task. const INFLIGHT: &str = "smg_scheduler_inflight"; @@ -84,6 +90,30 @@ pub fn describe() { PARTITION_ADMIT_TOTAL, "Priority-scheduler admission outcomes by partition, class, and outcome" ); + describe_counter!( + FAIR_SHARE_CHARGED_OUTPUT_TOKENS_TOTAL, + "Output tokens charged to the global fair-share ledger by tenant" + ); + describe_gauge!( + FAIR_SHARE_VIRTUAL_FINISH, + "Process-global active-set normalized virtual finish used for weighted dispatch" + ); + describe_gauge!( + FAIR_SHARE_RESERVED_OUTPUT_TOKENS, + "Provisional output-token charges for active requests by tenant" + ); + describe_histogram!( + FAIR_SHARE_QUEUE_WAIT_SECONDS, + "Fair-share queue wait by tenant and priority class" + ); + describe_counter!( + FAIR_SHARE_FALLBACK_TOTAL, + "Fair-share settlements that used an estimate instead of terminal usage" + ); + describe_counter!( + FAIR_SHARE_UNKNOWN_TENANT_TOTAL, + "Requests whose resolved tenant has no explicit fair-share weight" + ); describe_gauge!(INFLIGHT, "Current in-flight request count per class"); describe_gauge!(QUEUE_DEPTH, "Current queued waiter count per class"); describe_gauge!( @@ -183,6 +213,47 @@ pub fn record_starvation_promotion(class: Class) { counter!(STARVATION_PROMOTION_TOTAL, "class" => class.as_str()).increment(1); } +pub fn record_fair_share_charged_output_tokens(tenant: &str, tokens: u32) { + counter!( + FAIR_SHARE_CHARGED_OUTPUT_TOKENS_TOTAL, + "tenant" => intern_string(tenant) + ) + .increment(u64::from(tokens)); +} + +pub fn set_fair_share_virtual_finish(tenant: &str, virtual_finish: f64) { + gauge!(FAIR_SHARE_VIRTUAL_FINISH, "tenant" => intern_string(tenant)).set(virtual_finish); +} + +pub fn set_fair_share_reserved_output_tokens(tenant: &str, tokens: u64) { + gauge!( + FAIR_SHARE_RESERVED_OUTPUT_TOKENS, + "tenant" => intern_string(tenant) + ) + .set(tokens as f64); +} + +pub fn record_fair_share_queue_wait(tenant: &str, class: Class, wait: Duration) { + histogram!( + FAIR_SHARE_QUEUE_WAIT_SECONDS, + "tenant" => intern_string(tenant), + "class" => class.as_str() + ) + .record(wait.as_secs_f64()); +} + +pub fn record_fair_share_fallback(reason: &'static str) { + counter!(FAIR_SHARE_FALLBACK_TOTAL, "reason" => reason).increment(1); +} + +pub fn record_fair_share_unknown_tenant(tenant: &str) { + counter!( + FAIR_SHARE_UNKNOWN_TENANT_TOTAL, + "tenant" => intern_string(tenant) + ) + .increment(1); +} + /// Set the in-flight gauge for a class (sampler). pub fn set_inflight(class: Class, count: u16) { gauge!(INFLIGHT, "class" => class.as_str()).set(f64::from(count)); diff --git a/model_gateway/src/middleware/scheduler/mod.rs b/model_gateway/src/middleware/scheduler/mod.rs index ea74b5c20..ad4a6d90b 100644 --- a/model_gateway/src/middleware/scheduler/mod.rs +++ b/model_gateway/src/middleware/scheduler/mod.rs @@ -7,8 +7,10 @@ pub mod config; pub mod engine; pub mod error; pub mod extract; +pub mod fair_share; pub mod inflight; pub mod metrics; +mod output_tokens; pub mod policy; pub mod queue; pub mod slots; @@ -18,13 +20,14 @@ pub use admission::priority_admission_middleware; pub use body::SchedulerGuardBody; pub use class::{Class, PRIORITY_HEADER}; pub use config::{ - AdmissionPartitionConfig, ClassConfig, ClassRuntimeConfig, PrioritySchedulerYaml, - SchedulerSettings, SettingsValidationError, TenantPolicyConfig, + AdmissionPartitionConfig, ClassConfig, ClassRuntimeConfig, FairShareConfig, + PrioritySchedulerYaml, SchedulerSettings, SettingsValidationError, TenantPolicyConfig, }; pub use engine::{ AdmitOutcome, PriorityScheduler, RejectionReason, SchedulerInitError, SchedulerPermit, }; pub use error::{SchedulerError, HEADER_X_SMG_PREEMPTED}; pub use extract::PreemptionGuard; +pub use fair_share::{GlobalFairShare, SettlementKind, OUTPUT_TOKEN_ESTIMATE_HEADER}; pub use policy::{StaticTenantPolicyResolver, TenantPolicy, TenantPolicyResolver}; pub use state::{AdmissionMode, SchedulerState, ADMISSION_PARTITION_HEADER}; diff --git a/model_gateway/src/middleware/scheduler/output_tokens.rs b/model_gateway/src/middleware/scheduler/output_tokens.rs new file mode 100644 index 000000000..4763eec6d --- /dev/null +++ b/model_gateway/src/middleware/scheduler/output_tokens.rs @@ -0,0 +1,97 @@ +//! Terminal response-usage parsing for fair-share settlement. +//! +//! This parser is intentionally owned by the scheduler completion path. It +//! does not call or depend on adaptive admission's predictor. + +fn value_u32(value: Option<&serde_json::Value>) -> Option { + value?.as_u64().and_then(|value| u32::try_from(value).ok()) +} + +fn output_tokens_from_value(value: &serde_json::Value) -> Option { + [ + "/usage/completion_tokens", + "/usage/output_tokens", + "/meta_info/completion_tokens", + ] + .into_iter() + .find_map(|pointer| value_u32(value.pointer(pointer))) +} + +fn output_tokens_from_truncated_json_tail(body: &[u8]) -> Option { + [b"\"usage\"".as_slice(), b"\"meta_info\"".as_slice()] + .into_iter() + .find_map(|key| { + let offset = body.windows(key.len()).rposition(|window| window == key)? + key.len(); + let remainder = &body[offset..]; + let object = &remainder[remainder.iter().position(|byte| *byte == b':')? + 1..]; + let value = serde_json::Deserializer::from_slice(object) + .into_iter::() + .next()? + .ok()?; + value_u32(value.get("completion_tokens")) + .or_else(|| value_u32(value.get("output_tokens"))) + }) +} + +pub(crate) fn observed_output_tokens(body: &[u8]) -> Option { + if let Ok(value) = serde_json::from_slice::(body) { + if let Some(tokens) = output_tokens_from_value(&value) { + return Some(tokens); + } + } + if let Some(tokens) = output_tokens_from_truncated_json_tail(body) { + return Some(tokens); + } + + body.split(|byte| *byte == b'\n') + .filter_map(|line| { + let line = line + .strip_suffix(b"\r") + .unwrap_or(line) + .strip_prefix(b"data:")?; + let line = line.strip_prefix(b" ").unwrap_or(line); + if line == b"[DONE]" { + return None; + } + serde_json::from_slice::(line) + .ok() + .and_then(|value| output_tokens_from_value(&value)) + }) + .next_back() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_openai_anthropic_and_sglang_shapes() { + assert_eq!( + observed_output_tokens(br#"{"usage":{"completion_tokens":42}}"#), + Some(42) + ); + assert_eq!( + observed_output_tokens(br#"{"usage":{"output_tokens":13}}"#), + Some(13) + ); + assert_eq!( + observed_output_tokens(br#"{"meta_info":{"completion_tokens":7}}"#), + Some(7) + ); + } + + #[test] + fn parses_terminal_stream_usage() { + assert_eq!( + observed_output_tokens( + b"data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: {\"choices\":[],\"usage\":{\"completion_tokens\":17}}\n\ndata: [DONE]\n\n" + ), + Some(17) + ); + } + + #[test] + fn missing_usage_is_not_invented() { + assert_eq!(observed_output_tokens(br#"{"choices":[]}"#), None); + } +} diff --git a/model_gateway/src/middleware/scheduler/queue.rs b/model_gateway/src/middleware/scheduler/queue.rs index 51aea0944..dd2e54ffe 100644 --- a/model_gateway/src/middleware/scheduler/queue.rs +++ b/model_gateway/src/middleware/scheduler/queue.rs @@ -15,7 +15,12 @@ use smg_auth::RequestId; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; -use super::{engine::SchedulerPermit, Class}; +use super::{ + engine::SchedulerPermit, + fair_share::{FairShareCandidate, FairShareReservation, GlobalFairShare}, + Class, +}; +use crate::tenant::TenantKey; /// One work-conserving occupancy budget shared by all priority queues. /// @@ -90,6 +95,9 @@ pub struct Waiter { pub cancel: CancellationToken, pub request_id: RequestId, pub permit_tx: oneshot::Sender, + pub tenant: Option, + pub estimated_output_tokens: u32, + pub fair_share_reservation: Option, } impl Waiter { @@ -105,6 +113,29 @@ impl Waiter { cancel, request_id, permit_tx, + tenant: None, + estimated_output_tokens: 0, + fair_share_reservation: None, + } + } + + pub fn new_fair( + class: Class, + cancel: CancellationToken, + request_id: RequestId, + permit_tx: oneshot::Sender, + tenant: TenantKey, + estimated_output_tokens: u32, + ) -> Self { + Self { + class, + queued_at: Instant::now(), + cancel, + request_id, + permit_tx, + tenant: Some(tenant), + estimated_output_tokens: estimated_output_tokens.max(1), + fair_share_reservation: None, } } } @@ -205,6 +236,119 @@ impl ClassQueue for FifoClassQueue { } } +/// Work-conserving per-partition queue ordered by one shared token ledger. +/// +/// Only waiters in this concrete queue are candidates, so debt in another +/// non-fungible model pool can never idle local capacity. +pub struct FairClassQueue { + waiters: Mutex>, + soft_limit: usize, + budget: Arc, + ledger: Arc, + scope_id: u64, + class: Class, +} + +impl FairClassQueue { + pub fn with_shared_budget( + class: Class, + soft_limit: usize, + budget: Arc, + ledger: Arc, + scope_id: u64, + ) -> Self { + Self { + waiters: Mutex::new(VecDeque::with_capacity(soft_limit.min(64))), + soft_limit, + budget, + ledger, + scope_id, + class, + } + } +} + +impl ClassQueue for FairClassQueue { + fn try_enqueue(&self, waiter: Waiter) -> Result<(), Waiter> { + let Some(tenant) = waiter.tenant.as_ref() else { + return Err(waiter); + }; + if !self.budget.try_acquire() { + return Err(waiter); + } + let mut guard = self.waiters.lock(); + self.ledger + .register_waiter(self.scope_id, tenant, self.class); + guard.push_back(waiter); + Ok(()) + } + + fn pop_eligible(&self) -> Option { + let mut guard = self.waiters.lock(); + let selection = { + let candidates: Vec<_> = guard + .iter() + .enumerate() + .filter(|(_, waiter)| !waiter.cancel.is_cancelled()) + .filter_map(|(index, waiter)| { + waiter.tenant.as_ref().map(|tenant| FairShareCandidate { + index, + tenant, + estimated_output_tokens: waiter.estimated_output_tokens, + }) + }) + .collect(); + self.ledger + .reserve_local_candidate(self.scope_id, self.class, &candidates) + }?; + let Some(mut waiter) = guard.remove(selection.index) else { + selection.reservation.cancel(); + return None; + }; + self.budget.release(); + if let Some(tenant) = waiter.tenant.as_ref() { + self.ledger + .record_queue_wait(tenant, self.class, waiter.queued_at.elapsed()); + } + waiter.fair_share_reservation = Some(selection.reservation); + Some(waiter) + } + + fn head_age(&self) -> Option { + self.waiters + .lock() + .iter() + .map(|waiter| waiter.queued_at.elapsed()) + .max() + } + + fn depth(&self) -> usize { + self.waiters.lock().len() + } + + fn capacity(&self) -> usize { + self.soft_limit + } + + fn drop_cancelled_head(&self) { + let mut guard = self.waiters.lock(); + let mut index = 0; + while index < guard.len() { + if !guard[index].cancel.is_cancelled() { + index += 1; + continue; + } + let Some(waiter) = guard.remove(index) else { + break; + }; + self.budget.release(); + if let Some(tenant) = waiter.tenant.as_ref() { + self.ledger.remove_waiter(self.scope_id, tenant, self.class); + } + } + } +} + #[cfg(test)] mod tests { use std::time::Duration; diff --git a/model_gateway/src/middleware/scheduler/state.rs b/model_gateway/src/middleware/scheduler/state.rs index 89c257daf..a0549e594 100644 --- a/model_gateway/src/middleware/scheduler/state.rs +++ b/model_gateway/src/middleware/scheduler/state.rs @@ -12,7 +12,8 @@ use tokio::sync::{broadcast, watch}; use tracing::{error, info}; use super::{ - Class, PriorityScheduler, SchedulerSettings, StaticTenantPolicyResolver, TenantPolicyResolver, + Class, GlobalFairShare, PriorityScheduler, SchedulerSettings, StaticTenantPolicyResolver, + TenantPolicyResolver, }; use crate::{ config::types::RouterConfig, @@ -159,6 +160,7 @@ impl AdmissionMode { if yaml.is_none() { settings = settings.with_global_queue_budget(rc.queue_size); } + let fair_share = GlobalFairShare::from_settings(&settings).map(Arc::new); let resolver: Arc = Arc::new(StaticTenantPolicyResolver::from_settings(&settings)); @@ -177,8 +179,12 @@ impl AdmissionMode { // The atomic value covers any update that won the race before the // receiver subscribed; subsequent updates remain queued for the // dispatcher through `capacity_watch`. - let scheduler = PriorityScheduler::new(&settings, worker_capacity.current()) - .map_err(|e| e.to_string())?; + let scheduler = PriorityScheduler::new_with_fair_share( + &settings, + worker_capacity.current(), + fair_share, + ) + .map_err(|e| e.to_string())?; scheduler.spawn_dispatcher_retaining_capacity(capacity_watch, worker_capacity); scheduler.spawn_sampler(SAMPLER_INTERVAL); return Ok(Self::Priority(Arc::new(SchedulerState { @@ -235,8 +241,12 @@ impl AdmissionMode { name, initial_replicas.get(name).copied().unwrap_or(0), ); - let scheduler = PriorityScheduler::new(&partition_settings, capacity) - .map_err(|e| format!("partition {name}: {e}"))?; + let scheduler = PriorityScheduler::new_with_fair_share( + &partition_settings, + capacity, + fair_share.clone(), + ) + .map_err(|e| format!("partition {name}: {e}"))?; let (capacity_tx, capacity_rx) = watch::channel(capacity); scheduler.spawn_dispatcher(capacity_rx); capacity_senders.push((name.clone(), capacity_tx)); diff --git a/model_gateway/src/middleware/tenant_resolution.rs b/model_gateway/src/middleware/tenant_resolution.rs index 3285d025b..74a3f4bc7 100644 --- a/model_gateway/src/middleware/tenant_resolution.rs +++ b/model_gateway/src/middleware/tenant_resolution.rs @@ -19,6 +19,7 @@ use crate::{ #[derive(Clone)] pub struct TenantResolutionState { trust_tenant_header: bool, + prefer_trusted_tenant_header: bool, trusted_tenant_header_name: HeaderName, } @@ -32,12 +33,18 @@ impl TenantResolutionState { Ok(Self { trust_tenant_header: config.trust_tenant_header, + prefer_trusted_tenant_header: config.prefer_trusted_tenant_header, trusted_tenant_header_name, }) } } fn resolve_raw_tenant_key(state: &TenantResolutionState, request: &Request) -> TenantKey { + if state.trust_tenant_header && state.prefer_trusted_tenant_header { + if let Some(tenant_id) = extract_trusted_tenant_id(state, request.headers()) { + return canonical_tenant_key(TenantIdentity::Header(Arc::from(tenant_id))); + } + } if let Some(caller) = request.extensions().get::() { return caller.tenant_key().clone(); } @@ -142,6 +149,86 @@ mod tests { assert_eq!(request_meta.tenant_key().as_str(), "auth:b3c2"); } + #[tokio::test] + async fn trusted_header_does_not_override_auth_without_preference() { + let mut config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["http://worker1:8000".to_string()], + }, + PolicyConfig::Random, + ); + config.tenant_resolution.trust_tenant_header = true; + let state = TenantResolutionState::new(&config).unwrap(); + let mut request = Request::builder() + .uri("/") + .header(DEFAULT_TENANT_HEADER_NAME, "alice") + .body(Body::empty()) + .unwrap(); + request + .extensions_mut() + .insert(DataPlaneCaller::new(TenantKey::from("auth:proxy"))); + + let request_meta = resolve_route_request_meta(&state, &request); + assert_eq!(request_meta.tenant_key().as_str(), "auth:proxy"); + } + + #[tokio::test] + async fn preferred_trusted_header_distinguishes_users_behind_shared_proxy_auth() { + let mut config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["http://worker1:8000".to_string()], + }, + PolicyConfig::Random, + ); + config.tenant_resolution.trust_tenant_header = true; + config.tenant_resolution.prefer_trusted_tenant_header = true; + let state = TenantResolutionState::new(&config).unwrap(); + for user in ["alice", "bob"] { + let mut request = Request::builder() + .uri("/") + .header(DEFAULT_TENANT_HEADER_NAME, user) + .body(Body::empty()) + .unwrap(); + request + .extensions_mut() + .insert(DataPlaneCaller::new(TenantKey::from("auth:proxy"))); + + let request_meta = resolve_route_request_meta(&state, &request); + assert_eq!(request_meta.tenant_key().as_str(), format!("header:{user}")); + } + } + + #[tokio::test] + async fn unusable_preferred_header_falls_back_to_auth() { + for value in [ + None, + Some(HeaderValue::from_static(" ")), + Some(HeaderValue::from_bytes(b"\x80").unwrap()), + ] { + let mut config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["http://worker1:8000".to_string()], + }, + PolicyConfig::Random, + ); + config.tenant_resolution.trust_tenant_header = true; + config.tenant_resolution.prefer_trusted_tenant_header = true; + let state = TenantResolutionState::new(&config).unwrap(); + let mut request = Request::builder().uri("/").body(Body::empty()).unwrap(); + if let Some(value) = value { + request + .headers_mut() + .insert(DEFAULT_TENANT_HEADER_NAME, value); + } + request + .extensions_mut() + .insert(DataPlaneCaller::new(TenantKey::from("auth:proxy"))); + + let request_meta = resolve_route_request_meta(&state, &request); + assert_eq!(request_meta.tenant_key().as_str(), "auth:proxy"); + } + } + #[tokio::test] async fn request_meta_uses_trusted_header_when_enabled() { let mut config = RouterConfig::new( From c1ce061302ff00c617e3e6644ceccecbecdd1cbd Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:18:45 -0700 Subject: [PATCH 6/8] test(scheduler): satisfy strict lint Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- model_gateway/src/middleware/scheduler/config.rs | 16 +++++++++------- model_gateway/src/middleware/scheduler/engine.rs | 4 ++-- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/model_gateway/src/middleware/scheduler/config.rs b/model_gateway/src/middleware/scheduler/config.rs index 661b6af68..beb083a55 100644 --- a/model_gateway/src/middleware/scheduler/config.rs +++ b/model_gateway/src/middleware/scheduler/config.rs @@ -557,13 +557,15 @@ fair_share: #[test] fn test_invalid_fair_share_weights_are_rejected() { for weight in [0.0, -1.0, f64::INFINITY, f64::NAN] { - let mut yaml = PrioritySchedulerYaml::default(); - yaml.fair_share = Some(FairShareConfig { - default_weight: 1.0, - default_output_tokens: 128, - trust_output_token_estimate_header: false, - tenant_weights: HashMap::from([("header:alice".to_string(), weight)]), - }); + let yaml = PrioritySchedulerYaml { + fair_share: Some(FairShareConfig { + default_weight: 1.0, + default_output_tokens: 128, + trust_output_token_estimate_header: false, + tenant_weights: HashMap::from([("header:alice".to_string(), weight)]), + }), + ..Default::default() + }; assert!(matches!( SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&yaml)), Err(SettingsValidationError::InvalidFairShareTenantWeight { .. }) diff --git a/model_gateway/src/middleware/scheduler/engine.rs b/model_gateway/src/middleware/scheduler/engine.rs index 74ac06f40..8708fc0ea 100644 --- a/model_gateway/src/middleware/scheduler/engine.rs +++ b/model_gateway/src/middleware/scheduler/engine.rs @@ -1108,7 +1108,7 @@ mod tests { use std::{collections::HashMap, sync::Arc, time::Duration}; use super::*; - use crate::middleware::scheduler::{ClassConfig, PrioritySchedulerYaml}; + use crate::middleware::scheduler::{ClassConfig, FairShareConfig, PrioritySchedulerYaml}; fn default_settings() -> SchedulerSettings { SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, None).unwrap() @@ -2054,7 +2054,7 @@ mod tests { } let yaml = PrioritySchedulerYaml { classes, - fair_share: Some(crate::middleware::scheduler::FairShareConfig { + fair_share: Some(FairShareConfig { default_weight: 1.0, default_output_tokens: 10, trust_output_token_estimate_header: false, From e9feeeb4f4e001e6dfa34b1b536c4ac7bb86ba95 Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 21:29:26 -0700 Subject: [PATCH 7/8] feat(scheduler): add per-model tenant fair sharing Isolate scheduling debt by canonical model while retaining process-global output accounting. Use hierarchical named and other buckets without idling non-fungible pools. Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- docs/reference/priority-scheduler.md | 49 +- .../src/middleware/scheduler/admission.rs | 20 +- .../src/middleware/scheduler/body.rs | 2 + .../src/middleware/scheduler/config.rs | 125 ++- .../src/middleware/scheduler/engine.rs | 28 +- .../src/middleware/scheduler/fair_share.rs | 990 +++++++++++++++++- .../src/middleware/scheduler/metrics.rs | 27 +- model_gateway/src/middleware/scheduler/mod.rs | 7 +- .../src/middleware/scheduler/queue.rs | 132 ++- .../src/middleware/scheduler/state.rs | 94 +- 10 files changed, 1420 insertions(+), 54 deletions(-) diff --git a/docs/reference/priority-scheduler.md b/docs/reference/priority-scheduler.md index b97c0776c..2cf25b58a 100644 --- a/docs/reference/priority-scheduler.md +++ b/docs/reference/priority-scheduler.md @@ -114,6 +114,7 @@ fair_share: default_weight: 1 default_output_tokens: 256 trust_output_token_estimate_header: false + trust_request_model_header: false tenant_weights: "header:alice": 10 "header:bob": 5 @@ -172,7 +173,37 @@ Priority class selection remains the outer policy. | `default_weight` | `1.0` | Weight assigned to a resolved tenant absent from `tenant_weights`. Must be finite and greater than zero. | | `default_output_tokens` | `256` | Provisional output-token charge used when no trusted estimate is available. Must be greater than zero. | | `trust_output_token_estimate_header` | `false` | Honor `x-smg-output-token-estimate`. Enable only behind a proxy that strips client copies and injects a validated estimate. | +| `trust_request_model_header` | `false` | Honor `x-smg-request-model` for per-model profile selection. Model profiles require this setting. Enable only behind a proxy that strips client copies and injects the parsed request model. | | `tenant_weights` | `{}` | Relative weights keyed by canonical tenant key. Every value must be finite and greater than zero. | +| `model_profiles` | `{}` | Optional model-scoped hierarchical policies. Each profile contains explicit `tenant_weights` plus one positive aggregate `other_weight`. | + +Per-model profiles use two-level weighted fair queueing. Explicit tenants and +one aggregate `other` bucket contend at the outer level. Every unlisted real +tenant retains its own identity and shares the `other` bucket equally at the +inner level: + +```yaml +fair_share: + default_output_tokens: 256 + trust_output_token_estimate_header: true + trust_request_model_header: true + model_profiles: + deepseek-v4-flash: + tenant_weights: + "header:junu": 30 + "header:xuezhou": 40 + "header:zhenting": 10 + other_weight: 20 + kimi-k3: + tenant_weights: + "header:mukhesh": 80 + other_weight: 20 +``` + +The percentages apply while the corresponding buckets are simultaneously +backlogged for that model. Idle shares are borrowed, so a free model slot is +never held empty. Requests for models without a configured profile continue to +use the legacy flat process-wide weights. The scheduler reserves the estimate when a request is admitted. A trustworthy terminal usage record replaces that estimate with actual output tokens. If a @@ -186,10 +217,12 @@ virtual time, so idle tenants do not bank unlimited catch-up credit. The queue is work-conserving: a model partition with a free slot and an eligible local request never idles for an underserved tenant that can use only another model. -One `GlobalFairShare` instance aggregates actual-token metrics and canonical -tenant virtual finish across every partition built by one SMG process. Each -partition chooses among only the requests eligible for its non-fungible model -pool. Consequently: +One `GlobalFairShare` instance aggregates actual-token accounting across every +partition built by one SMG process. Flat configuration keeps one canonical +tenant virtual finish across partitions. Per-model profiles instead keep +scheduling debt and active-set virtual time separate by canonical model. A +partition that contains multiple models selects the group containing its +oldest eligible waiter before applying that model's policy. Consequently: - Batch and interactive requests resolve to the same tenant ledger when they enter the same SMG process with the same canonical tenant identity. Priority @@ -197,6 +230,9 @@ pool. Consequently: - Configured percentages converge when tenants are simultaneously backlogged for the same constrained pool. Exact fleet-wide percentages are not enforceable for tenants targeting disjoint pools without idling capacity. +- Per-model service never creates scheduling debt in another model, while + charged and reserved output-token accounting remains process-global by real + tenant. - The ledger does not coordinate separate M1 and M2 gateways, or overlapping old and new gateway processes during a rollout. Strict cross-gateway fairness requires a distributed ledger or one authoritative admission front door. @@ -257,9 +293,10 @@ The scheduler exposes these Prometheus metrics (see the [Metrics Reference](metr | `smg_scheduler_utilization` | Gauge | — | Total in-flight divided by backend capacity. | | `smg_scheduler_class_capacity_pressure` | Gauge | `class` | Normalized 0.0–1.0 pressure (worse of queue and slot pressure). | | `smg_fair_share_charged_output_tokens_total` | Counter | `tenant` | Actual or conservative fallback output tokens charged across the process-local ledger. | -| `smg_fair_share_virtual_finish` | Gauge | `tenant` | Process-global active-set normalized virtual finish used for local eligible-candidate dispatch. | +| `smg_fair_share_virtual_finish` | Gauge | `model`, `tenant` | Active-set normalized tenant virtual finish. Flat mode uses `model="global"`. | +| `smg_fair_share_other_bucket_virtual_finish` | Gauge | `model` | Outer virtual finish of a model profile's aggregate `other` bucket. | | `smg_fair_share_reserved_output_tokens` | Gauge | `tenant` | Provisional output-token charges held by active requests. | -| `smg_fair_share_queue_wait_seconds` | Histogram | `tenant`, `class` | Fair-share queue wait by tenant and priority class. | +| `smg_fair_share_queue_wait_seconds` | Histogram | `model`, `tenant`, `class` | Fair-share queue wait by model profile, tenant, and priority class. | | `smg_fair_share_fallback_total` | Counter | `reason` | Settlement or estimate paths that used a configured fallback. | | `smg_fair_share_unknown_tenant_total` | Counter | `tenant` | Requests whose canonical tenant has no explicit configured weight. | diff --git a/model_gateway/src/middleware/scheduler/admission.rs b/model_gateway/src/middleware/scheduler/admission.rs index fe63190e5..c06512a8f 100644 --- a/model_gateway/src/middleware/scheduler/admission.rs +++ b/model_gateway/src/middleware/scheduler/admission.rs @@ -25,9 +25,9 @@ use tokio_util::sync::CancellationToken; use tracing::trace; use super::{ - metrics as sched_metrics, state::SchedulerState, AdmitOutcome, Class, GlobalFairShare, - RejectionReason, SchedulerError, SchedulerGuardBody, HEADER_X_SMG_PREEMPTED, - OUTPUT_TOKEN_ESTIMATE_HEADER, PRIORITY_HEADER, + fair_share::FairShareProfile, metrics as sched_metrics, state::SchedulerState, AdmitOutcome, + Class, GlobalFairShare, RejectionReason, SchedulerError, SchedulerGuardBody, + HEADER_X_SMG_PREEMPTED, OUTPUT_TOKEN_ESTIMATE_HEADER, PRIORITY_HEADER, }; use crate::{ middleware::{ @@ -167,18 +167,22 @@ pub async fn priority_admission_middleware( // once admitted, releasing the slot. let cancel = CancellationToken::new(); - let estimated_output_tokens = match partition.scheduler.fair_share() { - Some(ledger) => output_token_estimate(req.headers(), ledger), - None => 1, + let (estimated_output_tokens, fair_share_profile) = match partition.scheduler.fair_share() { + Some(ledger) => ( + output_token_estimate(req.headers(), ledger), + state.fair_share_profile_for(req.headers(), ledger), + ), + None => (1, FairShareProfile::Global), }; match partition .scheduler - .admit_for_tenant( + .admit_for_tenant_profile( class, request_id, cancel, tenant.clone(), + fair_share_profile, estimated_output_tokens, ) .await @@ -257,7 +261,9 @@ mod tests { default_weight: 1.0, default_output_tokens: 256, trust_output_token_estimate_header: trust_header, + trust_request_model_header: false, tenant_weights: HashMap::new(), + model_profiles: HashMap::new(), }) } diff --git a/model_gateway/src/middleware/scheduler/body.rs b/model_gateway/src/middleware/scheduler/body.rs index b68dda204..e2f326028 100644 --- a/model_gateway/src/middleware/scheduler/body.rs +++ b/model_gateway/src/middleware/scheduler/body.rs @@ -239,7 +239,9 @@ mod tests { default_weight: 1.0, default_output_tokens, trust_output_token_estimate_header: false, + trust_request_model_header: false, tenant_weights: HashMap::from([("header:alice".to_string(), 1.0)]), + model_profiles: HashMap::new(), }), ..Default::default() }; diff --git a/model_gateway/src/middleware/scheduler/config.rs b/model_gateway/src/middleware/scheduler/config.rs index beb083a55..ce2f5c63a 100644 --- a/model_gateway/src/middleware/scheduler/config.rs +++ b/model_gateway/src/middleware/scheduler/config.rs @@ -127,7 +127,23 @@ fn default_fair_share_output_tokens() -> u32 { 256 } -/// Process-wide weighted sharing with per-partition eligibility boundaries. +/// One model's hierarchical weighted-sharing policy. +/// +/// Explicit tenants and the aggregate `other` bucket contend at the outer +/// level. Real tenant identities are retained inside `other` and share that +/// bucket equally, so settlement and accounting never collapse to a synthetic +/// tenant. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ModelFairShareConfig { + /// Per-tenant outer weights keyed by canonical `TenantKey` string. + #[serde(default)] + pub tenant_weights: HashMap, + /// Aggregate outer weight for every tenant absent from `tenant_weights`. + #[serde(default = "default_fair_share_weight")] + pub other_weight: f64, +} + +/// Flat process-wide sharing plus optional hierarchical per-model profiles. /// /// Weights are relative and need not sum to 100. For example, weights 10 and /// 5 give two continuously contending tenants a 2:1 output-token share. @@ -144,9 +160,18 @@ pub struct FairShareConfig { /// strips client copies and injects a validated value. #[serde(default)] pub trust_output_token_estimate_header: bool, + /// Honor `x-smg-request-model` when selecting a per-model profile. Keep + /// false unless a trusted proxy strips client copies and injects the + /// authenticated request body's model value. + #[serde(default)] + pub trust_request_model_header: bool, /// Per-tenant relative weights keyed by canonical `TenantKey` string. #[serde(default)] pub tenant_weights: HashMap, + /// Optional hierarchical policies keyed by canonical model id. Absence + /// preserves the original flat process-wide ledger exactly. + #[serde(default)] + pub model_profiles: HashMap, } /// Admission budget for one trusted upstream partition selector. @@ -244,6 +269,16 @@ pub enum SettingsValidationError { ZeroFairShareDefaultOutputTokens, #[error("fair_share.tenant_weights[{tenant:?}] must be finite and > 0")] InvalidFairShareTenantWeight { tenant: String }, + #[error("fair_share.model_profiles require trust_request_model_header=true")] + ModelProfilesRequireTrustedModelHeader, + #[error("fair_share.model_profiles contains an empty or untrimmed model key {model:?}")] + InvalidFairShareModelKey { model: String }, + #[error( + "fair_share.model_profiles[{model:?}].tenant_weights[{tenant:?}] must be finite and > 0" + )] + InvalidModelFairShareTenantWeight { model: String, tenant: String }, + #[error("fair_share.model_profiles[{model:?}].other_weight must be finite and > 0")] + InvalidModelFairShareOtherWeight { model: String }, } /// Runtime scheduler configuration assembled from CLI flags + the @@ -391,6 +426,29 @@ impl SchedulerSettings { }); } } + if !config.model_profiles.is_empty() && !config.trust_request_model_header { + return Err(SettingsValidationError::ModelProfilesRequireTrustedModelHeader); + } + for (model, profile) in &config.model_profiles { + if model.trim().is_empty() || model.trim() != model { + return Err(SettingsValidationError::InvalidFairShareModelKey { + model: model.clone(), + }); + } + if !profile.other_weight.is_finite() || profile.other_weight <= 0.0 { + return Err(SettingsValidationError::InvalidModelFairShareOtherWeight { + model: model.clone(), + }); + } + for (tenant, weight) in &profile.tenant_weights { + if !weight.is_finite() || *weight <= 0.0 { + return Err(SettingsValidationError::InvalidModelFairShareTenantWeight { + model: model.clone(), + tenant: tenant.clone(), + }); + } + } + } } let tenant_policies = yaml @@ -562,7 +620,9 @@ fair_share: default_weight: 1.0, default_output_tokens: 128, trust_output_token_estimate_header: false, + trust_request_model_header: false, tenant_weights: HashMap::from([("header:alice".to_string(), weight)]), + model_profiles: HashMap::new(), }), ..Default::default() }; @@ -573,6 +633,69 @@ fair_share: } } + #[test] + fn test_model_profiles_require_the_trusted_request_model_header() { + let yaml = PrioritySchedulerYaml { + fair_share: Some(FairShareConfig { + default_weight: 1.0, + default_output_tokens: 128, + trust_output_token_estimate_header: false, + trust_request_model_header: false, + tenant_weights: HashMap::new(), + model_profiles: HashMap::from([( + "kimi-k3".to_string(), + ModelFairShareConfig { + tenant_weights: HashMap::new(), + other_weight: 20.0, + }, + )]), + }), + ..Default::default() + }; + + assert!(matches!( + SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&yaml)), + Err(SettingsValidationError::ModelProfilesRequireTrustedModelHeader) + )); + } + + #[test] + fn test_per_model_hierarchical_weights_round_trip() { + let parsed: PrioritySchedulerYaml = serde_yaml::from_str( + r#" +fair_share: + trust_request_model_header: true + model_profiles: + deepseek-v4-flash: + tenant_weights: + "header:junu": 30 + "header:xuezhou": 40 + "header:zhenting": 10 + other_weight: 20 + kimi-k3: + tenant_weights: + "header:mukhesh": 80 + other_weight: 20 +"#, + ) + .unwrap(); + let fair_share = parsed.fair_share.as_ref().unwrap(); + assert!(fair_share.trust_request_model_header); + assert_eq!( + fair_share.model_profiles["deepseek-v4-flash"].tenant_weights["header:junu"], + 30.0 + ); + assert_eq!( + fair_share.model_profiles["deepseek-v4-flash"].other_weight, + 20.0 + ); + assert_eq!( + fair_share.model_profiles["kimi-k3"].tenant_weights["header:mukhesh"], + 80.0 + ); + SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&parsed)).unwrap(); + } + #[test] fn test_yaml_admission_partitions_round_trip() { let yaml = r#" diff --git a/model_gateway/src/middleware/scheduler/engine.rs b/model_gateway/src/middleware/scheduler/engine.rs index 8708fc0ea..2029e80f7 100644 --- a/model_gateway/src/middleware/scheduler/engine.rs +++ b/model_gateway/src/middleware/scheduler/engine.rs @@ -23,7 +23,7 @@ use tokio_util::sync::CancellationToken; use tracing::{info, warn}; use super::{ - fair_share::{FairShareReservation, GlobalFairShare, SettlementKind}, + fair_share::{FairShareProfile, FairShareReservation, GlobalFairShare, SettlementKind}, inflight::InflightHandle, queue::{ClassQueue, FairClassQueue, FifoClassQueue, QueueBudget, Waiter}, slots::SlotPool, @@ -411,6 +411,26 @@ impl PriorityScheduler { cancel: CancellationToken, tenant: TenantKey, estimated_output_tokens: u32, + ) -> AdmitOutcome { + self.admit_for_tenant_profile( + class, + request_id, + cancel, + tenant, + FairShareProfile::Global, + estimated_output_tokens, + ) + .await + } + + pub(crate) async fn admit_for_tenant_profile( + self: &Arc, + class: Class, + request_id: RequestId, + cancel: CancellationToken, + tenant: TenantKey, + profile: FairShareProfile, + estimated_output_tokens: u32, ) -> AdmitOutcome { if self.fair_share.is_none() { return self.admit(class, request_id, cancel).await; @@ -427,6 +447,7 @@ impl PriorityScheduler { request_id, tx, tenant, + profile, estimated_output_tokens, ); if self.class_queues[class as usize] @@ -681,6 +702,7 @@ impl PriorityScheduler { } continue; } + let fair_share_reservation = fair_share_reservation.map(|reservation| *reservation); let permit = self.register_inflight(class, request_id, fair_share_reservation); // If the receiver was dropped between is_closed() above and // send below (unlikely race window), cancel the provisional @@ -2058,10 +2080,12 @@ mod tests { default_weight: 1.0, default_output_tokens: 10, trust_output_token_estimate_header: false, + trust_request_model_header: false, tenant_weights: HashMap::from([ ("header:a".to_string(), 1.0), ("header:b".to_string(), 1.0), ]), + model_profiles: HashMap::new(), }), ..Default::default() }; @@ -2115,6 +2139,7 @@ mod tests { rid("b-queued"), b_tx, b, + FairShareProfile::Global, 10, )) .unwrap(); @@ -2147,6 +2172,7 @@ mod tests { rid("a-queued"), a_tx, a, + FairShareProfile::Global, 10, )) .unwrap(); diff --git a/model_gateway/src/middleware/scheduler/fair_share.rs b/model_gateway/src/middleware/scheduler/fair_share.rs index aaa450851..842374874 100644 --- a/model_gateway/src/middleware/scheduler/fair_share.rs +++ b/model_gateway/src/middleware/scheduler/fair_share.rs @@ -1,11 +1,13 @@ //! Process-wide output-token accounting and weighted service with local queues. //! //! Every priority-scheduler partition receives the same [`Arc`]. -//! Lifetime charged-token accounting and canonical tenant virtual finish span -//! that process. Each partition stores only local queued/reservation membership. -//! Queued work is registered before a partition asks for its next local waiter, -//! and the ledger chooses the globally least-served tenant among candidates -//! eligible for that partition. This keeps every model pool work-conserving: +//! Lifetime charged-token accounting spans that process. Flat configuration +//! also shares canonical tenant virtual finish across partitions. Optional +//! model profiles instead keep virtual service and active-set time model-local, +//! with a hierarchical outer named-tenant/other-bucket policy. Each partition +//! stores only local queued/reservation membership. Queued work is registered +//! before a partition asks for its next local waiter. This keeps every model +//! pool work-conserving: //! unrelated or non-fungible capacity is never idled to repay another tenant's //! debt. //! @@ -13,10 +15,9 @@ //! all workload types that resolve to the same tenant key in one SMG process; //! it does not coordinate separate gateways or overlapping blue/green //! processes. Because model pools are non-fungible, it also cannot guarantee -//! exact aggregate percentages when users target disjoint pools. It enforces -//! the configured ratios whenever weighted tenants contend for substitutable -//! eligible capacity. Service and contention debt can follow a tenant across -//! pools, but cannot force a disjoint pool to idle. +//! exact aggregate percentages when users target disjoint pools. Flat service +//! debt can follow a tenant across pools; model-profile debt never crosses +//! models. Neither mode can force a disjoint pool to idle. use std::{ collections::{HashMap, HashSet}, @@ -35,6 +36,30 @@ use crate::tenant::TenantKey; /// Trusted estimate injected by the authenticated Comet proxy. pub const OUTPUT_TOKEN_ESTIMATE_HEADER: &str = "x-smg-output-token-estimate"; +/// Trusted canonical-model selector injected by an authenticated proxy. +pub const REQUEST_MODEL_HEADER: &str = "x-smg-request-model"; + +/// Scheduling-debt scope selected for one request. +/// +/// `Global` preserves the original process-wide flat ledger. A configured +/// model profile has its own virtual clock and hierarchical service state, +/// while charged/reserved token accounting remains process-global. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) enum FairShareProfile { + Global, + Model(Arc), +} + +impl FairShareProfile { + #[must_use] + pub(crate) fn metric_label(&self) -> &str { + match self { + Self::Global => "global", + Self::Model(model) => model, + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SettlementKind { Observed, @@ -64,6 +89,34 @@ struct ScopeTenant { #[derive(Debug, Default)] struct ScopeLedger { tenants: HashMap, + model_tenants: HashMap, HashMap>, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +enum ModelBucket { + Tenant(TenantKey), + Other, +} + +#[derive(Debug)] +struct ServiceClock { + entries: HashMap, + system_virtual_time: f64, +} + +impl Default for ServiceClock { + fn default() -> Self { + Self { + entries: HashMap::new(), + system_virtual_time: 0.0, + } + } +} + +#[derive(Debug, Default)] +struct ModelServiceLedger { + outer: ServiceClock, + other: ServiceClock, } #[derive(Debug, Default)] @@ -72,6 +125,7 @@ struct LedgerState { service: HashMap, scopes: HashMap, system_virtual_time: f64, + model_service: HashMap, ModelServiceLedger>, } /// One local candidate offered by a partition queue. @@ -81,6 +135,13 @@ pub(crate) struct FairShareCandidate<'a> { pub estimated_output_tokens: u32, } +#[derive(Debug)] +struct ModelFairSharePolicy { + canonical_model: Arc, + tenant_weights: HashMap, + other_weight: f64, +} + /// The candidate selected by the shared ledger, plus its provisional charge. pub(crate) struct FairShareSelection { pub index: usize, @@ -92,7 +153,9 @@ pub struct GlobalFairShare { default_weight: f64, default_output_tokens: u32, trust_output_token_estimate_header: bool, + trust_request_model_header: bool, tenant_weights: HashMap, + model_profiles: HashMap, metric_tenants: HashSet, state: Mutex, next_scope: AtomicU64, @@ -107,7 +170,12 @@ impl std::fmt::Debug for GlobalFairShare { "trust_output_token_estimate_header", &self.trust_output_token_estimate_header, ) + .field( + "trust_request_model_header", + &self.trust_request_model_header, + ) .field("tenant_weights", &self.tenant_weights) + .field("model_profiles", &self.model_profiles) .finish_non_exhaustive() } } @@ -117,8 +185,19 @@ impl GlobalFairShare { pub fn from_settings(settings: &SchedulerSettings) -> Option { settings.fair_share_config().map(|config| { let mut ledger = Self::from_config(config); - let mut metric_tenants: Vec<_> = ledger.tenant_weights.keys().cloned().collect(); + let mut metric_tenants: Vec<_> = ledger + .tenant_weights + .keys() + .chain( + ledger + .model_profiles + .values() + .flat_map(|profile| profile.tenant_weights.keys()), + ) + .cloned() + .collect(); metric_tenants.sort_by(|left, right| left.as_str().cmp(right.as_str())); + metric_tenants.dedup(); metric_tenants.truncate(settings.tenant_metric_top_n as usize); ledger.metric_tenants = metric_tenants.into_iter().collect(); ledger @@ -132,12 +211,42 @@ impl GlobalFairShare { .iter() .map(|(tenant, weight)| (TenantKey::new(tenant), *weight)) .collect(); - let metric_tenants = tenant_weights.keys().cloned().collect(); + let model_profiles: HashMap = config + .model_profiles + .iter() + .map(|(model, profile)| { + let tenant_weights = profile + .tenant_weights + .iter() + .map(|(tenant, weight)| (TenantKey::new(tenant), *weight)) + .collect(); + let canonical_model = Arc::new(model.clone()); + ( + model.clone(), + ModelFairSharePolicy { + canonical_model, + tenant_weights, + other_weight: profile.other_weight, + }, + ) + }) + .collect(); + let metric_tenants = tenant_weights + .keys() + .chain( + model_profiles + .values() + .flat_map(|profile: &ModelFairSharePolicy| profile.tenant_weights.keys()), + ) + .cloned() + .collect(); Self { default_weight: config.default_weight, default_output_tokens: config.default_output_tokens, trust_output_token_estimate_header: config.trust_output_token_estimate_header, + trust_request_model_header: config.trust_request_model_header, tenant_weights, + model_profiles, metric_tenants, state: Mutex::new(LedgerState::default()), next_scope: AtomicU64::new(0), @@ -154,15 +263,88 @@ impl GlobalFairShare { self.trust_output_token_estimate_header } + #[must_use] + pub fn trusts_request_model_header(&self) -> bool { + self.trust_request_model_header + } + + #[must_use] + pub fn has_model_profiles(&self) -> bool { + !self.model_profiles.is_empty() + } + + #[must_use] + pub(crate) fn profile_for_model(&self, model: &str) -> FairShareProfile { + self.model_profiles + .get(model) + .map(|policy| FairShareProfile::Model(Arc::clone(&policy.canonical_model))) + .unwrap_or(FairShareProfile::Global) + } + pub(crate) fn new_scope(&self) -> u64 { self.next_scope.fetch_add(1, Ordering::Relaxed) } - pub(crate) fn record_queue_wait(&self, tenant: &TenantKey, class: Class, wait: Duration) { - super::metrics::record_fair_share_queue_wait(self.metric_tenant(tenant), class, wait); + pub(crate) fn record_queue_wait( + &self, + profile: &FairShareProfile, + tenant: &TenantKey, + class: Class, + wait: Duration, + ) { + super::metrics::record_fair_share_queue_wait( + profile.metric_label(), + self.metric_tenant(tenant), + class, + wait, + ); } + #[cfg(test)] pub(crate) fn register_waiter(&self, scope_id: u64, tenant: &TenantKey, class: Class) { + self.register_waiter_in_profile(scope_id, &FairShareProfile::Global, tenant, class); + } + + pub(crate) fn register_waiter_in_profile( + &self, + scope_id: u64, + profile: &FairShareProfile, + tenant: &TenantKey, + class: Class, + ) { + let FairShareProfile::Model(model) = profile else { + self.register_global_waiter(scope_id, tenant, class); + return; + }; + let Some(policy) = self.model_profiles.get(model.as_str()) else { + self.register_global_waiter(scope_id, tenant, class); + return; + }; + let bucket = if policy.tenant_weights.contains_key(tenant) { + ModelBucket::Tenant(tenant.clone()) + } else { + ModelBucket::Other + }; + let is_other = bucket == ModelBucket::Other; + let mut state = self.state.lock(); + let ledger = state.model_service.entry(Arc::clone(model)).or_default(); + Self::register_service(&mut ledger.outer, bucket); + if is_other { + Self::register_service(&mut ledger.other, tenant.clone()); + } + state + .scopes + .entry(scope_id) + .or_default() + .model_tenants + .entry(Arc::clone(model)) + .or_default() + .entry(tenant.clone()) + .or_default() + .queued[class as usize] += 1; + } + + fn register_global_waiter(&self, scope_id: u64, tenant: &TenantKey, class: Class) { let unknown = !self.tenant_weights.contains_key(tenant); let mut state = self.state.lock(); let system_virtual_time = state.system_virtual_time; @@ -189,7 +371,52 @@ impl GlobalFairShare { } } - pub(crate) fn remove_waiter(&self, scope_id: u64, tenant: &TenantKey, class: Class) { + pub(crate) fn remove_waiter_in_profile( + &self, + scope_id: u64, + profile: &FairShareProfile, + tenant: &TenantKey, + class: Class, + ) { + let FairShareProfile::Model(model) = profile else { + self.remove_global_waiter(scope_id, tenant, class); + return; + }; + let Some(policy) = self.model_profiles.get(model.as_str()) else { + self.remove_global_waiter(scope_id, tenant, class); + return; + }; + let bucket = if policy.tenant_weights.contains_key(tenant) { + ModelBucket::Tenant(tenant.clone()) + } else { + ModelBucket::Other + }; + let is_other = bucket == ModelBucket::Other; + let mut state = self.state.lock(); + let Some(membership) = state + .scopes + .get_mut(&scope_id) + .and_then(|scope| scope.model_tenants.get_mut(model)) + .and_then(|tenants| tenants.get_mut(tenant)) + else { + return; + }; + let queued = &mut membership.queued[class as usize]; + debug_assert!(*queued > 0, "model fair-share waiter removed below zero"); + if *queued == 0 { + return; + } + *queued -= 1; + let Some(ledger) = state.model_service.get_mut(model) else { + return; + }; + Self::remove_service_waiter(&mut ledger.outer, &bucket); + if is_other { + Self::remove_service_waiter(&mut ledger.other, tenant); + } + } + + fn remove_global_waiter(&self, scope_id: u64, tenant: &TenantKey, class: Class) { let mut state = self.state.lock(); { let Some(scope) = state.scopes.get_mut(&scope_id) else { @@ -217,11 +444,44 @@ impl GlobalFairShare { /// Selection never consults tenants that are queued only in another /// partition. That is the work-conserving boundary for non-fungible model /// pools: a local slot is never idled for work that cannot use it. + #[cfg(test)] pub(crate) fn reserve_local_candidate( self: &Arc, scope_id: u64, class: Class, candidates: &[FairShareCandidate<'_>], + ) -> Option { + self.reserve_local_candidate_in_profile( + scope_id, + &FairShareProfile::Global, + class, + candidates, + ) + } + + pub(crate) fn reserve_local_candidate_in_profile( + self: &Arc, + scope_id: u64, + profile: &FairShareProfile, + class: Class, + candidates: &[FairShareCandidate<'_>], + ) -> Option { + match profile { + FairShareProfile::Global => self.reserve_global_candidate(scope_id, class, candidates), + FairShareProfile::Model(model) if self.model_profiles.contains_key(model.as_str()) => { + self.reserve_model_candidate(scope_id, model, class, candidates) + } + FairShareProfile::Model(_) => { + self.reserve_global_candidate(scope_id, class, candidates) + } + } + } + + fn reserve_global_candidate( + self: &Arc, + scope_id: u64, + class: Class, + candidates: &[FairShareCandidate<'_>], ) -> Option { let mut state = self.state.lock(); let class_index = class as usize; @@ -265,7 +525,11 @@ impl GlobalFairShare { let reserved_output_tokens = accounting.reserved_output_tokens; drop(state); - super::metrics::set_fair_share_virtual_finish(self.metric_tenant(&tenant), virtual_service); + super::metrics::set_fair_share_virtual_finish( + "global", + self.metric_tenant(&tenant), + virtual_service, + ); super::metrics::set_fair_share_reserved_output_tokens( self.metric_tenant(&tenant), reserved_output_tokens, @@ -277,6 +541,144 @@ impl GlobalFairShare { ledger: Arc::clone(self), tenant, scope_id, + profile: FairShareProfile::Global, + estimated_output_tokens, + settled: false, + }, + }) + } + + fn reserve_model_candidate( + self: &Arc, + scope_id: u64, + model: &Arc, + class: Class, + candidates: &[FairShareCandidate<'_>], + ) -> Option { + let policy = self.model_profiles.get(model.as_str())?; + let class_index = class as usize; + let mut state = self.state.lock(); + let selected_bucket = { + let scope_tenants = state.scopes.get(&scope_id)?.model_tenants.get(model)?; + let ledger = state.model_service.get(model)?; + candidates + .iter() + .filter_map(|candidate| { + let membership = scope_tenants.get(candidate.tenant)?; + if membership.queued[class_index] == 0 { + return None; + } + let bucket = Self::model_bucket(policy, candidate.tenant); + let service = ledger.outer.entries.get(&bucket)?; + Some((bucket, service.virtual_finish)) + }) + .min_by(|(left_bucket, left_finish), (right_bucket, right_finish)| { + left_finish + .total_cmp(right_finish) + .then_with(|| Self::compare_model_buckets(left_bucket, right_bucket)) + })? + .0 + }; + + let selected = { + let scope_tenants = state.scopes.get(&scope_id)?.model_tenants.get(model)?; + let ledger = state.model_service.get(model)?; + candidates + .iter() + .filter(|candidate| { + scope_tenants + .get(candidate.tenant) + .is_some_and(|membership| membership.queued[class_index] > 0) + && Self::model_bucket(policy, candidate.tenant) == selected_bucket + }) + .min_by(|left, right| { + let left_finish = if selected_bucket == ModelBucket::Other { + ledger + .other + .entries + .get(left.tenant) + .map_or(f64::INFINITY, |service| service.virtual_finish) + } else { + 0.0 + }; + let right_finish = if selected_bucket == ModelBucket::Other { + ledger + .other + .entries + .get(right.tenant) + .map_or(f64::INFINITY, |service| service.virtual_finish) + } else { + 0.0 + }; + left_finish + .total_cmp(&right_finish) + .then_with(|| left.tenant.as_str().cmp(right.tenant.as_str())) + .then_with(|| left.index.cmp(&right.index)) + })? + }; + + let tenant = selected.tenant.clone(); + let selected_index = selected.index; + let estimated_output_tokens = selected.estimated_output_tokens.max(1); + let membership = state + .scopes + .get_mut(&scope_id)? + .model_tenants + .get_mut(model)? + .get_mut(&tenant)?; + debug_assert!(membership.queued[class_index] > 0); + membership.queued[class_index] = membership.queued[class_index].saturating_sub(1); + membership.active_reservations = membership.active_reservations.saturating_add(1); + + let ledger = state.model_service.get_mut(model)?; + let outer_weight = self.model_bucket_weight(policy, &selected_bucket); + let outer_service = ledger.outer.entries.get_mut(&selected_bucket)?; + outer_service.queued = outer_service.queued.saturating_sub(1); + outer_service.active_reservations = outer_service.active_reservations.saturating_add(1); + outer_service.virtual_finish += f64::from(estimated_output_tokens) / outer_weight; + let outer_virtual_finish = outer_service.virtual_finish; + + let inner_virtual_finish = if selected_bucket == ModelBucket::Other { + let service = ledger.other.entries.get_mut(&tenant)?; + service.queued = service.queued.saturating_sub(1); + service.active_reservations = service.active_reservations.saturating_add(1); + service.virtual_finish += f64::from(estimated_output_tokens) / self.default_weight; + Some(service.virtual_finish) + } else { + None + }; + Self::advance_service_clock(&mut ledger.outer); + if selected_bucket == ModelBucket::Other { + Self::advance_service_clock(&mut ledger.other); + } + + let accounting = state.accounting.entry(tenant.clone()).or_default(); + accounting.reserved_output_tokens = accounting + .reserved_output_tokens + .saturating_add(u64::from(estimated_output_tokens)); + let reserved_output_tokens = accounting.reserved_output_tokens; + drop(state); + + super::metrics::set_fair_share_virtual_finish( + model, + self.metric_tenant(&tenant), + inner_virtual_finish.unwrap_or(outer_virtual_finish), + ); + if selected_bucket == ModelBucket::Other { + super::metrics::set_fair_share_other_bucket_virtual_finish(model, outer_virtual_finish); + } + super::metrics::set_fair_share_reserved_output_tokens( + self.metric_tenant(&tenant), + reserved_output_tokens, + ); + + Some(FairShareSelection { + index: selected_index, + reservation: FairShareReservation { + ledger: Arc::clone(self), + tenant, + scope_id, + profile: FairShareProfile::Model(Arc::clone(model)), estimated_output_tokens, settled: false, }, @@ -290,10 +692,75 @@ impl GlobalFairShare { .unwrap_or(self.default_weight) } + fn model_bucket(policy: &ModelFairSharePolicy, tenant: &TenantKey) -> ModelBucket { + if policy.tenant_weights.contains_key(tenant) { + ModelBucket::Tenant(tenant.clone()) + } else { + ModelBucket::Other + } + } + + fn model_bucket_weight(&self, policy: &ModelFairSharePolicy, bucket: &ModelBucket) -> f64 { + match bucket { + ModelBucket::Tenant(tenant) => policy + .tenant_weights + .get(tenant) + .copied() + .unwrap_or(self.default_weight), + ModelBucket::Other => policy.other_weight, + } + } + + fn compare_model_buckets(left: &ModelBucket, right: &ModelBucket) -> std::cmp::Ordering { + match (left, right) { + (ModelBucket::Tenant(left), ModelBucket::Tenant(right)) => { + left.as_str().cmp(right.as_str()) + } + (ModelBucket::Tenant(_), ModelBucket::Other) => std::cmp::Ordering::Less, + (ModelBucket::Other, ModelBucket::Tenant(_)) => std::cmp::Ordering::Greater, + (ModelBucket::Other, ModelBucket::Other) => std::cmp::Ordering::Equal, + } + } + fn is_active(service: &TenantService) -> bool { service.active_reservations > 0 || service.queued > 0 } + fn register_service(clock: &mut ServiceClock, key: K) + where + K: Eq + std::hash::Hash, + { + let system_virtual_time = clock.system_virtual_time; + let service = clock.entries.entry(key).or_default(); + if !Self::is_active(service) { + service.virtual_finish = service.virtual_finish.max(system_virtual_time); + } + service.queued = service.queued.saturating_add(1); + Self::advance_service_clock(clock); + } + + fn remove_service_waiter(clock: &mut ServiceClock, key: &K) + where + K: Eq + std::hash::Hash, + { + if let Some(service) = clock.entries.get_mut(key) { + service.queued = service.queued.saturating_sub(1); + } + Self::advance_service_clock(clock); + } + + fn advance_service_clock(clock: &mut ServiceClock) { + if let Some(active_minimum) = clock + .entries + .values() + .filter(|service| Self::is_active(service)) + .map(|service| service.virtual_finish) + .min_by(f64::total_cmp) + { + clock.system_virtual_time = clock.system_virtual_time.max(active_minimum); + } + } + fn advance_system_virtual_time(state: &mut LedgerState) { if let Some(active_minimum) = state .service @@ -314,7 +781,32 @@ impl GlobalFairShare { } } - fn cancel_reservation(&self, scope_id: u64, tenant: &TenantKey, estimated_output_tokens: u32) { + fn cancel_reservation( + &self, + scope_id: u64, + profile: &FairShareProfile, + tenant: &TenantKey, + estimated_output_tokens: u32, + ) { + match profile { + FairShareProfile::Global => { + self.cancel_global_reservation(scope_id, tenant, estimated_output_tokens); + } + FairShareProfile::Model(model) if self.model_profiles.contains_key(model.as_str()) => { + self.cancel_model_reservation(scope_id, model, tenant, estimated_output_tokens); + } + FairShareProfile::Model(_) => { + self.cancel_global_reservation(scope_id, tenant, estimated_output_tokens); + } + } + } + + fn cancel_global_reservation( + &self, + scope_id: u64, + tenant: &TenantKey, + estimated_output_tokens: u32, + ) { let mut state = self.state.lock(); { let Some(membership) = state @@ -345,7 +837,88 @@ impl GlobalFairShare { .saturating_sub(u64::from(estimated_output_tokens)); let reserved_output_tokens = accounting.reserved_output_tokens; drop(state); - super::metrics::set_fair_share_virtual_finish(self.metric_tenant(tenant), virtual_service); + super::metrics::set_fair_share_virtual_finish( + "global", + self.metric_tenant(tenant), + virtual_service, + ); + super::metrics::set_fair_share_reserved_output_tokens( + self.metric_tenant(tenant), + reserved_output_tokens, + ); + } + + fn cancel_model_reservation( + &self, + scope_id: u64, + model: &Arc, + tenant: &TenantKey, + estimated_output_tokens: u32, + ) { + let Some(policy) = self.model_profiles.get(model.as_str()) else { + return; + }; + let bucket = Self::model_bucket(policy, tenant); + let is_other = bucket == ModelBucket::Other; + let mut state = self.state.lock(); + let Some(membership) = state + .scopes + .get_mut(&scope_id) + .and_then(|scope| scope.model_tenants.get_mut(model)) + .and_then(|tenants| tenants.get_mut(tenant)) + else { + return; + }; + if membership.active_reservations == 0 { + return; + } + membership.active_reservations -= 1; + + let Some(ledger) = state.model_service.get_mut(model) else { + return; + }; + let outer_weight = self.model_bucket_weight(policy, &bucket); + let outer_system_virtual_time = ledger.outer.system_virtual_time; + let Some(outer_service) = ledger.outer.entries.get_mut(&bucket) else { + return; + }; + outer_service.active_reservations = outer_service.active_reservations.saturating_sub(1); + outer_service.virtual_finish = (outer_service.virtual_finish + - f64::from(estimated_output_tokens) / outer_weight) + .max(outer_system_virtual_time); + let outer_virtual_finish = outer_service.virtual_finish; + let inner_virtual_finish = if is_other { + let inner_system_virtual_time = ledger.other.system_virtual_time; + let Some(service) = ledger.other.entries.get_mut(tenant) else { + return; + }; + service.active_reservations = service.active_reservations.saturating_sub(1); + service.virtual_finish = (service.virtual_finish + - f64::from(estimated_output_tokens) / self.default_weight) + .max(inner_system_virtual_time); + Some(service.virtual_finish) + } else { + None + }; + Self::advance_service_clock(&mut ledger.outer); + if is_other { + Self::advance_service_clock(&mut ledger.other); + } + let accounting = state.accounting.entry(tenant.clone()).or_default(); + accounting.reserved_output_tokens = accounting + .reserved_output_tokens + .saturating_sub(u64::from(estimated_output_tokens)); + let reserved_output_tokens = accounting.reserved_output_tokens; + drop(state); + + super::metrics::set_fair_share_virtual_finish( + model, + self.metric_tenant(tenant), + inner_virtual_finish.unwrap_or(outer_virtual_finish), + ); + if is_other { + super::metrics::set_fair_share_other_bucket_virtual_finish(model, outer_virtual_finish); + } super::metrics::set_fair_share_reserved_output_tokens( self.metric_tenant(tenant), reserved_output_tokens, @@ -353,6 +926,43 @@ impl GlobalFairShare { } fn settle_reservation( + &self, + scope_id: u64, + profile: &FairShareProfile, + tenant: &TenantKey, + estimated_output_tokens: u32, + observed_output_tokens: Option, + kind: SettlementKind, + ) { + match profile { + FairShareProfile::Global => self.settle_global_reservation( + scope_id, + tenant, + estimated_output_tokens, + observed_output_tokens, + kind, + ), + FairShareProfile::Model(model) if self.model_profiles.contains_key(model.as_str()) => { + self.settle_model_reservation( + scope_id, + model, + tenant, + estimated_output_tokens, + observed_output_tokens, + kind, + ); + } + FairShareProfile::Model(_) => self.settle_global_reservation( + scope_id, + tenant, + estimated_output_tokens, + observed_output_tokens, + kind, + ), + } + } + + fn settle_global_reservation( &self, scope_id: u64, tenant: &TenantKey, @@ -392,7 +1002,89 @@ impl GlobalFairShare { let metric_tenant = self.metric_tenant(tenant); super::metrics::record_fair_share_charged_output_tokens(metric_tenant, charged); - super::metrics::set_fair_share_virtual_finish(metric_tenant, virtual_service); + super::metrics::set_fair_share_virtual_finish("global", metric_tenant, virtual_service); + super::metrics::set_fair_share_reserved_output_tokens( + metric_tenant, + reserved_output_tokens, + ); + if kind != SettlementKind::Observed { + super::metrics::record_fair_share_fallback(kind.as_str()); + } + } + + fn settle_model_reservation( + &self, + scope_id: u64, + model: &Arc, + tenant: &TenantKey, + estimated_output_tokens: u32, + observed_output_tokens: Option, + kind: SettlementKind, + ) { + let Some(policy) = self.model_profiles.get(model.as_str()) else { + return; + }; + let charged = observed_output_tokens.unwrap_or(estimated_output_tokens); + let bucket = Self::model_bucket(policy, tenant); + let is_other = bucket == ModelBucket::Other; + let mut state = self.state.lock(); + let membership = state + .scopes + .entry(scope_id) + .or_default() + .model_tenants + .entry(Arc::clone(model)) + .or_default() + .entry(tenant.clone()) + .or_default(); + membership.active_reservations = membership.active_reservations.saturating_sub(1); + + let ledger = state.model_service.entry(Arc::clone(model)).or_default(); + let outer_weight = self.model_bucket_weight(policy, &bucket); + let outer_service = ledger.outer.entries.entry(bucket.clone()).or_default(); + outer_service.active_reservations = outer_service.active_reservations.saturating_sub(1); + let outer_correction = + (f64::from(charged) - f64::from(estimated_output_tokens)) / outer_weight; + outer_service.virtual_finish = + (outer_service.virtual_finish + outer_correction).max(ledger.outer.system_virtual_time); + let outer_virtual_finish = outer_service.virtual_finish; + + let inner_virtual_finish = if is_other { + let service = ledger.other.entries.entry(tenant.clone()).or_default(); + service.active_reservations = service.active_reservations.saturating_sub(1); + let correction = + (f64::from(charged) - f64::from(estimated_output_tokens)) / self.default_weight; + service.virtual_finish = + (service.virtual_finish + correction).max(ledger.other.system_virtual_time); + Some(service.virtual_finish) + } else { + None + }; + Self::advance_service_clock(&mut ledger.outer); + if is_other { + Self::advance_service_clock(&mut ledger.other); + } + + let accounting = state.accounting.entry(tenant.clone()).or_default(); + accounting.reserved_output_tokens = accounting + .reserved_output_tokens + .saturating_sub(u64::from(estimated_output_tokens)); + accounting.charged_output_tokens = accounting + .charged_output_tokens + .saturating_add(u64::from(charged)); + let reserved_output_tokens = accounting.reserved_output_tokens; + drop(state); + + let metric_tenant = self.metric_tenant(tenant); + super::metrics::record_fair_share_charged_output_tokens(metric_tenant, charged); + super::metrics::set_fair_share_virtual_finish( + model, + metric_tenant, + inner_virtual_finish.unwrap_or(outer_virtual_finish), + ); + if is_other { + super::metrics::set_fair_share_other_bucket_virtual_finish(model, outer_virtual_finish); + } super::metrics::set_fair_share_reserved_output_tokens( metric_tenant, reserved_output_tokens, @@ -426,6 +1118,38 @@ impl GlobalFairShare { ) } + #[cfg(test)] + fn model_snapshot( + &self, + scope_id: u64, + model: &str, + tenant: &TenantKey, + ) -> (u64, u64, u64, [usize; 4]) { + let state = self.state.lock(); + let accounting = state + .accounting + .get(tenant) + .expect("tenant accounting exists"); + let membership = state + .scopes + .get(&scope_id) + .and_then(|scope| { + scope + .model_tenants + .iter() + .find(|(configured, _)| configured.as_str() == model) + .map(|(_, tenants)| tenants) + }) + .and_then(|tenants| tenants.get(tenant)) + .expect("tenant model membership exists"); + ( + accounting.charged_output_tokens, + accounting.reserved_output_tokens, + membership.active_reservations, + membership.queued, + ) + } + #[cfg(test)] fn virtual_snapshot(&self, scope_id: u64, tenant: &TenantKey) -> (f64, f64) { let state = self.state.lock(); @@ -464,6 +1188,7 @@ pub struct FairShareReservation { ledger: Arc, tenant: TenantKey, scope_id: u64, + profile: FairShareProfile, estimated_output_tokens: u32, settled: bool, } @@ -473,6 +1198,7 @@ impl std::fmt::Debug for FairShareReservation { f.debug_struct("FairShareReservation") .field("tenant", &self.tenant) .field("scope_id", &self.scope_id) + .field("profile", &self.profile) .field("estimated_output_tokens", &self.estimated_output_tokens) .field("settled", &self.settled) .finish() @@ -483,6 +1209,7 @@ impl FairShareReservation { pub fn settle(mut self, observed_output_tokens: Option, kind: SettlementKind) { self.ledger.settle_reservation( self.scope_id, + &self.profile, &self.tenant, self.estimated_output_tokens, observed_output_tokens, @@ -492,8 +1219,12 @@ impl FairShareReservation { } pub fn cancel(mut self) { - self.ledger - .cancel_reservation(self.scope_id, &self.tenant, self.estimated_output_tokens); + self.ledger.cancel_reservation( + self.scope_id, + &self.profile, + &self.tenant, + self.estimated_output_tokens, + ); self.settled = true; } } @@ -505,6 +1236,7 @@ impl Drop for FairShareReservation { } self.ledger.settle_reservation( self.scope_id, + &self.profile, &self.tenant, self.estimated_output_tokens, None, @@ -517,16 +1249,19 @@ impl Drop for FairShareReservation { #[cfg(test)] mod tests { use super::*; + use crate::middleware::scheduler::ModelFairShareConfig; fn config(weights: &[(&str, f64)]) -> FairShareConfig { FairShareConfig { default_weight: 1.0, default_output_tokens: 10, trust_output_token_estimate_header: false, + trust_request_model_header: false, tenant_weights: weights .iter() .map(|(tenant, weight)| ((*tenant).to_string(), *weight)) .collect(), + model_profiles: HashMap::new(), } } @@ -551,6 +1286,55 @@ mod tests { .map(|selection| selection.reservation) } + fn model_config() -> FairShareConfig { + let mut config = config(&[]); + config.trust_request_model_header = true; + config.model_profiles = HashMap::from([ + ( + "deepseek-v4-flash".to_string(), + ModelFairShareConfig { + tenant_weights: HashMap::from([ + ("header:junu".to_string(), 30.0), + ("header:xuezhou".to_string(), 40.0), + ("header:zhenting".to_string(), 10.0), + ]), + other_weight: 20.0, + }, + ), + ( + "kimi-k3".to_string(), + ModelFairShareConfig { + tenant_weights: HashMap::from([("header:mukhesh".to_string(), 80.0)]), + other_weight: 20.0, + }, + ), + ]); + config + } + + fn reserve_one_in_profile( + ledger: &Arc, + scope_id: u64, + profile: &FairShareProfile, + class: Class, + tenant: &TenantKey, + estimated_output_tokens: u32, + ) -> Option { + ledger.register_waiter_in_profile(scope_id, profile, tenant, class); + ledger + .reserve_local_candidate_in_profile( + scope_id, + profile, + class, + &[FairShareCandidate { + index: 0, + tenant, + estimated_output_tokens, + }], + ) + .map(|selection| selection.reservation) + } + #[test] fn observed_tokens_replace_the_provisional_estimate_once() { let ledger = Arc::new(GlobalFairShare::from_config(&config(&[("header:a", 1.0)]))); @@ -891,4 +1675,170 @@ mod tests { selected.reservation.cancel(); assert_eq!(ledger.snapshot(scope_id, &a).0, 50); } + + #[test] + fn deepseek_profile_enforces_named_shares_and_one_aggregate_other_bucket() { + let ledger = Arc::new(GlobalFairShare::from_config(&model_config())); + let scope_id = ledger.new_scope(); + let profile = ledger.profile_for_model("deepseek-v4-flash"); + let tenants = [ + TenantKey::new("header:junu"), + TenantKey::new("header:xuezhou"), + TenantKey::new("header:zhenting"), + TenantKey::new("header:other-a"), + TenantKey::new("header:other-b"), + ]; + for tenant in &tenants { + for _ in 0..100 { + ledger.register_waiter_in_profile(scope_id, &profile, tenant, Class::Default); + } + } + let mut admitted = [0_u32; 5]; + for _ in 0..100 { + let candidates: Vec<_> = tenants + .iter() + .enumerate() + .map(|(index, tenant)| FairShareCandidate { + index, + tenant, + estimated_output_tokens: 10, + }) + .collect(); + let selection = ledger + .reserve_local_candidate_in_profile(scope_id, &profile, Class::Default, &candidates) + .expect("a deepseek contender must be selected"); + admitted[selection.index] += 1; + selection + .reservation + .settle(Some(10), SettlementKind::Observed); + } + + assert_eq!(admitted, [30, 40, 10, 10, 10]); + assert_eq!( + admitted[3] + admitted[4], + 20, + "all unlisted users share one fixed 20% outer bucket" + ); + } + + #[test] + fn kimi_profile_gives_mukhesh_eighty_and_other_users_twenty() { + let ledger = Arc::new(GlobalFairShare::from_config(&model_config())); + let scope_id = ledger.new_scope(); + let profile = ledger.profile_for_model("kimi-k3"); + let tenants = [ + TenantKey::new("header:mukhesh"), + TenantKey::new("header:other-a"), + TenantKey::new("header:other-b"), + ]; + for tenant in &tenants { + for _ in 0..100 { + ledger.register_waiter_in_profile(scope_id, &profile, tenant, Class::Default); + } + } + let mut admitted = [0_u32; 3]; + for _ in 0..100 { + let candidates: Vec<_> = tenants + .iter() + .enumerate() + .map(|(index, tenant)| FairShareCandidate { + index, + tenant, + estimated_output_tokens: 10, + }) + .collect(); + let selection = ledger + .reserve_local_candidate_in_profile(scope_id, &profile, Class::Default, &candidates) + .expect("a kimi contender must be selected"); + admitted[selection.index] += 1; + selection + .reservation + .settle(Some(10), SettlementKind::Observed); + } + + assert_eq!(admitted, [80, 10, 10]); + } + + #[test] + fn model_scoped_debt_is_independent_but_accounting_is_process_global() { + let mut config = config(&[]); + config.trust_request_model_header = true; + let equal = ModelFairShareConfig { + tenant_weights: HashMap::from([ + ("header:a".to_string(), 1.0), + ("header:b".to_string(), 1.0), + ]), + other_weight: 1.0, + }; + config.model_profiles = HashMap::from([ + ("deepseek".to_string(), equal.clone()), + ("kimi".to_string(), equal), + ]); + let ledger = Arc::new(GlobalFairShare::from_config(&config)); + let scope_id = ledger.new_scope(); + let deepseek = ledger.profile_for_model("deepseek"); + let kimi = ledger.profile_for_model("kimi"); + let a = TenantKey::new("header:a"); + let b = TenantKey::new("header:b"); + + ledger.register_waiter_in_profile(scope_id, &deepseek, &b, Class::Default); + reserve_one_in_profile(&ledger, scope_id, &deepseek, Class::Default, &a, 50) + .unwrap() + .settle(Some(50), SettlementKind::Observed); + + ledger.register_waiter_in_profile(scope_id, &kimi, &a, Class::Default); + ledger.register_waiter_in_profile(scope_id, &kimi, &b, Class::Default); + let candidates = [ + FairShareCandidate { + index: 0, + tenant: &a, + estimated_output_tokens: 10, + }, + FairShareCandidate { + index: 1, + tenant: &b, + estimated_output_tokens: 10, + }, + ]; + let selected = ledger + .reserve_local_candidate_in_profile(scope_id, &kimi, Class::Default, &candidates) + .unwrap(); + assert_eq!( + selected.index, 0, + "deepseek debt must not bias an independent kimi profile" + ); + selected + .reservation + .settle(Some(10), SettlementKind::Observed); + assert_eq!( + ledger.model_snapshot(scope_id, "kimi", &a).0, + 60, + "charged output accounting remains process-global across models" + ); + } + + #[test] + fn model_reservation_cancel_and_missing_usage_keep_existing_semantics() { + let ledger = Arc::new(GlobalFairShare::from_config(&model_config())); + let scope_id = ledger.new_scope(); + let profile = ledger.profile_for_model("kimi-k3"); + let tenant = TenantKey::new("header:other-a"); + + reserve_one_in_profile(&ledger, scope_id, &profile, Class::Default, &tenant, 100) + .unwrap() + .cancel(); + assert_eq!( + ledger.model_snapshot(scope_id, "kimi-k3", &tenant), + (0, 0, 0, [0; 4]) + ); + + drop( + reserve_one_in_profile(&ledger, scope_id, &profile, Class::Default, &tenant, 100) + .unwrap(), + ); + assert_eq!( + ledger.model_snapshot(scope_id, "kimi-k3", &tenant), + (100, 0, 0, [0; 4]) + ); + } } diff --git a/model_gateway/src/middleware/scheduler/metrics.rs b/model_gateway/src/middleware/scheduler/metrics.rs index 1ed9d56a2..4f69054d4 100644 --- a/model_gateway/src/middleware/scheduler/metrics.rs +++ b/model_gateway/src/middleware/scheduler/metrics.rs @@ -27,6 +27,7 @@ const STARVATION_PROMOTION_TOTAL: &str = "smg_scheduler_starvation_promotion_tot const PARTITION_ADMIT_TOTAL: &str = "smg_scheduler_partition_admit_total"; const FAIR_SHARE_CHARGED_OUTPUT_TOKENS_TOTAL: &str = "smg_fair_share_charged_output_tokens_total"; const FAIR_SHARE_VIRTUAL_FINISH: &str = "smg_fair_share_virtual_finish"; +const FAIR_SHARE_OTHER_BUCKET_VIRTUAL_FINISH: &str = "smg_fair_share_other_bucket_virtual_finish"; const FAIR_SHARE_RESERVED_OUTPUT_TOKENS: &str = "smg_fair_share_reserved_output_tokens"; const FAIR_SHARE_QUEUE_WAIT_SECONDS: &str = "smg_fair_share_queue_wait_seconds"; const FAIR_SHARE_FALLBACK_TOTAL: &str = "smg_fair_share_fallback_total"; @@ -96,7 +97,11 @@ pub fn describe() { ); describe_gauge!( FAIR_SHARE_VIRTUAL_FINISH, - "Process-global active-set normalized virtual finish used for weighted dispatch" + "Active-set normalized tenant virtual finish by model profile" + ); + describe_gauge!( + FAIR_SHARE_OTHER_BUCKET_VIRTUAL_FINISH, + "Outer aggregate other-bucket virtual finish by model profile" ); describe_gauge!( FAIR_SHARE_RESERVED_OUTPUT_TOKENS, @@ -221,8 +226,21 @@ pub fn record_fair_share_charged_output_tokens(tenant: &str, tokens: u32) { .increment(u64::from(tokens)); } -pub fn set_fair_share_virtual_finish(tenant: &str, virtual_finish: f64) { - gauge!(FAIR_SHARE_VIRTUAL_FINISH, "tenant" => intern_string(tenant)).set(virtual_finish); +pub fn set_fair_share_virtual_finish(model: &str, tenant: &str, virtual_finish: f64) { + gauge!( + FAIR_SHARE_VIRTUAL_FINISH, + "model" => intern_string(model), + "tenant" => intern_string(tenant) + ) + .set(virtual_finish); +} + +pub fn set_fair_share_other_bucket_virtual_finish(model: &str, virtual_finish: f64) { + gauge!( + FAIR_SHARE_OTHER_BUCKET_VIRTUAL_FINISH, + "model" => intern_string(model) + ) + .set(virtual_finish); } pub fn set_fair_share_reserved_output_tokens(tenant: &str, tokens: u64) { @@ -233,9 +251,10 @@ pub fn set_fair_share_reserved_output_tokens(tenant: &str, tokens: u64) { .set(tokens as f64); } -pub fn record_fair_share_queue_wait(tenant: &str, class: Class, wait: Duration) { +pub fn record_fair_share_queue_wait(model: &str, tenant: &str, class: Class, wait: Duration) { histogram!( FAIR_SHARE_QUEUE_WAIT_SECONDS, + "model" => intern_string(model), "tenant" => intern_string(tenant), "class" => class.as_str() ) diff --git a/model_gateway/src/middleware/scheduler/mod.rs b/model_gateway/src/middleware/scheduler/mod.rs index ad4a6d90b..20c6652f7 100644 --- a/model_gateway/src/middleware/scheduler/mod.rs +++ b/model_gateway/src/middleware/scheduler/mod.rs @@ -21,13 +21,16 @@ pub use body::SchedulerGuardBody; pub use class::{Class, PRIORITY_HEADER}; pub use config::{ AdmissionPartitionConfig, ClassConfig, ClassRuntimeConfig, FairShareConfig, - PrioritySchedulerYaml, SchedulerSettings, SettingsValidationError, TenantPolicyConfig, + ModelFairShareConfig, PrioritySchedulerYaml, SchedulerSettings, SettingsValidationError, + TenantPolicyConfig, }; pub use engine::{ AdmitOutcome, PriorityScheduler, RejectionReason, SchedulerInitError, SchedulerPermit, }; pub use error::{SchedulerError, HEADER_X_SMG_PREEMPTED}; pub use extract::PreemptionGuard; -pub use fair_share::{GlobalFairShare, SettlementKind, OUTPUT_TOKEN_ESTIMATE_HEADER}; +pub use fair_share::{ + GlobalFairShare, SettlementKind, OUTPUT_TOKEN_ESTIMATE_HEADER, REQUEST_MODEL_HEADER, +}; pub use policy::{StaticTenantPolicyResolver, TenantPolicy, TenantPolicyResolver}; pub use state::{AdmissionMode, SchedulerState, ADMISSION_PARTITION_HEADER}; diff --git a/model_gateway/src/middleware/scheduler/queue.rs b/model_gateway/src/middleware/scheduler/queue.rs index dd2e54ffe..6bccdc05c 100644 --- a/model_gateway/src/middleware/scheduler/queue.rs +++ b/model_gateway/src/middleware/scheduler/queue.rs @@ -17,7 +17,7 @@ use tokio_util::sync::CancellationToken; use super::{ engine::SchedulerPermit, - fair_share::{FairShareCandidate, FairShareReservation, GlobalFairShare}, + fair_share::{FairShareCandidate, FairShareProfile, FairShareReservation, GlobalFairShare}, Class, }; use crate::tenant::TenantKey; @@ -96,8 +96,9 @@ pub struct Waiter { pub request_id: RequestId, pub permit_tx: oneshot::Sender, pub tenant: Option, + pub(crate) fair_share_profile: FairShareProfile, pub estimated_output_tokens: u32, - pub fair_share_reservation: Option, + pub fair_share_reservation: Option>, } impl Waiter { @@ -114,17 +115,19 @@ impl Waiter { request_id, permit_tx, tenant: None, + fair_share_profile: FairShareProfile::Global, estimated_output_tokens: 0, fair_share_reservation: None, } } - pub fn new_fair( + pub(crate) fn new_fair( class: Class, cancel: CancellationToken, request_id: RequestId, permit_tx: oneshot::Sender, tenant: TenantKey, + fair_share_profile: FairShareProfile, estimated_output_tokens: u32, ) -> Self { Self { @@ -134,6 +137,7 @@ impl Waiter { request_id, permit_tx, tenant: Some(tenant), + fair_share_profile, estimated_output_tokens: estimated_output_tokens.max(1), fair_share_reservation: None, } @@ -277,19 +281,29 @@ impl ClassQueue for FairClassQueue { return Err(waiter); } let mut guard = self.waiters.lock(); - self.ledger - .register_waiter(self.scope_id, tenant, self.class); + self.ledger.register_waiter_in_profile( + self.scope_id, + &waiter.fair_share_profile, + tenant, + self.class, + ); guard.push_back(waiter); Ok(()) } fn pop_eligible(&self) -> Option { let mut guard = self.waiters.lock(); + let profile = guard + .iter() + .filter(|waiter| !waiter.cancel.is_cancelled()) + .min_by_key(|waiter| waiter.queued_at) + .map(|waiter| waiter.fair_share_profile.clone())?; let selection = { let candidates: Vec<_> = guard .iter() .enumerate() .filter(|(_, waiter)| !waiter.cancel.is_cancelled()) + .filter(|(_, waiter)| waiter.fair_share_profile == profile) .filter_map(|(index, waiter)| { waiter.tenant.as_ref().map(|tenant| FairShareCandidate { index, @@ -298,8 +312,12 @@ impl ClassQueue for FairClassQueue { }) }) .collect(); - self.ledger - .reserve_local_candidate(self.scope_id, self.class, &candidates) + self.ledger.reserve_local_candidate_in_profile( + self.scope_id, + &profile, + self.class, + &candidates, + ) }?; let Some(mut waiter) = guard.remove(selection.index) else { selection.reservation.cancel(); @@ -307,10 +325,14 @@ impl ClassQueue for FairClassQueue { }; self.budget.release(); if let Some(tenant) = waiter.tenant.as_ref() { - self.ledger - .record_queue_wait(tenant, self.class, waiter.queued_at.elapsed()); + self.ledger.record_queue_wait( + &waiter.fair_share_profile, + tenant, + self.class, + waiter.queued_at.elapsed(), + ); } - waiter.fair_share_reservation = Some(selection.reservation); + waiter.fair_share_reservation = Some(Box::new(selection.reservation)); Some(waiter) } @@ -343,7 +365,12 @@ impl ClassQueue for FairClassQueue { }; self.budget.release(); if let Some(tenant) = waiter.tenant.as_ref() { - self.ledger.remove_waiter(self.scope_id, tenant, self.class); + self.ledger.remove_waiter_in_profile( + self.scope_id, + &waiter.fair_share_profile, + tenant, + self.class, + ); } } } @@ -351,12 +378,13 @@ impl ClassQueue for FairClassQueue { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::time::Duration; use tokio_util::sync::CancellationToken; use super::*; - use crate::middleware::scheduler::Class; + use crate::middleware::scheduler::{Class, FairShareConfig, ModelFairShareConfig}; fn waiter(class: Class) -> Waiter { let (tx, _rx) = oneshot::channel(); @@ -368,6 +396,45 @@ mod tests { Waiter::new(class, cancel, RequestId("t".into()), tx) } + fn model_ledger() -> Arc { + Arc::new(GlobalFairShare::from_config(&FairShareConfig { + default_weight: 1.0, + default_output_tokens: 10, + trust_output_token_estimate_header: false, + trust_request_model_header: true, + tenant_weights: HashMap::new(), + model_profiles: HashMap::from([ + ( + "model-a".to_string(), + ModelFairShareConfig { + tenant_weights: HashMap::new(), + other_weight: 1.0, + }, + ), + ( + "model-b".to_string(), + ModelFairShareConfig { + tenant_weights: HashMap::new(), + other_weight: 1.0, + }, + ), + ]), + })) + } + + fn model_waiter(id: &str, tenant: &str, profile: FairShareProfile) -> Waiter { + let (tx, _rx) = oneshot::channel(); + Waiter::new_fair( + Class::Default, + CancellationToken::new(), + RequestId(id.into()), + tx, + TenantKey::new(tenant), + profile, + 10, + ) + } + #[test] fn test_try_enqueue_fills_to_capacity_then_rejects() { let q = FifoClassQueue::new(4); @@ -439,6 +506,47 @@ mod tests { ); } + #[test] + fn oldest_model_group_wins_before_model_local_virtual_service() { + let ledger = model_ledger(); + let scope_id = ledger.new_scope(); + let queue = FairClassQueue::with_shared_budget( + Class::Default, + 8, + Arc::new(QueueBudget::new(8)), + Arc::clone(&ledger), + scope_id, + ); + let mut older = model_waiter( + "older-model-a", + "header:a", + ledger.profile_for_model("model-a"), + ); + older.queued_at = Instant::now() - Duration::from_secs(1); + let newer = model_waiter( + "newer-model-b", + "header:b", + ledger.profile_for_model("model-b"), + ); + queue.try_enqueue(older).unwrap(); + queue.try_enqueue(newer).unwrap(); + + let mut first = queue.pop_eligible().expect("oldest model group dispatches"); + assert_eq!(first.request_id.0, "older-model-a"); + first + .fair_share_reservation + .take() + .expect("model waiter has reservation") + .cancel(); + let mut second = queue.pop_eligible().expect("other model remains eligible"); + assert_eq!(second.request_id.0, "newer-model-b"); + second + .fair_share_reservation + .take() + .expect("model waiter has reservation") + .cancel(); + } + #[test] fn test_drop_cancelled_head_leaves_live_head_alone() { let q = FifoClassQueue::new(4); diff --git a/model_gateway/src/middleware/scheduler/state.rs b/model_gateway/src/middleware/scheduler/state.rs index a0549e594..836285438 100644 --- a/model_gateway/src/middleware/scheduler/state.rs +++ b/model_gateway/src/middleware/scheduler/state.rs @@ -12,6 +12,7 @@ use tokio::sync::{broadcast, watch}; use tracing::{error, info}; use super::{ + fair_share::{FairShareProfile, REQUEST_MODEL_HEADER}, Class, GlobalFairShare, PriorityScheduler, SchedulerSettings, StaticTenantPolicyResolver, TenantPolicyResolver, }; @@ -49,6 +50,7 @@ pub struct SchedulerState { partitions: HashMap, default_partition: Arc, pub resolver: Arc, + model_registry: Arc, /// Per-second RPS sibling check, run before admission. Set only when an /// explicit `rate_limit_tokens_per_second` is configured; the bucket's /// concurrency-cap role is owned by the scheduler, so we must not consult @@ -75,6 +77,37 @@ impl SchedulerState { scheduler: Arc::clone(&self.scheduler), }) } + + /// Resolve the trusted request model to a configured fair-share profile. + /// Model aliases are canonicalized only for accounting; the outbound + /// request body remains untouched by the scheduler. + pub(crate) fn fair_share_profile_for( + &self, + headers: &HeaderMap, + ledger: &GlobalFairShare, + ) -> FairShareProfile { + if !ledger.has_model_profiles() { + return FairShareProfile::Global; + } + if !ledger.trusts_request_model_header() { + super::metrics::record_fair_share_fallback("untrusted_request_model"); + return FairShareProfile::Global; + } + let Some(raw_model) = headers + .get(REQUEST_MODEL_HEADER) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + super::metrics::record_fair_share_fallback("missing_request_model"); + return FairShareProfile::Global; + }; + let canonical = self + .model_registry + .resolve_model_alias(raw_model) + .unwrap_or_else(|| Arc::from(raw_model)); + ledger.profile_for_model(&canonical) + } } /// Which admission path the protected routes use. Chosen once at startup. @@ -161,6 +194,7 @@ impl AdmissionMode { settings = settings.with_global_queue_budget(rc.queue_size); } let fair_share = GlobalFairShare::from_settings(&settings).map(Arc::new); + let model_registry = Arc::clone(®istry); let resolver: Arc = Arc::new(StaticTenantPolicyResolver::from_settings(&settings)); @@ -192,6 +226,7 @@ impl AdmissionMode { partitions: HashMap::new(), default_partition: Arc::from("global"), resolver, + model_registry, rate_limiter, }))); }; @@ -295,6 +330,7 @@ impl AdmissionMode { partitions, default_partition: Arc::from(default_partition_name), resolver, + model_registry, rate_limiter, }))) } @@ -674,7 +710,63 @@ mod tests { use tokio::time::{sleep, Duration}; use super::*; - use crate::worker::BasicWorkerBuilder; + use crate::{ + middleware::scheduler::{FairShareConfig, ModelFairShareConfig, PrioritySchedulerYaml}, + worker::BasicWorkerBuilder, + }; + + #[test] + fn trusted_request_model_alias_selects_the_canonical_profile() { + let registry = Arc::new(WorkerRegistry::new()); + let worker = Arc::new( + BasicWorkerBuilder::new("http://alias-worker:8000") + .model( + openai_protocol::model_card::ModelCard::new("deepseek-v4-flash") + .with_alias("deepseek-flash"), + ) + .build(), + ); + registry.register(worker).expect("worker should register"); + let fair_config = FairShareConfig { + default_weight: 1.0, + default_output_tokens: 10, + trust_output_token_estimate_header: false, + trust_request_model_header: true, + tenant_weights: HashMap::new(), + model_profiles: HashMap::from([( + "deepseek-v4-flash".to_string(), + ModelFairShareConfig { + tenant_weights: HashMap::new(), + other_weight: 1.0, + }, + )]), + }; + let yaml = PrioritySchedulerYaml { + fair_share: Some(fair_config.clone()), + ..Default::default() + }; + let settings = + SchedulerSettings::from_cli_and_yaml(true, Class::Default, 32, Some(&yaml)).unwrap(); + let state = SchedulerState { + scheduler: PriorityScheduler::new(&settings, 1).unwrap(), + partitions: HashMap::new(), + default_partition: Arc::from("global"), + resolver: Arc::new(StaticTenantPolicyResolver::from_settings(&settings)), + model_registry: Arc::clone(®istry), + rate_limiter: None, + }; + let ledger = GlobalFairShare::from_config(&fair_config); + let mut headers = HeaderMap::new(); + headers.insert( + REQUEST_MODEL_HEADER, + HeaderValue::from_static("deepseek-flash"), + ); + + assert_eq!( + state.fair_share_profile_for(&headers, &ledger), + FairShareProfile::Model(Arc::new("deepseek-v4-flash".to_string())) + ); + } #[tokio::test] async fn priority_mode_applies_capacity_changes_after_startup() { From 37747d371da04fff2ff11d308f8e6dd0121624c9 Mon Sep 17 00:00:00 2001 From: David <12414531+DavidBellamy@users.noreply.github.com> Date: Tue, 4 Aug 2026 21:35:05 -0700 Subject: [PATCH 8/8] fix(scheduler): account for model-profile other tenants Signed-off-by: David <12414531+DavidBellamy@users.noreply.github.com> --- model_gateway/src/middleware/scheduler/fair_share.rs | 4 ++++ model_gateway/src/middleware/scheduler/queue.rs | 3 +-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/model_gateway/src/middleware/scheduler/fair_share.rs b/model_gateway/src/middleware/scheduler/fair_share.rs index 842374874..a84552344 100644 --- a/model_gateway/src/middleware/scheduler/fair_share.rs +++ b/model_gateway/src/middleware/scheduler/fair_share.rs @@ -342,6 +342,10 @@ impl GlobalFairShare { .entry(tenant.clone()) .or_default() .queued[class as usize] += 1; + drop(state); + if is_other { + super::metrics::record_fair_share_unknown_tenant(self.metric_tenant(tenant)); + } } fn register_global_waiter(&self, scope_id: u64, tenant: &TenantKey, class: Class) { diff --git a/model_gateway/src/middleware/scheduler/queue.rs b/model_gateway/src/middleware/scheduler/queue.rs index 6bccdc05c..695a02624 100644 --- a/model_gateway/src/middleware/scheduler/queue.rs +++ b/model_gateway/src/middleware/scheduler/queue.rs @@ -378,8 +378,7 @@ impl ClassQueue for FairClassQueue { #[cfg(test)] mod tests { - use std::collections::HashMap; - use std::time::Duration; + use std::{collections::HashMap, time::Duration}; use tokio_util::sync::CancellationToken;