scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): use INCRBYFLOAT + DB floor for global proxy spend

Browse source 
This switches to the same pattern the spend counters already use:

- Write path: Redis INCRBYFLOAT creates the counter from zero when missing
  (so post-reset increments are represented immediately) and mirrors the
  authoritative Redis value into local memory.

- Read path: read the shared Redis scalar directly, then floor to the
  authoritative DB spend. This prevents a stale-low Redis value (from a reset
  race, Redis restart, or key expiry) from underestimating current-period spend
  and letting requests leak past the global cap.

- Reset job: deletes the shared key. With atomic increments and the DB floor,
  this neither resurrects stale values nor clobbers concurrent increments.

Removes the now-unnecessary async_increment_if_exists helper and its tests.
Adds regression tests for the Redis-direct read and DB-floor repair.
This commit is contained in:
lostmartian 2026-08-15 02:29:41 +05:30
parent 16c8a15939
commit 955203729b
6 changed files with 97 additions and 151 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/caching/redis_cache.py
View file

@ -1033,41 +1033,6 @@ class RedisCache(BaseCache):
result = result.decode()
return float(result)
@_redis_circuit_breaker_guard
async def async_increment_if_exists(
self,
key: str,
value: float,
ttl: float | None = None,
) -> float | None:
"""Atomically increment ``key`` by ``value`` only when the key already
exists, returning the new value; an absent key is left missing and
``None`` is returned so callers can fall back to loading the
authoritative total. The GET/increment/SET runs in a single Lua call,
so it is race-free across callers and pods, and refreshing the TTL on
a present key keeps the accumulator alive while it is being used.
"""
_redis_client: Final = self.init_async_client()
_used_ttl: Final = self.get_ttl(ttl=ttl)
key = self.check_and_fix_namespace(key=key)
lua: Final = (
"local cur = redis.call('GET', KEYS[1]) "
"if cur == false then return nil end "
"local new = tonumber(cur) + tonumber(ARGV[1]) "
"redis.call('SET', KEYS[1], new) "
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end "
"return new"
)
result = cast( # cast-ok: redis eval returns an untyped scalar
"str | bytes | int | float | None",
await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))),
)
if result is None:
return None
if isinstance(result, bytes):
result = result.decode()
return float(result)
async def flush_cache_buffer(self):
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

41
litellm/proxy/auth/user_api_key_auth.py
View file

