diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..3b90cce211d 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -709,6 +709,11 @@ class RedisCache(BaseCache): def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]: return call() + @property + def circuit_breaker_closed(self) -> bool: + """Whether calls reach Redis without a recovery probe. Reading it, unlike making a call, never starts one.""" + return not self._circuit_breaker.enabled or self._circuit_breaker.state == RedisCircuitBreaker.CLOSED + def _setup_health_pings(self): """Setup async and sync health pings for Redis.""" # ASYNC HEALTH PING diff --git a/litellm/constants.py b/litellm/constants.py index 28e7b4a7fd7..199d7bc59bc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1924,6 +1924,7 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTER RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3)) RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000) +SPEND_COUNTER_COLD_SEED_JOIN_SECONDS: Final = 0.05 PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25)) PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0)) PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0)) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 1d9eb5f95b2..2da5f0128e1 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, ClassVar, Final, Optional from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger -from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE +from litellm.constants import SPEND_COUNTER_COLD_SEED_JOIN_SECONDS, SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate @@ -92,7 +92,8 @@ class _CachedSpendSeed: return in_flight def _settle(self, done: asyncio.Future[float | None]) -> None: - self.in_flight = None + if self.in_flight is done: + self.in_flight = None if not done.cancelled() and done.exception() is None: self.landed_at = spend_counter_tick() @@ -102,6 +103,11 @@ def _cached_spend_seed(spend_counter_cache: "DualCache", counter_key: str) -> _C return _CachedSpendSeed() +def _redis_breaker_closed(spend_counter_cache: "DualCache") -> bool: + redis_cache: Final = spend_counter_cache.redis_cache + return redis_cache is None or redis_cache.circuit_breaker_closed + + def _as_utc(value: datetime) -> datetime: return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) @@ -155,20 +161,35 @@ class SpendCounterReseed: 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 or after tick ``read_at``) before that add landed. A request joins an - add that started after its read and gets its result or error. 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 on another event loop than the add in flight, and each other process, still adds its own cached - spend.""" + 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 + ``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.""" seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) - earlier: Final = seed.in_flight - if earlier is not None and seed.started_at < read_at and seed.on_this_loop(): - await asyncio.wait((earlier,)) - if seed.in_flight is None and seed.landed_at > read_at: - forget_spend_counter(counter_key) - if await is_warm(): - return - added: Final = await (asyncio.shield(seed.in_flight or seed.start(add)) if seed.on_this_loop() else add()) + 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: + 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: + 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 + 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) + + added: Final = await add_unless_covered(read_at) 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 b43733286cd..c29ff0bee4a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4086,8 +4086,8 @@ async def ensure_spend_counter_initialized( counter_key: str, source_cache_key: str | list[str], ) -> None: - read_at: Final = spend_counter_read_tick() is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) + read_at: Final = spend_counter_read_tick(counter_key) if is_warm is False: # Shares the per-counter lock with get_current_spend. db_spend: Final = await SpendCounterReseed.coalesced( diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 49280448f75..4937fea039b 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -6,7 +6,7 @@ from collections.abc import Iterator, Mapping, Sequence from contextvars import ContextVar, Token from dataclasses import dataclass from types import MappingProxyType, TracebackType -from typing import Final +from typing import Final, TypeAlias from pydantic import TypeAdapter @@ -22,6 +22,8 @@ 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) @@ -53,20 +55,20 @@ class SpendCounterBatch: "_loaded", "_lock", "_open", + "_read_at", "_redis_cache", "_request_batch", - "opened_at", ) def __init__(self, redis_cache: RedisCache) -> None: self._redis_cache: Final = redis_cache - self.opened_at: Final = spend_counter_tick() self._lock: Final = asyncio.Lock() self._open = True self._keys: frozenset[str] = frozenset() self._fetched: frozenset[str] = frozenset() self._loaded: Mapping[str, float | None] = _NO_VALUES - self._inflight: Final[list[BatchResult[Mapping[str, object]]]] = [] # mutable-ok: drained by _load + self._read_at: Mapping[str, int] = _NO_TICKS + 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: self._request_batch.add_flush_hook(self._declare_pending) @@ -96,6 +98,10 @@ 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 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.""" if not self._open: @@ -111,6 +117,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}) async def _load(self) -> Mapping[str, float | None]: async with self._lock: @@ -121,10 +128,13 @@ class SpendCounterBatch: pending: Final = self._keys - self._fetched if pending: self._fetched = self._fetched | pending - fetched: Final = await self._fetch(pending) - self._loaded = MappingProxyType({**fetched, **self._loaded}) + self._merge(await self._fetch(pending), spend_counter_tick()) return self._loaded + def _merge(self, fetched: Mapping[str, float | None], read_at: int) -> None: + self._loaded = MappingProxyType({**fetched, **self._loaded}) + self._read_at = MappingProxyType({**dict.fromkeys(fetched, read_at), **self._read_at}) + def _declare_pending(self) -> None: """Flush hook: put every bound-but-unread counter on the request pipeline that is about to go out.""" if self._request_batch is None or not self._open: @@ -133,18 +143,21 @@ class SpendCounterBatch: if pending: self._fetched = self._fetched | pending with service_target(SPEND_COUNTERS_TARGET): - self._inflight.append(self._request_batch.mget(sorted(pending))) + 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)) async def _collect_inflight(self) -> None: - results: Final = tuple(self._inflight) + reads: Final = tuple(self._inflight) self._inflight.clear() - for result in results: + for read, returned_at in reads: try: - fetched: Mapping[str, float | None] = _CounterValues.validate_python(await result) + 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._loaded = MappingProxyType({**fetched, **self._loaded}) + self._merge(fetched, returned_at.result() if returned_at.done() else spend_counter_tick()) async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]: try: @@ -163,10 +176,11 @@ def active_spend_counter_batch() -> SpendCounterBatch | None: return _active_batch.get() -def spend_counter_read_tick() -> int: - """A tick no later than this context's next spend counter read: the open batch's opening tick, else a new one.""" +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.""" batch: Final = _active_batch.get() - return batch.opened_at if batch is not None and batch.is_open else spend_counter_tick() + 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() 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 ac5e3ff2388..06c8444f2d8 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -41,7 +41,8 @@ from litellm._internal_context import current_service_target from litellm._service_logger import ServiceLogging from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker, RedisCircuitBreakerOpenError from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.models.end_user import LiteLLM_EndUserTable @@ -1808,10 +1809,10 @@ class _FakeRedisRunningIncrementIfPresent(FakeAsyncRedis): class _RedisCacheOnClient(RedisCache): - def __init__(self, client: FakeAsyncRedis) -> None: + def __init__(self, client: FakeAsyncRedis, breaker: RedisCircuitBreaker | None = None) -> None: self._injected_client: Final = client self.namespace = None - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) + self._circuit_breaker = breaker or RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) self.service_logger_obj = ServiceLogging() self.default_ttl = None @@ -2026,6 +2027,8 @@ _COLD_SEED_SOURCE: Final = "end_user_id:customer-1" class _SharedRedis: """The counter commands a cold start sends, yielding like a network call so overlapping requests interleave.""" + circuit_breaker_closed: Final = True + def __init__(self) -> None: self.values: Final[dict[str, float]] = {} # mutable-ok: stands in for the Redis keyspace self.failing = False @@ -2051,7 +2054,9 @@ class _SharedRedis: self.values.pop(key, None) -def _process(monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | None, cached_spend: float) -> DualCache: +def _process( + monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | RedisCache | None, cached_spend: float +) -> DualCache: counters: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, redis)) objects: Final = DualCache(in_memory_cache=InMemoryCache()) objects.in_memory_cache.set_cache(key=_COLD_SEED_SOURCE, value={"spend": cached_spend}) @@ -2061,7 +2066,7 @@ def _process(monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | None, cached return counters -async def _cold_start(redis: _SharedRedis | None, counter: str) -> None: +async def _cold_start(redis: _SharedRedis | RedisCache | None, counter: str) -> None: with spend_counter_batch_scope(cast(RedisCache, redis) if redis is not None else None, frozenset((counter,))): await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) @@ -2172,7 +2177,7 @@ async def test_another_process_seeding_after_this_request_read_the_counter_still @pytest.mark.asyncio -async def test_a_failed_seed_fails_the_requests_that_joined_it_and_the_next_cold_start_seeds( +async def test_seeds_that_keep_failing_fail_every_cold_start_and_the_next_cold_start_seeds( monkeypatch: pytest.MonkeyPatch, ) -> None: counter: Final = "spend:end_user:failed-seed" @@ -2250,3 +2255,160 @@ def test_a_cold_start_on_another_event_loop_adds_on_its_own(monkeypatch: pytest. other_loop.join(5) assert redis.values[counter] == 12.0 + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_cold_after_this_process_seeded_it_still_adds_its_cached_spend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:read-after-seed" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + opened: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _late_reader() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + opened.set() + await seeded.wait() + del redis.values[counter] + await _read_cold(counter) + redis.values[counter] = 0.5 + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await opened.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_late_reader(), _first_request()) + + assert redis.values[counter] == 6.5 + + +@pytest.mark.asyncio +async def test_a_request_whose_pipelined_read_returned_before_this_process_seeded_the_counter_does_not_add_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:pipelined-reader" + client: Final = FakeAsyncRedis(server=FakeServer()) + redis: Final = _RedisCacheOnClient(client) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _pipelined_request() -> None: + with request_redis_batch_scope() as batches, spend_counter_batch_scope(redis, frozenset((counter,))): + await batches.batch(redis).flush() + read.set() + await seeded.wait() + 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() + + await asyncio.gather(_pipelined_request(), _first_request()) + + assert float(await client.get(counter) or 0) == 6.0 + + +class _DroppedFirstIncrementRedis(_SharedRedis): + """Resets the connection that carries the first increment and answers every later one.""" + + def __init__(self) -> None: + super().__init__() + self.dropped = False + + async def async_increment(self, key: str, value: float, **_: object) -> float: + if not self.dropped: + self.dropped = True + await asyncio.sleep(0) + raise ConnectionError("connection reset by peer") + return await super().async_increment(key, value) + + +@pytest.mark.asyncio +async def test_requests_that_joined_a_seed_whose_connection_reset_seed_again_and_only_its_starter_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:reset-seed" + redis: Final = _DroppedFirstIncrementRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + outcomes: Final = await asyncio.gather(*(_cold_start(redis, counter) for _ in range(4)), return_exceptions=True) + + failures: Final = [outcome for outcome in outcomes if outcome is not None] + assert [type(failure) for failure in failures] == [ConnectionError], outcomes + assert redis.values[counter] == 6.0 + + +class _StalledFirstIncrementRedis(_SharedRedis): + """Never answers the first increment until released, like a connection that stopped responding.""" + + def __init__(self) -> None: + super().__init__() + self.stalled: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_increment(self, key: str, value: float, **_: object) -> float: + if not self.stalled.is_set(): + self.stalled.set() + await self.release.wait() + raise TimeoutError("redis did not answer") + return await super().async_increment(key, value) + + +@pytest.mark.asyncio +async def test_cold_starts_stop_waiting_on_a_stalled_seed_and_add_their_own(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = "spend:end_user:stalled-seed" + redis: Final = _StalledFirstIncrementRedis() + _process(monkeypatch, redis, cached_spend=6.0) + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.stalled.wait() + + try: + await asyncio.wait_for(asyncio.gather(*(_cold_start(redis, counter) for _ in range(3))), timeout=10) + seeded_while_stalled: Final = redis.values.get(counter) + finally: + redis.release.set() + + with pytest.raises(TimeoutError): + await stalled + assert seeded_while_stalled == 18.0 + + +class _FakeRedisHoldingIncrements(FakeAsyncRedis): + def __init__(self, server: FakeServer) -> None: + super().__init__(server=server) + self.holding: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def incrbyfloat(self, name: str, amount: float = 1.0) -> float: + self.holding.set() + await self.release.wait() + return await super().incrbyfloat(name, amount) + + +@pytest.mark.asyncio +async def test_a_cold_start_after_the_redis_breaker_opened_fails_fast_instead_of_using_an_add_in_flight( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:breaker-open" + breaker: Final = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=30) + client: Final = _FakeRedisHoldingIncrements(FakeServer()) + redis: Final = _RedisCacheOnClient(client, breaker) + _process(monkeypatch, redis, cached_spend=6.0) + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await client.holding.wait() + breaker.record_failure() + + later: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(100): + await asyncio.sleep(0) + client.release.set() + + outcomes: Final = await asyncio.gather(first, later, return_exceptions=True) + + assert outcomes[0] is None and isinstance(outcomes[1], RedisCircuitBreakerOpenError), outcomes