mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-19 16:39:15 +00:00
feat(provider): half-cooldown probe so recovered slots return immediately
A slot that failed ten times landed in a 30-minute cooldown and was a hard
skip for the whole window: ModelAvailable() said no, the scheduler dropped the
candidate, and the ONLY way back was a success that could never happen. In
practice an upstream that blipped for a minute — or whose quota reset in
seconds — stayed unusable for half an hour.
Cooldown is now a window with a probe gate instead of a wall:
* backoffCap 30min -> 5min. 401/403 no longer implicitly "jump to the cap"
but get an explicit authCooldown (10min); 429 keeps its 30s fixed window.
* ModelState records cooldownFrom alongside cooldownUntil, so the window has
a measurable midpoint, plus a single-token probe permit (probeInFlight).
* TryProbe() hands out that permit only in the SECOND half of the window: a
freshly failed slot stays completely silent, and past the midpoint exactly
one request may go through as a probe.
* ProbeSlots() = clamp(max_concurrent/10, 1, 2), so a source configured for
max_concurrent=10 lets exactly one request probe.
* A probe is a REAL request: success runs RecordSuccess(), which clears the
cooldown and failCount on the spot and returns the slot to full rotation.
Failure reopens the window, pushing the next probe to the NEW midpoint
instead of retrying immediately.
Quota exhaustion is modelled separately (RecordQuotaExhausted): running out of
allowance is not the upstream being broken, so it must not feed the exponential
ladder or inflate failCount. It cools until the allowance window can plausibly
have reset (capped at 30min so a manual top-up is noticed) and is recognised
from the adapter-normalized error text via ReportStatusReason.
Scheduler side: collectCands() splits a tier into normal + probe candidates
with probes strictly LAST, so probe traffic is what a tier falls back to, never
what it prefers, and the round-robin cursor rotates over the normal head only.
releaseProbes() returns every permit on all exits (success, hard failure, ctx
cancel, tier fallthrough); the 90-line inlined busy-poll loop is extracted into
pollBusyTier() to keep those exits auditable. Direct Chat/ChatStream/Image take
the same gate.
Also fixes a real blacklist bug found on the way: a slot whose pref sank to
prefMin was refused by ModelAvailable forever, long after every cooldown had
expired. Such a slot is now probeable, and one success lifts it off the floor.
Verified live against a controllable upstream (tier1 flaky, tier2 healthy):
failure ladder 5s/9s/20s/40s/79s; after healing the upstream, 10 AUTO requests
across the silent half produced ZERO upstream calls; past the midpoint a single
AUTO request produced exactly ONE upstream call, cleared the remaining 39s of
cooldown, and the next three requests were served by the recovered slot.
This commit is contained in:
JianFeeeee
@ -13,6 +13,7 @@ import (
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@ -36,6 +37,16 @@ type ModelState struct {
|
||||
pref atomic.Int64 // +1 per success / -5 per failure, clamped
|
||||
failCount atomic.Int64
|
||||
cooldownUntil atomic.Int64 // unix seconds; 0 = schedulable
|
||||
// cooldownFrom is the unix second the current cooldown window opened. It
|
||||
// is the reference point for the half-cooldown probe gate: probing is
|
||||
// allowed only in the second half of [cooldownFrom, cooldownUntil), so a
|
||||
// freshly failed slot stays fully silent for a while instead of being
|
||||
// hammered immediately.
|
||||
cooldownFrom atomic.Int64
|
||||
// probeInFlight is the single-token probe permit. At most one probe
|
||||
// request may be in flight per (source, model) while cooling; the token is
|
||||
// released by ProbeDone regardless of outcome.
|
||||
probeInFlight atomic.Bool
|
||||
}
|
||||
|
||||
const (
|
||||
@ -43,11 +54,33 @@ const (
|
||||
prefMin = -20
|
||||
prefMax = 20
|
||||
backoffBase = 5 * time.Second
|
||||
backoffCap = 30 * time.Minute
|
||||
// backoffCap bounds the exponential schedule. It used to be 30min, which
|
||||
// in practice meant "a slot whose upstream blipped is unusable for half an
|
||||
// hour even though its quota recovered in seconds". Combined with the
|
||||
// half-cooldown probe gate below a much shorter cap is safe: a still-dead
|
||||
// upstream costs one probe request per cap/2, while a recovered one is back
|
||||
// in full rotation on its first probe success.
|
||||
backoffCap = 5 * time.Minute
|
||||
// failures at/above this count back off at the capped duration
|
||||
backoffCapN = 10
|
||||
)
|
||||
|
||||
// authCooldown is the cooldown applied to 401/403 credential failures. A bad
|
||||
// key does not fix itself in seconds, so it cools longer than a transport
|
||||
// blip — but it is still a self-healing cooldown (a rotated key is picked up by
|
||||
// the next probe), never a permanent blacklist.
|
||||
const authCooldown = 10 * time.Minute
|
||||
|
||||
// maxQuotaCooldown bounds a quota-exhausted cooldown. Even when the reported
|
||||
// reset window is huge (monthly quotas), the slot is re-probed at least this
|
||||
// often so a manually topped-up allowance is noticed.
|
||||
const maxQuotaCooldown = 30 * time.Minute
|
||||
|
||||
// quotaDefaultCooldown is used when the source gives no hint about its quota
|
||||
// window length. It is long enough not to hammer an exhausted allowance and
|
||||
// short enough that the half-cooldown probe retries within minutes.
|
||||
const quotaDefaultCooldown = 10 * time.Minute
|
||||
|
||||
// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit
|
||||
// rejection means "too fast", not "broken": unlike transport/5xx failures it
|
||||
// must NOT escalate exponentially to backoffCap, or a quota-rich source gets
|
||||
@ -81,25 +114,30 @@ func (s *ModelState) Available() bool {
|
||||
}
|
||||
|
||||
// RecordFailure counts one consecutive failure and schedules exponential
|
||||
// cooldown (5s, 10s, 20s … capped at 30min). auth marks 401/403 credential
|
||||
// failures: it jumps straight to the capped cooldown and doubles the
|
||||
// cooldown (5s, 10s, 20s … capped at backoffCap). auth marks 401/403
|
||||
// credential failures: they get the longer fixed authCooldown and a doubled
|
||||
// preference penalty, but the model still recovers when the cooldown expires.
|
||||
func (s *ModelState) RecordFailure(auth bool) {
|
||||
n := s.failCount.Add(1)
|
||||
if auth && n < backoffCapN {
|
||||
n = backoffCapN
|
||||
s.failCount.Store(n) // persist the cap so FailCount() reports it too
|
||||
}
|
||||
var cd time.Duration
|
||||
if n >= backoffCapN {
|
||||
switch {
|
||||
case auth:
|
||||
// Persist the cap so FailCount() reports "this is a hard problem" and a
|
||||
// following non-auth failure does not restart the ladder from the bottom.
|
||||
if n < backoffCapN {
|
||||
n = backoffCapN
|
||||
s.failCount.Store(n)
|
||||
}
|
||||
cd = authCooldown
|
||||
case n >= backoffCapN:
|
||||
cd = backoffCap
|
||||
} else {
|
||||
default:
|
||||
cd = backoffBase * time.Duration(1<<(n-1))
|
||||
if cd > backoffCap {
|
||||
cd = backoffCap
|
||||
}
|
||||
}
|
||||
s.cooldownUntil.Store(time.Now().Add(cd).Unix())
|
||||
s.openCooldown(cd)
|
||||
pen := int64(prefFailStep)
|
||||
if auth {
|
||||
pen *= 2
|
||||
@ -108,22 +146,57 @@ func (s *ModelState) RecordFailure(auth bool) {
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
|
||||
// openCooldown starts a fresh cooldown window of length cd. Recording both
|
||||
// ends of the window (not just the deadline) is what makes the half-cooldown
|
||||
// probe gate possible, and resetting cooldownFrom on every new failure means a
|
||||
// failed probe pushes the next probe to the middle of the NEW window instead of
|
||||
// retrying immediately.
|
||||
func (s *ModelState) openCooldown(cd time.Duration) {
|
||||
now := time.Now()
|
||||
s.cooldownFrom.Store(now.Unix())
|
||||
s.cooldownUntil.Store(now.Add(cd).Unix())
|
||||
}
|
||||
|
||||
// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the
|
||||
// exponential schedule, so a merely-throttled source recovers in seconds and
|
||||
// its remaining quota stays usable. The preference penalty still applies so
|
||||
// other slots are preferred while cooling.
|
||||
func (s *ModelState) RecordRateLimit() {
|
||||
s.failCount.Add(1)
|
||||
s.cooldownUntil.Store(time.Now().Add(rateLimitCooldown).Unix())
|
||||
s.openCooldown(rateLimitCooldown)
|
||||
s.pref.Add(-prefFailStep)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
|
||||
// RecordQuotaExhausted marks the model as out of allowance until its quota
|
||||
// window resets. Unlike a failure this is not "the upstream is broken", so it
|
||||
// must not feed the exponential ladder: the slot is unusable until the window
|
||||
// boundary and immediately usable again after it. resetIn is the time left in
|
||||
// the current quota window; non-positive values fall back to the short
|
||||
// rate-limit window (the caller could not determine a period).
|
||||
func (s *ModelState) RecordQuotaExhausted(resetIn time.Duration) {
|
||||
if resetIn <= 0 {
|
||||
resetIn = rateLimitCooldown
|
||||
}
|
||||
if resetIn > maxQuotaCooldown {
|
||||
resetIn = maxQuotaCooldown
|
||||
}
|
||||
s.openCooldown(resetIn)
|
||||
// No failCount bump: an exhausted quota is not an error streak, and
|
||||
// letting it inflate failCount would make the NEXT real failure jump
|
||||
// straight to the capped backoff.
|
||||
s.pref.Add(-prefFailStep)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
|
||||
// RecordSuccess resets the failure counter and cooldown and bumps the
|
||||
// preference score by one.
|
||||
// preference score by one. A single success is enough to lift a slot off the
|
||||
// prefMin floor (prefMin+1 > prefMin), which is what lets one successful probe
|
||||
// return a written-off slot to normal rotation.
|
||||
func (s *ModelState) RecordSuccess() {
|
||||
s.failCount.Store(0)
|
||||
s.cooldownUntil.Store(0)
|
||||
s.cooldownFrom.Store(0)
|
||||
s.pref.Add(1)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
@ -131,9 +204,53 @@ func (s *ModelState) RecordSuccess() {
|
||||
func (s *ModelState) reset() {
|
||||
s.failCount.Store(0)
|
||||
s.cooldownUntil.Store(0)
|
||||
s.cooldownFrom.Store(0)
|
||||
s.probeInFlight.Store(false)
|
||||
s.pref.Store(0)
|
||||
}
|
||||
|
||||
// TryProbe attempts to claim the single probe permit for a model that normal
|
||||
// scheduling refuses. It returns true only when a probe is both needed and due:
|
||||
//
|
||||
// - cooling: allowed once the window passed its midpoint (the first half stays
|
||||
// completely silent).
|
||||
// - cooldown expired but the preference score sank to prefMin: ModelAvailable
|
||||
// still refuses the slot, so a probe is the ONLY way back. Without this a
|
||||
// slot that failed prefMin/prefFailStep times in a row would be blacklisted
|
||||
// forever despite every cooldown having long expired.
|
||||
//
|
||||
// A healthy, non-cooling model returns false: the caller should schedule it
|
||||
// normally. On true the caller MUST call ProbeDone exactly once.
|
||||
func (s *ModelState) TryProbe() bool {
|
||||
now := time.Now().Unix()
|
||||
until := s.cooldownUntil.Load()
|
||||
if until <= now {
|
||||
if s.pref.Load() > prefMin {
|
||||
return false // schedulable normally, no probe needed
|
||||
}
|
||||
// preference floor: probe to earn the way out
|
||||
return s.probeInFlight.CompareAndSwap(false, true)
|
||||
}
|
||||
from := s.cooldownFrom.Load()
|
||||
if from <= 0 || from >= until {
|
||||
return false // no measurable window (legacy/degenerate state)
|
||||
}
|
||||
if now < from+(until-from)/2 {
|
||||
return false // first half of the window: stay silent
|
||||
}
|
||||
return s.probeInFlight.CompareAndSwap(false, true)
|
||||
}
|
||||
|
||||
// ProbeDone releases the probe permit claimed by TryProbe.
|
||||
func (s *ModelState) ProbeDone() { s.probeInFlight.Store(false) }
|
||||
|
||||
// Probing reports whether a probe request currently holds the permit.
|
||||
func (s *ModelState) Probing() bool { return s.probeInFlight.Load() }
|
||||
|
||||
// CooldownFrom is the unix timestamp the current cooldown window opened; 0
|
||||
// when the model is not cooling.
|
||||
func (s *ModelState) CooldownFrom() int64 { return s.cooldownFrom.Load() }
|
||||
|
||||
// Pref is the current preference score (higher = preferred).
|
||||
func (s *ModelState) Pref() int64 { return s.pref.Load() }
|
||||
|
||||
@ -155,15 +272,15 @@ type Provider struct {
|
||||
client *http.Client
|
||||
stream *http.Client
|
||||
|
||||
mu sync.Mutex
|
||||
sem chan struct{}
|
||||
states map[string]*ModelState // key = model id
|
||||
mu sync.Mutex
|
||||
sem chan struct{}
|
||||
states map[string]*ModelState // key = model id
|
||||
// proactive rate limiting: nextOK is the earliest unix-nano time the next
|
||||
// request may leave; throttle() spaces requests 60s/RPM apart. Zero when
|
||||
// cfg.RPM <= 0 (unlimited).
|
||||
rateMu sync.Mutex
|
||||
nextOK int64
|
||||
rpmGap time.Duration
|
||||
rateMu sync.Mutex
|
||||
nextOK int64
|
||||
rpmGap time.Duration
|
||||
lastProbe struct {
|
||||
ok bool
|
||||
err string
|
||||
@ -370,6 +487,55 @@ func (p *Provider) ModelAvailable(model string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// ModelSchedulable is the scheduling gate that replaces a bare ModelAvailable
|
||||
// check. It reports whether the model may take a request right now and whether
|
||||
// doing so is a cooldown PROBE rather than normal traffic:
|
||||
//
|
||||
// - (true, false): fully available, schedule normally.
|
||||
// - (true, true): cooling, but past the half-cooldown mark and this caller
|
||||
// claimed the probe permit. The caller MUST call ProbeDone(model) once the
|
||||
// attempt finished, and SHOULD prefer non-probe candidates first.
|
||||
// - (false, false): cooling and not probeable (first half of the window, the
|
||||
// permit is taken, or the slot is persistently bad).
|
||||
//
|
||||
// Probing is what stops a recovered upstream from sitting out a whole cooldown:
|
||||
// the probe is a real request, so success immediately clears the cooldown via
|
||||
// RecordSuccess and the slot is back in full rotation.
|
||||
func (p *Provider) ModelSchedulable(model string) (ok bool, isProbe bool) {
|
||||
p.mu.Lock()
|
||||
s, known := p.states[model]
|
||||
p.mu.Unlock()
|
||||
if !known {
|
||||
return true, false
|
||||
}
|
||||
if s.Available() && s.Pref() > prefMin {
|
||||
return true, false
|
||||
}
|
||||
if s.TryProbe() {
|
||||
return true, true
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
// ProbeDone releases a probe permit claimed through ModelSchedulable.
|
||||
func (p *Provider) ProbeDone(model string) { p.state(model).ProbeDone() }
|
||||
|
||||
// ProbeSlots is how many concurrent probe requests this source tolerates while
|
||||
// a model cools: one per 10 configured concurrency slots, at least 1 and never
|
||||
// more than 2. A source configured for max_concurrent=10 therefore lets exactly
|
||||
// one request through to probe, which is the intended "small-scale probe"
|
||||
// behaviour; a large source still cannot flood a cooling upstream.
|
||||
func (p *Provider) ProbeSlots() int {
|
||||
n := p.cfg.MaxConcurrent / 10
|
||||
if n < 1 {
|
||||
return 1
|
||||
}
|
||||
if n > 2 {
|
||||
return 2
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// Pref returns the adaptive preference score of a model (higher = preferred
|
||||
// within a priority tier). Used by the AUTO chain to order same-tier slots.
|
||||
func (p *Provider) Pref(model string) int64 {
|
||||
@ -393,6 +559,19 @@ func (p *Provider) ModelHealthInfo(model string) (pref, failCount, cooldownUntil
|
||||
return s.Pref(), s.FailCount(), s.CooldownUntil()
|
||||
}
|
||||
|
||||
// ModelProbeInfo exposes the probe-gate state for the web UI: when the current
|
||||
// cooldown window opened, when probing becomes allowed (its midpoint), and
|
||||
// whether a probe is in flight right now. All zero when not cooling.
|
||||
func (p *Provider) ModelProbeInfo(model string) (from, probeAfter int64, probing bool) {
|
||||
s := p.state(model)
|
||||
until := s.CooldownUntil()
|
||||
from = s.CooldownFrom()
|
||||
if until <= 0 || from <= 0 || from >= until {
|
||||
return 0, 0, s.Probing()
|
||||
}
|
||||
return from, from + (until-from)/2, s.Probing()
|
||||
}
|
||||
|
||||
// Probe performs a lightweight reachability + auth check against the source.
|
||||
// It first tries GET <base>/models (fast, ~1s for OpenAI-compatible upstreams)
|
||||
// and only falls back to a 1-token chat call when that endpoint is
|
||||
@ -518,17 +697,33 @@ func (p *Provider) RecordSuccess(model string) {
|
||||
}
|
||||
|
||||
// ReportStatus records an upstream HTTP status for the given model and drives
|
||||
// the (source, model) backoff state. 401/403 → capped self-healing cooldown
|
||||
// the (source, model) backoff state. 401/403 → long self-healing authCooldown
|
||||
// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed
|
||||
// cooldown (rateLimitCooldown): a rate rejection means "too fast", not
|
||||
// "broken", so the slot must re-enter rotation quickly instead of escalating
|
||||
// to a 30-minute lockout that wastes remaining quota.
|
||||
// to a long lockout that wastes remaining quota.
|
||||
//
|
||||
// Prefer ReportStatusReason when the upstream error text is available: a
|
||||
// quota-exhausted rejection is cooled to its window boundary instead of being
|
||||
// treated as a failure streak.
|
||||
func (p *Provider) ReportStatus(model string, code int) {
|
||||
p.ReportStatusReason(model, code, "")
|
||||
}
|
||||
|
||||
// ReportStatusReason is ReportStatus plus the upstream error text, which lets a
|
||||
// quota/allowance rejection be distinguished from a generic rate limit. Quota
|
||||
// exhaustion is not an error streak: it is cooled until the allowance window
|
||||
// can plausibly have reset, and never inflates the exponential ladder.
|
||||
func (p *Provider) ReportStatusReason(model string, code int, reason string) {
|
||||
if code == 401 || code == 403 {
|
||||
p.RecordFailure(model, code)
|
||||
return
|
||||
}
|
||||
if code == 429 {
|
||||
if code == 429 || code == 402 {
|
||||
if quotaExhaustedReason(reason) {
|
||||
p.state(model).RecordQuotaExhausted(p.quotaResetIn(model))
|
||||
return
|
||||
}
|
||||
p.state(model).RecordRateLimit()
|
||||
return
|
||||
}
|
||||
@ -537,6 +732,82 @@ func (p *Provider) ReportStatus(model string, code int) {
|
||||
}
|
||||
}
|
||||
|
||||
// quotaResetIn estimates how long the model's allowance window still has to
|
||||
// run. Upstreams do not report this, so it is taken from the source's optional
|
||||
// `quota_hours` meta hint when present (aligned to the epoch, matching the
|
||||
// gateway's own quota-window convention) and otherwise falls back to
|
||||
// quotaDefaultCooldown. Either way the value is capped by maxQuotaCooldown, so
|
||||
// a wrong hint can only delay recovery to that bound — and the half-cooldown
|
||||
// probe fires at half of it.
|
||||
func (p *Provider) quotaResetIn(model string) time.Duration {
|
||||
hours := p.quotaHours()
|
||||
if hours <= 0 {
|
||||
return quotaDefaultCooldown
|
||||
}
|
||||
win := int64(hours * float64(time.Hour/time.Second))
|
||||
if win <= 0 {
|
||||
return quotaDefaultCooldown
|
||||
}
|
||||
elapsed := time.Now().Unix() % win
|
||||
return time.Duration(win-elapsed) * time.Second
|
||||
}
|
||||
|
||||
// quotaHours reads the optional per-source `quota_hours` meta hint describing
|
||||
// how long its allowance window is. Absent or unparseable = 0 (unknown).
|
||||
func (p *Provider) quotaHours() float64 {
|
||||
v, ok := p.cfg.Meta["quota_hours"]
|
||||
if !ok {
|
||||
return 0
|
||||
}
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n
|
||||
case int:
|
||||
return float64(n)
|
||||
case int64:
|
||||
return float64(n)
|
||||
case string:
|
||||
f, err := strconv.ParseFloat(strings.TrimSpace(n), 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return f
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// quotaKeywords are the upstream phrases that mean "you are out of allowance"
|
||||
// rather than "you are going too fast". Matching is case-insensitive substring
|
||||
// matching against the adapter-normalized error reason.
|
||||
var quotaKeywords = []string{
|
||||
"insufficient_quota",
|
||||
"insufficient quota",
|
||||
"quota exceeded",
|
||||
"quota_exceeded",
|
||||
"exceeded your current quota",
|
||||
"out of credit",
|
||||
"insufficient balance",
|
||||
"insufficient_user_quota",
|
||||
"billing_hard_limit_reached",
|
||||
"credit balance is too low",
|
||||
"余额不足",
|
||||
"额度不足",
|
||||
"额度已用完",
|
||||
}
|
||||
|
||||
func quotaExhaustedReason(reason string) bool {
|
||||
if reason == "" {
|
||||
return false
|
||||
}
|
||||
low := strings.ToLower(reason)
|
||||
for _, k := range quotaKeywords {
|
||||
if strings.Contains(low, k) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ResetHealth resets the scheduling state of every model of this source
|
||||
// (cooldown and preference to zero), so the source becomes fully schedulable
|
||||
// again. Called after AUTO-chain edits and from the admin UI.
|
||||
@ -717,8 +988,9 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
|
||||
return nil, err
|
||||
}
|
||||
if status != 200 {
|
||||
p.ReportStatus(model, status)
|
||||
return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
|
||||
reason := p.apiErrReason(status, raw)
|
||||
p.ReportStatusReason(model, status, reason)
|
||||
return nil, fmt.Errorf("%s", reason)
|
||||
}
|
||||
unified, err := p.vm.Transform(p.adapter, "transform_response", raw)
|
||||
if err != nil {
|
||||
@ -792,9 +1064,10 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
if sel.resp.StatusCode != 200 {
|
||||
raw, _ := io.ReadAll(sel.resp.Body)
|
||||
sel.resp.Body.Close()
|
||||
p.ReportStatus(model, sel.resp.StatusCode)
|
||||
reason := p.apiErrReason(sel.resp.StatusCode, string(raw))
|
||||
p.ReportStatusReason(model, sel.resp.StatusCode, reason)
|
||||
p.Release()
|
||||
return nil, fmt.Errorf("%s", p.apiErrReason(sel.resp.StatusCode, string(raw)))
|
||||
return nil, fmt.Errorf("%s", reason)
|
||||
}
|
||||
go func() {
|
||||
defer p.Release()
|
||||
@ -1043,8 +1316,9 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
|
||||
return nil, err
|
||||
}
|
||||
if status != 200 {
|
||||
p.ReportStatus(model, status)
|
||||
return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
|
||||
reason := p.apiErrReason(status, raw)
|
||||
p.ReportStatusReason(model, status, reason)
|
||||
return nil, fmt.Errorf("%s", reason)
|
||||
}
|
||||
var out types.UnifiedResponse
|
||||
// try adapter transform_response; if missing, parse standard openai image format
|
||||
|
||||
@ -198,8 +198,8 @@ func TestProviderAuthFailureSelfHeals(t *testing.T) {
|
||||
if st.FailCount() != backoffCapN {
|
||||
t.Fatalf("auth failure must jump to capped count, got %d", st.FailCount())
|
||||
}
|
||||
if until := st.CooldownUntil(); until <= time.Now().Add(25*time.Minute).Unix() {
|
||||
t.Fatalf("auth failure must cool near the cap (until=%d)", until)
|
||||
if until := st.CooldownUntil(); until <= time.Now().Add(authCooldown-time.Minute).Unix() {
|
||||
t.Fatalf("auth failure must cool for authCooldown (until=%d)", until)
|
||||
}
|
||||
if st.Pref() != -2*int64(prefFailStep) {
|
||||
t.Fatalf("auth failure pref penalty must be doubled, got %d", st.Pref())
|
||||
@ -678,3 +678,369 @@ func TestToolIndexRemapper(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ---- half-cooldown probe gate (plan 阶段 1) ----
|
||||
|
||||
// TestProbeGateHalfCooldown pins the core rule: a cooling model is fully silent
|
||||
// in the first half of its cooldown window and probeable exactly once in the
|
||||
// second half.
|
||||
func TestProbeGateHalfCooldown(t *testing.T) {
|
||||
var st ModelState
|
||||
// open a 100s window centred so that "now" sits in the first half
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - 10)
|
||||
st.cooldownUntil.Store(now + 90)
|
||||
if st.TryProbe() {
|
||||
t.Fatal("first half of the cooldown window must stay silent")
|
||||
}
|
||||
|
||||
// move into the second half
|
||||
st.cooldownFrom.Store(now - 90)
|
||||
st.cooldownUntil.Store(now + 10)
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("second half of the cooldown window must allow one probe")
|
||||
}
|
||||
if st.TryProbe() {
|
||||
t.Fatal("probe permit must be single-token")
|
||||
}
|
||||
if !st.Probing() {
|
||||
t.Fatal("Probing() must report the claimed permit")
|
||||
}
|
||||
st.ProbeDone()
|
||||
if st.Probing() {
|
||||
t.Fatal("ProbeDone must release the permit")
|
||||
}
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("permit must be reusable after ProbeDone")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeGateNoProbeWhenHealthy: a healthy model needs no probe permit — the
|
||||
// caller schedules it normally.
|
||||
func TestProbeGateNoProbeWhenHealthy(t *testing.T) {
|
||||
var idle ModelState
|
||||
if idle.TryProbe() {
|
||||
t.Fatal("a healthy model must not be probed (schedule it normally)")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeGateRescuesPrefFloor documents a fix that comes with the probe gate:
|
||||
// a slot whose preference sank to prefMin used to be refused by ModelAvailable
|
||||
// forever, even long after every cooldown expired — a permanent blacklist. It is
|
||||
// now probeable, so one success brings it back.
|
||||
func TestProbeGateRescuesPrefFloor(t *testing.T) {
|
||||
var dead ModelState
|
||||
dead.pref.Store(prefMin)
|
||||
if !dead.TryProbe() {
|
||||
t.Fatal("a slot stuck at prefMin must be probeable, not blacklisted")
|
||||
}
|
||||
dead.RecordSuccess()
|
||||
dead.ProbeDone()
|
||||
if dead.Pref() <= prefMin {
|
||||
t.Fatalf("one probe success must lift the slot off the floor, pref=%d", dead.Pref())
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeSuccessRestoresImmediately is the user-visible payoff: a recovered
|
||||
// upstream must not sit out the rest of its cooldown.
|
||||
func TestProbeSuccessRestoresImmediately(t *testing.T) {
|
||||
var st ModelState
|
||||
for i := 0; i < backoffCapN; i++ {
|
||||
st.RecordFailure(false)
|
||||
}
|
||||
if cd := st.CooldownUntil() - time.Now().Unix(); cd < int64(backoffCap.Seconds())-2 {
|
||||
t.Fatalf("10 failures must cool for backoffCap, got %ds", cd)
|
||||
}
|
||||
// jump into the probe window: the midpoint of [from, until) must be in the
|
||||
// past, so push the window's start back past its full length
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - int64(backoffCap.Seconds()) - 10)
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("probe must be allowed past the midpoint")
|
||||
}
|
||||
if !st.Probing() {
|
||||
t.Fatal("probe permit must be held while the probe runs")
|
||||
}
|
||||
st.RecordSuccess()
|
||||
st.ProbeDone()
|
||||
if !st.Available() {
|
||||
t.Fatal("probe success must clear the cooldown immediately")
|
||||
}
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("probe success must reset failCount, got %d", st.FailCount())
|
||||
}
|
||||
if st.CooldownUntil() != 0 || st.CooldownFrom() != 0 {
|
||||
t.Fatal("probe success must clear both cooldown bounds")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeFailureDefersNextProbe: a failed probe must not immediately retry —
|
||||
// the new window's midpoint pushes the next probe out.
|
||||
func TestProbeFailureDefersNextProbe(t *testing.T) {
|
||||
var st ModelState
|
||||
st.RecordFailure(false) // 5s window
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - 100)
|
||||
st.cooldownUntil.Store(now + 10)
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("probe must be allowed past the midpoint")
|
||||
}
|
||||
st.RecordFailure(false) // probe failed: a fresh, longer window opens
|
||||
st.ProbeDone()
|
||||
if st.TryProbe() {
|
||||
t.Fatal("a failed probe must defer the next probe to the new midpoint")
|
||||
}
|
||||
if from := st.CooldownFrom(); from < now {
|
||||
t.Fatalf("failed probe must reopen the window (from=%d, now=%d)", from, now)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuthCooldownIsSeparateFromCap documents that credential failures cool on
|
||||
// their own (longer) schedule instead of reusing the exponential cap.
|
||||
func TestAuthCooldownIsSeparateFromCap(t *testing.T) {
|
||||
var st ModelState
|
||||
st.RecordFailure(true)
|
||||
cd := st.CooldownUntil() - time.Now().Unix()
|
||||
if cd < int64(authCooldown.Seconds())-2 || cd > int64(authCooldown.Seconds())+2 {
|
||||
t.Fatalf("auth failure must cool for authCooldown (%v), got %ds", authCooldown, cd)
|
||||
}
|
||||
if st.FailCount() != backoffCapN {
|
||||
t.Fatalf("auth failure must persist the capped count, got %d", st.FailCount())
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuotaExhaustedDoesNotEscalate is the second half of the user's complaint:
|
||||
// running out of allowance must not be treated as a failure streak.
|
||||
func TestQuotaExhaustedDoesNotEscalate(t *testing.T) {
|
||||
var st ModelState
|
||||
st.RecordQuotaExhausted(2 * time.Minute)
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("quota exhaustion must not bump failCount, got %d", st.FailCount())
|
||||
}
|
||||
cd := st.CooldownUntil() - time.Now().Unix()
|
||||
if cd < 118 || cd > 122 {
|
||||
t.Fatalf("quota cooldown must honour the reset window, got %ds", cd)
|
||||
}
|
||||
|
||||
// an unbounded window is clamped so a topped-up allowance is still noticed
|
||||
var big ModelState
|
||||
big.RecordQuotaExhausted(72 * time.Hour)
|
||||
if cd := big.CooldownUntil() - time.Now().Unix(); cd > int64(maxQuotaCooldown.Seconds())+2 {
|
||||
t.Fatalf("quota cooldown must be capped at maxQuotaCooldown, got %ds", cd)
|
||||
}
|
||||
|
||||
// no hint at all falls back to the short window
|
||||
var none ModelState
|
||||
none.RecordQuotaExhausted(0)
|
||||
if cd := none.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
|
||||
t.Fatalf("unknown reset window must fall back to rateLimitCooldown, got %ds", cd)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuotaReasonRouting checks that a 429 carrying a quota phrase is cooled as
|
||||
// quota exhaustion while a plain rate limit keeps the short fixed window.
|
||||
func TestQuotaReasonRouting(t *testing.T) {
|
||||
p := newTestProvider(t, src("q", "http://127.0.0.1:1", "openai", "m"))
|
||||
p.ReportStatusReason("m", 429, "429 insufficient_quota: you exceeded your current quota")
|
||||
st := p.state("m")
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("quota rejection must not bump failCount, got %d", st.FailCount())
|
||||
}
|
||||
if cd := st.CooldownUntil() - time.Now().Unix(); cd <= int64(rateLimitCooldown.Seconds()) {
|
||||
t.Fatalf("quota rejection must cool longer than a plain 429, got %ds", cd)
|
||||
}
|
||||
|
||||
p2 := newTestProvider(t, src("r", "http://127.0.0.1:1", "openai", "m"))
|
||||
p2.ReportStatusReason("m", 429, "rate limit reached, slow down")
|
||||
st2 := p2.state("m")
|
||||
if cd := st2.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
|
||||
t.Fatalf("plain 429 must keep the short window, got %ds", cd)
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelSchedulableProbeHandoff covers the provider-level gate the scheduler
|
||||
// consumes.
|
||||
func TestModelSchedulableProbeHandoff(t *testing.T) {
|
||||
p := newTestProvider(t, src("s", "http://127.0.0.1:1", "openai", "m"))
|
||||
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
|
||||
t.Fatalf("healthy model must be normally schedulable, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
st := p.state("m")
|
||||
st.RecordFailure(false)
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - 100)
|
||||
st.cooldownUntil.Store(now + 10)
|
||||
|
||||
ok, probe := p.ModelSchedulable("m")
|
||||
if !ok || !probe {
|
||||
t.Fatalf("cooling model past midpoint must be probeable, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
if ok2, _ := p.ModelSchedulable("m"); ok2 {
|
||||
t.Fatal("only one probe permit may be outstanding")
|
||||
}
|
||||
p.ProbeDone("m")
|
||||
if ok3, probe3 := p.ModelSchedulable("m"); !ok3 || !probe3 {
|
||||
t.Fatal("permit must be reclaimable after ProbeDone")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeSlotsFollowsMaxConcurrent pins the user's example: max_concurrent=10
|
||||
// yields exactly one probe slot.
|
||||
func TestProbeSlotsFollowsMaxConcurrent(t *testing.T) {
|
||||
for _, tc := range []struct{ mc, want int }{
|
||||
{0, 1}, {1, 1}, {4, 1}, {10, 1}, {20, 2}, {100, 2},
|
||||
} {
|
||||
s := src("p", "http://127.0.0.1:1", "openai", "m")
|
||||
s.MaxConcurrent = tc.mc
|
||||
p := newTestProvider(t, s)
|
||||
if got := p.ProbeSlots(); got != tc.want {
|
||||
t.Fatalf("max_concurrent=%d: probe slots = %d, want %d", tc.mc, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelProbeInfo checks the UI-facing view of the probe window.
|
||||
func TestModelProbeInfo(t *testing.T) {
|
||||
p := newTestProvider(t, src("i", "http://127.0.0.1:1", "openai", "m"))
|
||||
if from, after, probing := p.ModelProbeInfo("m"); from != 0 || after != 0 || probing {
|
||||
t.Fatalf("idle model must report no probe window, got %d %d %v", from, after, probing)
|
||||
}
|
||||
st := p.state("m")
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now)
|
||||
st.cooldownUntil.Store(now + 100)
|
||||
from, after, _ := p.ModelProbeInfo("m")
|
||||
if from != now || after != now+50 {
|
||||
t.Fatalf("probe window = (%d,%d), want (%d,%d)", from, after, now, now+50)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpstreamRecoveryWithoutWaitingOutCooldown is the end-to-end shape of the
|
||||
// user's complaint: an upstream that 500s ten times lands in the capped
|
||||
// cooldown, and used to be unusable for a full 30 minutes even after it
|
||||
// recovered. With the probe gate the slot is silent for the first half of the
|
||||
// window and then serves the very next request itself, clearing the cooldown.
|
||||
//
|
||||
// Real time is not waited out: the window bounds are moved so the test observes
|
||||
// the same state the scheduler would see at cap/2.
|
||||
func TestUpstreamRecoveryWithoutWaitingOutCooldown(t *testing.T) {
|
||||
var healthy atomic.Bool
|
||||
var hits atomic.Int64
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hits.Add(1)
|
||||
if !healthy.Load() {
|
||||
w.WriteHeader(500)
|
||||
return
|
||||
}
|
||||
fmt.Fprint(w, `{"choices":[{"message":{"content":"back"},"finish_reason":"stop"}]}`)
|
||||
}))
|
||||
defer up.Close()
|
||||
p := newTestProvider(t, src("recov", up.URL, "openai", "m"))
|
||||
req := func() *types.ChatRequest {
|
||||
return &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
}
|
||||
}
|
||||
|
||||
// 1. drive the slot into the capped cooldown
|
||||
for i := 0; i < backoffCapN; i++ {
|
||||
if _, err := p.Chat(context.Background(), req()); err == nil {
|
||||
t.Fatalf("attempt %d: expected upstream failure", i)
|
||||
}
|
||||
}
|
||||
st := p.state("m")
|
||||
if p.ModelAvailable("m") {
|
||||
t.Fatal("10 consecutive 5xx must take the slot out of normal rotation")
|
||||
}
|
||||
cd := st.CooldownUntil() - time.Now().Unix()
|
||||
if cd < int64(backoffCap.Seconds())-2 {
|
||||
t.Fatalf("cooldown = %ds, want ~%v", cd, backoffCap)
|
||||
}
|
||||
// the old 30-minute lockout is gone
|
||||
if cd > int64((6 * time.Minute).Seconds()) {
|
||||
t.Fatalf("cooldown %ds is too long: the cap must be minutes, not half an hour", cd)
|
||||
}
|
||||
|
||||
// 2. upstream recovers, but we are still in the FIRST half of the window:
|
||||
// the slot must stay completely silent (no probe traffic at all).
|
||||
healthy.Store(true)
|
||||
quiet := hits.Load()
|
||||
if ok, probe := p.ModelSchedulable("m"); ok || probe {
|
||||
t.Fatalf("first half of the window must stay silent, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
if hits.Load() != quiet {
|
||||
t.Fatal("no request may reach the upstream during the silent half")
|
||||
}
|
||||
|
||||
// 3. advance past the midpoint: exactly one probe is admitted.
|
||||
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
|
||||
ok, isProbe := p.ModelSchedulable("m")
|
||||
if !ok || !isProbe {
|
||||
t.Fatalf("past the midpoint one probe must be admitted, got ok=%v probe=%v", ok, isProbe)
|
||||
}
|
||||
if ok2, _ := p.ModelSchedulable("m"); ok2 {
|
||||
t.Fatal("only ONE probe may be in flight while cooling")
|
||||
}
|
||||
|
||||
// 4. the probe is a real request: its success restores the slot fully.
|
||||
resp, err := p.Chat(context.Background(), req())
|
||||
p.ProbeDone("m")
|
||||
if err != nil {
|
||||
t.Fatalf("probe request failed: %v", err)
|
||||
}
|
||||
if resp.Content != "back" {
|
||||
t.Fatalf("probe response = %q, want \"back\"", resp.Content)
|
||||
}
|
||||
if !p.ModelAvailable("m") {
|
||||
t.Fatal("a successful probe must return the slot to normal rotation immediately")
|
||||
}
|
||||
if st.CooldownUntil() != 0 || st.FailCount() != 0 {
|
||||
t.Fatalf("probe success must clear cooldown/failCount, got until=%d fails=%d",
|
||||
st.CooldownUntil(), st.FailCount())
|
||||
}
|
||||
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
|
||||
t.Fatalf("recovered slot must schedule normally, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRecoveryProbeBudgetIsBounded quantifies the cost of the shorter cap
|
||||
// against a still-dead upstream: one probe per window, not a retry storm.
|
||||
func TestRecoveryProbeBudgetIsBounded(t *testing.T) {
|
||||
var hits atomic.Int64
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hits.Add(1)
|
||||
w.WriteHeader(500)
|
||||
}))
|
||||
defer up.Close()
|
||||
p := newTestProvider(t, src("dead", up.URL, "openai", "m"))
|
||||
req := &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
}
|
||||
for i := 0; i < backoffCapN; i++ {
|
||||
_, _ = p.Chat(context.Background(), req)
|
||||
}
|
||||
baseline := hits.Load()
|
||||
st := p.state("m")
|
||||
|
||||
// simulate three consecutive probe windows against a dead upstream
|
||||
for i := 0; i < 3; i++ {
|
||||
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
|
||||
ok, isProbe := p.ModelSchedulable("m")
|
||||
if !ok || !isProbe {
|
||||
t.Fatalf("window %d: expected a probe permit", i)
|
||||
}
|
||||
// a second caller in the same window gets nothing
|
||||
if ok2, _ := p.ModelSchedulable("m"); ok2 {
|
||||
t.Fatalf("window %d: more than one probe admitted", i)
|
||||
}
|
||||
if _, err := p.Chat(context.Background(), req); err == nil {
|
||||
t.Fatalf("window %d: expected the probe to fail", i)
|
||||
}
|
||||
p.ProbeDone("m")
|
||||
}
|
||||
if extra := hits.Load() - baseline; extra != 3 {
|
||||
t.Fatalf("3 probe windows must cost exactly 3 upstream requests, got %d", extra)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user