diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e6b0430..0ed7bcb 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,4 +1,5 @@ -# Unit tests for the Go components. +# Unit tests: the Go components (bodylog-listener) and the Lua engine modules +# (test/utest_*.lua, run with the resty CLI by test/run_lua_unit_tests.sh). # # Separate from release.yml so a pull request is tested without going anywhere # near the publish path. release.yml runs the same tests as a gate before it @@ -7,13 +8,17 @@ name: test on: pull_request: - paths: &go_paths + paths: &test_paths - 'bodylog-listener-go/**' - 'Dockerfile.bodylog' + - 'lua/**' + - 'test/utest_*.lua' + - 'test/run_lua_unit_tests.sh' + - 'Dockerfile' - '.github/workflows/test.yml' push: branches: [main] - paths: *go_paths + paths: *test_paths workflow_dispatch: permissions: @@ -37,3 +42,23 @@ jobs: env: CGO_ENABLED: 1 run: go test ./... + + lua: + # jammy, so the OpenResty pinned in the Dockerfile (OPENRESTY_VER) installs as is + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + + - name: install openresty + resty CLI (the Dockerfile's version) + run: | + set -eux + ver=$(sed -n 's/^ARG OPENRESTY_VER=//p' Dockerfile) + test -n "$ver" + wget -qO - https://openresty.org/package/pubkey.gpg | sudo gpg --dearmor -o /usr/share/keyrings/openresty.gpg + echo "deb [arch=amd64 signed-by=/usr/share/keyrings/openresty.gpg] http://openresty.org/package/ubuntu jammy main" \ + | sudo tee /etc/apt/sources.list.d/openresty.list + sudo apt-get update + sudo apt-get install -y --no-install-recommends "openresty=$ver" "openresty-resty=$ver" + + - name: lua unit tests + run: bash test/run_lua_unit_tests.sh /usr/local/openresty/bin/resty diff --git a/README.md b/README.md index 6ac8648..da0d6f5 100644 --- a/README.md +++ b/README.md @@ -238,6 +238,42 @@ pairs. The fix for a small pool is to *lower* `session_prio_r_new`, or supply a also that under adaptive concurrency the limit itself is a float that moves over time, so the effective number of distinct ceilings moves with it. +## Response post-processing + +A route can rewrite upstream responses before they reach the client, to hide backend +quirks. It is off unless the route names a handler (`lua/postprocess.lua`): + +```yaml +# ModelRoute +spec: + nginx: + values: + postprocess: strip_actual_usage +``` + +(in a route conf: `postprocess = "strip_actual_usage"`; `_G.POSTPROCESS_DEFAULT` in +`session_base.conf` sets one for every route that names none.) + +| handler | what it does | +|---|---| +| `strip_actual_usage` | removes the top-level `actual_usage` some sglang builds return next to `usage` (a second token count, e.g. with different `cached_tokens`). Nothing else is touched. | +| `vllm_format` | normalizes vLLM's response shape: drops vLLM's placeholder fields when they are null (`prompt_token_ids`, `prompt_logprobs`, `logprobs`, `token_ids`, `refusal`, `annotations`, ...), always drops `stop_reason`, `service_tier`, `system_fingerprint`, `kv_transfer_params`, `ec_transfer_params` and `actual_usage`, and renames `reasoning` to `reasoning_content` and `created_cache_tokens` to `cache_write_tokens`. | + +How it behaves: + +- One handler per route. +- Only 2xx responses, and only `text/event-stream` (each `data:` frame) and + `application/json` bodies; anything else passes through. +- A frame or body with nothing to change is sent byte-identical. A changed one is + re-encoded, so its key order and `/` escaping can differ from the backend's. +- `Content-Length` is dropped on the route's 2xx responses (they are sent chunked), + and a non-stream JSON body is held until it is complete (up to 8 MiB; a larger + one passes through unchanged). Streams are not held: frames go out as they arrive. +- Fail-open: a frame that does not parse, or a handler error, leaves the original bytes. +- The bodylog sees the raw upstream response (it runs first). The listener keeps + `usage` and `actual_usage` in the record's `resp_meta`, so a stripped + `actual_usage` is still in the bodylog. + ## Authentication Keys are read from a file, default `/etc/openresty/api-keys/keys`, overridable with @@ -305,15 +341,17 @@ never added to the real 429 counters. Content rules (`reject_rules`) have their switch and are not covered. Three layers, highest first: - the runtime override `POST //_rate_limit_toggle?on=0|1` (127.0.0.1 only) takes effect without a reload and survives one; `?reset=1` removes it; -- the route option `rate_limit` (route conf, or `nginx.values` on a ModelRoute); -- the global `_G.RATE_LIMIT_ENABLED` in `session_base.conf`. +- the route option `load_rate_limit` (route conf, or `nginx.values` on a ModelRoute). Named so, + not `rate_limit`: on a ModelRoute that key is a video route's download speed limit, and the CRD + rejects it on llm routes; +- the global `_G.LOAD_RATE_LIMIT_ENABLED` in `session_base.conf`. Things to watch when switching it off at runtime. The override lives in this openresty instance's shared memory: with several replicas (e.g. two openresty pods) send it to every one of them, and a restarted instance comes back without it. It survives reloads and lasts until `?reset=1` or a full restart, so remove it once the incident is over (`/_rate_limit_status` shows `source: runtime` while it is in place). To keep limiting off for longer than an incident, set the -route option (`rate_limit: false` in the ModelRoute's `nginx.values`) instead: it is part of the +route option (`load_rate_limit: false` in the ModelRoute's `nginx.values`) instead: it is part of the deployed configuration and applies to every instance. And the adaptive limit keeps moving on traffic it no longer caps -- down under overload, up under pressure -- so switching back on can bring a burst of 429s until it settles. Likewise `session_prio`'s turn-1 hold flag is neither set nor released while @@ -387,12 +425,16 @@ cd test ./test_lazy_init_timing.sh # lazy route registration across workers ``` -Pure-Lua unit tests need only `resty`: +Pure-Lua unit tests (`test/utest_*.lua`) need only `resty`. Run them all from the +repository root, as CI does on every pull request that touches `lua/` or a unit test: ```bash -resty utest_surgical_normalize.lua ../lua/reqtransform.lua +bash test/run_lua_unit_tests.sh [path/to/resty] ``` +It exits non-zero if any test fails, or if it finds none. A test that needs a special +invocation (a shared dict, another working directory) is listed in the script. + A note on how these are written: several assertions exist specifically to fail when the thing under test is removed. A check that only confirms "no error appeared" also passes when the feature never ran, so the harnesses pair each such check with a diff --git a/bodylog-listener-go/main.go b/bodylog-listener-go/main.go index bc3928d..ec3c15a 100644 --- a/bodylog-listener-go/main.go +++ b/bodylog-listener-go/main.go @@ -272,6 +272,7 @@ type respFields struct { finishReason string model string usage any + actualUsage any // some sglang builds: a second token count next to usage (openresty may strip it from the client response) toolCalls []any stopReason string // Anthropic errorObj any @@ -365,6 +366,9 @@ func extractFromObj(obj map[string]any, rf *respFields) { if u, ok := obj["usage"].(map[string]any); ok && u != nil { rf.usage = u } + if u, ok := obj["actual_usage"].(map[string]any); ok && u != nil { + rf.actualUsage = u + } if m, ok := obj["model"].(string); ok && m != "" { rf.model = m } @@ -490,6 +494,9 @@ func buildRespMeta(rf respFields) map[string]any { if rf.usage != nil { m["usage"] = rf.usage } + if rf.actualUsage != nil { + m["actual_usage"] = rf.actualUsage + } if len(rf.toolCalls) > 0 { m["tool_calls"] = reassembleToolCalls(rf.toolCalls) } diff --git a/bodylog-listener-go/resp_meta_actual_usage_test.go b/bodylog-listener-go/resp_meta_actual_usage_test.go new file mode 100644 index 0000000..f5d52fa --- /dev/null +++ b/bodylog-listener-go/resp_meta_actual_usage_test.go @@ -0,0 +1,60 @@ +package main + +import ( + "encoding/json" + "testing" +) + +// Some sglang builds return a top-level actual_usage next to usage. openresty can strip it +// from the client response (postprocess strip_actual_usage), so the bodylog is the only place +// it survives: resp_meta must carry it, as it carries usage. +func TestRespMetaKeepsActualUsage(t *testing.T) { + usage := `{"prompt_tokens":4096,"completion_tokens":128,"total_tokens":4224,"prompt_tokens_details":{"cached_tokens":2048,"cache_write_tokens":1792}}` + actual := `{"prompt_tokens":4096,"completion_tokens":128,"total_tokens":4224,"prompt_tokens_details":{"cached_tokens":3072}}` + cases := map[string]string{ + "stream": `data: {"id":"c1","choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}` + "\n\n" + + `data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + "\n\n" + + `data: {"id":"c1","choices":[],"usage":` + usage + `,"actual_usage":` + actual + `}` + "\n\n" + + "data: [DONE]\n\n", + "json": `{"id":"c2","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":` + + usage + `,"actual_usage":` + actual + `}`, + } + for name, body := range cases { + meta := buildRespMeta(extractResp(body)) + if meta == nil { + t.Fatalf("%s: no resp_meta", name) + } + b, _ := json.Marshal(meta) + var got struct { + Usage struct { + Details struct { + Cached int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"usage"` + Actual *struct { + Prompt int `json:"prompt_tokens"` + Details struct { + Cached int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + } `json:"actual_usage"` + } + if err := json.Unmarshal(b, &got); err != nil { + t.Fatalf("%s: %v", name, err) + } + if got.Usage.Details.Cached != 2048 { + t.Errorf("%s: usage.cached_tokens = %d, want 2048 (%s)", name, got.Usage.Details.Cached, b) + } + if got.Actual == nil { + t.Fatalf("%s: resp_meta has no actual_usage: %s", name, b) + } + if got.Actual.Details.Cached != 3072 || got.Actual.Prompt != 4096 { + t.Errorf("%s: actual_usage = %+v, want cached 3072 / prompt 4096", name, *got.Actual) + } + } + + // absent upstream -> absent in resp_meta (no empty key) + meta := buildRespMeta(extractResp(`{"choices":[{"message":{"content":"x"},"finish_reason":"stop"}],"usage":` + usage + `}`)) + if _, ok := meta["actual_usage"]; ok { + t.Errorf("actual_usage present though upstream sent none: %v", meta) + } +} diff --git a/lua/debug_endpoints.lua b/lua/debug_endpoints.lua index d74a0dd..dcd05cd 100644 --- a/lua/debug_endpoints.lua +++ b/lua/debug_endpoints.lua @@ -502,8 +502,8 @@ local function rate_limit_view(opts) enabled = enabled, source = source, -- runtime | route | global runtime_override = override, - route_value = opts.rate_limit, -- route conf, or the global when unset - global_default = _G.RATE_LIMIT_ENABLED ~= false, + route_value = opts.load_rate_limit, -- route conf, or the global when unset + global_default = _G.LOAD_RATE_LIMIT_ENABLED ~= false, covers = route.RATE_LIMIT_REASONS, -- While the switch is off: per gate, the requests that tripped it and were let through. -- A trip, not a refusal: a probe slot or a warm turn in session_prio's soft band might have diff --git a/lua/directives.lua b/lua/directives.lua index ea89bdf..908c88b 100644 --- a/lua/directives.lua +++ b/lua/directives.lua @@ -62,7 +62,7 @@ M.DIRECTIVES = { { key = "default_max", label = "Default peer maximum", group = "Capacity & failover", kind = "number", detail = "In-flight requests allowed per peer when the peer entry names no maximum of its own. The static pool maximum -- what the adaptive limit is bounded by -- is the largest per-tier sum of peer maximums." }, - { key = "rate_limit", label = "Rate limiting", group = "Capacity & failover", kind = "boolean", + { key = "load_rate_limit", label = "Load-based rate limiting", group = "Capacity & failover", kind = "boolean", detail = "Master switch for the load-based 429s: concurrency (the static pool limit or the adaptive one, session_prio's per-turn shares included), TTFT (session_prio's graded gate included) and decode rate. Off lets those requests through while everything behind them -- the measurements, the adaptive limit, health bans and the 503 for an all-banned pool -- keeps running, and /_rate_limit_status counts what would have been refused. Content rules (reject_rules) are not covered. A runtime override from /_rate_limit_toggle wins over this value. Accepts true/false (also yes/no, on/off, 1/0); anything else is logged and the default stands." }, { hidden = true, key = "rt_limit_factor", label = "Runtime limit factor", group = "Capacity & failover", kind = "number", detail = "Scales the pool limit at which requests are turned away for concurrency: a request is refused once in-flight requests reach limit x factor. Must be over 0; anything else is logged and the global default stands." }, diff --git a/lua/postprocess.lua b/lua/postprocess.lua index 3a1f785..9392e27 100644 --- a/lua/postprocess.lua +++ b/lua/postprocess.lua @@ -44,12 +44,20 @@ -- exceeds it, the framework gives up buffering, flushes what it has and -- passes the rest through un-transformed (fail-open, logged). -- --- Route opt: opts.postprocess = "" selects the handler for that route. --- Global fallback: _G.POSTPROCESS_DEFAULT = "". nil in both = off. +-- Route opt: opts.postprocess = "" selects the handler for that route (one +-- handler per route). Global fallback: _G.POSTPROCESS_DEFAULT = "". nil in +-- both = off. Built-in handlers (bottom of this file): +-- vllm_format -- normalize vLLM's response quirks (includes dropping actual_usage) +-- strip_actual_usage -- only drop the top-level actual_usage some sglang builds add local M = {} -local cjson = require "cjson.safe" -- decode/encode return nil on error (no throw) +-- decode/encode return nil on error (no throw). A private instance, so the setting +-- below does not leak into other modules: decoded arrays keep an array marker and an +-- empty one re-encodes as [] -- by default cjson writes every empty table as {}, which +-- would turn the final usage chunk's "choices":[] into "choices":{} on a rewrite. +local cjson = require("cjson.safe").new() +cjson.decode_array_with_array_mt(true) local json_null = (require "cjson").null -- Cap on how much of a response may be accumulated in memory, so a large @@ -206,6 +214,7 @@ local vllm_choice_null_fields = { local vllm_message_null_fields = { "refusal", "annotations", "audio", "function_call" } local vllm_top_drop_fields = { "service_tier", "system_fingerprint", "kv_transfer_params", "ec_transfer_params", + "actual_usage", -- some sglang builds; see strip_actual_usage below } local function vllm_drop_null_fields(obj, fields) @@ -281,72 +290,95 @@ local function vllm_rewrite_json(payload) return enc or payload -- fail open on an encode error end --- Rewrite one SSE data frame. Leave unchanged frames byte-identical. -local function vllm_rewrite_sse_line(line) +-- Rewrite one SSE data frame with rewrite_json. Leave unchanged frames byte-identical. +local function rewrite_sse_line(line, rewrite_json) local payload = line:match("^data:%s*(.-)%s*\r?\n?$") if not payload or payload == "" or payload == "[DONE]" then return line -- comment/heartbeat/terminator/non-data end - local rewritten = vllm_rewrite_json(payload) + local rewritten = rewrite_json(payload) if rewritten == payload then return line end return "data: " .. rewritten .. (line:match("\r?\n$") or "") end --- vllm_format: normalize SSE frames and complete JSON responses. SSE is handled --- line by line; non-stream JSON is capped and buffered until eof. -M.register("vllm_format", { - buffered = false, - rewrites_length = true, - accept_status = function(status) return status >= 200 and status < 300 end, - header = function(ctx, opts) - local content_type = (ngx.header.content_type or ""):lower() - ctx.pp_sse = content_type:find("text/event-stream", 1, true) ~= nil - ctx.pp_json = not ctx.pp_sse and content_type:find("application/json", 1, true) ~= nil - end, - body = function(chunk, eof, ctx, opts) - if ctx.pp_json then - if ctx.pp_json_overflow then return nil end - local buf = ctx.pp_json_buf - if not buf then - buf = {} - ctx.pp_json_buf = buf - ctx.pp_json_size = 0 - end - if chunk and chunk ~= "" then - buf[#buf + 1] = chunk - ctx.pp_json_size = ctx.pp_json_size + #chunk - end - if ctx.pp_json_size > M.DEFAULT_MAX_BUFFER then - ngx.log(ngx.ERR, "postprocess: JSON body over ", M.DEFAULT_MAX_BUFFER, - " bytes; passing through un-transformed") - ctx.pp_json_overflow = true +-- A handler that rewrites OpenAI-style responses one JSON document at a time: +-- each SSE data frame, or a whole non-stream JSON body (capped and buffered until +-- eof). rewrite_json(payload) -> payload' returns its input unchanged (the same +-- string) when there is nothing to do, so untouched frames stay byte-identical. +local function json_rewrite_handler(rewrite_json) + return { + buffered = false, + rewrites_length = true, + accept_status = function(status) return status >= 200 and status < 300 end, + header = function(ctx, opts) + local content_type = (ngx.header.content_type or ""):lower() + ctx.pp_sse = content_type:find("text/event-stream", 1, true) ~= nil + ctx.pp_json = not ctx.pp_sse and content_type:find("application/json", 1, true) ~= nil + end, + body = function(chunk, eof, ctx, opts) + if ctx.pp_json then + if ctx.pp_json_overflow then return nil end + local buf = ctx.pp_json_buf + if not buf then + buf = {} + ctx.pp_json_buf = buf + ctx.pp_json_size = 0 + end + if chunk and chunk ~= "" then + buf[#buf + 1] = chunk + ctx.pp_json_size = ctx.pp_json_size + #chunk + end + if ctx.pp_json_size > M.DEFAULT_MAX_BUFFER then + ngx.log(ngx.ERR, "postprocess: JSON body over ", M.DEFAULT_MAX_BUFFER, + " bytes; passing through un-transformed") + ctx.pp_json_overflow = true + ctx.pp_json_buf = nil + return table.concat(buf) + end + if not eof then return "" end ctx.pp_json_buf = nil - return table.concat(buf) + return rewrite_json(table.concat(buf)) end - if not eof then return "" end - ctx.pp_json_buf = nil - return vllm_rewrite_json(table.concat(buf)) - end - if not ctx.pp_sse then return nil end -- other media types pass through - local pending = (ctx.pp_leftover or "") .. (chunk or "") - ctx.pp_leftover = nil - local out, pos = {}, 1 - while true do - local nl = pending:find("\n", pos, true) - if not nl then - local rest = pending:sub(pos) -- incomplete trailing line - if eof then - out[#out + 1] = vllm_rewrite_sse_line(rest) - else - ctx.pp_leftover = rest -- stash; the next chunk completes it + if not ctx.pp_sse then return nil end -- other media types pass through + local pending = (ctx.pp_leftover or "") .. (chunk or "") + ctx.pp_leftover = nil + local out, pos = {}, 1 + while true do + local nl = pending:find("\n", pos, true) + if not nl then + local rest = pending:sub(pos) -- incomplete trailing line + if eof then + out[#out + 1] = rewrite_sse_line(rest, rewrite_json) + else + ctx.pp_leftover = rest -- stash; the next chunk completes it + end + break end - break + out[#out + 1] = rewrite_sse_line(pending:sub(pos, nl), rewrite_json) + pos = nl + 1 end - out[#out + 1] = vllm_rewrite_sse_line(pending:sub(pos, nl)) - pos = nl + 1 - end - return table.concat(out) - end, -}) + return table.concat(out) + end, + } +end + +-- vllm_format: normalize vLLM's SSE frames and complete JSON responses. +M.register("vllm_format", json_rewrite_handler(vllm_rewrite_json)) + +-- strip_actual_usage: drop the top-level "actual_usage" object some sglang builds add +-- next to "usage" (a second token count that clients and billing should not see; the +-- bodylog still records it, since it runs before post-processing). Nothing else is +-- touched. A plain substring test runs first, so only the frame that carries the +-- field -- the final usage chunk of a stream -- is decoded at all. +local function strip_actual_usage_json(payload) + if not payload:find('"actual_usage"', 1, true) then return payload end + local obj = cjson.decode(payload) + if type(obj) ~= "table" or obj.actual_usage == nil then return payload end + obj.actual_usage = nil + return cjson.encode(obj) or payload -- fail open on an encode error +end +M.strip_actual_usage_json = strip_actual_usage_json -- exposed for unit tests + +M.register("strip_actual_usage", json_rewrite_handler(strip_actual_usage_json)) return M diff --git a/lua/route.lua b/lua/route.lua index 1adc19a..770ddbc 100644 --- a/lua/route.lua +++ b/lua/route.lua @@ -68,7 +68,7 @@ M.DEFAULTS = { default_max = global("MAX_CONCURRENCY_PER_PEER"), rt_limit_factor = global("RT_LIMIT_FACTOR"), -- `~= false`: a session_base.conf older than this switch has no such global, and limiting must stay on - rate_limit = function() return _G.RATE_LIMIT_ENABLED ~= false end, + load_rate_limit = function() return _G.LOAD_RATE_LIMIT_ENABLED ~= false end, max_more_tries = 2, cross_tier_fallback = true, -- health checks / routing @@ -223,12 +223,14 @@ function _G.register_route(name, opts_factory) " is invalid (need a positive number) - falling back to ", D("rt_limit_factor")) opts.rt_limit_factor = D("rt_limit_factor") end - -- rate_limit: master switch for the load-based 429s (concurrency -- static or adaptive --, TTFT, TPS). + -- load_rate_limit: master switch for the load-based 429s (concurrency -- static or adaptive --, TTFT, TPS). + -- Not named rate_limit: on a ModelRoute that key is a video route's download speed limit, and the + -- ModelRoute CRD rejects it on llm routes. -- Off lets those requests through; measurement, AIMD and the 503 for an all-banned pool are unchanged. -- A runtime override (/_rate_limit_toggle) wins over this; see M.rate_limit_state. -- An unparsable value falls back to the global default, so it must not be reported as source=route. - opts._rate_limit_from_route = opts.rate_limit ~= nil and util.parse_bool(opts.rate_limit) ~= nil - util.bool_opt(opts, "rate_limit", D("rate_limit"), name) + opts._load_rate_limit_from_route = opts.load_rate_limit ~= nil and util.parse_bool(opts.load_rate_limit) ~= nil + util.bool_opt(opts, "load_rate_limit", D("load_rate_limit"), name) util.bool_opt(opts, "cch_default_enabled", _G.CCH_STRIP_DEFAULT_ENABLED == true, name) -- Kimi tool_call_id normalize 每路由默认（factory 显式设 false 即永久关，重启也不恢复；缺省回落全局 _G.KIMI_NORMALIZE_DEFAULT_ENABLED） util.bool_opt(opts, "kimi_normalize_default_enabled", _G.KIMI_NORMALIZE_DEFAULT_ENABLED == true, name) @@ -720,8 +722,8 @@ end -- Three layers, highest first: -- runtime ngx.shared[opts.cch_ctl_dict] "rate_limit_enabled" = 1/0, set by /_rate_limit_toggle; -- no reload, survives one; ?reset=1 removes it --- route opts.rate_limit (route conf / ModelRoute nginx.values), parsed in register_route --- global _G.RATE_LIMIT_ENABLED (session_base.conf), the route option's default +-- route opts.load_rate_limit (route conf / ModelRoute nginx.values), parsed in register_route +-- global _G.LOAD_RATE_LIMIT_ENABLED (session_base.conf), the route option's default -- Returns enabled, source ("runtime" | "route" | "global"). local RL_KEY = "rate_limit_enabled" M.RATE_LIMIT_REASONS = { "concurrency", "ttft", "tps" } @@ -730,8 +732,8 @@ function M.rate_limit_state(opts) local ctl = ngx.shared[opts.cch_ctl_dict] local ovr = ctl and ctl:get(RL_KEY) if ovr ~= nil then return ovr == 1, "runtime" end - local src = opts._rate_limit_from_route and "route" or "global" - return opts.rate_limit ~= false, src + local src = opts._load_rate_limit_from_route and "route" or "global" + return opts.load_rate_limit ~= false, src end -- With the switch off, a request that trips a gate is let through and counted here instead: per diff --git a/session_base.conf b/session_base.conf index 6d725d1..327a2d0 100644 --- a/session_base.conf +++ b/session_base.conf @@ -143,10 +143,10 @@ init_by_lua_block { -- 注:hit_avg(avg_5min 那路)已暂时禁用,只用 realtime 判过载(见 do_route)。 _G.RT_LIMIT_FACTOR = 1 - -- Rate-limit master switch, the default for every route's `rate_limit` option. false lets the + -- Rate-limit master switch, the default for every route's `load_rate_limit` option. false lets the -- load-based 429s (concurrency, TTFT, TPS) through on all routes; measurement keeps running. -- Per route: `rate_limit = false` in the route conf, or POST //_rate_limit_toggle?on=0. - _G.RATE_LIMIT_ENABLED = true + _G.LOAD_RATE_LIMIT_ENABLED = true -- ── TTFT 限流开关 + 全局阈值 ── -- _G.TTFT_ENABLED:总开关(config 级,reload 生效)。true=对所有路由生效(用下方全局阈值, diff --git a/test/mock_vllm_sse.py b/test/mock_vllm_sse.py index fb1ec92..1a55535 100644 --- a/test/mock_vllm_sse.py +++ b/test/mock_vllm_sse.py @@ -27,6 +27,8 @@ {"prefill_delay_ms": 0} -> override prefill delay {"emit_reasoning": true}-> 用 reasoning 字段（Kimi-K2.5 思考链）而非 content {"emit_tool_call": true}-> 末尾追加一个 tool_call（chat 接口） + {"emit_actual_usage": true} -> chat usage also carries a top-level "actual_usage" (as some + sglang builds do; same as --actual-usage for every request) 非流式 finish_reason / stop_reason 默认 "length"（mock 总是发完 max_tokens 才停）， 和真 vllm 在 max_tokens 触发时一致。 @@ -103,9 +105,19 @@ def _chat_tool_call_chunk(cid: str, model: str, created: int, name: str, args: s })) +def _actual_usage(usage: dict) -> dict: + """A second, different token count next to "usage", as some sglang builds return it: + same totals, but its own prompt_tokens_details (cached_tokens + 1024, capped at the + prompt) so a test can tell the two apart.""" + cached = ((usage.get("prompt_tokens_details") or {}).get("cached_tokens") or 0) + 1024 + actual = dict(usage) + actual["prompt_tokens_details"] = {"cached_tokens": min(cached, usage["prompt_tokens"])} + return actual + + def _chat_usage_chunk(cid: str, model: str, created: int, prompt_tokens: int, completion_tokens: int, - cached_tokens=None) -> bytes: + cached_tokens=None, actual_usage=False) -> bytes: """include_usage=true 时的最终 chunk（choices=[] + usage）。""" usage = {"prompt_tokens": prompt_tokens, "total_tokens": prompt_tokens + completion_tokens, "completion_tokens": completion_tokens} @@ -114,11 +126,14 @@ def _chat_usage_chunk(cid: str, model: str, created: int, # default wire is byte-identical for every other harness. if cached_tokens is not None: usage["prompt_tokens_details"] = {"cached_tokens": cached_tokens} - return _sse(json.dumps({ + chunk = { "id": cid, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [], "usage": usage, - })) + } + if actual_usage: + chunk["actual_usage"] = _actual_usage(usage) + return _sse(json.dumps(chunk)) # ─── /v1/completions: chunk builders（legacy text completion）───────────── @@ -301,6 +316,8 @@ async def chat_completions(request: web.Request): "prompt_logprobs": None, "prompt_token_ids": None, "kv_transfer_params": None, } + if app["actual_usage"] or body.get("emit_actual_usage"): + payload["actual_usage"] = _actual_usage(payload["usage"]) h = {"X-Mock-Peer": app["name"]}; h.update(extra_hdr) return web.json_response(payload, headers=h) @@ -354,7 +371,9 @@ async def chat_completions(request: web.Request): # usage chunk（仅当 stream_options.include_usage=true） if (body.get("stream_options") or {}).get("include_usage"): await resp.write(_chat_usage_chunk(cid, model, created, prompt_tokens, out_len, - cached_tokens=app.get("cached_tokens"))) + cached_tokens=app.get("cached_tokens"), + actual_usage=app["actual_usage"] + or bool(body.get("emit_actual_usage")))) await resp.write(_DONE) await resp.write_eof() @@ -490,6 +509,7 @@ def make_app(args): app["models_status"] = args.models_status app["chat_status"] = args.chat_status app["cached_tokens"] = args.cached_tokens + app["actual_usage"] = args.actual_usage app["stats"] = {"req_count": 0, "in_flight": 0} app.router.add_get("/health", health) app.router.add_get("/v1/models", models_list) @@ -512,12 +532,16 @@ def main(): p.add_argument("--cached-tokens", type=int, default=None, help="when set, usage chunks also carry prompt_tokens_details.cached_tokens (prefix cache); " "unset keeps the wire byte-identical to before") + p.add_argument("--actual-usage", action="store_true", + help="chat usage also carries a top-level actual_usage (some sglang builds); " + "per request: {\"emit_actual_usage\": true}") + p.add_argument("--host", default="127.0.0.1", help="listen address (0.0.0.0 inside a pod)") args = p.parse_args() app = make_app(args) print(f"[mock-sse] {args.name} on :{args.port} out={args.output_len} " f"chunk={args.chunk_delay_ms}ms prefill={args.prefill_delay_ms}ms", flush=True) - web.run_app(app, host="127.0.0.1", port=args.port, access_log=None, print=lambda *_: None) + web.run_app(app, host=args.host, port=args.port, access_log=None, print=lambda *_: None) if __name__ == "__main__": diff --git a/test/run_lua_unit_tests.sh b/test/run_lua_unit_tests.sh new file mode 100644 index 0000000..b2bb757 --- /dev/null +++ b/test/run_lua_unit_tests.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Run every Lua unit test (test/utest_*.lua) with the resty CLI; used by CI (.github/workflows/test.yml) +# and locally: +# bash test/run_lua_unit_tests.sh [resty] default resty: $RESTY, else `resty` on PATH +# Each test exits non-zero on failure (os.exit(1), or an assert error). Exits non-zero if any test fails, +# or if no test was found -- a runner that finds nothing must not pass. +# Most tests run from the repo root with `-I lua`; the ones below need their own invocation. +set -uo pipefail +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +RESTY=${1:-${RESTY:-resty}} +command -v "$RESTY" >/dev/null || { echo "no resty CLI ($RESTY)"; exit 2; } +cd "$ROOT" +pass=0 fail=0 failed="" +for t in test/utest_*.lua; do + [ -f "$t" ] || continue + case $t in + test/utest_bodylog_gate.lua) out=$("$RESTY" -I lua --shdict 'bodylog_ctl 1m' "$t" 2>&1) ;; + test/utest_window_stat.lua) out=$(cd test && "$RESTY" -I ../lua "${t#test/}" 2>&1) ;; # loads ../lua/util.lua + *) out=$("$RESTY" -I lua "$t" 2>&1) ;; + esac + rc=$? + # resty prints a _G write-guard warning with a traceback for each module that sets a global; not a failure + summary=$(printf '%s\n' "$out" | grep -vE "_G write guard|stack traceback|in main chunk|file_gen|init_worker_by_lua|xpcall|^[[:space:]]*\[C\]|^[[:space:]]*(test|lua)/" | grep . | tail -1) + if [ $rc = 0 ]; then + pass=$((pass + 1)); printf 'ok %-46s %s\n' "$t" "$summary" + else + fail=$((fail + 1)); failed="$failed $t" + printf 'FAIL %-46s rc=%s\n' "$t" "$rc"; printf '%s\n' "$out" | sed 's/^/ /' | tail -30 + fi +done +echo "lua unit tests: $pass passed, $fail failed${failed:+ (}${failed# }${failed:+)}" +[ $((pass + fail)) -gt 0 ] || { echo "no test/utest_*.lua found"; exit 1; } +[ $fail = 0 ] diff --git a/test/test_rate_limit_switch.sh b/test/test_rate_limit_switch.sh index b6332dd..1e1cbef 100755 --- a/test/test_rate_limit_switch.sh +++ b/test/test_rate_limit_switch.sh @@ -4,7 +4,7 @@ # Covers the three load-based 429 exits and what must NOT change with the switch off: # S1-S4 static concurrency limit: on -> 429; runtime off -> all 200, counted as would-reject and NOT # as real 429s; ?reset=1 -> back to the route/global value; bad ?on -> 400 -# S5-S6 route conf `rate_limit = false` -> source=route, all 200; runtime on=1 wins over it +# S5-S6 route conf `load_rate_limit = false` -> source=route, all 200; runtime on=1 wins over it # S7 adaptive concurrency route: off -> all 200 # S8 the runtime override survives `openresty -s reload` # S9 TTFT 429: on -> 429; off -> all 200, would_reject.ttft > 0 @@ -64,14 +64,14 @@ cat <<'EOF' _G.TPS_BUCKETS = {5,10,20,40,80,150,300} _G.ADAPTIVE_CC_HOLD_WHEN_IDLE = false _G.TPS_LIMIT_TPS = 20 -- pinned, not inherited from session_base.conf - _G.RATE_LIMIT_ENABLED = true -- pinned likewise + _G.LOAD_RATE_LIMIT_ENABLED = true -- pinned likewise local tight = {{"127.0.0.1",28951,"m1",0,1},{"127.0.0.1",28952,"m2",0,1}} -- static max = 2 local wide = {{"127.0.0.1",28951,"m1"},{"127.0.0.1",28952,"m2"},{"127.0.0.1",28953,"m3"}} local C = {default_max=50, bodylog_default_enabled=false, health_check_interval=2} local function R(x) local t={} for k,v in pairs(C) do t[k]=v end for k,v in pairs(x) do t[k]=v end return t end _G.register_route("st", function() return R({peers=tight, rt_limit_factor=1, adaptive_cc=false}) end) _G.register_route("ac", function() return R({peers=tight, rt_limit_factor=1, adaptive_cc=true}) end) - _G.register_route("rf", function() return R({peers=tight, rt_limit_factor=1, adaptive_cc=false, rate_limit=false}) end) + _G.register_route("rf", function() return R({peers=tight, rt_limit_factor=1, adaptive_cc=false, load_rate_limit=false}) end) _G.register_route("tt", function() return R({peers=wide, adaptive_cc=false, ttft_limit_ms=400, ttft_window=3, ttft_ttl=8, ttft_probe_window=3, ttft_probe_per_window=2}) end) _G.register_route("tp", function() return R({peers=wide, adaptive_cc=false, tps_limit_tps=50, tps_window=3, tps_ttl=8, @@ -86,7 +86,7 @@ cat <<'EOF' ttft_window=3, ttft_ttl=30, ttft_probe_window=3, ttft_probe_per_window=2}) end) _G.register_route("sp", function() return R({peers=wide, adaptive_cc=false, session_prio=true, tps_limit_tps=50, tps_window=3, tps_ttl=8, tps_probe_window=3, tps_probe_per_window=2, tps_min_decode_s=0.3}) end) - _G.register_route("rv", function() return R({peers=tight, rt_limit_factor=1, adaptive_cc=false, rate_limit="maybe"}) end) + _G.register_route("rv", function() return R({peers=tight, rt_limit_factor=1, adaptive_cc=false, load_rate_limit="maybe"}) end) } EOF for r in $ROUTES; do p=$(port_of $r); cat < 400" || no "S4 bad on -> $code" o=$(tog st reset_counts=1); [ "$(echo "$o" | js "d['would_reject']['concurrency']")" = "0" ] && ok "S4 reset_counts=1 zeroes the counts" || no "S4 reset_counts: $o" -# ── route conf rate_limit=false ── +# ── route conf load_rate_limit=false ── [ "$(st rf "d['enabled']")" = "False" ] && [ "$(st rf "d['source']")" = "route" ] && ok "S5 rf: rate_limit=false -> enabled=false, source=route" || no "S5 rf: $(curl -s $(U rf)/_rate_limit_status)" r=$(burst rf 20 1500 5 5); allok "$r" 20 && ok "S5 rf: all 200 ($r)" || no "S5 rf: $r" tog rf on=1 >/dev/null; r=$(burst rf 20 1500 5 5); limited "$r" && ok "S6 runtime on=1 beats route false: 429 ($r)" || no "S6 override not applied: $r" @@ -186,7 +186,7 @@ code=$(fire dd 20 5 5); [ "$code" = 503 ] && ok "S11 all peers banned -> 503 wit v=$(curl -s "$(U st)/_tps_status" | js "(d['rate_limit_enabled'], d['rate_limit_source'])"); [ "$v" = "(True, 'global')" ] && ok "S12 /_tps_status reports the switch $v" || no "S12 tps_status: $v" v=$(curl -s "$(U st)/_429_status?all=1" | js "sorted(d['rate_limit'].keys())"); echo "$v" | grep -q "'tt'" && ok "S12 /_429_status?all=1 lists every route's switch" || no "S12 429_status: $v" v=$(curl -s "$(U st)/_429_status" | js "list(d['by_reason'].keys())"); echo "$v" | grep -q would && no "S12 would-reject leaked into by_reason: $v" || ok "S12 by_reason holds real 429s only" -v=$(curl -s "$(U rf)/_route_opts" | js "[x for x in d.get('directives',[]) if x['key']=='rate_limit'][0]['value']"); [ "$v" = "False" ] && ok "S12 /_route_opts: rate_limit value=False on rf" || no "S12 route_opts: $v" +v=$(curl -s "$(U rf)/_route_opts" | js "[x for x in d.get('directives',[]) if x['key']=='load_rate_limit'][0]['value']"); [ "$v" = "False" ] && ok "S12 /_route_opts: load_rate_limit value=False on rf" || no "S12 route_opts: $v" # ── session_prio: per-turn concurrency share ── (pool of 10; turn 1 capped at 8.5) # 10 concurrent: the 10th sees 9 in flight -- over turn 1's 8.5, under the pool's 10. diff --git a/test/utest_postprocess_strip_actual_usage.lua b/test/utest_postprocess_strip_actual_usage.lua new file mode 100644 index 0000000..22a4e30 --- /dev/null +++ b/test/utest_postprocess_strip_actual_usage.lua @@ -0,0 +1,101 @@ +-- Run: resty test/utest_postprocess_strip_actual_usage.lua [lua/postprocess.lua] +local cjson = require "cjson.safe" +local pp = assert(loadfile(arg[1] or "lua/postprocess.lua"))() + +-- Drive the handler through the real header/body phase entry points. +local function run(name, content_type, pieces, status) + local ctx = {} + local previous_ngx = _G.ngx + _G.ngx = { + ctx = ctx, status = status or 200, + header = { content_type = content_type, content_length = 999 }, + arg = {}, ERR = "error", log = function() end, + } + local opts = { postprocess = name } + _G.do_postprocess_header(opts) + local out = {} + for i, piece in ipairs(pieces) do + ngx.arg[1], ngx.arg[2] = piece, i == #pieces + _G.do_postprocess_body(opts) + out[#out + 1] = ngx.arg[1] + end + local content_length = ngx.header.content_length + _G.ngx = previous_ngx + return table.concat(out), content_length +end + +local function pieces_of(s, n) -- split s into n roughly equal pieces + local out, step = {}, math.max(1, math.floor(#s / n)) + local i = 1 + while i <= #s do out[#out + 1] = s:sub(i, i + step - 1); i = i + step end + return out +end + +local usage = '{"prompt_tokens":4096,"completion_tokens":128,"total_tokens":4224,' + .. '"prompt_tokens_details":{"cached_tokens":2048,"cache_write_tokens":1792}}' +local actual = '{"prompt_tokens":4096,"completion_tokens":128,"total_tokens":4224,' + .. '"prompt_tokens_details":{"cached_tokens":3072}}' + +-- 1. Stream: content frames stay byte-identical; the final usage frame loses +-- actual_usage, keeps usage untouched and keeps "choices":[] an array. +local content = 'data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}\n\n' +local final = 'data: {"id":"c1","object":"chat.completion.chunk","choices":[],"usage":' .. usage + .. ',"actual_usage":' .. actual .. '}\n\n' +local stream = content .. final .. "data: [DONE]\n\n" +for _, n in ipairs({ 1, 3, 7, #stream }) do -- whole, and split mid-line / per byte + local out, cl = run("strip_actual_usage", "text/event-stream", pieces_of(stream, n)) + assert(cl == nil, "Content-Length must be dropped") + assert(out:sub(1, #content) == content, "content frame changed (n=" .. n .. ")") + assert(not out:find("actual_usage", 1, true), "actual_usage left (n=" .. n .. ")") + assert(out:sub(-14) == "data: [DONE]\n\n", "terminator changed (n=" .. n .. ")") + local last + for json in out:gmatch("data: ({[^\n]*})\n") do last = json end -- the last JSON frame + assert(last) + assert(last:find('"choices":[]', 1, true), "choices no longer an empty array: " .. last) + local obj = assert(cjson.decode(last)) + assert(obj.usage.prompt_tokens_details.cached_tokens == 2048) + assert(obj.usage.prompt_tokens_details.cache_write_tokens == 1792) + assert(obj.usage.total_tokens == 4224 and obj.id == "c1") +end + +-- 2. Non-stream JSON body, split across chunks. +local body = '{"id":"c2","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":' + .. usage .. ',"actual_usage":' .. actual .. '}' +local out, cl = run("strip_actual_usage", "application/json; charset=utf-8", pieces_of(body, 5)) +assert(cl == nil) +local obj = assert(cjson.decode(out)) +assert(obj.actual_usage == nil and obj.usage.prompt_tokens_details.cached_tokens == 2048) +assert(obj.choices[1].message.content == "ok") + +-- 3. Nothing to strip: byte-identical, stream and JSON. +local plain = 'data: {"choices":[],"usage":' .. usage .. '}\n\n' +assert(run("strip_actual_usage", "text/event-stream", { plain }) == plain) +local plain_json = '{ "choices": [], "usage": ' .. usage .. ' }' +assert(run("strip_actual_usage", "application/json", { plain_json }) == plain_json) +-- the name inside a string value: its quotes are escaped there, so the substring test +-- does not match and the frame is not even decoded +local in_text = 'data: {"choices":[{"delta":{"content":"\\"actual_usage\\" is a field"}}]}\n' +assert(run("strip_actual_usage", "text/event-stream", { in_text }) == in_text) +-- nested actual_usage is not the top-level field: left alone +local nested = '{"choices":[],"usage":{"actual_usage":1}}' +assert(pp.strip_actual_usage_json(nested) == nested) + +-- 4. Gates: non-2xx and other media types pass through untouched. +assert(run("strip_actual_usage", "application/json", { body }, 500) == body) +assert(run("strip_actual_usage", "text/plain", { body }) == body) +-- malformed JSON: fail open +local broken = 'data: {"actual_usage": \n' +assert(run("strip_actual_usage", "text/event-stream", { broken }) == broken) + +-- 5. vllm_format drops it too, and keeps "choices":[] an array on the usage frame. +local v_out = run("vllm_format", "text/event-stream", + { 'data: {"choices":[],"system_fingerprint":"fp","usage":' .. usage .. ',"actual_usage":' .. actual .. '}\n' }) +assert(not v_out:find("actual_usage", 1, true) and not v_out:find("system_fingerprint", 1, true)) +assert(v_out:find('"choices":[]', 1, true), "vllm_format turned choices into an object: " .. v_out) + +-- 6. Each route gets only the handler it names. +assert(pp.resolve({ postprocess = "strip_actual_usage" }) == pp.handlers.strip_actual_usage) +assert(pp.resolve({ postprocess = "vllm_format" }) == pp.handlers.vllm_format) +assert(pp.resolve({}) == nil) + +print("strip_actual_usage SSE and JSON rewrite: PASS")