jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-20 00:48:00 +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

233
internal/scheduler/scheduler.go
View File

@ -42,6 +42,12 @@ type Provider interface {
Name() string
ModelFor(reqModel string) string
ModelAvailable(model string) bool
// ModelSchedulable is the probe-aware availability gate: ok reports
// whether the model may take a request now, isProbe marks that it is only
// allowed as a cooldown probe (the caller must release the permit with
// ProbeDone once the attempt finished).
ModelSchedulable(model string) (ok bool, isProbe bool)
ProbeDone(model string)
Pref(model string) int64
Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error)
ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error)
@ -176,18 +182,84 @@ type tierResult struct {
hard []TierError // hard failures seen in this pass (nil = none)
}
// candidate is one schedulable slot in a tier pass. probe marks a slot that is
// still cooling but past its half-cooldown mark and holding the probe permit:
// it is tried only after every normal slot, so probe traffic is what the tier
// falls back to instead of what it prefers.
type candidate struct {
slot *Slot
probe bool
}
// collectCands partitions a tier's slots into normal and probe candidates,
// normal first. Quota-exhausted slots are dropped outright. Every probe
// candidate returned holds a probe permit, so the caller MUST call
// releaseProbes on the result exactly once.
func collectCands(slots []*Slot, exhausted func(*Slot) bool) []candidate {
var normal, probes []candidate
for _, sl := range slots {
if exhausted != nil && exhausted(sl) {
continue
}
ok, isProbe := sl.Prov.ModelSchedulable(sl.Model)
if !ok {
continue
}
if isProbe {
probes = append(probes, candidate{slot: sl, probe: true})
continue
}
normal = append(normal, candidate{slot: sl})
}
return append(normal, probes...)
}
// releaseProbes hands every claimed probe permit back, whether or not the probe
// slot was actually used.
func releaseProbes(cands []candidate) {
for _, c := range cands {
if c.probe {
c.slot.Prov.ProbeDone(c.slot.Model)
}
}
}
// normalCount is how many leading candidates are normal (non-probe). The
// round-robin cursor rotates only over those: probe slots are a strictly
// ordered tail, never a rotation target.
func normalCount(cands []candidate) int {
for i, c := range cands {
if c.probe {
return i
}
}
return len(cands)
}
// runTier executes one tier pass starting at the round-robin base index.
// Cooldown is the only hard skip (re-verified per slot); a busy slot is
// skipped without any penalty; a hard failure is recorded and the pass moves
// on to the next slot (plan 2.3: "单请求内不重试已失败槽" — the failed slot is
// not retried, the others still are). hard == nil and no success means every
// candidate was merely busy/cooling, so the caller may wait a bounded time.
func runTier(ctx context.Context, tn *TierNode, cands []*Slot, base int64, req *types.ChatRequest, stream bool) tierResult {
//
// Normal candidates rotate by base; probe candidates form a fixed tail tried
// only after every normal slot failed or was busy.
func runTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool) tierResult {
n := len(cands)
norm := normalCount(cands)
var hard []TierError
for i := 0; i < n; i++ {
sl := cands[(int(base)+i)%n]
if !sl.Prov.ModelAvailable(sl.Model) {
var c candidate
if i < norm {
c = cands[(int(base)+i)%norm] // rotate within the normal head
} else {
c = cands[i] // probe tail keeps its order
}
sl := c.slot
// A probe candidate is intentionally NOT re-checked here: it is cooling
// by definition, and its permit was already claimed.
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
continue
}
r := *req
@ -233,75 +305,40 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
}
var ce ChainErr
for _, tn := range chain.Tiers {
// initial filter: quota-exhausted and cooling slots are dropped
var cands []*Slot
for _, sl := range tn.Slots {
if exhausted != nil && exhausted(sl) {
continue
}
if !sl.Prov.ModelAvailable(sl.Model) {
continue
}
cands = append(cands, sl)
}
// initial filter: quota-exhausted slots are dropped, cooling slots are
// dropped unless they qualify as half-cooldown probes (appended last).
cands := collectCands(tn.Slots, exhausted)
if len(cands) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling or quota exhausted)", tn.Tier))
continue
}
// No Pref sort: load balancing is done by round-robin cursor.
// Persistently failing slots are excluded by ModelAvailable
// Persistently failing slots are excluded by ModelSchedulable
// (which checks Pref > prefMin).
base := tn.NextStart()
res := runTier(ctx, tn, cands, base, req, stream)
if res.resp != nil || res.chunks != nil {
releaseProbes(cands)
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
releaseProbes(cands)
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
releaseProbes(cands)
continue // hard failures: fall through to the next tier, no waiting
}
// every candidate was busy or cooling: bounded poll before downgrading
deadline := time.Now().Add(busyWait)
timer := time.NewTimer(busyPoll)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return nil, nil, "", "", ctx.Err()
case <-timer.C:
if err := s.pollBusyTier(ctx, tn, cands, base, req, stream, &ce); err != nil {
releaseProbes(cands)
if r, ok := err.(*tierSuccess); ok {
return r.res.resp, r.res.chunks, r.res.src, r.res.model, nil
}
done := time.Now().After(deadline)
if done {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
// refresh candidates: cooldowns may have expired meanwhile
var again []*Slot
for _, sl := range cands {
if sl.Prov.ModelAvailable(sl.Model) {
again = append(again, sl)
}
}
if len(again) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
res = runTier(ctx, tn, again, base, req, stream)
if res.resp != nil || res.chunks != nil {
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
break // hard failure while waiting: stop waiting, fall through
}
timer.Reset(busyPoll)
return nil, nil, "", "", err
}
releaseProbes(cands)
}
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
return nil, nil, "", "", fmt.Errorf("no auto slot configured")
@ -309,6 +346,57 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
return nil, nil, "", "", &ce
}
// tierSuccess carries a successful result out of pollBusyTier through the error
// return. It is never surfaced to callers of chainDrive.
type tierSuccess struct{ res tierResult }
func (t *tierSuccess) Error() string { return "tier success" }
// pollBusyTier waits a bounded time for a fully-busy tier to free a slot,
// retrying the pass while cooldowns expire. It returns nil when the tier should
// be abandoned (caller falls through to the next tier), a *tierSuccess when a
// retry succeeded, or a context error.
func (s *Scheduler) pollBusyTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool, ce *ChainErr) error {
deadline := time.Now().Add(busyWait)
timer := time.NewTimer(busyPoll)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
}
if time.Now().After(deadline) {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
return nil
}
// refresh candidates: cooldowns may have expired meanwhile. Probe
// candidates keep their already-claimed permit and stay eligible.
var again []candidate
for _, c := range cands {
if c.probe || c.slot.Prov.ModelAvailable(c.slot.Model) {
again = append(again, c)
}
}
if len(again) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
return nil
}
res := runTier(ctx, tn, again, base, req, stream)
if res.resp != nil || res.chunks != nil {
return &tierSuccess{res: res}
}
if ctx.Err() != nil {
return ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
return nil // hard failure while waiting: stop waiting, fall through
}
timer.Reset(busyPoll)
}
}
// ChainChat runs a non-streaming AUTO request down the chain. exhausted, when
// non-nil, decides slot token-quota exhaustion. Returns the response, the
// serving source and the exact model id used; on total failure a *ChainErr
@ -338,32 +426,35 @@ func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.Ima
}
var ce ChainErr
for _, tn := range chain.Tiers {
var cands []*Slot
for _, sl := range tn.Slots {
if !sl.Prov.ModelAvailable(sl.Model) {
continue
}
cands = append(cands, sl)
}
cands := collectCands(tn.Slots, nil)
if len(cands) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling)", tn.Tier))
continue
}
// No Pref sort: load balancing is done by round-robin cursor.
base := tn.NextStart()
norm := normalCount(cands)
var hard []TierError
for i := 0; i < len(cands); i++ {
sl := cands[(int(base)+i)%len(cands)]
if !sl.Prov.ModelAvailable(sl.Model) {
var c candidate
if i < norm {
c = cands[(int(base)+i)%norm]
} else {
c = cands[i]
}
sl := c.slot
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
continue
}
r := *req
r.Model = sl.Model
resp, err := sl.Prov.Image(ctx, &r)
if ctx.Err() != nil {
releaseProbes(cands)
return nil, "", "", ctx.Err()
}
if err == nil {
releaseProbes(cands)
return resp, sl.Source, sl.Model, nil
}
if errors.Is(err, types.ErrBusy) {
@ -371,6 +462,7 @@ func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.Ima
}
hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
}
releaseProbes(cands)
if len(hard) > 0 {
ce.Tiers = append(ce.Tiers, hard...)
}
@ -399,11 +491,15 @@ func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatR
p := cands[i]
r := *req
r.Model = p.ModelFor(req.Model)
if !p.ModelAvailable(r.Model) {
ok, isProbe := p.ModelSchedulable(r.Model)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
continue
}
resp, err := p.Chat(ctx, &r)
if isProbe {
p.ProbeDone(r.Model)
}
if ctx.Err() != nil {
return nil, "", "", ctx.Err()
}
@ -431,11 +527,15 @@ func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types
p := cands[i]
r := *req
r.Model = p.ModelFor(req.Model)
if !p.ModelAvailable(r.Model) {
ok, isProbe := p.ModelSchedulable(r.Model)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
continue
}
resp, err := p.ChatStream(ctx, &r)
if isProbe {
p.ProbeDone(r.Model)
}
if err == nil {
return resp, p.Name(), r.Model, nil
}
@ -454,11 +554,16 @@ func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.Imag
var lastErr error
for i := 0; i < attempts && i < len(cands); i++ {
p := cands[i]
if !p.ModelAvailable(p.ModelFor(req.Model)) {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), p.ModelFor(req.Model))
im := p.ModelFor(req.Model)
ok, isProbe := p.ModelSchedulable(im)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), im)
continue
}
resp, err := p.Image(ctx, req)
if isProbe {
p.ProbeDone(im)
}
if err == nil {
return resp, p.Name(), nil
}