@ -494,6 +494,10 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use
# Reusable coordinator for global spend to prevent cache stampede
_global_spend_coordinator: Final = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEND]")
# Cache the authoritative DB spend read for a few seconds per worker so the
# global-proxy budget floor check does not query the DB on every request.
_GLOBAL_PROXY_SPEND_DB_FLOOR_TTL_SECONDS: Final = 5
async def _fetch_global_spend_with_event_coordination(
cache_key: str,
@ -522,14 +526,49 @@ async def _fetch_global_spend_with_event_coordination(
)
return float(proxy_budget_row.spend) if proxy_budget_row is not None else None
async def _authoritative_db_spend() -> float | None:
# Cached per-worker for a few seconds to avoid a DB query per request.
marker_key: Final = f"global_proxy_spend_db_floor:{cache_key}"
cached = user_api_key_cache.in_memory_cache.get_cache(key=marker_key)
if cached is not None:
return float(cached)
db_spend = await _load_global_spend()
if db_spend is not None:
user_api_key_cache.in_memory_cache.set_cache(
key=marker_key,
value=db_spend,
ttl=_GLOBAL_PROXY_SPEND_DB_FLOOR_TTL_SECONDS,
)
return db_spend
spend_cache: Final = (
user_api_key_cache.redis_cache if user_api_key_cache.redis_cache is not None else user_api_key_cache
)
return await _global_spend_coordinator.get_or_load(
value = await _global_spend_coordinator.get_or_load(
cache_key=cache_key,
cache=spend_cache, # pyright: ignore[reportArgumentType]
load_fn=_load_global_spend,
)
if value is None:
return None
try:
numeric_value = float(value)
except (TypeError, ValueError):
numeric_value = 0.0
# Floor to the authoritative DB spend when Redis is in use. A Redis counter
# that was clobbered by a reset race, restarted from an old RDB snapshot,
# or recreated from zero after expiry can under-report current-period spend;
# capping the value to the DB row prevents requests from leaking past the
# global cap. Without Redis the coordinator already loaded the DB row.
if user_api_key_cache.redis_cache is not None:
db_spend = await _authoritative_db_spend()
if db_spend is not None and db_spend > numeric_value:
await spend_cache.async_set_cache(key=cache_key, value=db_spend)
numeric_value = db_spend
return numeric_value
async def get_global_proxy_spend(

24
litellm/proxy/proxy_server.py
View file

@ -3126,25 +3126,21 @@ async def update_cache(
if response_cost is None:
return
try:
# Atomically accumulate the shared global-proxy scalar without the
# read-modify-write race with ResetBudgetJob. A plain INCRBYFLOAT
# would recreate a missing key from zero, so increment-if-present
# leaves an absent key missing and lets the auth-time loader
# restore the authoritative DB spend instead.
# Atomically accumulate the shared global-proxy scalar in Redis.
# The returned value is authoritative across workers, so mirror it
# into this worker's in-memory tier to avoid accumulating a stale
# pre-reset local value.
if user_api_key_cache.redis_cache is not None:
current = await user_api_key_cache.redis_cache.async_increment_if_exists(
current = await user_api_key_cache.redis_cache.async_increment(
key=GLOBAL_PROXY_SPEND_CACHE_KEY,
value=response_cost,
ttl=get_management_object_ttl(user_api_key_cache),
refresh_ttl=True,
)
await user_api_key_cache.in_memory_cache.async_set_cache(
key=GLOBAL_PROXY_SPEND_CACHE_KEY,
value=current,
)
if current is not None:
# Mirror the authoritative Redis value into this worker's
# in-memory tier so a stale pre-reset value is never
# accumulated or read back.
await user_api_key_cache.in_memory_cache.async_set_cache(
key=GLOBAL_PROXY_SPEND_CACHE_KEY,
value=current,
)
else:
# No Redis: single-process, so an in-memory increment is coherent.
await user_api_key_cache.async_increment_cache(

60
tests/test_litellm/caching/test_redis_cache.py
View file

@ -435,66 +435,6 @@ def test_increment_cache_namespaces_key(
mock_client.incr.assert_called_once_with(name=expected, amount=1)
@pytest.mark.asyncio
async def test_async_increment_if_exists_increments_present_key(
monkeypatch, redis_no_ping
):
"""A present key is atomically incremented and its TTL refreshed."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=None)
mock_redis_instance = AsyncMock()
mock_redis_instance.eval = AsyncMock(return_value="5.0")
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
result = await redis_cache.async_increment_if_exists(key="k", value=2.0, ttl=60)
assert result == 5.0
lua, num_keys, key, value, ttl = mock_redis_instance.eval.await_args.args
assert num_keys == 1
assert key == "k"
assert value == "2.0"
assert ttl == "60"
assert "if cur == false then return nil end" in lua
@pytest.mark.asyncio
async def test_async_increment_if_exists_decodes_bytes_result(
monkeypatch, redis_no_ping
):
"""Redis returns bytes; the method decodes them to a float."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=None)
mock_redis_instance = AsyncMock()
mock_redis_instance.eval = AsyncMock(return_value=b"7.0")
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
result = await redis_cache.async_increment_if_exists(key="k", value=2.0)
assert result == 7.0
@pytest.mark.asyncio
async def test_async_increment_if_exists_leaves_missing_key_absent(
monkeypatch, redis_no_ping
):
"""An absent key stays missing: the method returns None and does not create it."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache(namespace=None)
mock_redis_instance = AsyncMock()
mock_redis_instance.eval = AsyncMock(return_value=None)
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
result = await redis_cache.async_increment_if_exists(key="k", value=2.0)
assert result is None
@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")])
def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")

42
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -6134,11 +6134,14 @@ async def test_global_proxy_spend_reads_shared_redis_scalar_when_configured():
cache = UserApiKeyCache()
redis_cache = MagicMock()
redis_cache.async_get_cache = AsyncMock(return_value=8.5)
redis_cache.async_set_cache = AsyncMock()
cache.redis_cache = redis_cache
proxy_budget_row = MagicMock()
proxy_budget_row.spend = 5.0
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
side_effect=AssertionError("a Redis hit must not hit the DB")
return_value=proxy_budget_row
)
result = await _fetch_global_spend_with_event_coordination(
@ -6149,6 +6152,43 @@ async def test_global_proxy_spend_reads_shared_redis_scalar_when_configured():
assert result == 8.5
redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend")
redis_cache.async_set_cache.assert_not_called()
@pytest.mark.asyncio
async def test_global_proxy_spend_floors_stale_low_redis_to_db():
"""A Redis counter that is stale-low (e.g., recreated from zero after a
reset or Redis restart) must be floored to the authoritative DB spend so
requests cannot leak past the global cap."""
from litellm.proxy.auth.user_api_key_auth import (
_fetch_global_spend_with_event_coordination,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
cache = UserApiKeyCache()
redis_cache = MagicMock()
redis_cache.async_get_cache = AsyncMock(return_value="2.0")
redis_cache.async_set_cache = AsyncMock()
cache.redis_cache = redis_cache
proxy_budget_row = MagicMock()
proxy_budget_row.spend = 9.0
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=proxy_budget_row
)
result = await _fetch_global_spend_with_event_coordination(
cache_key="default_user_id:spend",
user_api_key_cache=cache,
prisma_client=prisma_client,
)
assert result == 9.0
redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend")
redis_cache.async_set_cache.assert_awaited_once_with(
key="default_user_id:spend", value=9.0
)
@pytest.mark.asyncio

46
tests/test_litellm/proxy/proxy_server/test_spend_counters.py
View file

@ -1334,40 +1334,6 @@ async def test_update_cache_global_proxy_spend_concurrent_increments_after_reset
assert final_value == pytest.approx(0.5)
@pytest.mark.asyncio
async def test_update_cache_global_proxy_spend_leaves_missing_key_absent(monkeypatch):
"""An absent global spend key must not be recreated from zero by an
increment; the auth-time loader restores the authoritative DB spend."""
from unittest.mock import AsyncMock, MagicMock
from litellm.caching.caching import DualCache
redis_cache = MagicMock()
redis_cache.async_increment_if_exists = AsyncMock(return_value=None)
cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300)
monkeypatch.setattr(ps, "user_api_key_cache", cache)
await cache.async_set_cache(
key="user-missing",
value={"user_id": "user-missing", "spend": 0.0},
)
await ps.update_cache(
token=None,
user_id="user-missing",
end_user_id=None,
team_id=None,
response_cost=1.0,
parent_otel_span=None,
)
redis_cache.async_increment_if_exists.assert_awaited_once_with(
key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY,
value=1.0,
ttl=ps.get_management_object_ttl(cache),
)
assert await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) is None
@pytest.mark.asyncio
async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monkeypatch):
"""In-memory mirrors the authoritative Redis result instead of accruing a
@ -1377,7 +1343,7 @@ async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monke
from litellm.caching.caching import DualCache
redis_cache = MagicMock()
redis_cache.async_increment_if_exists = AsyncMock(return_value=3.5)
redis_cache.async_increment = AsyncMock(return_value=3.5)
cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300)
monkeypatch.setattr(ps, "user_api_key_cache", cache)
@ -1394,10 +1360,11 @@ async def test_update_cache_global_proxy_spend_mirrors_authoritative_value(monke
parent_otel_span=None,
)
redis_cache.async_increment_if_exists.assert_awaited_once_with(
redis_cache.async_increment.assert_awaited_once_with(
key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY,
value=1.5,
ttl=ps.get_management_object_ttl(cache),
refresh_ttl=True,
)
assert await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) == 3.5
@ -1410,9 +1377,7 @@ async def test_update_cache_global_proxy_spend_swallows_redis_failure(monkeypatc
from litellm.caching.caching import DualCache
redis_cache = MagicMock()
redis_cache.async_increment_if_exists = AsyncMock(
side_effect=RuntimeError("redis down")
)
redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down"))
cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300)
monkeypatch.setattr(ps, "user_api_key_cache", cache)
@ -1429,8 +1394,9 @@ async def test_update_cache_global_proxy_spend_swallows_redis_failure(monkeypatc
parent_otel_span=None,
)
redis_cache.async_increment_if_exists.assert_awaited_once_with(
redis_cache.async_increment.assert_awaited_once_with(
key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY,
value=1.0,
ttl=ps.get_management_object_ttl(cache),
refresh_ttl=True,
)