fix(proxy): compare cold seeds to when a read went out and came back

A cold read now carries two ticks, when it was sent and when its reply came back
A request adds its cached spend directly only when this process's last add landed before the read was sent, and otherwise re-reads the counter
A re-read is stamped before its GET goes out, and a pipelined read keeps the tick from before its MGET, so a seed that lands while a read is in flight is never added twice
A request joins an add in flight only when that add started after its read came back

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)
 

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

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:

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