diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 079cbd163dc..222c5547b37 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -8,12 +8,15 @@ deployments. This module reseeds from the authoritative DB instead. A per-counter singleflight lock collapses concurrent reseeds on the same pod -to one DB query per cold-cache window. The lock dict is bounded LRU to cap -memory in long-lived deployments. +to one DB query per cold-cache window. The lock registry retains active locks +and bounds idle entries with LRU eviction in long-lived deployments. """ import asyncio from collections import OrderedDict +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, ClassVar, Optional @@ -36,6 +39,12 @@ from litellm.proxy.utils import PrismaClient +@dataclass +class _CounterLock: + lock: asyncio.Lock + users: int = 0 + + class SpendCounterReseed: """ Reseeds spend counters from the authoritative DB and warms the cache, @@ -53,23 +62,60 @@ class SpendCounterReseed: and get_tag_objects_batch(); callers pass those values as fallback_spend. """ - _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() + _locks: ClassVar["OrderedDict[str, _CounterLock]"] = OrderedDict() _registry_lock: ClassVar[Optional[asyncio.Lock]] = None @staticmethod - async def _get_lock(counter_key: str) -> asyncio.Lock: + def _get_registry_lock() -> asyncio.Lock: if SpendCounterReseed._registry_lock is None: SpendCounterReseed._registry_lock = asyncio.Lock() - async with SpendCounterReseed._registry_lock: - lock = SpendCounterReseed._locks.get(counter_key) - if lock is not None: - SpendCounterReseed._locks.move_to_end(counter_key) - return lock - lock = asyncio.Lock() - SpendCounterReseed._locks[counter_key] = lock - if len(SpendCounterReseed._locks) > SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: - SpendCounterReseed._locks.popitem(last=False) - return lock + return SpendCounterReseed._registry_lock + + @staticmethod + def _get_or_create_lock(counter_key: str) -> _CounterLock: + counter_lock = SpendCounterReseed._locks.get(counter_key) + if counter_lock is None: + counter_lock = _CounterLock(lock=asyncio.Lock()) + SpendCounterReseed._locks[counter_key] = counter_lock + SpendCounterReseed._locks.move_to_end(counter_key) + return counter_lock + + @staticmethod + def _prune_idle_locks() -> None: + while len(SpendCounterReseed._locks) > SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: + idle_key = next( + ( + key + for key, counter_lock in SpendCounterReseed._locks.items() + if counter_lock.users == 0 and not counter_lock.lock.locked() + ), + None, + ) + if idle_key is None: + return + SpendCounterReseed._locks.pop(idle_key) + + @staticmethod + async def _get_lock(counter_key: str) -> asyncio.Lock: + async with SpendCounterReseed._get_registry_lock(): + counter_lock = SpendCounterReseed._get_or_create_lock(counter_key) + SpendCounterReseed._prune_idle_locks() + return counter_lock.lock + + @staticmethod + @asynccontextmanager + async def _counter_lock(counter_key: str) -> AsyncIterator[None]: + async with SpendCounterReseed._get_registry_lock(): + counter_lock = SpendCounterReseed._get_or_create_lock(counter_key) + counter_lock.users += 1 + try: + lock = await SpendCounterReseed._get_lock(counter_key) + async with lock: + yield + finally: + async with SpendCounterReseed._get_registry_lock(): + counter_lock.users -= 1 + SpendCounterReseed._prune_idle_locks() @staticmethod async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> Optional[float]: @@ -152,8 +198,7 @@ async def coalesced( Returns the spend value (including 0.0 from a fresh budget reset) when the DB read succeeds, or None when the DB is unavailable. """ - lock = await SpendCounterReseed._get_lock(counter_key) - async with lock: + async with SpendCounterReseed._counter_lock(counter_key): # Re-check after acquiring the lock. Skip in-memory on a clean # Redis miss - in-memory is per-pod-stale. redis_clean_miss = False @@ -197,7 +242,15 @@ async def coalesced( value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True) + cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if cached is not None: + current_value = float(cached) + else: + current_value = await spend_counter_cache.async_increment_cache( + key=counter_key, + value=db_spend, + refresh_ttl=True, + ) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -262,8 +315,7 @@ async def coalesced_window( entity_id: str, window_start: datetime, ) -> Optional[float]: - lock = await SpendCounterReseed._get_lock(counter_key) - async with lock: + async with SpendCounterReseed._counter_lock(counter_key): redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 18a927e7a44..a49e2df1656 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2248,24 +2248,25 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None in-memory copy is guarded by a read-compare-write with no await in between, so it is atomic within the worker. """ - cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - needs_update = True - if cached is not None: - try: - needs_update = float(cached) < db_spend - except (TypeError, ValueError): - needs_update = True - if needs_update: - spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) - if spend_counter_cache.redis_cache is not None: - try: - await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend) - except Exception: - verbose_proxy_logger.debug( - "Unable to repair stale spend counter %s in Redis", - counter_key, - exc_info=True, - ) + async with SpendCounterReseed._counter_lock(counter_key): + cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + needs_update = True + if cached is not None: + try: + needs_update = float(cached) < db_spend + except (TypeError, ValueError): + needs_update = True + if needs_update: + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend) + except Exception: + verbose_proxy_logger.debug( + "Unable to repair stale spend counter %s in Redis", + counter_key, + exc_info=True, + ) async def reseed_spend_counter_from_db(counter_key: str) -> None: @@ -2638,11 +2639,6 @@ async def _init_and_increment_spend_counter( 2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`. Falls back to the cached object's `.spend` via user_api_key_cache only if prisma is unavailable, since that value can lag the flusher. - 3. Seed counter via async_increment_cache (not async_set_cache) to avoid a - check-then-set race: if two pods cold-start simultaneously, both may see - the counter as absent and seed it. Using increment means the worst case - is over-counting (conservative, blocks slightly early) rather than - under-counting (would allow overspend). 4. Increment atomically (both in-memory + Redis) """ await _ensure_spend_counter_initialized( @@ -2778,11 +2774,12 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float): ) return current_value - return await spend_counter_cache.async_increment_cache( - key=counter_key, - value=increment, - refresh_ttl=True, - ) + async with SpendCounterReseed._counter_lock(counter_key): + return await spend_counter_cache.async_increment_cache( + key=counter_key, + value=increment, + refresh_ttl=True, + ) async def _invalidate_spend_counter(counter_key: str): diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 51980342a1d..8efc679957f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -1179,6 +1179,28 @@ async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( assert fake_cache.redis_cache.async_delete_cache.called is True +@pytest.mark.asyncio +async def test_increment_spend_counter_cache_no_redis_uses_in_memory_increment( + monkeypatch, +): + fake_cache = _make_spend_counter_cache( + redis_increment_value=44.0, + with_redis=False, + ) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + + result = await ps._increment_spend_counter_cache( + counter_key="spend:key:k", increment=4.0 + ) + + assert result == 44.0 + fake_cache.async_increment_cache.assert_awaited_once_with( + key="spend:key:k", + value=4.0, + refresh_ttl=True, + ) + + # --------------------------------------------------------------------------- # _invalidate_spend_counter # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 62f5ced7a39..a2eec2be566 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8053,6 +8053,189 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): ps.prisma_client = orig_prisma +@pytest.mark.asyncio +async def test_cold_reseed_serializes_with_direct_counter_repair(): + """A repair cannot write between a cold reseed's DB read and seed.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + counter_key = "spend:team:team-repair-race" + db_spend = 100.0 + db_read_started = asyncio.Event() + release_db_read = asyncio.Event() + + async def from_db(*_args, **_kwargs): + db_read_started.set() + await release_db_read.wait() + return db_spend + + import litellm.proxy.proxy_server as ps + + original_counter_cache = ps.spend_counter_cache + original_prisma_client = ps.prisma_client + ps.spend_counter_cache = DualCache() + ps.prisma_client = MagicMock() + try: + with patch.object(SpendCounterReseed, "from_db", side_effect=from_db): + cold_reseed = asyncio.create_task( + SpendCounterReseed.coalesced( + prisma_client=ps.prisma_client, + spend_counter_cache=ps.spend_counter_cache, + counter_key=counter_key, + ) + ) + await db_read_started.wait() + repair = asyncio.create_task( + ps._repair_stale_spend_counter( + counter_key=counter_key, + db_spend=db_spend, + ) + ) + await asyncio.sleep(0) + release_db_read.set() + await asyncio.gather(cold_reseed, repair) + + assert ps.spend_counter_cache.in_memory_cache.get_cache( + key=counter_key + ) == pytest.approx(db_spend) + finally: + ps.spend_counter_cache = original_counter_cache + ps.prisma_client = original_prisma_client + + +@pytest.mark.asyncio +async def test_cold_reseed_serializes_with_reservation_reconciliation_and_increment(): + """Reconciliation reseeds and request increments retain the DB floor.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + counter_key = "spend:team:team-reservation-race" + db_spend = 100.0 + request_cost = 3.5 + db_read_started = asyncio.Event() + reconciliation_db_read = asyncio.Event() + release_db_read = asyncio.Event() + db_read_count = 0 + + async def from_db(*_args, **_kwargs): + nonlocal db_read_count + db_read_count += 1 + if db_read_count == 1: + db_read_started.set() + await release_db_read.wait() + else: + reconciliation_db_read.set() + return db_spend + + budget_reservation = { + "reserved_cost": 1.0, + "entries": [ + { + "counter_key": counter_key, + "reserved_cost": 1.0, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + import litellm.proxy.proxy_server as ps + + original_counter_cache = ps.spend_counter_cache + original_prisma_client = ps.prisma_client + ps.spend_counter_cache = DualCache() + ps.prisma_client = MagicMock() + try: + with patch.object(SpendCounterReseed, "from_db", side_effect=from_db): + cold_reseed = asyncio.create_task( + SpendCounterReseed.coalesced( + prisma_client=ps.prisma_client, + spend_counter_cache=ps.spend_counter_cache, + counter_key=counter_key, + ) + ) + await db_read_started.wait() + reconciliation = asyncio.create_task( + ps.increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.5, + budget_reservation=budget_reservation, + ) + ) + await reconciliation_db_read.wait() + request_increment = asyncio.create_task( + ps._increment_spend_counter_cache( + counter_key=counter_key, + increment=request_cost, + ) + ) + await asyncio.sleep(0) + release_db_read.set() + await asyncio.gather(cold_reseed, reconciliation, request_increment) + + assert budget_reservation["finalized"] is True + assert ps.spend_counter_cache.in_memory_cache.get_cache( + key=counter_key + ) == pytest.approx(db_spend + request_cost) + finally: + ps.spend_counter_cache = original_counter_cache + ps.prisma_client = original_prisma_client + + +@pytest.mark.asyncio +async def test_request_increment_waits_for_cold_reseed_and_preserves_db_floor(): + """A request increment after a cold counter miss starts from DB spend.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + counter_key = "spend:team:team-increment-race" + db_spend = 100.0 + request_cost = 3.5 + db_read_started = asyncio.Event() + release_db_read = asyncio.Event() + + async def from_db(*_args, **_kwargs): + db_read_started.set() + await release_db_read.wait() + return db_spend + + import litellm.proxy.proxy_server as ps + + original_counter_cache = ps.spend_counter_cache + original_prisma_client = ps.prisma_client + ps.spend_counter_cache = DualCache() + ps.prisma_client = MagicMock() + try: + with patch.object(SpendCounterReseed, "from_db", side_effect=from_db): + cold_reseed = asyncio.create_task( + SpendCounterReseed.coalesced( + prisma_client=ps.prisma_client, + spend_counter_cache=ps.spend_counter_cache, + counter_key=counter_key, + ) + ) + await db_read_started.wait() + request_increment = asyncio.create_task( + ps._increment_spend_counter_cache( + counter_key=counter_key, + increment=request_cost, + ) + ) + await asyncio.sleep(0) + assert request_increment.done() is False + release_db_read.set() + await asyncio.gather(cold_reseed, request_increment) + + assert ps.spend_counter_cache.in_memory_cache.get_cache( + key=counter_key + ) == pytest.approx(db_spend + request_cost) + finally: + ps.spend_counter_cache = original_counter_cache + ps.prisma_client = original_prisma_client + + @pytest.mark.asyncio async def test_concurrent_read_and_write_paths_share_one_db_query(): """ @@ -8182,6 +8365,91 @@ async def test_reseed_locks_dict_is_bounded(): SpendCounterReseed._locks.update(orig_locks) +@pytest.mark.asyncio +async def test_active_reseed_lock_survives_registry_pressure(monkeypatch): + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + import litellm.proxy.db.spend_counter_reseed as scr + + target_key = "spend:key:active-lock" + original_locks = SpendCounterReseed._locks.copy() + SpendCounterReseed._locks.clear() + monkeypatch.setattr(scr, "SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 2) + holder_entered = asyncio.Event() + release_holder = asyncio.Event() + second_entered = asyncio.Event() + + async def hold_target_lock(): + async with SpendCounterReseed._counter_lock(target_key): + holder_entered.set() + await release_holder.wait() + + async def acquire_target_lock(): + async with SpendCounterReseed._counter_lock(target_key): + second_entered.set() + + holder = asyncio.create_task(hold_target_lock()) + second = None + try: + await holder_entered.wait() + first_lock = await SpendCounterReseed._get_lock(target_key) + for index in range(3): + await SpendCounterReseed._get_lock(f"spend:key:pressure-{index}") + second_lock = await SpendCounterReseed._get_lock(target_key) + + assert second_lock is first_lock + second = asyncio.create_task(acquire_target_lock()) + await asyncio.sleep(0) + assert second_entered.is_set() is False + finally: + release_holder.set() + await holder + if second is not None: + await second + SpendCounterReseed._locks.clear() + SpendCounterReseed._locks.update(original_locks) + + assert second_entered.is_set() is True + + +@pytest.mark.asyncio +async def test_reseed_registry_exceeds_limit_only_while_locks_are_active(monkeypatch): + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + import litellm.proxy.db.spend_counter_reseed as scr + + original_locks = SpendCounterReseed._locks.copy() + SpendCounterReseed._locks.clear() + monkeypatch.setattr(scr, "SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 1) + first_entered = asyncio.Event() + second_entered = asyncio.Event() + release_locks = asyncio.Event() + remaining_lock_count: int | None = None + + async def hold_lock(counter_key: str, entered: asyncio.Event): + async with SpendCounterReseed._counter_lock(counter_key): + entered.set() + await release_locks.wait() + + first = asyncio.create_task(hold_lock("spend:key:first", first_entered)) + second = None + try: + await first_entered.wait() + second = asyncio.create_task(hold_lock("spend:key:second", second_entered)) + await second_entered.wait() + + assert len(SpendCounterReseed._locks) == 2 + finally: + release_locks.set() + await first + if second is not None: + await second + remaining_lock_count = len(SpendCounterReseed._locks) + SpendCounterReseed._locks.clear() + SpendCounterReseed._locks.update(original_locks) + + assert remaining_lock_count is not None + assert remaining_lock_count <= 1 + + @pytest.mark.asyncio async def test_reseed_warms_cache_even_on_zero_db_spend(): """