From c06c0a2ce08613f71d60ffcdb081aba56f7995af Mon Sep 17 00:00:00 2001 From: ishan Date: Thu, 6 Aug 2026 23:12:14 -0700 Subject: [PATCH 1/4] fix(gateway): warn when priority scheduler runs a multi-model fleet WorkerCapacity derives one fleet-wide admission budget across all healthy workers with no model dimension, so with multiple models an idle model's slots mask a saturated model's queue (#2069). The real fix (per-model capacity) is an interface change; this lands the interim mitigations from the issue: - one-time warn! at startup or on the first event that makes the healthy fleet multi-model, pointing at mitigations (--worker-capacity-override, one gateway per model) - document the limitation on priority_scheduler_enabled - replace the dangling .claude/priority-scheduling design-doc reference in capacity.rs with the actual fleet-wide semantics Refs: #2069 Signed-off-by: ishan --- model_gateway/src/config/types.rs | 8 ++ model_gateway/src/worker/capacity.rs | 114 ++++++++++++++++++++++++++- 2 files changed, 119 insertions(+), 3 deletions(-) diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index fc56e7bbc..fc955832c 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -96,6 +96,14 @@ pub struct RouterConfig { /// Enable the priority-aware admission scheduler. When false (default), /// the legacy concurrency-limit middleware stays wired — zero behavior /// change for existing deployments. + /// + /// Limitation: admission capacity is a single fleet-wide scalar with no + /// model dimension, so multi-model deployments get incorrect admission + /// limits (an idle model's capacity masks a saturated model's queue). + /// SMG logs a one-time warning when the healthy fleet reports more than + /// one model id. Until per-model capacity lands, prefer + /// `--worker-capacity-override` or one gateway per model. See + /// smg-project/smg#2069. #[serde(default)] pub priority_scheduler_enabled: bool, /// Max priority class applied to tenants not listed in the scheduler diff --git a/model_gateway/src/worker/capacity.rs b/model_gateway/src/worker/capacity.rs index 0fccd6d9d..56b56914f 100644 --- a/model_gateway/src/worker/capacity.rs +++ b/model_gateway/src/worker/capacity.rs @@ -1,10 +1,14 @@ //! Aggregate backend capacity tracking. //! -//! See `.claude/priority-scheduling/01-worker-capacity-design.md` for the -//! full design rationale. +//! Capacity is a single fleet-wide scalar (summed across every healthy +//! worker, with no model dimension). That is correct only for +//! single-model fleets: with multiple models, an idle model's slots +//! inflate other models' admission budgets (smg-project/smg#2069). The +//! tracker logs a one-time warning when the healthy fleet reports more +//! than one model id. use std::sync::{ - atomic::{AtomicU16, AtomicU8, Ordering}, + atomic::{AtomicBool, AtomicU16, AtomicU8, Ordering}, Arc, Weak, }; @@ -113,6 +117,8 @@ pub struct WorkerCapacity { capacity: AtomicU16, source: AtomicU8, watch_tx: watch::Sender, + /// One-shot latch for the multi-model warning (#2069). + warned_multi_model: AtomicBool, } impl WorkerCapacity { @@ -141,9 +147,32 @@ impl WorkerCapacity { capacity: AtomicU16::new(capacity), source: AtomicU8::new(source as u8), watch_tx: tx, + warned_multi_model: AtomicBool::new(false), }) } + /// Log once when the healthy fleet serves more than one model id. + /// + /// Capacity is a fleet-wide scalar, so multi-model admission limits are + /// incorrect (#2069): an idle model's slots mask a saturated model's + /// queue. The warning fires at most once per tracker, at startup or on + /// the first event that makes the fleet multi-model. + fn warn_once_if_multi_model(&self, workers: &[Arc]) { + if self.warned_multi_model.load(Ordering::Relaxed) { + return; + } + let models = distinct_model_count(workers); + if models > 1 { + self.warned_multi_model.store(true, Ordering::Relaxed); + tracing::warn!( + models, + "priority scheduler capacity is fleet-wide: an idle model's slots \ + inflate other models' admission budgets (#2069). Mitigations: \ + --worker-capacity-override, or one gateway per model." + ); + } + } + /// Construct a `WorkerCapacity`, compute the initial value /// synchronously, and spawn the supervised event-loop task. /// @@ -163,7 +192,9 @@ impl WorkerCapacity { capacity: AtomicU16::new(initial_capacity), source: AtomicU8::new(initial_source as u8), watch_tx, + warned_multi_model: AtomicBool::new(false), }); + this.warn_once_if_multi_model(&workers); // Supervised loop: any panic in the inner future restarts the task // after a 1s backoff. Graceful exit (Ok) breaks out. @@ -216,6 +247,19 @@ fn healthy_workers(registry: &WorkerRegistry) -> Vec> { .collect() } +/// Count distinct model ids across the given workers. +/// +/// Used by the multi-model warning: capacity is a fleet-wide scalar, so a +/// fleet reporting more than one model id gets incorrect per-model +/// admission limits (#2069). +pub(super) fn distinct_model_count(workers: &[Arc]) -> usize { + let mut seen = std::collections::HashSet::new(); + for w in workers { + seen.insert(w.model_id()); + } + seen.len() +} + async fn run_event_loop( tracker: Weak, registry: Weak, @@ -245,6 +289,7 @@ async fn run_event_loop( let workers = healthy_workers(&r); let (new_capacity, new_source) = recompute(&settings, &workers); + t.warn_once_if_multi_model(&workers); let old_capacity = t.capacity.swap(new_capacity, Ordering::AcqRel); let old_source_raw = t.source.swap(new_source as u8, Ordering::AcqRel); @@ -619,4 +664,67 @@ mod tests { let rx = tracker.watch(); assert_eq!(*rx.borrow(), 42); } + + fn worker_with_model(url: &str, model: &str) -> Arc { + Arc::new( + BasicWorkerBuilder::new(url) + .model(openai_protocol::model_card::ModelCard::new(model)) + .build(), + ) + } + + #[test] + fn distinct_model_count_counts_unique_models() { + let workers = vec![ + worker_with_model("http://w1", "llama-7b"), + worker_with_model("http://w2", "llama-7b"), + worker_with_model("http://w3", "llama-70b"), + ]; + assert_eq!(distinct_model_count(&workers), 2); + } + + #[test] + fn distinct_model_count_single_model_fleet() { + let workers = vec![ + worker_with_model("http://w1", "llama-7b"), + worker_with_model("http://w2", "llama-7b"), + ]; + assert_eq!(distinct_model_count(&workers), 1); + } + + #[test] + fn warn_once_if_multi_model_sets_flag_once() { + let tracker = WorkerCapacity::for_test_with_value(64, CapacitySource::Mixed); + let workers = vec![ + worker_with_model("http://w1", "llama-7b"), + worker_with_model("http://w2", "llama-70b"), + ]; + assert!(!tracker.warned_multi_model.load(Ordering::Relaxed)); + tracker.warn_once_if_multi_model(&workers); + assert!(tracker.warned_multi_model.load(Ordering::Relaxed)); + // Second call is a no-op regardless of input. + tracker.warn_once_if_multi_model(&[]); + assert!(tracker.warned_multi_model.load(Ordering::Relaxed)); + } + + #[tracing_test::traced_test] + #[test] + fn warn_once_if_multi_model_logs_for_multi_model_fleet() { + let tracker = WorkerCapacity::for_test_with_value(64, CapacitySource::Mixed); + let workers = vec![ + worker_with_model("http://w1", "llama-7b"), + worker_with_model("http://w2", "llama-70b"), + ]; + tracker.warn_once_if_multi_model(&workers); + assert!(logs_contain("fleet-wide")); + } + + #[tracing_test::traced_test] + #[test] + fn warn_once_if_multi_model_quiet_for_single_model() { + let tracker = WorkerCapacity::for_test_with_value(64, CapacitySource::Mixed); + let workers = vec![worker_with_model("http://w1", "llama-7b")]; + tracker.warn_once_if_multi_model(&workers); + assert!(!logs_contain("fleet-wide")); + } } From b4f53454afd3ce32013d1186150e392fa8503c92 Mon Sep 17 00:00:00 2001 From: ishan Date: Sun, 9 Aug 2026 15:05:35 -0700 Subject: [PATCH 2/4] fix(gateway): startup-only multi-model warning; count all model cards Review on #2071: - slin1237: the per-event warn check ran an O(fleet) distinct-model count in the worker-event loop just for a log. The warning is now startup-only (spawn), keeping the event loop O(1). - CodeRabbit: distinct_model_count used only the primary model-card id; workers advertising multiple model cards could suppress the warning. Count every card from Worker::models() instead. Adds a spawn-path latch test and a multi-card counting test. Refs: #2069 Signed-off-by: ishan --- model_gateway/src/worker/capacity.rs | 51 +++++++++++++++++++++++----- 1 file changed, 42 insertions(+), 9 deletions(-) diff --git a/model_gateway/src/worker/capacity.rs b/model_gateway/src/worker/capacity.rs index 56b56914f..fe5751339 100644 --- a/model_gateway/src/worker/capacity.rs +++ b/model_gateway/src/worker/capacity.rs @@ -151,12 +151,13 @@ impl WorkerCapacity { }) } - /// Log once when the healthy fleet serves more than one model id. + /// Log once at startup when the healthy fleet serves more than one + /// model id. /// /// Capacity is a fleet-wide scalar, so multi-model admission limits are /// incorrect (#2069): an idle model's slots mask a saturated model's - /// queue. The warning fires at most once per tracker, at startup or on - /// the first event that makes the fleet multi-model. + /// queue. Startup-only by design — the check is O(fleet size), which is + /// not acceptable in the worker-event loop. fn warn_once_if_multi_model(&self, workers: &[Arc]) { if self.warned_multi_model.load(Ordering::Relaxed) { return; @@ -247,15 +248,18 @@ fn healthy_workers(registry: &WorkerRegistry) -> Vec> { .collect() } -/// Count distinct model ids across the given workers. +/// Count distinct advertised model ids across the given workers. /// -/// Used by the multi-model warning: capacity is a fleet-wide scalar, so a -/// fleet reporting more than one model id gets incorrect per-model -/// admission limits (#2069). +/// Every model card counts, not just the primary — a worker serving +/// several models must not read as single-model. Used by the multi-model +/// warning: capacity is a fleet-wide scalar, so a fleet reporting more +/// than one model id gets incorrect per-model admission limits (#2069). pub(super) fn distinct_model_count(workers: &[Arc]) -> usize { let mut seen = std::collections::HashSet::new(); for w in workers { - seen.insert(w.model_id()); + for card in w.models() { + seen.insert(card.id); + } } seen.len() } @@ -289,7 +293,6 @@ async fn run_event_loop( let workers = healthy_workers(&r); let (new_capacity, new_source) = recompute(&settings, &workers); - t.warn_once_if_multi_model(&workers); let old_capacity = t.capacity.swap(new_capacity, Ordering::AcqRel); let old_source_raw = t.source.swap(new_source as u8, Ordering::AcqRel); @@ -692,6 +695,36 @@ mod tests { assert_eq!(distinct_model_count(&workers), 1); } + #[test] + fn distinct_model_count_counts_every_advertised_model_card() { + // One worker can advertise multiple model cards; counting only the + // primary id would undercount and suppress the warning. + let spec: openai_protocol::worker::WorkerSpec = serde_json::from_value(serde_json::json!({ + "url": "http://w1", + "models": [{"id": "llama-7b"}, {"id": "llama-70b"}], + })) + .expect("multi-model spec"); + let workers = + vec![Arc::new(BasicWorkerBuilder::from_spec(spec).build()) as Arc]; + assert_eq!(distinct_model_count(&workers), 2); + } + + #[tokio::test] + async fn spawn_latches_warning_for_multi_model_registry() { + let registry = Arc::new(WorkerRegistry::new()); + let id1 = registry + .register(worker_with_model("http://w1", "llama-7b")) + .expect("registered"); + registry.transition_status(&id1, WorkerStatus::Ready); + let id2 = registry + .register(worker_with_model("http://w2", "llama-70b")) + .expect("registered"); + registry.transition_status(&id2, WorkerStatus::Ready); + + let tracker = WorkerCapacity::spawn(registry, CapacityTrackerSettings::default()); + assert!(tracker.warned_multi_model.load(Ordering::Relaxed)); + } + #[test] fn warn_once_if_multi_model_sets_flag_once() { let tracker = WorkerCapacity::for_test_with_value(64, CapacitySource::Mixed); From 6b4f282aa4f304ab07bf7c6c27b4ffa38b31f8aa Mon Sep 17 00:00:00 2001 From: ishan Date: Tue, 18 Aug 2026 16:02:10 -0700 Subject: [PATCH 3/4] chore(ci): retrigger PR checks after runner socket hang-up Signed-off-by: ishan From 52b5ccf1dc3b56079de10c121b6f771842c36681 Mon Sep 17 00:00:00 2001 From: ishan Date: Tue, 18 Aug 2026 16:11:19 -0700 Subject: [PATCH 4/4] fix(gateway): count label-only workers in the multi-model warning CodeRabbit on #2071: distinct_model_count ignored workers with an empty model-card list, so a fleet with distinct primary model ids via the model_id label counted as zero models and suppressed the startup warning. Fall back to worker.model_id() when models() is empty, matching WorkerRegistry::worker_model_ids(). Refs: #2069 Signed-off-by: ishan --- model_gateway/src/worker/capacity.rs | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/model_gateway/src/worker/capacity.rs b/model_gateway/src/worker/capacity.rs index fe5751339..fab2ae1e4 100644 --- a/model_gateway/src/worker/capacity.rs +++ b/model_gateway/src/worker/capacity.rs @@ -257,8 +257,13 @@ fn healthy_workers(registry: &WorkerRegistry) -> Vec> { pub(super) fn distinct_model_count(workers: &[Arc]) -> usize { let mut seen = std::collections::HashSet::new(); for w in workers { - for card in w.models() { - seen.insert(card.id); + let cards = w.models(); + if cards.is_empty() { + // No model cards (e.g. wildcard or label-only worker) — fall back + // to the primary id, matching WorkerRegistry::worker_model_ids(). + seen.insert(w.model_id().to_string()); + } else { + seen.extend(cards.into_iter().map(|card| card.id)); } } seen.len() @@ -709,6 +714,22 @@ mod tests { assert_eq!(distinct_model_count(&workers), 2); } + #[test] + fn distinct_model_count_falls_back_to_primary_model_id_when_no_cards() { + // Label-only workers (no model cards) must still count — matches the + // WorkerRegistry::worker_model_ids() fallback. + let with_label = |url: &str, model: &str| { + let mut labels = HashMap::new(); + labels.insert("model_id".to_string(), model.to_string()); + Arc::new(BasicWorkerBuilder::new(url).labels(labels).build()) as Arc + }; + let workers = vec![ + with_label("http://w1", "llama-7b"), + with_label("http://w2", "llama-70b"), + ]; + assert_eq!(distinct_model_count(&workers), 2); + } + #[tokio::test] async fn spawn_latches_warning_for_multi_model_registry() { let registry = Arc::new(WorkerRegistry::new());