Repository navigation
fix(caching): set the counter TTL atomically in async_increment - #40715
songkuan-zheng wants to merge 4 commits into
Conversation
RedisCache.async_increment sent INCRBYFLOAT and then EXPIRE (or TTL plus EXPIRE) as separate awaits. A task cancelled between the two, which happens whenever the client disconnects mid request, left the counter without any expiry, so a spend or rate limit counter could live forever in Redis The increment and the TTL decision now run in one Lua call, following the pattern async_increment_with_floor and async_set_max already use in this file. The refresh_ttl semantics are unchanged: the flag re-arms the TTL on every increment, otherwise only a key with no TTL gets one Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com>
Greptile SummaryThis PR makes Redis counter increment and TTL handling atomic by executing both operations in one Lua call, preventing cancellation between separate Redis commands from leaving persistent counters.
Confidence Score: 5/5The PR appears safe to merge; the atomic Lua operation fixes the cancellation window without leaving any outstanding review finding. The previously reported non-positive-TTL regression is fixed because literal TTL values, including zero, are passed to EXPIRE, and the separate test-structure concern was withdrawn after the repository’s mocked-test convention and real Redis evidence were clarified. The staging merge did not alter this PR’s production caching implementation.
|
| Filename | Overview |
|---|---|
| litellm/caching/redis_cache.py | Replaces the multi-await increment/expiry sequence with one atomic Redis Lua evaluation while preserving existing TTL semantics. |
| tests/test_litellm/caching/test_redis_cache.py | Adds mocked Redis boundary coverage for script arguments, result handling, TTL edge cases, namespacing, and round-trip behavior. |
Reviews (4): Last reviewed commit: "Merge upstream litellm_internal_staging ..." | Re-trigger Greptile
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
The Lua script skipped EXPIRE for any non-positive TTL, while the old code only skipped it when get_ttl returned None and otherwise passed the value through, so EXPIRE 0 still deleted the key. The TTL now travels as an empty string for None and as the literal value otherwise, and the script only branches on that emptiness Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com>
|
Ready for maintainer review, Greptile 5/5 on c12a3db |
…rement_ttl Upstream (BerriAI#40834, BerriAI#40841) moved the increment into a _incrbyfloat_with_ttl helper that pipelines INCRBYFLOAT with EXPIRE or TTL and still issues a separate EXPIRE when the key had no TTL. The helper now runs the single Lua script instead, so the first increment of a fresh key cannot be split by a cancellation either. Upstream's round-trip test keeps its fake client, which gained an eval that applies the script's semantics, and now expects one round trip per increment in every case Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com>
|
Merged litellm_internal_staging (fd41183) to resolve the conflict with #40834 and #40841, which introduced _incrbyfloat_with_ttl. That helper pipelines INCRBYFLOAT with EXPIRE or TTL but still sends a separate EXPIRE when the key had no TTL, so the first increment of a fresh key kept the cancellation window. The helper body is now the single Lua script; the new round-trip test keeps its fake client (with an eval that applies the script) and expects one round trip in every case. Proof re-run at fd41183: cancelled increment leaves TTL 59, ttl 0 still deletes the key, the proxy issues one EVAL per counter and the counter reads TTL 50 ten seconds later |
|
Ready for maintainer review, Greptile 5/5 on fd41183 (rebased onto current litellm_internal_staging via merge) |
…rement_ttl Test-file only: upstream added pool-timeout breaker tests next to the spy tests of this branch. Both sets are kept, production code is unchanged Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com>
TLDR
Problem this solves:
RedisCache.async_incrementsends INCRBYFLOAT, then EXPIRE, as two separate awaitsHow it solves it:
async_increment_with_floorandasync_set_maxrefresh_ttlsemantics unchanged: re-arm always, or only arm a missing TTLUser Flow
Before: a proxy operator who runs several replicas with Redis coordination sees rate limit and spend counters that never expire
router_settings.routing_strategy: usage-based-routing-v2withredis_host/redis_portand restart the proxyTTLon it answers-1redis-cli --scan --pattern '*:tpm:*'weeks later and finds thousands of stale per-minute keys that memory eviction has to clean upAfter: the same disconnect leaves a counter that expires on schedule
router_settingsand restart the proxy on the fixed buildTTLon the usage counter answers60, the counter disappears one minute later like every other oneredis-cli --scan --pattern '*:tpm:*'only ever shows the current minute's keysRelevant issues
No existing issue.
RedisCache.async_increment(litellm/caching/redis_cache.py) awaitsincrbyfloatand then awaitsexpire(orttlfollowed byexpire). asyncio delivers a cancellation at whichever await is pending, so a task cancelled after the first round trip has already incremented the key but never sets its TTL. Every caller that routes throughDualCache.async_increment_cachewith a Redis backend is exposed: the usage-based routing counters, the parallel request limiter refund path, the spend counter warm-up, the MCP DCR rate counterLinear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*,make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
The tip 39a95d6 is a merge of litellm_internal_staging that only touched the test file (upstream added breaker tests next to these), so the After run captured at fd41183 still describes the production change byte for byte
A curl on its own cannot show this bug: the happy path writes the TTL on both sides, the defect only appears when the client goes away between the two Redis round trips. So the proof has two cases. Case 1 is the live proxy on port 4105 with
redis-cli monitorattached, showing what the proxy actually sends to Redis for one real chat completion. Case 2 reproduces the disconnect with the realRedisCacheclass against a real Redis 7 behind a tiny TCP relay that delays every reply by 300ms, then cancels the increment task while the reply is in flight. No mocks anywhere; the relay only adds latencyShared setup
/tmp/proof/config.yaml(the model is a real OpenAI compatible deployment, key and base come from the environment)/tmp/proof/cancel_probe.pyBefore (9a715df)
Case 1: what the proxy sends to Redis for one chat completion
uv run python litellm/proxy/proxy_cli.py --config /tmp/proof/config.yaml --port 4105docker exec up-redis-atomic-redis redis-cli monitor > /tmp/monitor-before.log &curl -s http://localhost:4105/v1/chat/completions -H "Authorization: Bearer sk-proof-1234" -H "Content-Type: application/json" -d '{"model":"proof-model","messages":[{"role":"user","content":"Reply with the single word ok"}],"max_tokens":5}'grep tpm /tmp/monitor-before.logshows the deployment usage counter written as three separate commands, each its own round tripdocker exec up-redis-atomic-redis redis-cli TTL "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:11-00"answers57, because nothing interrupted this requestCase 2: the client disconnects while the reply is in flight
uv run python /tmp/proof/cancel_probe.py before-9a715df212The increment reached Redis, the expiry never did: the counter now lives forever
After (fd41183)
Case 1: what the proxy sends to Redis for one chat completion
uv run python litellm/proxy/proxy_cli.py --config /tmp/proof/config.yaml --port 4105docker exec up-redis-atomic-redis redis-cli monitor > /tmp/monitor-after.log &Same curl as before
grep tpm /tmp/monitor-after.logshows the same counter written in one server side step (MONITOR also echoes the script's own INCRBYFLOAT and EXPIRE calls, tagged as the lua client, after the single EVAL)docker exec up-redis-atomic-redis redis-cli TTL "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:16-58"answers50(read ten seconds after the request), same happy path as beforeCase 2: the client disconnects while the reply is in flight
uv run python /tmp/proof/cancel_probe.py after-fd4118392fThe same cancellation now leaves a counter that expires in a minute
uv run python /tmp/proof/cancel_probe.py after-ttl0-fd4118392f 0confirms a zero TTL still deletes the key, as the two step version did with EXPIRE 0Type
🐛 Bug Fix
Caveats (if any)
Low
async_increment_pipelinestill queues INCRBYFLOAT and EXPIRE as two pipeline commands, but a pipeline ships both in one write so a cancellation cannot split them; left untouchedincrement_cachekeeps its INCR then EXPIRE shape; a sync call has no await to be cancelled atredis-cli monitorcapture also shows aglobal_router:counter written by the pipeline path, unchanged by this PRFinal Attestation