173
internal/scheduler/scheduler_test.go
View File

@ -22,6 +22,14 @@ type fakeProvider struct {
busy atomic.Bool
fail atomic.Bool
chatHits atomic.Int64
// probeable makes an unavailable provider eligible as a half-cooldown
// probe candidate. probePermit is the single-token permit; probeClaims
// counts how many times it was handed out and probeDones how many times it
// was returned, so tests can assert the permit is always released.
probeable atomic.Bool
probePermit atomic.Bool
probeClaims atomic.Int64
probeDones atomic.Int64
}
func fakeProv(name, model string) *fakeProvider {
@ -36,6 +44,25 @@ func (f *fakeProvider) ModelFor(reqModel string) string { return f.model }
func (f *fakeProvider) ModelAvailable(model string) bool { return f.available.Load() }
func (f *fakeProvider) ModelSchedulable(model string) (bool, bool) {
if f.available.Load() {
return true, false
}
if !f.probeable.Load() {
return false, false
}
if f.probePermit.CompareAndSwap(false, true) {
f.probeClaims.Add(1)
return true, true
}
return false, false
}
func (f *fakeProvider) ProbeDone(model string) {
f.probeDones.Add(1)
f.probePermit.Store(false)
}
func (f *fakeProvider) Pref(model string) int64 { return f.pref.Load() }
func (f *fakeProvider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
@ -322,3 +349,149 @@ func TestChainResetCooldownAfterSwap(t *testing.T) {
t.Fatalf("first start = %d, want 0", base)
}
}
// ---- half-cooldown probe candidates (plan 阶段 1.3) ----
// TestChainProbeIsLastResort pins the ordering rule: a probe candidate must not
// steal traffic from a healthy slot in the same tier.
func TestChainProbeIsLastResort(t *testing.T) {
healthy := fakeProv("healthy", "h")
cooling := fakeProv("cooling", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
ch := BuildChain([]Rule{
{Tier: 0, Model: "c", Source: "cooling"}, // configured FIRST on purpose
{Tier: 0, Model: "h", Source: "healthy"},
}, bySource(healthy, cooling))
s := New(0)
for i := 0; i < 3; i++ {
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("iter %d: %v", i, err)
}
if src != "healthy" {
t.Fatalf("iter %d served by %q; a probe slot must never outrank a healthy one", i, src)
}
}
if cooling.chatHits.Load() != 0 {
t.Fatalf("probe slot was hit %d times while a healthy slot existed", cooling.chatHits.Load())
}
if got := cooling.probeDones.Load(); got != cooling.probeClaims.Load() {
t.Fatalf("probe permits leaked: claims=%d dones=%d", cooling.probeClaims.Load(), got)
}
}
// TestChainProbeServesWhenNothingElseCan is the recovery path: with every normal
// slot gone, the cooling slot's probe carries the request instead of the tier
// failing outright.
func TestChainProbeServesWhenNothingElseCan(t *testing.T) {
cooling := fakeProv("cooling", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
s := New(0)
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("probe must serve the request: %v", err)
}
if src != "cooling" || model != "c" {
t.Fatalf("served by src=%q model=%q, want cooling/c", src, model)
}
if resp == nil || resp.Content != "cooling" {
t.Fatalf("bad response: %#v", resp)
}
if cooling.chatHits.Load() != 1 {
t.Fatalf("probe must issue exactly one request, got %d", cooling.chatHits.Load())
}
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
}
}
// TestChainProbePermitReleasedOnFailure: a failed probe must still hand its
// permit back, otherwise the slot could never be probed again.
func TestChainProbePermitReleasedOnFailure(t *testing.T) {
cooling := fakeProv("cooling", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
cooling.fail.Store(true)
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
s := New(0)
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
t.Fatal("expected the failing probe to surface an error")
}
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
}
// permit is free again for the next attempt
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
t.Fatal("expected the second probe to fail too")
}
if cooling.probeClaims.Load() != 2 {
t.Fatalf("permit must be reclaimable, claims=%d", cooling.probeClaims.Load())
}
}
// TestChainProbeDoesNotBlockTierFallthrough: a cooling tier-1 slot that is not
// probeable must still let the request fall through to tier 2.
func TestChainProbeDoesNotBlockTierFallthrough(t *testing.T) {
cold := fakeProv("cold", "c")
cold.available.Store(false) // cooling and NOT probeable
backup := fakeProv("backup", "b")
ch := BuildChain([]Rule{
{Tier: 1, Model: "c", Source: "cold"},
{Tier: 2, Model: "b", Source: "backup"},
}, bySource(cold, backup))
s := New(0)
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil || src != "backup" {
t.Fatalf("want fallthrough to backup, got src=%q err=%v", src, err)
}
if cold.chatHits.Load() != 0 {
t.Fatal("non-probeable cooling slot must not be hit")
}
}
// TestDirectProbeReleasesPermit covers the non-AUTO paths (Chat/ChatStream/Image).
func TestDirectProbeReleasesPermit(t *testing.T) {
cooling := fakeProv("cooling", "m")
cooling.available.Store(false)
cooling.probeable.Store(true)
s := New(0)
_, src, _, err := s.Chat(context.Background(), []Provider{cooling}, chatReq())
if err != nil || src != "cooling" {
t.Fatalf("direct probe must serve: src=%q err=%v", src, err)
}
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
t.Fatalf("permit accounting: claims=%d dones=%d", cooling.probeClaims.Load(), cooling.probeDones.Load())
}
}
// TestChainProbeRoundRobinUnaffected: adding a probe tail must not disturb the
// round-robin rotation over the healthy head.
func TestChainProbeRoundRobinUnaffected(t *testing.T) {
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
cooling := fakeProv("s3", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
ch := BuildChain([]Rule{
{Tier: 0, Model: "a", Source: "s1"},
{Tier: 0, Model: "b", Source: "s2"},
{Tier: 0, Model: "c", Source: "s3"},
}, bySource(a, b, cooling))
s := New(0)
var got []string
for i := 0; i < 4; i++ {
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("iter %d: %v", i, err)
}
got = append(got, src)
}
want := []string{"s1", "s2", "s1", "s2"}
for i := range want {
if got[i] != want[i] {
t.Fatalf("rr order = %v, want %v (probe tail must not join the rotation)", got, want)
}
}
}