diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 2da5f0128e1..4faca60779d 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -29,6 +29,7 @@ from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate from litellm.proxy.spend_tracking.spend_counter_batch import ( SPEND_COUNTERS_TARGET, + SpendCounterRead, forget_spend_counter, read_batched_spend_counter, record_spend_counter_value, @@ -156,15 +157,15 @@ class SpendCounterReseed: async def add_cached_spend_once( spend_counter_cache: "DualCache", counter_key: str, - read_at: int, + read: SpendCounterRead, add: Callable[[], Awaitable[float | None]], is_warm: Callable[[], Awaitable[bool]], ) -> None: """Run ``add``, this process adding its cached spend to a cold counter, once for all of this process's - requests that read the counter cold (at tick ``read_at``) before that add landed. A request joins an add that - started after its read and gets its result, and starts or joins another add when that one fails. When an add - landed after its read, or one that started before its read finishes, the request re-reads the counter and adds - again only if the counter is gone. A request waits on other requests' adds for at most + requests that read the counter cold before that add landed. A request joins an add that started after its read + came back and gets its result, and starts or joins another add when that one fails. When an add landed after + its read went out, or one that started before its read came back finishes, the request re-reads the counter + and adds again only if the counter is gone. A request waits on other requests' adds for at most ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` in all and then adds its own. A request on another event loop than the add in flight, one that finds the Redis circuit breaker not closed, and each other process, still add their own cached spend.""" @@ -172,24 +173,27 @@ class SpendCounterReseed: loop: Final = asyncio.get_running_loop() join_by: Final = loop.time() + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS - async def add_unless_covered(cold_at: int) -> float | None: + async def add_unless_covered(cold: SpendCounterRead) -> float | None: in_flight: Final = seed.in_flight if not seed.on_this_loop() or not _redis_breaker_closed(spend_counter_cache): return await add() if in_flight is None: - if seed.landed_at <= cold_at: + if seed.landed_at < cold.sent_at: return await asyncio.shield(seed.start(add)) forget_spend_counter(counter_key) - return None if await is_warm() else await add_unless_covered(spend_counter_tick()) - covers: Final = seed.started_at > cold_at + sent_at: Final = spend_counter_tick() + if await is_warm(): + return None + return await add_unless_covered(SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick())) + covers: Final = seed.started_at > cold.returned_at await asyncio.wait((in_flight,), timeout=max(0.0, join_by - loop.time())) if not in_flight.done(): return await asyncio.shield(seed.start(add)) if covers and not in_flight.cancelled() and in_flight.exception() is None: return in_flight.result() - return await add_unless_covered(cold_at) + return await add_unless_covered(cold) - added: Final = await add_unless_covered(read_at) + added: Final = await add_unless_covered(read) if added is not None: record_spend_counter_value(counter_key, added) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c29ff0bee4a..5a6e9756abb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -936,7 +936,8 @@ from litellm.proxy.spend_tracking.spend_counter_batch import ( read_batched_spend_counter, record_spend_counter_value, spend_counter_batch_scope, - spend_counter_read_tick, + spend_counter_read, + spend_counter_tick, ) from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, @@ -4086,8 +4087,9 @@ async def ensure_spend_counter_initialized( counter_key: str, source_cache_key: str | list[str], ) -> None: + sent_at: Final = spend_counter_tick() is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) - read_at: Final = spend_counter_read_tick(counter_key) + read: Final = spend_counter_read(counter_key, sent_at) if is_warm is False: # Shares the per-counter lock with get_current_spend. db_spend: Final = await SpendCounterReseed.coalesced( @@ -4103,7 +4105,7 @@ async def ensure_spend_counter_initialized( await SpendCounterReseed.add_cached_spend_once( spend_counter_cache=spend_counter_cache, counter_key=counter_key, - read_at=read_at, + read=read, add=lambda: increment_spend_counter_cache(counter_key=counter_key, increment=base_spend), is_warm=lambda: _is_spend_counter_cache_warm(counter_key=counter_key), ) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 4937fea039b..e15d3a62122 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -22,8 +22,6 @@ from litellm.proxy.common_utils.user_api_key_cache import ( _CounterValues: Final = TypeAdapter(dict[str, float | None]) _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({}) -_NO_TICKS: Final[Mapping[str, int]] = MappingProxyType({}) -_PendingMGet: TypeAlias = tuple[BatchResult[Mapping[str, object]], "asyncio.Future[int]"] SPEND_COUNTERS_TARGET: Final = "spend_counters" _ticks: Final = itertools.count(1) @@ -33,6 +31,18 @@ def spend_counter_tick() -> int: return next(_ticks) +@dataclass(frozen=True, slots=True) +class SpendCounterRead: + """When a read of a spend counter went out and when its reply came back, as spend counter ticks.""" + + sent_at: int + returned_at: int + + +_NO_READS: Final[Mapping[str, SpendCounterRead]] = MappingProxyType({}) +_PendingMGet: TypeAlias = tuple[BatchResult[Mapping[str, object]], int, "asyncio.Future[int]"] + + @dataclass(frozen=True, slots=True) class PendingSpendIncrement: counter_key: str @@ -55,7 +65,7 @@ class SpendCounterBatch: "_loaded", "_lock", "_open", - "_read_at", + "_reads", "_redis_cache", "_request_batch", ) @@ -67,7 +77,7 @@ class SpendCounterBatch: self._keys: frozenset[str] = frozenset() self._fetched: frozenset[str] = frozenset() self._loaded: Mapping[str, float | None] = _NO_VALUES - self._read_at: Mapping[str, int] = _NO_TICKS + self._reads: Mapping[str, SpendCounterRead] = _NO_READS self._inflight: Final[list[_PendingMGet]] = [] # mutable-ok: drained by _load self._request_batch: Final[RedisBatch | None] = active_request_redis_batch(redis_cache) if self._request_batch is not None: @@ -98,9 +108,9 @@ class SpendCounterBatch: return None return loaded[counter_key], True - def read_tick(self, counter_key: str) -> int | None: - """The tick at which the MGET behind this batch's value for the counter returned, None without that value.""" - return self._read_at.get(counter_key) if self._open and counter_key in self._loaded else None + def counter_read(self, counter_key: str) -> SpendCounterRead | None: + """When the MGET behind this batch's value for the counter went out and came back, None without that value.""" + return self._reads.get(counter_key) if self._open and counter_key in self._loaded else None def record(self, counter_key: str, value: float) -> None: """A write returned the counter's new value; later reads in this scope see it instead of the MGET value.""" @@ -117,7 +127,7 @@ class SpendCounterBatch: self._keys = self._keys - key self._fetched = self._fetched - key self._loaded = MappingProxyType({k: v for k, v in self._loaded.items() if k != counter_key}) - self._read_at = MappingProxyType({k: v for k, v in self._read_at.items() if k != counter_key}) + self._reads = MappingProxyType({k: v for k, v in self._reads.items() if k != counter_key}) async def _load(self) -> Mapping[str, float | None]: async with self._lock: @@ -128,12 +138,14 @@ class SpendCounterBatch: pending: Final = self._keys - self._fetched if pending: self._fetched = self._fetched | pending - self._merge(await self._fetch(pending), spend_counter_tick()) + sent_at: Final = spend_counter_tick() + fetched: Final = await self._fetch(pending) + self._merge(fetched, SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick())) return self._loaded - def _merge(self, fetched: Mapping[str, float | None], read_at: int) -> None: + def _merge(self, fetched: Mapping[str, float | None], read: SpendCounterRead) -> None: self._loaded = MappingProxyType({**fetched, **self._loaded}) - self._read_at = MappingProxyType({**dict.fromkeys(fetched, read_at), **self._read_at}) + self._reads = MappingProxyType({**dict.fromkeys(fetched, read), **self._reads}) def _declare_pending(self) -> None: """Flush hook: put every bound-but-unread counter on the request pipeline that is about to go out.""" @@ -142,22 +154,28 @@ class SpendCounterBatch: pending: Final = self._keys - self._fetched if pending: self._fetched = self._fetched | pending + sent_at: Final = spend_counter_tick() with service_target(SPEND_COUNTERS_TARGET): read: Final = self._request_batch.mget(sorted(pending)) returned_at: Final[asyncio.Future[int]] = asyncio.get_running_loop().create_future() read.on_settled(lambda _: returned_at.set_result(spend_counter_tick())) - self._inflight.append((read, returned_at)) + self._inflight.append((read, sent_at, returned_at)) async def _collect_inflight(self) -> None: reads: Final = tuple(self._inflight) self._inflight.clear() - for read, returned_at in reads: + for read, sent_at, returned_at in reads: try: fetched: Mapping[str, float | None] = _CounterValues.validate_python(await read) except Exception as e: # noqa: BLE001 # per-key reads take over and apply their own Redis fallback verbose_proxy_logger.debug("spend counter batch read failed, falling back to per-key reads: %s", e) continue - self._merge(fetched, returned_at.result() if returned_at.done() else spend_counter_tick()) + self._merge( + fetched, + SpendCounterRead( + sent_at=sent_at, returned_at=returned_at.result() if returned_at.done() else spend_counter_tick() + ), + ) async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]: try: @@ -176,11 +194,12 @@ def active_spend_counter_batch() -> SpendCounterBatch | None: return _active_batch.get() -def spend_counter_read_tick(counter_key: str) -> int: - """The tick this context last read the counter at: when the open batch's MGET for it returned, else now.""" +def spend_counter_read(counter_key: str, sent_at: int) -> SpendCounterRead: + """When this context's last read of the counter went out and came back: the open batch's MGET for it, else a read + that went out at ``sent_at`` and came back now.""" batch: Final = _active_batch.get() - read_at: Final = batch.read_tick(counter_key) if batch is not None else None - return read_at if read_at is not None else spend_counter_tick() + batched: Final = batch.counter_read(counter_key) if batch is not None else None + return batched if batched is not None else SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick()) class spend_counter_batch_scope: diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 06c8444f2d8..f8e0dd16ea1 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -2412,3 +2412,163 @@ async def test_a_cold_start_after_the_redis_breaker_opened_fails_fast_instead_of outcomes: Final = await asyncio.gather(first, later, return_exceptions=True) assert outcomes[0] is None and isinstance(outcomes[1], RedisCircuitBreakerOpenError), outcomes + + +class _HeldFirstMGetRedis(_SharedRedis): + """Reads the first MGET's keys at once but holds its reply until released, like a slow network round trip.""" + + def __init__(self) -> None: + super().__init__() + self.mget_read: Final = asyncio.Event() + self.release_mget: Final = asyncio.Event() + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + values: Final = {key: self.values.get(key) for key in key_list} + if not self.mget_read.is_set(): + self.mget_read.set() + await self.release_mget.wait() + return values + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_while_a_cold_read_is_on_its_way_back_is_not_added_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:seed-during-read" + redis: Final = _HeldFirstMGetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + async def _first_request() -> None: + await redis.mget_read.wait() + await _cold_start(redis, counter) + redis.release_mget.set() + + await asyncio.gather(_cold_start(redis, counter), _first_request()) + + assert redis.values[counter] == 6.0 + + +class _HeldReadRedis(_SharedRedis): + """Answers an MGET with the values its keys had when it arrived, holding the reply while ``hold_reads`` is set.""" + + def __init__(self) -> None: + super().__init__() + self.hold_reads = False + self.read_arrived: Final = asyncio.Event() + self.release_read: Final = asyncio.Event() + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + values: Final = {key: self.values.get(key) for key in key_list} + if self.hold_reads: + self.read_arrived.set() + await self.release_read.wait() + return values + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_while_a_request_re_reads_the_counter_is_not_added_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:seed-during-reread" + redis: Final = _HeldReadRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + del redis.values[counter] + redis.hold_reads = True + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + async def _request_seeding_during_the_re_read() -> None: + await redis.read_arrived.wait() + redis.hold_reads = False + await _cold_start(redis, counter) + redis.release_read.set() + + await asyncio.gather(_stale_request(), _first_request(), _request_seeding_during_the_re_read()) + + assert redis.values[counter] == 6.0 + + +class _HeldGetRedis(_SharedRedis): + """Answers the first GET with the value its key had when it arrived, holding the reply until released.""" + + def __init__(self) -> None: + super().__init__() + self.get_arrived: Final = asyncio.Event() + self.release_get: Final = asyncio.Event() + + async def async_get_cache(self, key: str, **_: object) -> float | None: + value: Final = self.values.get(key) + if not self.get_arrived.is_set(): + self.get_arrived.set() + await self.release_get.wait() + return value + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_while_an_unbatched_cold_read_is_on_its_way_back_is_not_added_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:seed-during-unbatched-read" + redis: Final = _HeldGetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + async def _first_request() -> None: + await redis.get_arrived.wait() + await _cold_start(redis, counter) + redis.release_get.set() + + await asyncio.gather( + ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE), + _first_request(), + ) + + assert redis.values[counter] == 6.0 + + +class _QueuedReadRedis(_HeldReplyRedis): + """Holds the first MGET before Redis runs it, while increments apply at once and hold their replies.""" + + def __init__(self) -> None: + super().__init__() + self.read_sent: Final = asyncio.Event() + self.run_read: Final = asyncio.Event() + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + if not self.read_sent.is_set(): + self.read_sent.set() + await self.run_read.wait() + return {key: self.values.get(key) for key in key_list} + + +@pytest.mark.asyncio +async def test_a_cold_read_sent_before_an_add_started_does_not_join_it_when_the_counter_was_deleted_after_it_landed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:deleted-while-read-queued" + redis: Final = _QueuedReadRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + reader: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.read_sent.wait() + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.applied.wait() + del redis.values[counter] + redis.run_read.set() + for _ in range(20): + await asyncio.sleep(0) + redis.release.set() + await asyncio.gather(reader, first) + + assert redis.values.get(counter) == 6.0