Skip to content

fix(caching): set the counter TTL atomically in async_increment - #40715

Closed
songkuan-zheng wants to merge 4 commits into
BerriAI:litellm_internal_stagingfrom
GhishaDev:litellm_redis_atomic_increment_ttl
Closed

songkuan-zheng wants to merge 4 commits into
BerriAI:litellm_internal_stagingfrom
GhishaDev:litellm_redis_atomic_increment_ttl

Conversation

@songkuan-zheng

@songkuan-zheng songkuan-zheng commented Sep 11, 2026 •

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • RedisCache.async_increment sends INCRBYFLOAT, then EXPIRE, as two separate awaits
  • A cancellation between them leaves the counter in Redis with no TTL
  • Cancelled requests are routine (client disconnects), so counters slowly accumulate forever

How it solves it:

  • Increment and TTL decision run in one Lua call, one Redis round trip
  • Same pattern this file already uses for async_increment_with_floor and async_set_max
  • refresh_ttl semantics unchanged: re-arm always, or only arm a missing TTL

User Flow

Before: a proxy operator who runs several replicas with Redis coordination sees rate limit and spend counters that never expire

  1. They set router_settings.routing_strategy: usage-based-routing-v2 with redis_host / redis_port and restart the proxy
  2. A developer sends POST https://litellm-domain/v1/chat/completions and drops the connection while Redis is slow to answer (a mobile client, a load balancer timeout, a cancelled agent run)
  3. The per-minute usage counter for that deployment was incremented in Redis but never received its expiry, so TTL on it answers -1
  4. The operator runs redis-cli --scan --pattern '*:tpm:*' weeks later and finds thousands of stale per-minute keys that memory eviction has to clean up

After: the same disconnect leaves a counter that expires on schedule

  1. They set the same router_settings and restart the proxy on the fixed build
  2. The developer sends the same POST https://litellm-domain/v1/chat/completions and drops the connection at the same moment
  3. TTL on the usage counter answers 60, the counter disappears one minute later like every other one
  4. redis-cli --scan --pattern '*:tpm:*' only ever shows the current minute's keys

Relevant issues

No existing issue. RedisCache.async_increment (litellm/caching/redis_cache.py) awaits incrbyfloat and then awaits expire (or ttl followed by expire). 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 through DualCache.async_increment_cache with 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 counter

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • The handful of test files covering my change pass locally, e.g. 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
  • My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to 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 monitor attached, showing what the proxy actually sends to Redis for one real chat completion. Case 2 reproduces the disconnect with the real RedisCache class 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 latency

Shared setup

docker run -d --name up-redis-atomic-redis -p 6395:6379 redis:7

/tmp/proof/config.yaml (the model is a real OpenAI compatible deployment, key and base come from the environment)

model_list:
  - model_name: proof-model
    litellm_params:
      model: openai/openai-a/gpt-5.6-luna
      api_key: os.environ/OPENAI_API_KEY
      api_base: os.environ/OPENAI_API_BASE
      rpm: 1000
router_settings:
  routing_strategy: usage-based-routing-v2
  redis_host: localhost
  redis_port: 6395
general_settings:
  master_key: sk-proof-1234

/tmp/proof/cancel_probe.py

import asyncio
import sys

from redis.asyncio import Redis

from litellm.caching.redis_cache import RedisCache

REDIS_PORT = 6395
RELAY_PORT = 6396
REPLY_DELAY = 0.3


async def relay(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    up_reader, up_writer = await asyncio.open_connection("localhost", REDIS_PORT)

    async def client_to_server() -> None:
        while data := await reader.read(65536):
            up_writer.write(data)
            await up_writer.drain()

    async def server_to_client() -> None:
        while data := await up_reader.read(65536):
            await asyncio.sleep(REPLY_DELAY)
            writer.write(data)
            await writer.drain()

    await asyncio.gather(client_to_server(), server_to_client(), return_exceptions=True)


async def main() -> None:
    server = await asyncio.start_server(relay, "localhost", RELAY_PORT)
    cache = RedisCache(host="localhost", port=RELAY_PORT, namespace="probe")
    key = f"spend:key:{sys.argv[1]}"
    await cache.init_async_client().ping()
    ttl = int(sys.argv[2]) if len(sys.argv) > 2 else 60
    task = asyncio.ensure_future(cache.async_increment(key=key, value=0.25, ttl=ttl, refresh_ttl=True))
    await asyncio.sleep(0.15)
    task.cancel()
    try:
        await task
        print("increment returned", task.result())
    except asyncio.CancelledError:
        print("client cancelled the increment while waiting for the Redis reply")
    await asyncio.sleep(0.5)
    admin = Redis(host="localhost", port=REDIS_PORT)
    print("GET", key, "=", await admin.get(f"probe:{key}"))
    print("TTL", key, "=", await admin.ttl(f"probe:{key}"))
    await admin.aclose()
    server.close()


asyncio.run(main())

Before (9a715df)

Case 1: what the proxy sends to Redis for one chat completion

  1. uv run python litellm/proxy/proxy_cli.py --config /tmp/proof/config.yaml --port 4105

  2. docker exec up-redis-atomic-redis redis-cli monitor > /tmp/monitor-before.log &

  3. 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}'

    {"id":"chatcmpl-EMtHv8SdDa2i3sqCvTOL0SYBJJOxi","created":1789124451,"model":"proof-model","object":"chat.completion","choices":[{"finish_reason":"stop","index":0,"message":{"content":"ok","role":"assistant",...}}],"usage":{"completion_tokens":4,"prompt_tokens":12,"total_tokens":16,...}}
    
  4. grep tpm /tmp/monitor-before.log shows the deployment usage counter written as three separate commands, each its own round trip

    "INCRBYFLOAT" "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:11-00" "16"
    "TTL" "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:11-00"
    "EXPIRE" "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:11-00" "60"
    
  5. docker exec up-redis-atomic-redis redis-cli TTL "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:11-00" answers 57, because nothing interrupted this request

