diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 5d94003..2003309 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -308,9 +308,8 @@ impl DagNode for SelectAccount { .protocol() .ok_or_else(|| GatewayError::internal("select_account before resolve_model"))?; // post-variant: the served model's knobs ride the provider lookup - let conf = ctx - .cfg - .find_model(served_model(ctx.request.model_param_v2.as_ref())); + let model = served_model(ctx.request.model_param_v2.as_ref()); + let conf = ctx.cfg.find_model(model); ctx.request.prompt_cache = conf.is_some_and(|m| m.prompt_cache); let provider = conf.and_then(|m| m.provider.as_deref()); let latency = ctx @@ -321,7 +320,7 @@ impl DagNode for SelectAccount { let account = ctx .state .pool - .select_healthy(mt, provider, &[], ctx.state.health.as_ref(), latency) + .select_healthy(mt, model, provider, &[], ctx.state.health.as_ref(), latency) .await; let Some(account) = account else { // unsampled, an exhausted pool would read no_data forever @@ -516,6 +515,7 @@ impl DagNode for CallEngine { .pool .select_healthy( mt, + served_model(ctx.request.model_param_v2.as_ref()), provider, std::slice::from_ref(&failed.name), ctx.state.health.as_ref(), @@ -599,7 +599,11 @@ async fn note_engine_outcome( .avail .record(requested_model(ctx.request.model_param_v2.as_ref()), true); if let Some(account) = ctx.request.account.as_ref() { - ctx.state.health.record_success(&account.name).await; + let model = served_model(ctx.request.model_param_v2.as_ref()); + ctx.state + .health + .record_success(&gw_state::health_key(&account.name, model)) + .await; if let Some(started) = started { ctx.state.latency.record(&account.name, started.elapsed()); } @@ -630,23 +634,24 @@ fn account_fault(class: gw_consts::ErrClass) -> bool { async fn note_failure(ctx: &mut DagContext, account: &str) { let threshold = ctx.cfg.stability.failure_threshold; let secs = ctx.cfg.stability.cooldown_seconds; + let key = gw_state::health_key(account, served_model(ctx.request.model_param_v2.as_ref())); if !ctx .state .health - .record_failure(account, threshold, std::time::Duration::from_secs(secs)) + .record_failure(&key, threshold, std::time::Duration::from_secs(secs)) .await { return; } + ctx.decide( + "account_health", + format!("{key} entered cooldown ({secs}s)"), + ); ctx.state.alerts.emit( "account_cooldown", - account.to_owned(), + key, format!("{threshold} consecutive failures; cooling {secs}s"), ); - ctx.decide( - "account_health", - format!("{account} entered cooldown ({secs}s)"), - ); } /// post_process/common_usage: RawUsageJSON -> CommonUsage. diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 1e627c7..793f7dc 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2361,6 +2361,44 @@ mod tests { } } + #[derive(Debug)] + struct RefusesOneModel { + model: &'static str, + in_band: bool, + } + + #[async_trait::async_trait] + impl gw_engines::transport::Transport for RefusesOneModel { + async fn send( + &self, + req: gw_engines::transport::UpstreamRequest, + ) -> GResult { + let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap_or_default(); + if body["model"] != self.model { + return gw_engines::MockTransport.send(req).await; + } + let refusal = r#"{"error":{"message":"high demand: exceeds the maximum usage size allowed during peak load","code":503}}"#; + let (status, body) = if self.in_band { + ( + 200, + gw_engines::transport::UpstreamBody::Sse( + format!("data: {refusal}\n\n").into_bytes(), + ), + ) + } else { + ( + 503, + gw_engines::transport::UpstreamBody::Json(bytes::Bytes::from(refusal)), + ) + }; + Ok(gw_engines::transport::UpstreamResponse { + status, + body, + headers: Default::default(), + }) + } + } + #[derive(Debug)] struct PausedStream { release: Arc, @@ -2485,7 +2523,10 @@ mod tests { let outcome = ctx.outcome.expect("outcome"); assert!(outcome.terminal_error.is_some(), "{code} native={native}"); assert_eq!( - h.state().health.available("a1").await, + h.state() + .health + .available(&gw_state::health_key("a1", "m")) + .await, !faults, "{code} native={native}" ); @@ -2526,21 +2567,75 @@ mod tests { .map(|error| error.class), Some(gw_consts::ErrClass::ModelStreamError) ); - assert!(!h.state().health.available("a-down").await); - assert!(!h.state().health.available("a-up").await); + let health = &h.state().health; + assert!(!health.available(&gw_state::health_key("a-down", "m")).await); + assert!(!health.available(&gw_state::health_key("a-up", "m")).await); let first = alerts.try_recv().expect("first account cooldown alert"); let second = alerts.try_recv().expect("retry account cooldown alert"); assert_eq!( (first.kind, first.subject.as_str()), - ("account_cooldown", "a-down") + ("account_cooldown", "a-down/m") ); assert_eq!( (second.kind, second.subject.as_str()), - ("account_cooldown", "a-up") + ("account_cooldown", "a-up/m") ); assert!(alerts.try_recv().is_err(), "one alert per failed account"); } + #[tokio::test] + async fn a_model_cooling_down_leaves_its_only_account_serving_other_models() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: pub, protocol: openai-chat, variants: [{model: busy, weight: 1}]}, {name: busy, protocol: openai-chat}, {name: calm, protocol: openai-chat}]\naccounts: [{name: a1, provider: p, protocols: [openai-chat]}]\nstability: {failure_threshold: 1, cooldown_seconds: 30}\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; + for (in_band, requested) in [(false, "busy"), (true, "pub")] { + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(RefusesOneModel { + model: "busy", + in_band, + }), + ); + let key = h.state().auth.get(&access_key_id("k1")).await.unwrap(); + let refused = h + .run(drained_stream_req(requested, "hi"), key.clone()) + .await; + assert!( + refused.map_or(true, |ctx| ctx + .outcome + .is_some_and(|o| o.terminal_error.is_some())), + "in_band={in_band}" + ); + let health = &h.state().health; + assert!( + !health.available(&gw_state::health_key("a1", "busy")).await, + "in_band={in_band}" + ); + assert!( + health.available(&gw_state::health_key("a1", "calm")).await, + "in_band={in_band}" + ); + + let ctx = h.run(chat_req("calm", "hi"), key.clone()).await.unwrap(); + let outcome = ctx.outcome.expect("calm is served"); + assert!(outcome.terminal_error.is_none(), "in_band={in_band}"); + assert_eq!( + ctx.request.account.as_ref().map(|a| a.name.as_str()), + Some("a1") + ); + + let cooling = h + .run(chat_req(requested, "hi"), key) + .await + .err() + .expect("busy is cooling"); + assert!( + cooling.message.contains("no healthy upstream"), + "in_band={in_band}: {cooling:?}" + ); + } + } + #[tokio::test] async fn client_disconnect_after_commit_stays_health_neutral() { let yaml = "listen: {host: h, port: 1}\nsecurity: {dlp_redact: false}\nmodels: [{name: m, protocol: openai-chat}]\naccounts: [{name: a1, provider: p, protocols: [openai-chat]}]\nstability: {failure_threshold: 1, cooldown_seconds: 30}\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; @@ -2575,7 +2670,12 @@ mod tests { avail.flush().await; let minute = gw_state::epoch_secs() / 60; assert_eq!(avail.window("m", minute - 5, minute).await, (0, 0)); - assert!(h.state().health.available("a1").await); + assert!( + h.state() + .health + .available(&gw_state::health_key("a1", "m")) + .await + ); assert!(alerts.try_recv().is_err()); } diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index 8ab8f4f..670f32f 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -4224,12 +4224,12 @@ async fn account_cooldown_and_recovery() { .unwrap() .iter() .find(|a| a["name"] == "mock-spark-down") - .unwrap()["health"] - .as_str() .unwrap() - .to_owned() + .clone() }; - assert_eq!(spark_health(&app).await, "cooling"); + let spark = spark_health(&app).await; + assert_eq!(spark["health"], "cooling"); + assert_eq!(spark["cooling_models"], serde_json::json!(["spark-lite"])); let r = app .clone() .oneshot(post("/v1/chat/completions", Some("ak-demo-123"), body)) @@ -4245,7 +4245,7 @@ async fn account_cooldown_and_recovery() { let mut recovered = false; for _ in 0..40 { tokio::time::sleep(std::time::Duration::from_millis(200)).await; - if spark_health(&app).await == "ok" { + if spark_health(&app).await["health"] == "ok" { recovered = true; break; } diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 21e3fce..12e706f 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -456,20 +456,22 @@ impl AccountPool { .map(Arc::clone) } - /// PTU accounts are preferred over paygo; `excluded` (failed-over) and - /// cooldown accounts are skipped; within a tier, pick by priority then - /// round-robin. A model bound to a `provider` only uses that provider's - /// accounts, and only those are health-checked. + /// PTU accounts are preferred over paygo; `excluded` (failed-over) accounts + /// and accounts cooling down for `model` are skipped; within a tier, pick by + /// priority then round-robin. A model bound to a `provider` only uses that + /// provider's accounts, and only those are health-checked. pub async fn select_healthy( &self, p: Protocol, + model: &str, provider: Option<&str>, excluded: &[String], health: &dyn HealthStore, latency: Option<&latency::Latency>, ) -> Option> { let serving = || self.accounts.iter().filter(|a| serves(a, p, provider)); - let checks = futures::future::join_all(serving().map(|a| health.available(&a.name))).await; + let keys: Vec = serving().map(|a| health_key(&a.name, model)).collect(); + let checks = futures::future::join_all(keys.iter().map(|k| health.available(k))).await; let unhealthy: Vec<&str> = serving() .zip(checks) .filter(|(_, ok)| !ok) @@ -988,6 +990,10 @@ pub fn epoch_millis() -> i64 { .unwrap_or(0) } +pub fn health_key(account: &str, model: &str) -> String { + format!("{account}/{model}") +} + /// Recovers a poisoned lock instead of panicking: every critical section here is infallible. pub(crate) fn lock(m: &std::sync::Mutex) -> std::sync::MutexGuard<'_, T> { m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) @@ -1133,6 +1139,7 @@ mod tests { s.pool .select_healthy( Protocol::OpenaiChat, + "m", Some("p"), &[], s.health.as_ref(), @@ -1525,19 +1532,19 @@ mod tests { let h = s.health.as_ref(); let a = s .pool - .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h, None) + .select_healthy(Protocol::OpenaiChat, "m", Some("openai"), &[], h, None) .await .unwrap(); let b = s .pool - .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h, None) + .select_healthy(Protocol::OpenaiChat, "m", Some("openai"), &[], h, None) .await .unwrap(); assert_eq!(a.name, "mock-openai-1"); assert_eq!(b.name, "mock-openai-1"); assert_eq!( s.pool - .select_healthy(Protocol::AnthropicMessages, None, &[], h, None) + .select_healthy(Protocol::AnthropicMessages, "m", None, &[], h, None) .await .unwrap() .name, @@ -1545,12 +1552,39 @@ mod tests { ); assert!( s.pool - .select_healthy(Protocol::Video, Some("nonexistent"), &[], h, None) + .select_healthy(Protocol::Video, "m", Some("nonexistent"), &[], h, None) .await .is_none() ); } + #[tokio::test] + async fn a_cooldown_benches_the_account_for_one_model_only() { + let yaml = "listen: {host: h, port: 1}\naccounts: [{name: only, provider: p, protocols: ['responses']}]"; + let s = GatewayState::from_config(&GatewayConfig::from_yaml(yaml).unwrap()); + let cd = Duration::from_secs(30); + assert!( + s.health + .record_failure(&health_key("only", "busy"), 1, cd) + .await + ); + let pick = |model| { + s.pool.select_healthy( + Protocol::Responses, + model, + None, + &[], + s.health.as_ref(), + None, + ) + }; + assert!(pick("busy").await.is_none()); + assert_eq!( + pick("calm").await.map(|a| a.name.clone()).as_deref(), + Some("only") + ); + } + #[test] fn failure_streak_decays_when_idle() { let h = AccountHealth::default(); @@ -1585,7 +1619,7 @@ mod tests { let h = s.health.as_ref(); let first = s .pool - .select_healthy(Protocol::OpenaiChat, Some("tencent"), &[], h, None) + .select_healthy(Protocol::OpenaiChat, "m", Some("tencent"), &[], h, None) .await .unwrap(); assert_eq!(first.name, "mock-hunyuan-ptu-down"); @@ -1594,6 +1628,7 @@ mod tests { .pool .select_healthy( Protocol::OpenaiChat, + "m", Some("tencent"), &["mock-hunyuan-ptu-down".into()], h, diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 08eafb4..8ab4b9e 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -452,6 +452,7 @@ async fn realtime_ws( .pool .select_healthy( mt, + &served, served_conf.and_then(|m| m.provider.as_deref()), &[], snap.state.health.as_ref(), @@ -1320,12 +1321,33 @@ async fn accounts(State(s): State, _: GlobalAdmin) -> Response { let health = &s.handler.state().health; let mut data: Vec = Vec::with_capacity(cfg.accounts.len()); for a in &cfg.accounts { + let served: Vec<&str> = cfg + .models + .iter() + .filter(|m| { + a.protocols.contains(&m.protocol) + && m.provider.as_deref().is_none_or(|p| p == a.provider) + }) + .map(|m| m.name.as_str()) + .collect(); + let checks = futures::future::join_all( + served + .iter() + .map(|m| async move { health.available(&gw_state::health_key(&a.name, m)).await }), + ) + .await; + let cooling: Vec<&str> = served + .into_iter() + .zip(checks) + .filter_map(|(m, ok)| (!ok).then_some(m)) + .collect(); data.push(json!({ "name": a.name, "provider": a.provider, "priority": a.priority, "tier": if a.tier.is_empty() { gw_consts::account_tier::PAYGO } else { a.tier.as_str() }, - "health": if health.available(&a.name).await { "ok" } else { "cooling" }, + "health": if cooling.is_empty() { "ok" } else { "cooling" }, + "cooling_models": cooling, "protocols": a.protocols, })); } diff --git a/docs/api.md b/docs/api.md index e1ed524..587c56f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -371,7 +371,7 @@ and admits on `response.create`. | GET | `/health` | liveness | | GET | `/metrics` | Prometheus registry (see [Observability](observability.md)) | | GET | `/internal/ledger` | billing records; `?limit=N` returns the N most recent, default 100 (oldest-first within the page; `count` is the total); global admin token only | -| GET | `/internal/accounts` | account pool view with health; global admin token only | +| GET | `/internal/accounts` | account pool view with health (`cooling` while any served model cools down, listed in `cooling_models`); global admin token only | `/internal/*` is an operator surface: it answers only to the global admin bearer (`admin.token_env`; 404 while no admin token at all is configured, 403 to a tenant token), and the raw rows diff --git a/docs/architecture.md b/docs/architecture.md index 991c341..ce85d35 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -75,11 +75,12 @@ default, so the whole pipeline is testable offline: - **`UserBudgetStore`** — per-user budget overrides inside a tenant. In-process by default; `PostgresUserBudgets` with `storage.postgres_url`, cached and invalidated the same way as keys. -- **`HealthStore`** — account cooldown/recovery. In-process breaker by - default; `RedisHealth` with `storage.redis_url` — a tripped account is - skipped by every instance. Only upstream faults count (5xx, timeouts, - broken streams, a 401/402/403 credential or billing refusal), never a - client-caused or rate-limit error. +- **`HealthStore`** — account cooldown/recovery, per account and served + model, so one model's refusals never bench the account for the others. + In-process breaker by default; `RedisHealth` with `storage.redis_url` — a + tripped account is skipped for that model by every instance. Only upstream + faults count (5xx, timeouts, broken streams, a 401/402/403 credential or + billing refusal), never a client-caused or rate-limit error. - **`Governance`** — rate/quota/TPM counters. In-process by default; `RedisGovernance` shares them (including pooled tenant QPS) fleet-wide. - **`Moderator`** — the external content review behind `security.moderate`; diff --git a/docs/configuration.md b/docs/configuration.md index 40a3f3c..8800d1f 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -287,7 +287,7 @@ moderation: # the external moderator behind security.moderate timeout_seconds: 10 stability: - failure_threshold: 3 # consecutive failures (5xx, timeouts, a 401/402/403 refusal) before an account cools down + failure_threshold: 3 # consecutive failures (5xx, timeouts, a 401/402/403 refusal) before an account cools down for that model cooldown_seconds: 30 availability_window_minutes: 5 # /admin/models/status judgment window (max 60) unstable_error_rate: 0.1 # window error rate that reports `unstable` diff --git a/docs/multi-instance.md b/docs/multi-instance.md index ad3dbb8..f241f29 100644 --- a/docs/multi-instance.md +++ b/docs/multi-instance.md @@ -8,7 +8,7 @@ shared, what stays local, and what the LB needs to do. | State | Backend | Shared across the fleet? | |-------|---------|--------------------------| | Rate limits / quotas / TPM (`Governance`) | Redis (`storage.redis_url`) | ✅ when Redis is set (includes pooled tenant QPS) | -| Account health / cooldown (`HealthStore`) | Redis (`storage.redis_url`) | ✅ when Redis is set — one instance's cooldown benches the account for all | +| Account health / cooldown (`HealthStore`) | Redis (`storage.redis_url`) | ✅ when Redis is set — one instance's cooldown benches the account for that model on all | | Per-model availability counts (`AvailStore`, behind `/admin/models/status`) | Redis (`storage.redis_url`) | ✅ when Redis is set — every instance's samples land in the same minute buckets (one hour retained); in-process otherwise | | Config: keys/models/providers/tenants (`ConfigStore`) | Postgres (`storage.postgres_url`) | ✅ when Postgres is set — versioned documents + a change feed | | Access-key table (`KeyStore`) | Postgres (`storage.postgres_url`) | ✅ when Postgres is set — admin key CRUD is fleet-wide in under a second and survives restarts; a key's MCP servers and tool allowlists are stored with it | diff --git a/docs/providers.md b/docs/providers.md index 43cde29..e965e5f 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -198,8 +198,9 @@ the traffic while an account never or not recently (60 s) sampled ranks first and gets probed; equal ranks stay round-robin. The realtime surface keeps round-robin. On an upstream 5xx or a 401/402/403 credential or billing refusal the failed account is excluded and another is tried once (a PTU→paygo switch is flagged `ptu_spillover`). Consecutive failures put -an account into cooldown (`stability.failure_threshold` / `cooldown_seconds`), -and it auto-recovers on expiry. A streaming response that already sent bytes to +an account into cooldown for the served model (`stability.failure_threshold` / +`cooldown_seconds`), so it keeps serving its other models, and it +auto-recovers on expiry. A streaming response that already sent bytes to the client is never failed over, but a provider error that breaks such a stream still counts against the account's health and the model's availability; a plain client disconnect counts as neither.