jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-20 08:57:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(provider): half-cooldown probe so recovered slots return immediately

Browse Source 
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
2026-08-30 08:04:50 +08:00
parent 2ed1f0ecde
commit 21cd0429a2
4 changed files with 1011 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

370
internal/provider/provider_test.go
View File

@ -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)
}
}