diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index e35b1c8c82b..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,7 +254,10 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True) + 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", @@ -438,11 +450,16 @@ class SpendCounterReseed: value=current_value, ) else: - await spend_counter_cache.async_increment_cache(key=counter_key, value=window_spend) + 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", counter_key, ) raise - return window_spend + return current_value diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c8fd7edfed6..ba1e2632489 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3370,10 +3370,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 ) @@ -3396,10 +3394,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 0361d5cfe8f..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 @@ -15,6 +16,7 @@ from typing import Final import pytest from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed @@ -78,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 @@ -270,6 +305,63 @@ async def test_coalesced_window_seeds_a_cold_counter_from_the_row(): assert prisma.db.litellm_spendlogs.call_count == 0 +@pytest.mark.asyncio +@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( + window: bool, + concurrent_spend: float, +) -> 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)) + + 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) + + 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)) @@ -304,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(