mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-10-03 23:54:06 +00:00
feat(plugin): AUTO 调度轨迹可见(chain_step stage)
被问"还有 auto 调度相关 stage 呢?"问出来的真实缺口。
## 问题
chainDrive 只返回 (resp, src, model, err),调用方只知道**最终哪个槽位赢了**。
遍历过程中算出来又丢掉的东西——哪些档被跳过、为什么跳过、哪些槽位硬失败、
哪档全忙——一律不可见。ChainErr 里其实有这些,但**只在全部失败时**才填,
而它是 error 返回值不是记录。于是:
"tier 1 冷却所以降级到 tier 3" == "tier 1 正常接单"
对插件而言 tier 只是个常量 -2("resolved by the chain"),信息量为零。而这
恰恰是优先级链存在的全部理由,也是"我那个贵模型为什么没被用"的答案。
## 做法(scheduler 侧零新依赖)
新增 TraceEvent / TraceSink,chainDrive 多一个可选 sink 参数:
- TraceEvent 是本包的普通 struct,sink 是 func 参数 ⇒ **不新增 import**,
scheduler 仍然可独立测试
- sink 为 nil 时每次 emit 只多一次 nil 判断;没有插件的网关在 AUTO 热路径上
零开销(gateway 的 chainTraceSink 直接返回 nil)
- 事件是纯观测:scheduler 不基于它做任何分支,gateway 也不把它喂回路由/
冷却/配额
四种 kind:tier_skip / slot_fail / tier_busy / selected,selected 每次成功
遍历恰好一次且是最后一步。顺序保证所有 step 在 routed 之前。
## 暴露给插件
新增 chain_step stage(逐个步骤),并在 request_end 载荷里加三个便于做报表的
字段:chain_walk(上限 12 步,防审计记录膨胀)、degraded、tier_served。
## ★ 计费口径(我按推荐的做,已写进文档,需要你确认)
**按实际服务的模型计费**:降级到 tier 3 仍按 tier 3 的价算,轨迹只作观测。
理由与 §7.5 的边界一致——插件只报表不执法,两套口径混在一起会引出"降级该不该
多收钱"这种无法从代码判断的争议。若要改成"按本该用的档计价",需要在 models
价目里允许按 tier 定价,这我没做,因为那是个产品决策。
## 计费插件同步消费
by_tier_served / skip_reasons / degraded_reqs 三个新维度。skip_reasons 的等待
时长做了归一(`no free slot within <wait>`),否则 busy-wait 文案一变就多一行。
降级次数在 request_end 里计而不是在 chain_step 里计:一次降级的请求要走多步,
按步计会重复计数。
## 判据(346 个测试全绿,新增 15 个)
scheduler 6 个:正常路径只发一个 selected / 跳档+降级可见 / 硬失败与跳档
严格区分(不可混为一谈,否则抖动上游看起来像空闲上游)/
nil sink 安全 / 全失败时轨迹与 ChainErr 并存且不互相破坏 /
空链不发事件
gateway 1 个端到端:tier 1 全 500 → 插件收到 slot_fail(tier 1) +
selected(tier 2),request_end 的 tier_served=2 且 degraded=true
lua 2 个:降级计数与按实际模型计价 / 跳过原因归一聚合
lua 1 个:chain_step 是真 stage 且顺序正确
3 个变异都红:去掉 slot_fail(3 个判据红)/ 去掉 tier_skip(1 个)/
去掉 degraded 字段(1 个)。
This commit is contained in:
@ -135,10 +135,43 @@ core.New
|
||||
|
||||
| stage | 代码位置 | 说明 |
|
||||
|---|---|---|
|
||||
| `request_start` | `gateway/chat.go` `handleChat` | 每个 chat 请求一次 |
|
||||
| `request_start` | `gateway/chat.go` `handleChat` / `fireImageStart` | 每个请求一次(chat 与生图各一条),在配额闸门**之前** |
|
||||
| `chain_step` | `gateway/chat.go` `chainTraceSink` | **仅 AUTO 路径**,每步一次 |
|
||||
| `routed` | `singleChat` / `streamChat` / `singleChatAuto` / `streamChatAuto` | 成功选定源之后,各一次 |
|
||||
| `request_end` | `gateway/chat.go` `writeRec` | **所有出口的唯一汇合点**,每个请求一次 |
|
||||
|
||||
### 3.1b 为什么单独有 `chain_step`
|
||||
|
||||
`routed` 只在**遍历结束后**触发一次,只带最终胜出的槽位。所以
|
||||
"tier 1 冷却所以降级到 tier 3"和"tier 1 正常接单"在它眼里**完全一样**——
|
||||
而这恰恰是优先级链存在的全部理由。
|
||||
|
||||
`chain_step` 补上这条信息,四种 `kind`:
|
||||
|
||||
| kind | 含义 | 何时产生 |
|
||||
|---|---|---|
|
||||
| `tier_skip` | 整档被跳过 | 该档所有槽位冷却中/配额用尽 |
|
||||
| `slot_fail` | 某个槽位硬失败 | 上游报错 / 适配器输出不可用 |
|
||||
| `tier_busy` | 整档全忙且有界等待超时 | 2s 内没等到空位 |
|
||||
| `selected` | 这个槽位接了单 | 每次成功遍历**恰好一次**,且是最后一步 |
|
||||
|
||||
顺序保证:所有 `chain_step` 都在 `routed` 之前,`selected` 是最后一步。
|
||||
所以只订阅 `request_end` 的插件也能拿到轨迹摘要(见下)。
|
||||
|
||||
### 3.1c `request_end` 里的轨迹摘要
|
||||
|
||||
除了逐个 `chain_step`,`request_end` 还带三个便于做报表的字段:
|
||||
|
||||
| 字段 | 含义 |
|
||||
|---|---|
|
||||
| `chain_walk` | 整个遍历的步骤数组(上限 12 步,超出截断) |
|
||||
| `degraded` | 布尔。`true` = 有过跳过/失败,即**发生了降级** |
|
||||
| `tier_served` | 实际服务的那一档;直连或全失败时为 `-1` |
|
||||
|
||||
> **计费口径**:按**实际服务的模型**计费。降级到 tier 3 仍按 tier 3 的价算,
|
||||
> `chain_step` / `degraded` / `tier_served` 只作**观测**,不参与计价。
|
||||
> 理由见 §7.5。
|
||||
|
||||
`request_end` 放在 `writeRec` 是因为四条入口路径(直连/AUTO × 流式/非流式)都
|
||||
经过它,既不会漏(流式的 token 数只有流结束才知道),也不会重复。
|
||||
|
||||
@ -345,6 +378,17 @@ curl -H "Authorization: Bearer $ADMIN_KEY" \
|
||||
`total` / `by_source` / `by_model` / `by_key` / `by_day`(`YYYY-MM-DD` UTC),
|
||||
每项含 `cost`、`requests`、`prompt_tokens`、`completion_tokens`、`failures`。
|
||||
|
||||
另有两个**降级观测**维度(来自 `chain_step`):
|
||||
|
||||
| 字段 | 含义 |
|
||||
|---|---|
|
||||
| `degraded_reqs` | 发生过降级的请求数 |
|
||||
| `by_tier_served` | 各档实际接单数(`{"1": 812, "2": 37}`) |
|
||||
| `skip_reasons` | 跳过原因计数,等待时长已归一(`no free slot within <wait>`) |
|
||||
|
||||
这三项是"网关是不是在悄悄降级"的核心指标:一个持续降级的网关,账单结构和健康
|
||||
网关看起来一模一样——**除非**单独统计降级次数。
|
||||
|
||||
### 7.4 计费策略:失败的请求怎么算
|
||||
|
||||
**保留 token 费用,丢弃固定费用。** 理由:上游在生成后才 500,token 确实被消耗
|
||||
|
||||
@ -930,6 +930,76 @@ func (g *Gateway) fireRouted(ctx context.Context, kind, source, model string, ti
|
||||
})
|
||||
}
|
||||
|
||||
// chainTraceSink adapts a scheduler TraceSink into the plugin chain_step stage.
|
||||
//
|
||||
// It returns nil when no plugin is loaded, so the scheduler's emit() does a
|
||||
// single nil check per event and the AUTO hot path pays nothing on a gateway
|
||||
// with no plugins.
|
||||
//
|
||||
// The events are also accumulated into walk so request_end can carry a compact
|
||||
// summary: a plugin that only listens to request_end still learns that a
|
||||
// degradation happened, which is the common case for a dashboard that does not
|
||||
// want to subscribe to a high-frequency stage.
|
||||
func (g *Gateway) chainTraceSink(ctx context.Context, kind string, walk *[]map[string]interface{}) scheduler.TraceSink {
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil || ps.Count() == 0 {
|
||||
return nil
|
||||
}
|
||||
key := keyID(reqKey(ctx))
|
||||
return func(ev scheduler.TraceEvent) {
|
||||
payload := map[string]interface{}{
|
||||
"stage": string(lua.StageChainStep),
|
||||
"kind": string(ev.Kind),
|
||||
"type": kind,
|
||||
"key": key,
|
||||
"tier": ev.Tier,
|
||||
"attempt": ev.Attempt,
|
||||
}
|
||||
if ev.Source != "" {
|
||||
payload["source"] = ev.Source
|
||||
}
|
||||
if ev.Model != "" {
|
||||
payload["model"] = ev.Model
|
||||
}
|
||||
if ev.Reason != "" {
|
||||
payload["reason"] = ev.Reason
|
||||
}
|
||||
if ev.Err != "" {
|
||||
payload["error"] = ev.Err
|
||||
}
|
||||
if walk != nil {
|
||||
// Keep the summary bounded: a pathological chain could emit many
|
||||
// steps, and request_end's payload is written to the audit trail.
|
||||
if len(*walk) < maxWalkSummary {
|
||||
*walk = append(*walk, map[string]interface{}{
|
||||
"kind": string(ev.Kind), "tier": ev.Tier,
|
||||
"source": ev.Source, "model": ev.Model, "reason": ev.Reason,
|
||||
})
|
||||
}
|
||||
}
|
||||
ps.Fire(lua.StageChainStep, payload)
|
||||
}
|
||||
}
|
||||
|
||||
// maxWalkSummary caps how many chain steps request_end carries, so a long
|
||||
// degradation cannot inflate every audit record.
|
||||
const maxWalkSummary = 12
|
||||
|
||||
// tierServed returns the AUTO tier that actually served the request, or -1 when
|
||||
// the walk is empty (a direct request) or ended without a selection (total
|
||||
// failure). It is the single most useful number for "why did my expensive tier
|
||||
// not get used".
|
||||
func tierServed(walk []map[string]interface{}) int {
|
||||
for i := len(walk) - 1; i >= 0; i-- {
|
||||
if k, _ := walk[i]["kind"].(string); k == string(scheduler.TraceSelected) {
|
||||
if t, ok := walk[i]["tier"].(int); ok {
|
||||
return t
|
||||
}
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// fireEnd dispatches the plugin request_end stage for one finished request.
|
||||
func (g *Gateway) fireEnd(rec *Req) {
|
||||
ps := g.core.Plugins()
|
||||
@ -953,6 +1023,13 @@ func (g *Gateway) fireEnd(rec *Req) {
|
||||
"image_count": rec.ImageCount,
|
||||
"error": rec.Err,
|
||||
"time": rec.Time,
|
||||
// chain_walk: the AUTO tier-by-tier trace, when the request went
|
||||
// through the chain. Empty for a direct request and for a gateway with
|
||||
// no plugins loaded. Absent rather than empty so a plugin can tell
|
||||
// "no chain" from "chain with no degradation".
|
||||
"degraded": len(rec.Walk) > 1,
|
||||
"chain_walk": rec.Walk,
|
||||
"tier_served": tierServed(rec.Walk),
|
||||
}
|
||||
// The merged result is intentionally discarded: request_end is the last
|
||||
// stage, so there is nobody downstream to read a plugin's additions. Plugins
|
||||
@ -1179,8 +1256,11 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, chain *scheduler.Chain, req *types.ChatRequest, rec *Req, quotaExhausted func(*scheduler.Slot) bool) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
|
||||
var walk []map[string]interface{}
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted,
|
||||
g.chainTraceSink(ctx, "chat", &walk))
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
rec.Walk = walk
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
g.writeRec(rec)
|
||||
@ -1213,7 +1293,10 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
|
||||
var walk []map[string]interface{}
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted,
|
||||
g.chainTraceSink(ctx, "stream", &walk))
|
||||
rec.Walk = walk
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
|
||||
@ -500,3 +500,127 @@ func TestRejectedChatStillFiresRequestStart(t *testing.T) {
|
||||
"plugin can count real traffic, not just served traffic", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChainStepReachesPluginOnDegradation is the end-to-end proof for the
|
||||
// AUTO trace: a request that had to drop from tier 1 to tier 2 must be visible
|
||||
// to a plugin as a tier_skip followed by a selected, and request_end must carry
|
||||
// tier_served=2.
|
||||
//
|
||||
// Before the trace existed the plugin saw only tier=-2 ("resolved by the
|
||||
// chain") and could not tell a degradation from a clean tier-1 hit — which is
|
||||
// the whole question a priority chain exists to answer.
|
||||
func TestChainStepReachesPluginOnDegradation(t *testing.T) {
|
||||
// tier 1's source always fails, so the walk must drop to tier 2.
|
||||
bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"error":"boom"}`))
|
||||
}))
|
||||
defer bad.Close()
|
||||
good := mockUpstream()
|
||||
defer good.Close()
|
||||
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
os.WriteFile(cfgPath, []byte("listen: :0\n"+
|
||||
"adapter_dir: "+filepath.Join(dir, "adapters")+"\n"+
|
||||
"plugin_dir: "+filepath.Join(dir, "plugins")+"\n"+
|
||||
"runtime_file: "+filepath.Join(dir, "runtime.json")+"\n"+
|
||||
"gateway_keys:\n - sk-test\n"), 0600)
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources = []config.Source{
|
||||
{Name: "t1", BaseURL: bad.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "hi-tier", Kind: "chat"}}},
|
||||
{Name: "t2", BaseURL: good.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "lo-tier", Kind: "chat"}}},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
|
||||
// A spy that records chain_step events too.
|
||||
sp := `
|
||||
local plugin = { name = "walker", version = "1.0.0" }
|
||||
plugin.state = { steps = {}, ends = {} }
|
||||
plugin.hooks = { chain_step = "step", request_end = "fin" }
|
||||
function plugin.step(p)
|
||||
table.insert(plugin.state.steps, { kind = p.kind, tier = p.tier, source = p.source, model = p.model, reason = p.reason })
|
||||
return nil
|
||||
end
|
||||
function plugin.fin(p)
|
||||
plugin.state.ends[#plugin.state.ends + 1] = {
|
||||
tier_served = p.tier_served, degraded = p.degraded,
|
||||
walk = p.chain_walk, source = p.source, model = p.model,
|
||||
}
|
||||
return nil
|
||||
end
|
||||
return plugin
|
||||
`
|
||||
if err := c.Plugins().LoadSource("walker", sp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Two tiers, both in the chain.
|
||||
if rr := doReq(t, g, http.MethodPut, "/api/auto",
|
||||
`{"rules":[{"model":"hi-tier","source":"t1","tier":1},{"model":"lo-tier","source":"t2","tier":2}]}`); rr.Code != 200 {
|
||||
t.Fatalf("save auto: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("chat = %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
raw := c.Plugins().State("walker")
|
||||
b, _ := json.Marshal(raw)
|
||||
var st struct {
|
||||
Steps []struct {
|
||||
Kind string `json:"kind"`
|
||||
Tier int `json:"tier"`
|
||||
Source string `json:"source"`
|
||||
Model string `json:"model"`
|
||||
} `json:"steps"`
|
||||
Ends []struct {
|
||||
TierServed int `json:"tier_served"`
|
||||
Degraded bool `json:"degraded"`
|
||||
Source string `json:"source"`
|
||||
Model string `json:"model"`
|
||||
} `json:"ends"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &st); err != nil {
|
||||
t.Fatalf("decode: %v (%s)", err, string(b))
|
||||
}
|
||||
if len(st.Steps) < 2 {
|
||||
t.Fatalf("chain_step events = %+v, want at least a slot_fail and a selected", st.Steps)
|
||||
}
|
||||
if st.Steps[0].Kind != "slot_fail" || st.Steps[0].Tier != 1 {
|
||||
t.Errorf("first step = %+v, want slot_fail on tier 1", st.Steps[0])
|
||||
}
|
||||
last := st.Steps[len(st.Steps)-1]
|
||||
if last.Kind != "selected" || last.Tier != 2 {
|
||||
t.Errorf("last step = %+v, want selected on tier 2", last)
|
||||
}
|
||||
if len(st.Ends) != 1 {
|
||||
t.Fatalf("request_end count = %d, want 1", len(st.Ends))
|
||||
}
|
||||
if st.Ends[0].TierServed != 2 {
|
||||
t.Errorf("tier_served = %d, want 2", st.Ends[0].TierServed)
|
||||
}
|
||||
if !st.Ends[0].Degraded {
|
||||
t.Error("degraded = false, but the request dropped a tier")
|
||||
}
|
||||
if st.Ends[0].Model != "lo-tier" {
|
||||
t.Errorf("served model = %q, want lo-tier", st.Ends[0].Model)
|
||||
}
|
||||
}
|
||||
|
||||
@ -52,6 +52,23 @@ type Req struct {
|
||||
// Kept separate from Compl/Prompt: image generation has no token concept,
|
||||
// so counting images as "completion tokens" would corrupt the token totals.
|
||||
ImageCount int `json:"image_count,omitempty"`
|
||||
|
||||
// Walk is the AUTO chain's step-by-step trace for this request: which tiers
|
||||
// were skipped and why, which slots hard-failed, which one served it. It is
|
||||
// the only way a consumer can tell "tier 1 served this" from "tier 1 was
|
||||
// cooling so we dropped to tier 3" — a distinction that is the entire point
|
||||
// of a priority chain.
|
||||
//
|
||||
// json:"-" — deliberately NOT persisted. The audit file is a hot append and
|
||||
// this is observational detail: on a degraded gateway every request would
|
||||
// carry a multi-element array, and the audit trail's own retention (16 files
|
||||
// x 16 MB) is already the largest thing on the box. A plugin that wants the
|
||||
// walk sees it live at request_end; an operator post-mortem reads it from the
|
||||
// plugin's own accumulated state or from /api/auto slot health.
|
||||
//
|
||||
// Only populated when a plugin is loaded (chainTraceSink returns nil
|
||||
// otherwise), so a gateway with no plugins allocates nothing for it.
|
||||
Walk []map[string]interface{} `json:"-"`
|
||||
}
|
||||
|
||||
// Stat aggregates counters for one dimension row.
|
||||
|
||||
@ -294,3 +294,85 @@ func TestBillingPluginLoadedByDefault(t *testing.T) {
|
||||
t.Errorf("billing.lua was not written to the plugin dir: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingCountsDegradations: the plugin must distinguish a request that had
|
||||
// to drop below the top tier from one the top tier served. Without the chain
|
||||
// trace these were identical in the accounts, so a quietly degraded gateway
|
||||
// looked healthy while spending more per request.
|
||||
func TestBillingCountsDegradations(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
_ = ps.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"models": map[string]interface{}{
|
||||
"hi-tier": map[string]interface{}{"prompt": 1e-5, "completion": 1e-5},
|
||||
"lo-tier": map[string]interface{}{"prompt": 1e-6, "completion": 1e-6},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
// Request 1: degraded. tier 1 hard-failed, tier 2 served it.
|
||||
ps.Fire(StageChainStep, map[string]interface{}{
|
||||
"kind": "slot_fail", "tier": 1, "source": "t1", "model": "hi-tier",
|
||||
})
|
||||
ps.Fire(StageChainStep, map[string]interface{}{
|
||||
"kind": "selected", "tier": 2, "source": "t2", "model": "lo-tier",
|
||||
})
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "lo-tier", "source": "t2", "key": "***d1", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 1000,
|
||||
"degraded": true, "tier_served": 2, "time": 1750000000000,
|
||||
})
|
||||
|
||||
// Request 2: clean, served by the top tier.
|
||||
ps.Fire(StageChainStep, map[string]interface{}{
|
||||
"kind": "selected", "tier": 1, "source": "t1", "model": "hi-tier",
|
||||
})
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "hi-tier", "source": "t1", "key": "***d1", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 1000,
|
||||
"degraded": false, "tier_served": 1, "time": 1750000000000,
|
||||
})
|
||||
|
||||
st := stateOf(t, ps)
|
||||
if got := st["degraded_reqs"].(float64); got != 1 {
|
||||
t.Errorf("degraded_reqs = %v, want 1 (one of the two requests dropped a tier)", got)
|
||||
}
|
||||
tiers := st["by_tier_served"].(map[string]interface{})
|
||||
if tiers["2"].(float64) != 1 {
|
||||
t.Errorf("by_tier_served[2] = %v, want 1", tiers["2"])
|
||||
}
|
||||
if tiers["1"].(float64) != 1 {
|
||||
t.Errorf("by_tier_served[1] = %v, want 1", tiers["1"])
|
||||
}
|
||||
// Cost reflects the model actually served, not the one that should have been.
|
||||
// 1000*1e-6*2 = 0.002 for the degraded one, 1000*1e-5*2 = 0.02 for the clean one.
|
||||
approx(t, "total", st["total"].(map[string]interface{})["cost"].(float64), 0.022)
|
||||
}
|
||||
|
||||
// TestBillingAggregatesSkipReasons: skip reasons are the actionable diagnostic
|
||||
// ("no schedulable slot (cooling or quota exhausted)"), so they must be
|
||||
// counted. The wait time is normalised, otherwise a fresh row per request would
|
||||
// appear whenever the busy-wait text varies.
|
||||
func TestBillingAggregatesSkipReasons(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
ps.Fire(StageChainStep, map[string]interface{}{
|
||||
"kind": "tier_skip", "tier": 1, "reason": "no schedulable slot (cooling or quota exhausted)",
|
||||
})
|
||||
ps.Fire(StageChainStep, map[string]interface{}{
|
||||
"kind": "tier_busy", "tier": 2, "reason": "no free slot within 2s",
|
||||
})
|
||||
ps.Fire(StageChainStep, map[string]interface{}{
|
||||
"kind": "tier_busy", "tier": 3, "reason": "no free slot within 2.0001s",
|
||||
})
|
||||
st := stateOf(t, ps)
|
||||
reasons := st["skip_reasons"].(map[string]interface{})
|
||||
if len(reasons) != 2 {
|
||||
t.Errorf("skip_reasons = %v, want 2 (the two variable waits must collapse to one)", reasons)
|
||||
}
|
||||
busy, ok := reasons["no free slot within <wait>"]
|
||||
if !ok {
|
||||
t.Errorf("busy reason missing; got %v", reasons)
|
||||
} else if busy.(float64) != 2 {
|
||||
t.Errorf("busy count = %v, want 2 (two different wait texts, one cause)", busy)
|
||||
}
|
||||
}
|
||||
|
||||
@ -64,6 +64,31 @@ const (
|
||||
// tier (AUTO tier, -1 for the direct path), stream.
|
||||
StageRouted Stage = "routed"
|
||||
|
||||
// StageChainStep fires ONCE PER STEP of an AUTO chain walk, and only on the
|
||||
// AUTO path (a direct request has no chain and therefore emits nothing).
|
||||
//
|
||||
// This exists because StageRouted cannot express degradation: it fires once,
|
||||
// after the walk, with the slot that finally won. "tier 1 was cooling so we
|
||||
// dropped to tier 3" and "tier 1 served it" were indistinguishable. That
|
||||
// distinction is the whole point of a priority chain, and it is what an
|
||||
// operator debugging "why did my expensive model not get used" needs.
|
||||
//
|
||||
// payload:
|
||||
// kind "tier_skip" | "slot_fail" | "tier_busy" | "selected"
|
||||
// tier the AUTO tier this step belongs to (1 = highest priority)
|
||||
// source / model set for slot_fail and selected
|
||||
// reason human-readable cause, for tier_skip and tier_busy
|
||||
// error the underlying error text, for slot_fail
|
||||
// attempt 1-based slot attempt within this walk
|
||||
//
|
||||
// Ordering: every step precedes StageRouted, and the "selected" step is the
|
||||
// last one. A plugin accumulating the walk therefore has the full picture
|
||||
// by the time request_end arrives.
|
||||
//
|
||||
// These events are OBSERVATION ONLY — see the accounting note in
|
||||
// docs/plugins.md: nothing here feeds back into routing, cooldown or quota.
|
||||
StageChainStep Stage = "chain_step"
|
||||
|
||||
// StageRequestEnd fires exactly once per request, after the client response
|
||||
// has been produced (or after a failure was recorded). payload adds:
|
||||
// source, model, ok, status, latency_ms, first_byte_ms, prompt_tokens,
|
||||
@ -78,6 +103,7 @@ const (
|
||||
// AllStages is the firing order, used by the docs and by the hook listing.
|
||||
var AllStages = []Stage{
|
||||
StageRequestStart,
|
||||
StageChainStep,
|
||||
StageRouted,
|
||||
StageRequestEnd,
|
||||
}
|
||||
|
||||
@ -194,9 +194,43 @@ end
|
||||
-- ---------- hooks ----------
|
||||
|
||||
plugin.hooks = {
|
||||
-- chain_step gives the per-tier walk; request_end gives the final accounting.
|
||||
-- Subscribing to chain_step is OPTIONAL here: the totals are driven by
|
||||
-- request_end alone, and the degradation counters below are pure observation.
|
||||
-- A gateway with thousands of requests can drop this hook to save the
|
||||
-- per-step Lua call without losing a single billed request.
|
||||
chain_step = "on_chain_step",
|
||||
request_end = "on_request_end",
|
||||
}
|
||||
|
||||
-- Tracks how often a request had to drop below the top tier, and which tier
|
||||
-- actually served it. Without this, "tier 1 was cooling" and "tier 1 served it"
|
||||
-- are indistinguishable in the accounts, and a quietly degraded gateway looks
|
||||
-- exactly like a healthy one.
|
||||
plugin.state.degraded_reqs = 0
|
||||
plugin.state.by_tier_served = {}
|
||||
plugin.state.skip_reasons = {}
|
||||
|
||||
function plugin.on_chain_step(payload)
|
||||
if payload == nil then return nil end
|
||||
local s = plugin.state
|
||||
if s == nil then return nil end
|
||||
if s.by_tier_served == nil then s.by_tier_served = {} end
|
||||
if s.skip_reasons == nil then s.skip_reasons = {} end
|
||||
|
||||
if payload.kind == "selected" then
|
||||
local t = tostring(payload.tier or "?")
|
||||
s.by_tier_served[t] = (s.by_tier_served[t] or 0) + 1
|
||||
elseif payload.kind == "tier_skip" or payload.kind == "tier_busy" then
|
||||
-- reason text is the ACTIONABLE part; normalise the volatile bits so the
|
||||
-- same cause aggregates instead of creating a new row per request.
|
||||
local r = tostring(payload.reason or payload.kind or "unknown")
|
||||
r = string.gsub(r, "within [%d%.%a]+", "within <wait>")
|
||||
s.skip_reasons[r] = (s.skip_reasons[r] or 0) + 1
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
function plugin.on_request_end(payload)
|
||||
if payload == nil then return nil end
|
||||
local prompt = tonumber(payload.prompt_tokens) or 0
|
||||
@ -216,6 +250,11 @@ function plugin.on_request_end(payload)
|
||||
if s.by_key == nil then s.by_key = {} end
|
||||
if s.by_day == nil then s.by_day = {} end
|
||||
if s.started == nil then s.started = payload.time or 0 end
|
||||
if s.degraded_reqs == nil then s.degraded_reqs = 0 end
|
||||
-- Degradation is counted here rather than in the chain_step hook because
|
||||
-- request_end sees the whole walk at once: one degraded request must count
|
||||
-- once, whereas the walk may contain several skipped tiers.
|
||||
if payload.degraded then s.degraded_reqs = s.degraded_reqs + 1 end
|
||||
|
||||
add(s.total, cost, prompt, completion, ok)
|
||||
if payload.source ~= nil and payload.source ~= "" then
|
||||
@ -313,6 +352,7 @@ plugin.ui = {
|
||||
document.getElementById("billing-kpis").innerHTML = [
|
||||
["Total", money(t.cost, cur)],
|
||||
["Requests", t.requests || 0],
|
||||
["Degraded", s.degraded_reqs || 0],
|
||||
["Prompt tokens", t.prompt_tokens || 0],
|
||||
["Completion tokens", t.completion_tokens || 0],
|
||||
["Failures", t.failures || 0]
|
||||
|
||||
@ -311,3 +311,63 @@ func TestPluginFireWithNoPluginsIsNoop(t *testing.T) {
|
||||
t.Errorf("empty registry misbehaved: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChainStepIsARealStage: chain_step is documented as a distinct stage that
|
||||
// fires once per step of an AUTO walk. If it were only a field on `routed`, a
|
||||
// plugin author following the docs would silently get one event instead of the
|
||||
// whole walk.
|
||||
func TestChainStepIsARealStage(t *testing.T) {
|
||||
found := false
|
||||
for _, s := range AllStages {
|
||||
if s == StageChainStep {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("StageChainStep is not in AllStages, so the dispatcher never registers it")
|
||||
}
|
||||
// Ordering: it must sit between request_start and routed, which is what
|
||||
// docs/plugins.md promises.
|
||||
var iStart, iStep, iRouted = -1, -1, -1
|
||||
for i, s := range AllStages {
|
||||
switch s {
|
||||
case StageRequestStart:
|
||||
iStart = i
|
||||
case StageChainStep:
|
||||
iStep = i
|
||||
case StageRouted:
|
||||
iRouted = i
|
||||
}
|
||||
}
|
||||
if !(iStart < iStep && iStep < iRouted) {
|
||||
t.Errorf("stage order = %v, want request_start < chain_step < routed", AllStages)
|
||||
}
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = { name = "stepper" }
|
||||
p.hooks = { chain_step = "s" }
|
||||
function p.s(payload)
|
||||
payload.kinds = (payload.kinds or "")
|
||||
return nil
|
||||
end
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "stepper", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
// It must actually dispatch.
|
||||
seen := false
|
||||
for _, row := range ps.List() {
|
||||
if row["name"] == "stepper" {
|
||||
hooks := row["hooks"].([]string)
|
||||
for _, h := range hooks {
|
||||
if h == string(StageChainStep) {
|
||||
seen = true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if !seen {
|
||||
t.Error("a plugin registered for chain_step is not reported as such")
|
||||
}
|
||||
}
|
||||
|
||||
@ -293,13 +293,75 @@ func runTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, r
|
||||
return tierResult{hard: hard}
|
||||
}
|
||||
|
||||
// TraceKind classifies one step of an AUTO chain walk.
|
||||
type TraceKind string
|
||||
|
||||
const (
|
||||
// TraceTierSkip: the whole tier was skipped — every slot was cooling,
|
||||
// quota-exhausted, or none was schedulable. Reason says which.
|
||||
TraceTierSkip TraceKind = "tier_skip"
|
||||
// TraceSlotFail: one slot failed hard (upstream error / bad adapter). The
|
||||
// walk continues to the next slot or tier.
|
||||
TraceSlotFail TraceKind = "slot_fail"
|
||||
// TraceTierBusy: the tier was fully busy and the bounded wait expired.
|
||||
TraceTierBusy TraceKind = "tier_busy"
|
||||
// TraceSelected: this slot served the request. Exactly one per successful
|
||||
// chain walk, and the last event emitted.
|
||||
TraceSelected TraceKind = "selected"
|
||||
)
|
||||
|
||||
// TraceEvent is one observable step of an AUTO chain walk.
|
||||
//
|
||||
// WHY THIS EXISTS: chainDrive's return value is (resp, src, model, err), so a
|
||||
// caller learns only which slot finally served the request. Everything the
|
||||
// scheduler decided on the way there — which tiers it skipped and WHY, which
|
||||
// slots hard-failed, whether a tier was merely busy — was computed and then
|
||||
// discarded. That is invisible to operators and to plugins: "tier 1 was cooling
|
||||
// so we degraded to tier 3" looked exactly like "tier 1 served it".
|
||||
//
|
||||
// The walk already accumulates this in ChainErr, but ONLY on total failure, and
|
||||
// ChainErr is an error return, not a record. Emitting a trace as it happens
|
||||
// covers the far more common case: a request that SUCCEEDED after degrading.
|
||||
//
|
||||
// Design constraints:
|
||||
// - scheduler stays dependency-free and independently testable. A TraceEvent
|
||||
// is a plain struct in this package and the sink is a func parameter, so no
|
||||
// import is added and no test has to change to observe a walk.
|
||||
// - The sink is optional (nil = emit nothing). The overhead on the hot path
|
||||
// is one nil check per event.
|
||||
// - Events are OBSERVATION ONLY. Nothing in the scheduler branches on them,
|
||||
// and the gateway does not feed them back into routing, cooldown or quota —
|
||||
// see docs/plugins.md for why accounting and enforcement are kept apart.
|
||||
type TraceEvent struct {
|
||||
Kind TraceKind
|
||||
Tier int
|
||||
Source string
|
||||
Model string
|
||||
Reason string // human-readable, for TraceTierSkip / TraceSlotFail
|
||||
Err string // the underlying error text, for TraceSlotFail
|
||||
// Attempt counts the 1-based slot attempt within the whole walk.
|
||||
Attempt int
|
||||
}
|
||||
|
||||
// TraceSink receives chain-walk events. It must not block: it is called from the
|
||||
// request path, and a slow sink slows the request.
|
||||
type TraceSink func(TraceEvent)
|
||||
|
||||
// chainDrive runs a request down the chain (plan 2.3): tiers ascending (tier
|
||||
// 1, the highest priority, first), per-tier round-robin starting at the tier
|
||||
// cursor, same-tier runs ordered by preference (negative prefs sink but stay
|
||||
// reachable). Quota-exhausted and cooling slots are filtered up front; a
|
||||
// fully busy tier is polled for a bounded time before falling through.
|
||||
// Failures are summarized in *ChainErr for the caller to map to HTTP 503.
|
||||
func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool, stream bool) (*types.UnifiedResponse, <-chan types.UnifiedChunk, string, string, error) {
|
||||
//
|
||||
// trace may be nil; when set it receives one event per observable step.
|
||||
func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool, stream bool, trace TraceSink) (*types.UnifiedResponse, <-chan types.UnifiedChunk, string, string, error) {
|
||||
emit := func(ev TraceEvent) {
|
||||
if trace != nil {
|
||||
trace(ev)
|
||||
}
|
||||
}
|
||||
attempt := 0
|
||||
if chain == nil || len(chain.Tiers) == 0 {
|
||||
return nil, nil, "", "", fmt.Errorf("no auto slot configured")
|
||||
}
|
||||
@ -309,7 +371,9 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
|
||||
// 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))
|
||||
reason := "no schedulable slot (cooling or quota exhausted)"
|
||||
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: %s", tn.Tier, reason))
|
||||
emit(TraceEvent{Kind: TraceTierSkip, Tier: tn.Tier, Reason: reason})
|
||||
continue
|
||||
}
|
||||
// No Pref sort: load balancing is done by round-robin cursor.
|
||||
@ -318,6 +382,8 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
|
||||
base := tn.NextStart()
|
||||
res := runTier(ctx, tn, cands, base, req, stream)
|
||||
if res.resp != nil || res.chunks != nil {
|
||||
attempt++
|
||||
emit(TraceEvent{Kind: TraceSelected, Tier: tn.Tier, Source: res.src, Model: res.model, Attempt: attempt})
|
||||
releaseProbes(cands)
|
||||
return res.resp, res.chunks, res.src, res.model, nil
|
||||
}
|
||||
@ -327,6 +393,13 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
|
||||
}
|
||||
if len(res.hard) > 0 {
|
||||
ce.Tiers = append(ce.Tiers, res.hard...)
|
||||
for _, h := range res.hard {
|
||||
attempt++
|
||||
emit(TraceEvent{
|
||||
Kind: TraceSlotFail, Tier: tn.Tier, Source: h.Source, Model: h.Model,
|
||||
Err: types.OneLine(h.Err.Error(), 200), Attempt: attempt,
|
||||
})
|
||||
}
|
||||
releaseProbes(cands)
|
||||
continue // hard failures: fall through to the next tier, no waiting
|
||||
}
|
||||
@ -334,10 +407,13 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
|
||||
if err := s.pollBusyTier(ctx, tn, cands, base, req, stream, &ce); err != nil {
|
||||
releaseProbes(cands)
|
||||
if r, ok := err.(*tierSuccess); ok {
|
||||
attempt++
|
||||
emit(TraceEvent{Kind: TraceSelected, Tier: tn.Tier, Source: r.res.src, Model: r.res.model, Attempt: attempt})
|
||||
return r.res.resp, r.res.chunks, r.res.src, r.res.model, nil
|
||||
}
|
||||
return nil, nil, "", "", err
|
||||
}
|
||||
emit(TraceEvent{Kind: TraceTierBusy, Tier: tn.Tier, Reason: fmt.Sprintf("no free slot within %v", busyWait)})
|
||||
releaseProbes(cands)
|
||||
}
|
||||
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
|
||||
@ -401,16 +477,16 @@ func (s *Scheduler) pollBusyTier(ctx context.Context, tn *TierNode, cands []cand
|
||||
// non-nil, decides slot token-quota exhaustion. Returns the response, the
|
||||
// serving source and the exact model id used; on total failure a *ChainErr
|
||||
// summarizing every tier.
|
||||
func (s *Scheduler) ChainChat(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool) (*types.UnifiedResponse, string, string, error) {
|
||||
resp, _, src, model, err := s.chainDrive(ctx, chain, req, exhausted, false)
|
||||
func (s *Scheduler) ChainChat(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool, trace TraceSink) (*types.UnifiedResponse, string, string, error) {
|
||||
resp, _, src, model, err := s.chainDrive(ctx, chain, req, exhausted, false, trace)
|
||||
return resp, src, model, err
|
||||
}
|
||||
|
||||
// ChainChatStream runs a streaming AUTO request down the chain. A slot is
|
||||
// abandoned only on connect failures / busy (before its first chunk); after a
|
||||
// stream starts it is pinned. Same return contract as ChainChat.
|
||||
func (s *Scheduler) ChainChatStream(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool) (<-chan types.UnifiedChunk, string, string, error) {
|
||||
_, chunks, src, model, err := s.chainDrive(ctx, chain, req, exhausted, true)
|
||||
func (s *Scheduler) ChainChatStream(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool, trace TraceSink) (<-chan types.UnifiedChunk, string, string, error) {
|
||||
_, chunks, src, model, err := s.chainDrive(ctx, chain, req, exhausted, true, trace)
|
||||
return chunks, src, model, err
|
||||
}
|
||||
|
||||
|
||||
@ -176,7 +176,7 @@ func TestChainRoundRobin(t *testing.T) {
|
||||
s := New(0)
|
||||
var got []string
|
||||
for i := 0; i < 4; i++ {
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("iter %d: %v", i, err)
|
||||
}
|
||||
@ -200,7 +200,7 @@ func TestChainPreferenceSinksButStaysReachable(t *testing.T) {
|
||||
{Tier: 0, Model: "g", Source: "good"},
|
||||
}, bySource(neg, good))
|
||||
s := New(0)
|
||||
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
}
|
||||
@ -213,7 +213,7 @@ func TestChainPreferenceSinksButStaysReachable(t *testing.T) {
|
||||
t.Fatal("neg was tried first and succeeded; good must not be attempted")
|
||||
}
|
||||
// Second request: cursor advances. neg wins again (good hard-fails).
|
||||
resp2, src2, _, err2 := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
resp2, src2, _, err2 := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err2 != nil {
|
||||
t.Fatalf("second chain: %v", err2)
|
||||
}
|
||||
@ -230,7 +230,7 @@ func TestChainBusySkipsWithoutPenalty(t *testing.T) {
|
||||
{Tier: 0, Model: "b", Source: "s2"},
|
||||
}, bySource(a, b))
|
||||
s := New(0)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
}
|
||||
@ -256,7 +256,7 @@ func TestChainAllBusyBoundedWaitThenNextTier(t *testing.T) {
|
||||
}, bySource(a, b, c))
|
||||
s := New(0)
|
||||
t0 := time.Now()
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
el := time.Since(t0)
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
@ -277,7 +277,7 @@ func TestChainQuotaExhausted(t *testing.T) {
|
||||
}, bySource(a, b))
|
||||
s := New(0)
|
||||
exhausted := func(sl *Slot) bool { return sl.Source == "s1" && sl.Quota > 0 }
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), exhausted)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), exhausted, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
}
|
||||
@ -300,7 +300,7 @@ func TestChainErrSummary(t *testing.T) {
|
||||
{Tier: 1, Model: "c", Source: "s3"},
|
||||
}, bySource(a, b, c))
|
||||
s := New(0)
|
||||
_, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
_, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
var ce *ChainErr
|
||||
if !errors.As(err, &ce) {
|
||||
t.Fatalf("err = %v, want *ChainErr", err)
|
||||
@ -325,7 +325,7 @@ func TestChainStreamFallsBackBeforeFirstChunk(t *testing.T) {
|
||||
{Tier: 0, Model: "b", Source: "s2"},
|
||||
}, bySource(a, b))
|
||||
s := New(0)
|
||||
chunks, src, model, err := s.ChainChatStream(context.Background(), ch, chatReq(), nil)
|
||||
chunks, src, model, err := s.ChainChatStream(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("chain stream: %v", err)
|
||||
}
|
||||
@ -365,7 +365,7 @@ func TestChainProbeIsLastResort(t *testing.T) {
|
||||
}, bySource(healthy, cooling))
|
||||
s := New(0)
|
||||
for i := 0; i < 3; i++ {
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("iter %d: %v", i, err)
|
||||
}
|
||||
@ -390,7 +390,7 @@ func TestChainProbeServesWhenNothingElseCan(t *testing.T) {
|
||||
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)
|
||||
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("probe must serve the request: %v", err)
|
||||
}
|
||||
@ -417,14 +417,14 @@ func TestChainProbePermitReleasedOnFailure(t *testing.T) {
|
||||
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 {
|
||||
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, 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 {
|
||||
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil); err == nil {
|
||||
t.Fatal("expected the second probe to fail too")
|
||||
}
|
||||
if cooling.probeClaims.Load() != 2 {
|
||||
@ -443,7 +443,7 @@ func TestChainProbeDoesNotBlockTierFallthrough(t *testing.T) {
|
||||
{Tier: 2, Model: "b", Source: "backup"},
|
||||
}, bySource(cold, backup))
|
||||
s := New(0)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil || src != "backup" {
|
||||
t.Fatalf("want fallthrough to backup, got src=%q err=%v", src, err)
|
||||
}
|
||||
@ -482,7 +482,7 @@ func TestChainProbeRoundRobinUnaffected(t *testing.T) {
|
||||
s := New(0)
|
||||
var got []string
|
||||
for i := 0; i < 4; i++ {
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("iter %d: %v", i, err)
|
||||
}
|
||||
|
||||
191
internal/scheduler/trace_test.go
Normal file
191
internal/scheduler/trace_test.go
Normal file
@ -0,0 +1,191 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The chain trace is the only way a caller learns that a request was DEGRADED
|
||||
// — served by a lower tier than the one that should have taken it. chainDrive's
|
||||
// return value carries only the winner, so without these events "tier 1 was
|
||||
// cooling and we dropped to tier 2" is indistinguishable from "tier 1 served
|
||||
// it", which is the exact question a priority chain exists to answer.
|
||||
|
||||
// recorder collects trace events for assertions.
|
||||
type recorder struct{ events []TraceEvent }
|
||||
|
||||
func (r *recorder) sink(ev TraceEvent) { r.events = append(r.events, ev) }
|
||||
|
||||
func (r *recorder) kinds() []TraceKind {
|
||||
out := make([]TraceKind, 0, len(r.events))
|
||||
for _, e := range r.events {
|
||||
out = append(out, e.Kind)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (r *recorder) find(k TraceKind) *TraceEvent {
|
||||
for i := range r.events {
|
||||
if r.events[i].Kind == k {
|
||||
return &r.events[i]
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestTraceSelectedOnlyOnHappyPath: a clean walk emits exactly one event.
|
||||
func TestTraceSelectedOnlyOnHappyPath(t *testing.T) {
|
||||
p1 := fakeProv("p1", "m1")
|
||||
ch := BuildChain([]Rule{{Model: "m1", Source: "p1", Tier: 1}},
|
||||
func(model, source string) Provider { return p1 })
|
||||
var rec recorder
|
||||
_, _, _, err := New(3).ChainChat(context.Background(), ch, chatReq(), nil, rec.sink)
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
}
|
||||
if got := rec.kinds(); len(got) != 1 || got[0] != TraceSelected {
|
||||
t.Errorf("events = %v, want a single selected", got)
|
||||
}
|
||||
e := rec.find(TraceSelected)
|
||||
if e.Tier != 1 || e.Source != "p1" || e.Model != "m1" {
|
||||
t.Errorf("selected event = %+v, want tier 1 / p1 / m1", e)
|
||||
}
|
||||
if e.Attempt != 1 {
|
||||
t.Errorf("Attempt = %d, want 1", e.Attempt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTraceRecordsTierSkipAndDegradation is the core case: tier 1 is
|
||||
// unschedulable, tier 2 answers. The trace must show the skip AND the eventual
|
||||
// selection, so a consumer can see the request was served one tier down.
|
||||
func TestTraceRecordsTierSkipAndDegradation(t *testing.T) {
|
||||
// p1 is unavailable (not probeable), so tier 1 yields no candidates.
|
||||
p1 := fakeProv("p1", "m1")
|
||||
p1.available.Store(false)
|
||||
p2 := fakeProv("p2", "m2")
|
||||
ch := BuildChain([]Rule{
|
||||
{Model: "m1", Source: "p1", Tier: 1},
|
||||
{Model: "m2", Source: "p2", Tier: 2},
|
||||
}, bySource(p1, p2))
|
||||
var rec recorder
|
||||
_, src, model, err := New(3).ChainChat(context.Background(), ch, chatReq(), nil, rec.sink)
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
}
|
||||
if src != "p2" || model != "m2" {
|
||||
t.Fatalf("served by %s/%s, want p2/m2", src, model)
|
||||
}
|
||||
skip := rec.find(TraceTierSkip)
|
||||
if skip == nil {
|
||||
t.Fatalf("no tier_skip event; events = %v", rec.kinds())
|
||||
}
|
||||
if skip.Tier != 1 {
|
||||
t.Errorf("skip tier = %d, want 1", skip.Tier)
|
||||
}
|
||||
if !strings.Contains(skip.Reason, "cooling") {
|
||||
t.Errorf("skip reason = %q, want it to mention cooling", skip.Reason)
|
||||
}
|
||||
sel := rec.find(TraceSelected)
|
||||
if sel == nil || sel.Tier != 2 {
|
||||
t.Errorf("selected = %+v, want tier 2", sel)
|
||||
}
|
||||
// The order matters: the skip must be observable BEFORE the selection.
|
||||
if rec.events[0].Kind != TraceTierSkip || rec.events[len(rec.events)-1].Kind != TraceSelected {
|
||||
t.Errorf("event order = %v, want skip first and selected last", rec.kinds())
|
||||
}
|
||||
}
|
||||
|
||||
// TestTraceRecordsHardSlotFailures: a slot that returns an upstream error is a
|
||||
// different event from a skip — the request tried it and it failed. Losing that
|
||||
// distinction makes a flaky upstream look like an idle one.
|
||||
func TestTraceRecordsHardSlotFailures(t *testing.T) {
|
||||
p1 := fakeProv("p1", "m1")
|
||||
p1.fail.Store(true) // Chat returns "upstream error"
|
||||
p2 := fakeProv("p2", "m2")
|
||||
ch := BuildChain([]Rule{
|
||||
{Model: "m1", Source: "p1", Tier: 1},
|
||||
{Model: "m2", Source: "p2", Tier: 2},
|
||||
}, bySource(p1, p2))
|
||||
var rec recorder
|
||||
_, src, _, err := New(3).ChainChat(context.Background(), ch, chatReq(), nil, rec.sink)
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
}
|
||||
if src != "p2" {
|
||||
t.Fatalf("served by %s, want p2", src)
|
||||
}
|
||||
fail := rec.find(TraceSlotFail)
|
||||
if fail == nil {
|
||||
t.Fatalf("no slot_fail event; events = %v", rec.kinds())
|
||||
}
|
||||
if fail.Tier != 1 || fail.Source != "p1" || fail.Model != "m1" {
|
||||
t.Errorf("slot_fail = %+v, want tier 1 / p1 / m1", fail)
|
||||
}
|
||||
if !strings.Contains(fail.Err, "upstream error") {
|
||||
t.Errorf("slot_fail error = %q, want the upstream text", fail.Err)
|
||||
}
|
||||
// A hard failure must NOT be reported as a skip.
|
||||
if rec.find(TraceTierSkip) != nil {
|
||||
t.Error("a hard failure was also reported as a tier_skip")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTraceNilSinkIsSafe: the gateway passes nil when no plugin is loaded, so
|
||||
// every emit path must tolerate it. This is the "plugins are optional" property
|
||||
// on the scheduler side.
|
||||
func TestTraceNilSinkIsSafe(t *testing.T) {
|
||||
p1 := fakeProv("p1", "m1")
|
||||
p1.fail.Store(true)
|
||||
p2 := fakeProv("p2", "m2")
|
||||
ch := BuildChain([]Rule{
|
||||
{Model: "m1", Source: "p1", Tier: 1},
|
||||
{Model: "m2", Source: "p2", Tier: 2},
|
||||
}, bySource(p1, p2))
|
||||
if _, _, _, err := New(3).ChainChat(context.Background(), ch, chatReq(), nil, nil); err != nil {
|
||||
t.Fatalf("a nil trace sink broke the walk: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTraceOnTotalFailure: when every tier fails, the walk still emits its
|
||||
// per-step events AND returns the ChainErr. The trace is additive — it must not
|
||||
// replace or disturb the error contract callers depend on for the 503.
|
||||
func TestTraceOnTotalFailure(t *testing.T) {
|
||||
p1 := fakeProv("p1", "m1")
|
||||
p1.fail.Store(true)
|
||||
p2 := fakeProv("p2", "m2")
|
||||
p2.fail.Store(true)
|
||||
ch := BuildChain([]Rule{
|
||||
{Model: "m1", Source: "p1", Tier: 1},
|
||||
{Model: "m2", Source: "p2", Tier: 2},
|
||||
}, bySource(p1, p2))
|
||||
var rec recorder
|
||||
_, _, _, err := New(3).ChainChat(context.Background(), ch, chatReq(), nil, rec.sink)
|
||||
var ce *ChainErr
|
||||
if !errors.As(err, &ce) {
|
||||
t.Fatalf("err = %v, want a *ChainErr so the gateway can answer 503", err)
|
||||
}
|
||||
if len(ce.Tiers) != 2 {
|
||||
t.Errorf("ChainErr.Tiers = %d, want 2 (the error contract must be unchanged)", len(ce.Tiers))
|
||||
}
|
||||
if n := len(rec.kinds()); n != 2 {
|
||||
t.Errorf("events = %v, want two slot_fail and no selection", rec.kinds())
|
||||
}
|
||||
if rec.find(TraceSelected) != nil {
|
||||
t.Error("a selected event was emitted for a walk that served nothing")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTraceSkipsEmptyChain: no chain configured must not emit anything; the
|
||||
// gateway answers 503 before scheduling in that case anyway.
|
||||
func TestTraceSkipsEmptyChain(t *testing.T) {
|
||||
var rec recorder
|
||||
_, _, _, err := New(3).ChainChat(context.Background(), &Chain{}, chatReq(), nil, rec.sink)
|
||||
if err == nil {
|
||||
t.Fatal("expected an error for an empty chain")
|
||||
}
|
||||
if len(rec.events) != 0 {
|
||||
t.Errorf("events = %v, want none", rec.kinds())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user