Case 2: the client disconnects while the reply is in flight

  1. uv run python /tmp/proof/cancel_probe.py before-9a715df212

    client cancelled the increment while waiting for the Redis reply
    GET spend:key:before-9a715df212 = b'0.25'
    TTL spend:key:before-9a715df212 = -1
    
  2. The 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

  1. uv run python litellm/proxy/proxy_cli.py --config /tmp/proof/config.yaml --port 4105

  2. docker exec up-redis-atomic-redis redis-cli monitor > /tmp/monitor-after.log &

  3. Same curl as before

    {"id":"chatcmpl-ENLLr3fn9vzo8XsYmSphxnvN6nDCb","created":1789127609,"model":"proof-model","object":"chat.completion","choices":[{"finish_reason":"stop","index":0,"message":{"content":"ok","role":"assistant",...}}],"usage":{"completion_tokens":4,"prompt_tokens":12,"total_tokens":16,...}}
    
  4. grep tpm /tmp/monitor-after.log shows 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)

    "EVAL" "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]) if ARGV[2] ~= '' and (ARGV[3] == '1' or redis.call('TTL', KEYS[1]) == -1) then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2])) end return value" "1" "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:16-58" "16" "60" "0"
    
  5. docker exec up-redis-atomic-redis redis-cli TTL "01666b58...:openai/openai-a/gpt-5.6-luna:tpm:16-58" answers 50 (read ten seconds after the request), same happy path as before

Case 2: the client disconnects while the reply is in flight

  1. uv run python /tmp/proof/cancel_probe.py after-fd4118392f

    client cancelled the increment while waiting for the Redis reply
    GET spend:key:after-fd4118392f = b'0.25'
    TTL spend:key:after-fd4118392f = 59
    
  2. The same cancellation now leaves a counter that expires in a minute

  3. uv run python /tmp/proof/cancel_probe.py after-ttl0-fd4118392f 0 confirms a zero TTL still deletes the key, as the two step version did with EXPIRE 0

    client cancelled the increment while waiting for the Redis reply
    GET spend:key:after-ttl0-fd4118392f = None
    TTL spend:key:after-ttl0-fd4118392f = -2
    

Type

🐛 Bug Fix

Caveats (if any)

Low

  • async_increment_pipeline still queues INCRBYFLOAT and EXPIRE as two pipeline commands, but a pipeline ships both in one write so a cancellation cannot split them; left untouched
  • The synchronous increment_cache keeps its INCR then EXPIRE shape; a sync call has no await to be cancelled at
  • The proxy redis-cli monitor capture also shows a global_router: counter written by the pipeline path, unchanged by this PR

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

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>
@songkuan-zheng
songkuan-zheng requested a review from a team September 11, 2026 11:35
@codspeed

codspeed Bot commented Sep 11, 2026 •

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing GhishaDev:litellm_redis_atomic_increment_ttl (39a95d6) with litellm_internal_staging (426e675)

Open in CodSpeed

@greptile-apps

greptile-apps Bot commented Sep 11, 2026 •

Copy link
Copy Markdown
Contributor

Greptile Summary

This 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.

  • Adds a Lua script that increments the counter and either refreshes or conditionally initializes its TTL.
  • Preserves None, zero, and refresh_ttl behavior while validating the Lua result as a float.
  • Adds boundary tests covering namespaced keys, TTL forwarding, refresh behavior, return conversion, and single-round-trip execution.

Confidence Score: 5/5

The 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.

Important Files Changed

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

Comment thread litellm/caching/redis_cache.py Outdated
Comment thread tests/test_litellm/caching/test_redis_cache.py Outdated
@codecov

codecov Bot commented Sep 11, 2026

Copy link
Copy Markdown

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>
@songkuan-zheng

Copy link
Copy Markdown
Contributor Author

@greptileai

@songkuan-zheng

Copy link
Copy Markdown
Contributor Author

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>
@songkuan-zheng

Copy link
Copy Markdown
Contributor Author

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

@greptileai

@songkuan-zheng

Copy link
Copy Markdown
Contributor Author

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>
@songkuan-zheng

Copy link
Copy Markdown
Contributor Author

Merged litellm_internal_staging again (39a95d6) to clear a test-file conflict with the new pool-timeout breaker tests; production code is unchanged from fd41183, all 84 tests in the file pass

@greptileai

@yuneng-berri
yuneng-berri deleted the branch BerriAI:litellm_internal_staging September 13, 2026 04:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants