diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index cc81530a2f8..1d9eb5f95b2 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/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) diff --git a/tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py b/tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py index d51c36a9c64..2662518e941 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py +++ b/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