A token budget that only gets checked after the response comes back gets checked too late. By then the model has already run and the money is spent. If many requests from one tenant arrive in the same second, each of them reads the same "remaining budget" value, each of them passes, and the overspend shows up in tomorrow's cost report.
Release v2.5.0 of ai-runtime-platform adds an Inference Budget Guard to the gateway's chat completions path. It treats spend as admission control. Before a request reaches the model, the gateway reserves an estimated cost against a per-tenant or per-API-key window. When the request finishes, the reservation is settled against what actually happened.
Where it sits in the request path
The guard runs inside the gateway process. From docs/budget-guard.md:
drain check
→ global admission
→ tenant / fleet concurrency admission
→ inference budget reserve
→ tenant attribution pre-reserve (signals)
→ Control Plane governance
→ route / execute
→ settle budget + reconcile attribution
Putting the reserve after concurrency admission means a request that was going to be rejected for capacity never touches the budget store. If the reserve is denied, the gateway returns 429 (inference_budget_exceeded) and the backend is never called.
One invariant, checked atomically
All amounts are integers in micro-USD. Estimates and actual charges round up, the budget cap rounds down. Per scope and window, the store keeps committed_micros (settled spend) and reserved_micros (in-flight estimates). The rule the guard enforces is short:
committed + reserved + estimate ≤ budget
With REDIS_URL set, that check and the increment happen inside one Lua script, so replicas share the budget and cannot race each other between read and write. The core of the reserve script, from app/gateway/budget_guard.py:
if (committed + reserved + estimate) > budget then
-- persist window state, return deny
return {0, committed, reserved, committed + reserved, expired_count}
end
reserved = reserved + estimate
redis.call('ZADD', KEYS[2], now + ttl, res_id)
redis.call('HSET', KEYS[3], res_id, string.format('%d', estimate))
Each reservation also lands in a sorted set keyed by expiry. If a gateway pod dies mid-request, the reservation is reclaimed after reservationTtlSeconds instead of holding budget forever. Window rollover resets only committed spend. In-flight reservations carry over, so a new window cannot admit new spend on top of requests that are still running.
Without REDIS_URL, the guard falls back to an in-process lock. The design doc is explicit that budgets are then per gateway process, which is fine for a single replica and wrong for anything behind a Service with more than one pod.
API keys never appear raw in Redis. The scope key is key:<tenant>:<sha256(apiKey)[:16]>, and the same rule applies to metric labels and error bodies.
Where the estimate comes from
The estimate is deliberately simple. Input tokens are approximated as message length divided by four. Output tokens come from max_completion_tokens or max_tokens if the client sent one, otherwise from the smaller of the input estimate and 512. Prices come from the policy:
{
"version": "budget-guard-1.0",
"enabled": true,
"requireStore": true,
"default": {
"budgetUsd": 10.0,
"windowSeconds": 3600,
"softWarningRatio": 0.8,
"reservationTtlSeconds": 120
},
"prices": {
"qwen": { "inputPerMillion": 0.5, "outputPerMillion": 1.0 },
"*": { "inputPerMillion": 1.0, "outputPerMillion": 2.0 }
}
}
That is the example policy shipped in docs/budget-guard.policy.json. It also defines a finance tenant with a tighter budget and a per-key override. Above softWarningRatio, responses carry X-AI-Budget-Warning: soft_threshold, which gives clients a signal before they start getting 429s.
Enabling it is a few environment variables on the gateway:
INFERENCE_BUDGET_ENABLED=true
INFERENCE_BUDGET_POLICY_FILE=/etc/gateway/budget-guard.policy.json
REDIS_URL=redis://redis:6379/0
INFERENCE_BUDGET_STORE_TIMEOUT_SECONDS=1.0
With INFERENCE_BUDGET_ENABLED unset, nothing changes: no reservation and no budget headers.
Settlement is where the edge cases live
Reserving is the easy half. Every reservation has to be settled exactly once, and the right amount depends on how far the request got:
- upstream success with a
usageblock: charged at actual priced usage - upstream success without
usage, which covers most streams: charged at the estimate - failure before the model was called (governance deny, unknown model, connection refused): released, charged zero
- failure after the model may have run (read timeout, unparsable body): charged at the estimate
- settle arriving after the TTL expired: actual spend still committed once, via a tombstone
- duplicate settle: no-op
The timeout case is the easiest one to get wrong. A read timeout does not mean the model did nothing. The GPU may have generated the full answer and the gateway simply stopped waiting. Charging the estimate there is conservative, and it is hard to argue for anything lower when the real number is unknown.
Fail-closed returns 503
When the guard is on and it cannot make a decision, it denies. Store down, store timeout, a corrupt payload in Redis and a malformed policy all return 503 without calling the model backend. Being strict here costs availability. If Redis goes away, chat completions stop. For a guard that exists to cap spend, a silent allow during a Redis outage would make the cap meaningless at exactly the moment nobody is watching it.
What the evidence says
The figures come from docs/EVIDENCE.md at v2.5.0. These rows are marked "Tested locally; re-run in CI" and cite local commit fee974a. On this setup:
- real Redis server, 8 OS processes × 25 concurrent reserves × 10 rounds
- allows per round
[142], expected 142, withoverspend_events=0 - max reserved
994000of1000000micro-USD over 10 rounds - concurrent settle/release/abandon exact in 10/10 rounds,
leaked_reserved_micros=0 -
store_down,store_timeout,store_corruptandmalformed_policyeach denied with 503, backend not called
The 142 matters more than the zero. A guard that denies everything also has zero overspend. Hitting the exact expected allow count every round shows the invariant is enforced without leaving budget on the table. The suite also includes a false_deny_count_legitimate test for within-budget requests.
Limits
The repo states these directly, and they are worth repeating before anyone wires this into a billing conversation:
- The pre-request estimate can differ from actual usage. Committed spend can exceed the budget by at most the in-flight estimate error.
- If Redis fails after the upstream call succeeds, settlement lags. The reservation is released by TTL and that request's actual spend is not committed.
- Estimate prices come from the budget policy, settlement prices from the gateway's model targets. If those two tables drift apart, the guard drifts with them.
- Reservation expiry uses each replica's wall clock, so clock skew shifts TTL release.
- It is one Redis. Not multi-region billing, not a chargeback ledger. EVIDENCE.md calls it "an admission control, not an invoice system".
That last line is the right mental model. The guard answers "may this request start" with a number that holds under concurrency. Invoicing stays with whatever system reads the usage events afterwards.
Code and evidence: github.com/justrunme/ai-runtime-platform.
Top comments (0)