fix(proxy): re-read a cold counter after an earlier seed lands, and seed on its own from another event loop

A request that read the counter cold after this process's add started now waits
for that add and re-reads the counter, so a counter deleted before the add
returned is seeded again. A request on another event loop than the add in
flight adds on its own, as before.

litellm/proxy/db/spend_counter_reseed.py

@@ -74,12 +74,18 @@ _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
 
 @dataclass(slots=True)
 class _CachedSpendSeed:
-    """This process adding its cached spend to one cold counter: the add in flight, and the tick the last one landed."""
+    """This process adding its cached spend to one cold counter: the add in flight, the tick it started at, and the
+    tick the last one landed at."""
 
+    started_at: int = 0
     landed_at: int = 0
     in_flight: asyncio.Task[float | None] | None = None
 
+    def on_this_loop(self) -> bool:
+        return self.in_flight is None or self.in_flight.get_loop() is asyncio.get_running_loop()
+
     def start(self, add: Callable[[], Awaitable[float | None]]) -> asyncio.Task[float | None]:
+        self.started_at = spend_counter_tick()
         in_flight: Final = asyncio.ensure_future(add())
         in_flight.add_done_callback(self._settle)
         self.in_flight = in_flight
@@ -149,15 +155,20 @@ 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 the
-        add in flight and gets its result or error. When the add already landed, the request re-reads the counter and
-        adds again only if the counter is gone. Each other process still adds its own cached spend."""
+        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."""
         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))
+        added: Final = await (asyncio.shield(seed.in_flight or seed.start(add)) if seed.on_this_loop() else add())
         if added is not None:
             record_spend_counter_value(counter_key, added)
 

tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py

@@ -1,6 +1,7 @@
 """Overlapping cold starts of one spend counter inside one process add that process's cached spend once."""
 
 import asyncio
+import threading
 from typing import Final, cast
 
 import pytest
@@ -177,3 +178,67 @@ async def test_a_failed_seed_fails_the_requests_that_joined_it_and_the_next_cold
 
     assert [type(outcome) for outcome in outcomes] == [ConnectionError] * 3
     assert redis.values[counter] == 6.0
+
+
+class _HeldReplyRedis(_SharedRedis):
+    """Applies an increment at once but holds its reply until released, like a slow network round trip."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.applied: Final = asyncio.Event()
+        self.release: Final = asyncio.Event()
+
+    async def async_increment(self, key: str, value: float, **_: object) -> float:
+        self.values[key] = self.values.get(key, 0.0) + value
+        total: Final = self.values[key]
+        self.applied.set()
+        await self.release.wait()
+        return total
+
+
+@pytest.mark.asyncio
+async def test_a_counter_deleted_before_an_earlier_add_returns_is_seeded_by_a_later_cold_reader(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    counter: Final = "spend:end_user:deleted-mid-add"
+    redis: Final = _HeldReplyRedis()
+    _process(monkeypatch, redis, cached_spend=6.0)
+
+    first: Final = asyncio.ensure_future(_cold_start(redis, counter))
+    await redis.applied.wait()
+    del redis.values[counter]
+    later: Final = asyncio.ensure_future(_cold_start(redis, counter))
+    for _ in range(20):
+        await asyncio.sleep(0)
+    redis.release.set()
+    await asyncio.gather(first, later)
+
+    assert redis.values[counter] == 6.0
+
+
+def test_a_cold_start_on_another_event_loop_adds_on_its_own(monkeypatch: pytest.MonkeyPatch) -> None:
+    counter: Final = "spend:end_user:two-loops"
+    held: Final = threading.Event()
+    release: Final = threading.Event()
+
+    class _ThreadHeldRedis(_SharedRedis):
+        async def async_increment(self, key: str, value: float, **_: object) -> float:
+            if threading.current_thread() is not threading.main_thread():
+                held.set()
+                while not release.is_set():
+                    await asyncio.sleep(0.001)
+            return await super().async_increment(key, value)
+
+    redis: Final = _ThreadHeldRedis()
+    _process(monkeypatch, redis, cached_spend=6.0)
+    other_loop: Final = threading.Thread(target=lambda: asyncio.run(_cold_start(redis, counter)), daemon=True)
+    other_loop.start()
+    assert held.wait(5)
+
+    try:
+        asyncio.run(_cold_start(redis, counter))
+    finally:
+        release.set()
+        other_loop.join(5)
+
+    assert redis.values[counter] == 12.0