diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index c6060a47f84..0131a67db8b 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -104,6 +104,15 @@ class SpendCounterReseed: SpendCounterReseed._locks.popitem(last=False) return lock + @staticmethod + async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None: + """Apply local deltas after an in-flight reseed establishes the spend balance.""" + lock: Final = await SpendCounterReseed._get_lock(counter_key) + async with lock: + return await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, local_only=True, refresh_ttl=True + ) + @staticmethod async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: """ @@ -245,16 +254,10 @@ class SpendCounterReseed: value=current_value, ) else: - # Repair/reservations can populate the counter during the DB read. - # Seed a floor without adding the database balance again. - # No await between read/compare/write: atomic within this worker. - cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - current_value = float(db_spend) - if cached is not None: - current_value = max(current_value, float(cached)) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=current_value - ) + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", @@ -447,16 +450,12 @@ class SpendCounterReseed: value=current_value, ) else: - # Repair/reservations can populate the counter during the DB read. - # Seed a floor without adding the database balance again. - # No await between read/compare/write: atomic within this worker. - cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - current_value = float(window_spend) - if cached is not None: - current_value = max(current_value, float(cached)) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=current_value + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = ( + max(window_spend, float(cached_spend)) if cached_spend is not None else window_spend ) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + return seeded_spend except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced_window: failed to warm counter %s", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9269fd48e6c..e8608661c48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3330,10 +3330,8 @@ 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, + return await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment ) @@ -3356,10 +3354,8 @@ async def _apply_spend_counter_increments(pending: Sequence[_PendingSpendIncreme redis_cache: Final = spend_counter_cache.redis_cache if redis_cache is None: for item in pending: - await spend_counter_cache.async_increment_cache( - key=item.counter_key, - value=item.increment, - refresh_ttl=True, + await SpendCounterReseed.increment_in_memory( + spend_counter_cache=spend_counter_cache, counter_key=item.counter_key, increment=item.increment ) return ttl: Final = redis_cache.get_ttl() diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index a0cf9128cc9..69ec174903b 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -8,6 +8,7 @@ allowed to run: only when the row is missing or belongs to an older window. from __future__ import annotations import asyncio +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final @@ -79,6 +80,39 @@ def _row(window_start: datetime, spend: float) -> SimpleNamespace: return SimpleNamespace(window_start=window_start, spend=spend) +class _PausedSpendTable: + def __init__(self, spend: float) -> None: + self.spend: Final = spend + self.read_started: Final = asyncio.Event() + self.resume_read: Final = asyncio.Event() + + async def find_unique(self, where: Mapping[str, object]) -> SimpleNamespace: + self.read_started.set() + await self.resume_read.wait() + return _row(WINDOW_START, self.spend) + + +async def _reseed_with_paused_table( + table: _PausedSpendTable, cache: DualCache, counter_key: str, window: bool +) -> float | None: + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_usertable=table, litellm_budgetwindowspend=table)) + if window: + return await SpendCounterReseed.coalesced_window( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + entity_type="Team", + entity_id="team-1", + window_duration="1d", + window_start=WINDOW_START, + ) + return await SpendCounterReseed.coalesced( + prisma_client=prisma, + spend_counter_cache=cache, + counter_key=counter_key, + ) + + @pytest.mark.asyncio async def test_window_from_table_reads_row_by_primary_key(): """The lookup must use the table's own entity_type values ("key"), not the @@ -275,49 +309,59 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): @pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) @pytest.mark.parametrize("concurrent_spend", [989.01459411, 995.0, 900.0]) async def test_cold_reseed_does_not_add_database_spend_to_concurrent_cache( - monkeypatch: pytest.MonkeyPatch, window: bool, concurrent_spend: float, -): - """A repair/reservation write may populate the counter while DB read runs. +) -> None: + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = "spend:team:team-1:window:1d" if window else "spend:user:user-1" + db_spend: Final = 989.01459411 + table: Final = _PausedSpendTable(db_spend) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) - Reseeding must establish the larger value, not increment the concurrent - value by the same authoritative spend a second time. - """ - cache = DualCache(in_memory_cache=InMemoryCache()) - counter_key = ( - "spend:team:team-1:window:1d" if window else "spend:user:user-1" - ) - db_spend = 989.01459411 + await asyncio.wait_for(table.read_started.wait(), timeout=5) + cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) + table.resume_read.set() + result: Final = await asyncio.wait_for(reseed_task, timeout=5) - async def read_db(*args, **kwargs): - cache.in_memory_cache.set_cache(key=counter_key, value=concurrent_spend) - return db_spend - - if window: - monkeypatch.setattr(SpendCounterReseed, "window_from_db", staticmethod(read_db)) - result = await SpendCounterReseed.coalesced_window( - prisma_client=None, - spend_counter_cache=cache, - counter_key=counter_key, - entity_type="Team", - entity_id="team-1", - window_duration="1d", - window_start=WINDOW_START, - ) - else: - monkeypatch.setattr(SpendCounterReseed, "from_db", staticmethod(read_db)) - result = await SpendCounterReseed.coalesced( - prisma_client=None, - spend_counter_cache=cache, - counter_key=counter_key, - ) - - expected = max(db_spend, concurrent_spend) + expected: Final = max(db_spend, concurrent_spend) assert cache.in_memory_cache.get_cache(key=counter_key) == expected assert result == expected +@pytest.mark.asyncio +@pytest.mark.parametrize("window", [False, True], ids=["primary", "window"]) +@pytest.mark.parametrize("batch", [False, True], ids=["single_increment", "batch_increment"]) +@pytest.mark.parametrize("increment", [5.0, -5.0], ids=["charge", "refund"]) +async def test_cold_reseed_preserves_concurrent_local_increment( + monkeypatch: pytest.MonkeyPatch, window: bool, batch: bool, increment: float +) -> None: + from litellm.proxy import proxy_server + + cache: Final = DualCache(in_memory_cache=InMemoryCache()) + counter_key: Final = ( + f"spend:team:concurrent-{batch}-{increment}:window:1d" + if window + else f"spend:user:concurrent-{batch}-{increment}" + ) + table: Final = _PausedSpendTable(100.0) + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + reseed_task: Final = asyncio.create_task(_reseed_with_paused_table(table, cache, counter_key, window)) + await asyncio.wait_for(table.read_started.wait(), timeout=5) + + increment_task: Final = asyncio.create_task( + proxy_server._apply_spend_counter_increments( + pending=(proxy_server._PendingSpendIncrement(counter_key=counter_key, increment=increment),) + ) + if batch + else proxy_server._increment_spend_counter_cache(counter_key=counter_key, increment=increment) + ) + await asyncio.sleep(0) + table.resume_read.set() + await asyncio.wait_for(asyncio.gather(reseed_task, increment_task), timeout=5) + + assert cache.in_memory_cache.get_cache(key=counter_key) == 100.0 + increment + + @pytest.mark.asyncio async def test_end_user_from_db_reads_the_end_user_row_by_user_id(): prisma: Final = _FakePrismaClient(end_user_row=SimpleNamespace(user_id="customer-42", spend=0.0)) @@ -352,8 +396,7 @@ async def test_end_user_from_db_ignores_other_counter_kinds_without_touching_the @pytest.mark.asyncio async def test_end_user_from_db_returns_none_without_a_row_a_client_or_on_db_error(): assert ( - await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") - is None + await SpendCounterReseed.end_user_from_db(prisma_client=None, counter_key="spend:end_user:customer-42") is None ) assert ( await SpendCounterReseed.end_user_from_db(