fix(proxy): retry a failed shared cold seed, stamp the real read tick, and skip sharing when the Redis breaker is not closed

A request that joined another request's cold seed starts or joins a new seed when that one fails, so one dropped
connection fails only the request that owned it. A request waits on other requests' seeds for at most 50 ms and then
adds its own. A request that finds the Redis circuit breaker not closed adds on its own, so it fails fast as on main.
The batch now stamps each counter read with the tick its MGET returned, so a request that read the counter cold after
this process's seed landed still adds its cached spend

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

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))

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)
 

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(

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:

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