Add adaptive token-work admission

model_gateway/src/app_context.rs

@@ -22,6 +22,7 @@ use crate::{
     rate_limit::RateLimitManager,
     routers::{
         common::{openai_bridge::FormatRegistry, realtime::RealtimeRegistry},
+        grpc::adaptive_admission::AdaptiveAdmissionController,
         grpc::multimodal::MultimodalConfigRegistry,
         router_manager::RouterManager,
     },
@@ -65,6 +66,7 @@ pub struct AppContext {
     pub conversation_storage: Arc<dyn ConversationStorage>,
     pub conversation_item_storage: Arc<dyn ConversationItemStorage>,
     pub worker_monitor: Option<Arc<WorkerMonitor>>,
+    pub(crate) adaptive_admission: Option<Arc<AdaptiveAdmissionController>>,
     pub configured_reasoning_parser: Option<String>,
     pub configured_tool_parser: Option<String>,
     pub worker_job_queue: Arc<OnceLock<Arc<JobQueue>>>,
@@ -338,6 +340,20 @@ impl AppContextBuilder {
         let worker_job_queue = self
             .worker_job_queue
             .ok_or(AppContextBuildError::MissingField("worker_job_queue"))?;
+        let worker_monitor = self.worker_monitor;
+        let adaptive_admission =
+            if router_config.adaptive_admission.mode == crate::config::AdaptiveAdmissionMode::Off {
+                None
+            } else {
+                let controller = AdaptiveAdmissionController::new(
+                    router_config.adaptive_admission.clone(),
+                    worker_registry.clone(),
+                );
+                if let Some(monitor) = &worker_monitor {
+                    controller.start_load_updates(monitor.subscribe());
+                }
+                Some(controller)
+            };
 
         // Create WorkerService from the already-built components
         let worker_service = Arc::new(WorkerService::new(
@@ -373,7 +389,8 @@ impl AppContextBuilder {
             conversation_item_storage: self.conversation_item_storage.ok_or(
                 AppContextBuildError::MissingField("conversation_item_storage"),
             )?,
-            worker_monitor: self.worker_monitor,
+            worker_monitor,
+            adaptive_admission,
             configured_reasoning_parser,
             configured_tool_parser,
             worker_job_queue,
@@ -605,7 +622,8 @@ impl AppContextBuilder {
                 .clone(),
             client.clone(),
             config.load_monitor_interval_secs,
-            config.engine_metrics,
+            config.engine_metrics
+                || config.adaptive_admission.mode != crate::config::AdaptiveAdmissionMode::Off,
         )));
         Ok(self)
     }

model_gateway/src/config/builder.rs

@@ -289,6 +289,11 @@ impl RouterConfigBuilder {
         self
     }
 
+    pub fn adaptive_admission(mut self, config: super::types::AdaptiveAdmissionConfig) -> Self {
+        self.config.adaptive_admission = config;
+        self
+    }
+
     // ==================== Tenant Rate Limit ====================
 
     pub fn tenant_rate_limit_enabled(mut self, enabled: bool) -> Self {

model_gateway/src/config/types.rs

@@ -11,6 +11,100 @@ pub use smg_data_connector::{
 use super::{validation::ConfigValidator, ConfigResult};
 use crate::{tenant::DEFAULT_TENANT_HEADER_NAME, worker::ConnectionMode};
 
+/// Runtime mode for predictive, token-work admission.
+///
+/// `shadow` learns from completed requests and records the decision that would
+/// have been made, but never delays or rejects a request. `enforce` is kept as
+/// an explicit operator action so a newly deployed estimator cannot change
+/// serving behavior before its calibration is measured.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum AdaptiveAdmissionMode {
+    #[default]
+    Off,
+    Shadow,
+    Enforce,
+}
+
+impl std::str::FromStr for AdaptiveAdmissionMode {
+    type Err = String;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        match value.trim().to_ascii_lowercase().as_str() {
+            "off" => Ok(Self::Off),
+            "shadow" => Ok(Self::Shadow),
+            "enforce" => Ok(Self::Enforce),
+            _ => Err(format!(
+                "adaptive admission mode must be one of off, shadow, enforce; got {value:?}"
+            )),
+        }
+    }
+}
+
+fn default_adaptive_work_horizon_secs() -> f64 {
+    30.0
+}
+
+fn default_adaptive_estimator_half_life_secs() -> f64 {
+    900.0
+}
+
+fn default_adaptive_prior_observations() -> f64 {
+    20.0
+}
+
+fn default_adaptive_max_segments() -> usize {
+    50_000
+}
+
+fn default_adaptive_min_load_coverage() -> f64 {
+    0.8
+}
+
+fn default_adaptive_cold_start_output_tokens() -> u32 {
+    4096
+}
+
+/// Predictive token-work admission settings.
+///
+/// The work horizon is an operator-facing latency objective rather than a
+/// request-concurrency guess: recent per-replica generation capacity multiplied
+/// by healthy replica count and this horizon gives the maximum predicted
+/// outstanding decode work. Router reservations are reconciled with engine
+/// running and waiting counts. Missing or insufficient telemetry fails open to
+/// the existing priority scheduler.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct AdaptiveAdmissionConfig {
+    #[serde(default)]
+    pub mode: AdaptiveAdmissionMode,
+    #[serde(default = "default_adaptive_work_horizon_secs")]
+    pub work_horizon_secs: f64,
+    #[serde(default = "default_adaptive_estimator_half_life_secs")]
+    pub estimator_half_life_secs: f64,
+    #[serde(default = "default_adaptive_prior_observations")]
+    pub prior_observations: f64,
+    #[serde(default = "default_adaptive_max_segments")]
+    pub max_segments: usize,
+    #[serde(default = "default_adaptive_min_load_coverage")]
+    pub min_load_coverage: f64,
+    #[serde(default = "default_adaptive_cold_start_output_tokens")]
+    pub cold_start_output_tokens: u32,
+}
+
+impl Default for AdaptiveAdmissionConfig {
+    fn default() -> Self {
+        Self {
+            mode: AdaptiveAdmissionMode::Off,
+            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(),
+        }
+    }
+}
+
 /// Main router configuration
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct RouterConfig {
@@ -104,6 +198,10 @@ pub struct RouterConfig {
     /// by inflight; the remainder bucket under `tenant="other"`).
     #[serde(default = "default_priority_scheduler_tenant_metric_top_n")]
     pub priority_scheduler_tenant_metric_top_n: u32,
+    /// Optional predictive token-work admission. Off by default. Shadow mode
+    /// is behavior-preserving and is the required first deployment state.
+    #[serde(default)]
+    pub adaptive_admission: AdaptiveAdmissionConfig,
     /// Enable per-tenant LLM token/request rate limiting. When false
     /// (default), no rate limiter is constructed — zero behavior change
     /// for existing deployments.
@@ -842,6 +940,7 @@ impl Default for RouterConfig {
             priority_scheduler_config: None,
             priority_scheduler_tenant_metric_top_n: default_priority_scheduler_tenant_metric_top_n(
             ),
+            adaptive_admission: AdaptiveAdmissionConfig::default(),
             tenant_rate_limit_enabled: false,
             tenant_rate_limit_config: None,
             cors_allowed_origins: vec![],

model_gateway/src/config/validation.rs

@@ -724,6 +724,55 @@ impl ConfigValidator {
             });
         }
 
+        let adaptive = &config.adaptive_admission;
+        if !adaptive.work_horizon_secs.is_finite() || adaptive.work_horizon_secs <= 0.0 {
+            return Err(ConfigError::InvalidValue {
+                field: "adaptive_admission.work_horizon_secs".to_string(),
+                value: adaptive.work_horizon_secs.to_string(),
+                reason: "Must be finite and > 0".to_string(),
+            });
+        }
+        if !adaptive.estimator_half_life_secs.is_finite()
+            || adaptive.estimator_half_life_secs <= 0.0
+        {
+            return Err(ConfigError::InvalidValue {
+                field: "adaptive_admission.estimator_half_life_secs".to_string(),
+                value: adaptive.estimator_half_life_secs.to_string(),
+                reason: "Must be finite and > 0".to_string(),
+            });
+        }
+        if !adaptive.prior_observations.is_finite() || adaptive.prior_observations < 0.0 {
+            return Err(ConfigError::InvalidValue {
+                field: "adaptive_admission.prior_observations".to_string(),
+                value: adaptive.prior_observations.to_string(),
+                reason: "Must be finite and >= 0".to_string(),
+            });
+        }
+        if adaptive.max_segments < 4 {
+            return Err(ConfigError::InvalidValue {
+                field: "adaptive_admission.max_segments".to_string(),
+                value: adaptive.max_segments.to_string(),
+                reason: "Must be >= 4".to_string(),
+            });
+        }
+        if !adaptive.min_load_coverage.is_finite()
+            || adaptive.min_load_coverage <= 0.0
+            || adaptive.min_load_coverage > 1.0
+        {
+            return Err(ConfigError::InvalidValue {
+                field: "adaptive_admission.min_load_coverage".to_string(),
+                value: adaptive.min_load_coverage.to_string(),
+                reason: "Must be finite and in (0, 1]".to_string(),
+            });
+        }
+        if adaptive.cold_start_output_tokens == 0 {
+            return Err(ConfigError::InvalidValue {
+                field: "adaptive_admission.cold_start_output_tokens".to_string(),
+                value: adaptive.cold_start_output_tokens.to_string(),
+                reason: "Must be > 0".to_string(),
+            });
+        }
+
         Ok(())
     }
 

model_gateway/src/main.rs

@@ -5,11 +5,11 @@ use openai_protocol::worker::TransportMode;
 use rand::{distr::Alphanumeric, RngExt};
 use smg::{
     config::{
-        validate_mesh_server_name, CircuitBreakerConfig, ConfigError, ConfigResult,
-        DiscoveryConfig, HealthCheckConfig, HistoryBackend, ManualAssignmentMode, MetricsConfig,
-        OracleConfig, PolicyConfig, PostgresConfig, RedisConfig, RetryConfig, RouterConfig,
-        RoutingKeyOverrideConfig, RoutingMode, SchemaConfig, TenantApiKeyEntry,
-        TokenizerCacheConfig, TraceConfig,
+        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,
     },
     observability::{
         metrics::PrometheusConfig,
@@ -97,6 +97,34 @@ impl std::fmt::Display for Backend {
     }
 }
 
+#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, ValueEnum)]
+enum AdaptiveAdmissionCliMode {
+    #[default]
+    Off,
+    Shadow,
+    Enforce,
+}
+
+impl std::fmt::Display for AdaptiveAdmissionCliMode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(match self {
+            Self::Off => "off",
+            Self::Shadow => "shadow",
+            Self::Enforce => "enforce",
+        })
+    }
+}
+
+impl From<AdaptiveAdmissionCliMode> for AdaptiveAdmissionMode {
+    fn from(value: AdaptiveAdmissionCliMode) -> Self {
+        match value {
+            AdaptiveAdmissionCliMode::Off => Self::Off,
+            AdaptiveAdmissionCliMode::Shadow => Self::Shadow,
+            AdaptiveAdmissionCliMode::Enforce => Self::Enforce,
+        }
+    }
+}
+
 #[derive(Parser, Debug)]
 #[command(name = "shepherd-model-gateway", alias = "smg", alias = "amg")]
 #[command(about = "Shepherd Model Gateway - High-performance inference gateway")]
@@ -488,6 +516,42 @@ struct CliArgs {
     #[arg(long, default_value_t = 32, help_heading = "Priority Scheduler")]
     priority_scheduler_tenant_metric_top_n: u32,
 
+    // ==================== Adaptive Admission ====================
+    /// Predictive token-work admission mode. Shadow mode learns and records
+    /// hypothetical decisions without delaying or rejecting requests.
+    #[arg(
+        long,
+        value_enum,
+        default_value_t = AdaptiveAdmissionCliMode::Off,
+        help_heading = "Adaptive Admission"
+    )]
+    adaptive_admission_mode: AdaptiveAdmissionCliMode,
+
+    /// 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,
+
+    /// Half-life for recency weighting of observed output lengths.
+    #[arg(long, default_value_t = 900.0, help_heading = "Adaptive Admission")]
+    adaptive_admission_estimator_half_life_secs: f64,
+
+    /// Hierarchical shrinkage strength in effective observations.
+    #[arg(long, default_value_t = 20.0, help_heading = "Adaptive Admission")]
+    adaptive_admission_prior_observations: f64,
+
+    /// Maximum in-memory predictor segments before stale-segment eviction.
+    #[arg(long, default_value_t = 50000, help_heading = "Adaptive Admission")]
+    adaptive_admission_max_segments: usize,
+
+    /// Minimum fraction of healthy replicas with fresh load telemetry before
+    /// a token-work decision is considered usable.
+    #[arg(long, default_value_t = 0.8, help_heading = "Adaptive Admission")]
+    adaptive_admission_min_load_coverage: f64,
+
+    /// Cold-start output-token prediction before a model has observations.
+    #[arg(long, default_value_t = 4096, help_heading = "Adaptive Admission")]
+    adaptive_admission_cold_start_output_tokens: u32,
+
     // ==================== Tenant Rate Limit ====================
     /// Enable per-tenant LLM token/request rate limiting. When unset
     /// (default), no rate limiter is constructed.
@@ -1490,6 +1554,15 @@ impl CliArgs {
             .priority_scheduler_default_max_class(self.priority_scheduler_default_max_class.clone())
             .priority_scheduler_config(self.priority_scheduler_config.clone())
             .priority_scheduler_tenant_metric_top_n(self.priority_scheduler_tenant_metric_top_n)
+            .adaptive_admission(AdaptiveAdmissionConfig {
+                mode: self.adaptive_admission_mode.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,
+            })
             .tenant_rate_limit_enabled(self.tenant_rate_limit_enabled)
             .tenant_rate_limit_config(self.tenant_rate_limit_config.clone())
             .cors_allowed_origins(self.cors_allowed_origins.clone())

model_gateway/src/observability/metrics.rs

@@ -480,7 +480,9 @@ 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;
     scheduler_metrics::describe();
+    adaptive_admission_metrics::describe_metrics();
 }
 
 #[expect(