From 2f4090725925a809d87039f0b91ce819295f6c88 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 05:53:03 -0700 Subject: [PATCH 01/18] fix(proxy): add a process's cached spend to a cold spend counter once Overlapping cold starts of one spend counter inside one proxy process each added the process's whole cached spend, so N concurrent requests left N times the spend and refused a customer still under budget. Requests that read the counter cold before this process's own add landed now join that add, or re-read the counter and add again only when it is gone. Other processes still add their own cached spend as before. Refs #47062 --- litellm/proxy/db/spend_counter_reseed.py | 51 ++++- litellm/proxy/proxy_server.py | 10 +- .../spend_tracking/spend_counter_batch.py | 26 ++- .../spend/test_end_user_budget_enforcement.py | 135 ++++++++++++- .../test_spend_counter_cold_seed.py | 179 ++++++++++++++++++ 5 files changed, 397 insertions(+), 4 deletions(-) create mode 100644 tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 721999b1868..cc81530a2f8 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -13,8 +13,10 @@ memory in long-lived deployments. """ import asyncio +import functools from collections import OrderedDict -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Optional @@ -27,8 +29,10 @@ 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, + forget_spend_counter, read_batched_spend_counter, record_spend_counter_value, + spend_counter_tick, ) from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository @@ -68,6 +72,30 @@ _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.""" + + landed_at: int = 0 + in_flight: asyncio.Task[float | None] | None = None + + def start(self, add: Callable[[], Awaitable[float | None]]) -> asyncio.Task[float | None]: + in_flight: Final = asyncio.ensure_future(add()) + in_flight.add_done_callback(self._settle) + self.in_flight = in_flight + return in_flight + + def _settle(self, done: asyncio.Future[float | None]) -> None: + self.in_flight = None + if not done.cancelled() and done.exception() is None: + self.landed_at = spend_counter_tick() + + +@functools.lru_cache(maxsize=SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE) +def _cached_spend_seed(spend_counter_cache: "DualCache", counter_key: str) -> _CachedSpendSeed: + return _CachedSpendSeed() + + def _as_utc(value: datetime) -> datetime: return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) @@ -112,6 +140,27 @@ class SpendCounterReseed: SpendCounterReseed._locks.popitem(last=False) return lock + @staticmethod + async def add_cached_spend_once( + spend_counter_cache: "DualCache", + counter_key: str, + read_at: int, + 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 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.""" + seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) + 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 added is not None: + record_spend_counter_value(counter_key, added) + @staticmethod @with_service_target(SPEND_COUNTERS_TARGET) async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a767f5423ed..b43733286cd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -936,6 +936,7 @@ 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, ) from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, @@ -4085,6 +4086,7 @@ 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) if is_warm is False: # Shares the per-counter lock with get_current_spend. @@ -4098,7 +4100,13 @@ async def ensure_spend_counter_initialized( # DB unavailable - fall back to in-process cache (may be stale). base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key) if base_spend > 0: - await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend) + await SpendCounterReseed.add_cached_spend_once( + spend_counter_cache=spend_counter_cache, + counter_key=counter_key, + read_at=read_at, + 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), + ) _ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7e0821e091e..49280448f75 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -1,6 +1,7 @@ """One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each.""" import asyncio +import itertools from collections.abc import Iterator, Mapping, Sequence from contextvars import ContextVar, Token from dataclasses import dataclass @@ -22,6 +23,12 @@ 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({}) SPEND_COUNTERS_TARGET: Final = "spend_counters" +_ticks: Final = itertools.count(1) + + +def spend_counter_tick() -> int: + """The next number in a per-process sequence that orders spend counter reads and cold seeds.""" + return next(_ticks) @dataclass(frozen=True, slots=True) @@ -39,10 +46,21 @@ class SpendCounterBatch: hook declares whatever is bound but unread, so whoever flushes first (the auth object prefetch, usually) carries the spend counters in the same round trip.""" - __slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_redis_cache", "_request_batch") + __slots__ = ( + "_fetched", + "_inflight", + "_keys", + "_loaded", + "_lock", + "_open", + "_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() @@ -145,6 +163,12 @@ 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.""" + batch: Final = _active_batch.get() + return batch.opened_at if batch is not None and batch.is_open else spend_counter_tick() + + class spend_counter_batch_scope: """Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it. Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET.""" diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py index 02c8aee8f5e..9ba53f133ec 100644 --- a/tests/integration/spend/test_end_user_budget_enforcement.py +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -1,17 +1,20 @@ import json +import os import re import uuid from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from typing import Final import httpx import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, write_rows from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import BaseModel, JsonValue +from redis import Redis PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key" CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint." @@ -639,3 +642,133 @@ def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) -> assert "data:" not in denied.text, denied.text _assert_budget_refusal(denied, end_user, 0.05) assert _observed_requests(upstream) == [] + + +_COLD_BURST: Final = 40 +_COLD_SEED_SPEND: Final = 6.0 +_COLD_BUDGET: Final = 10.0 +_CHAT_COST: Final = 20 * 0.001 + 20 * 0.002 + + +def _serves_model(gateway: Gateway, model: str) -> bool: + response: Final = gateway.request("GET", "/model/info") + return response.status_code == 200 and any( + object_value(entry)["model_name"] == model for entry in object_value(response.json())["data"] + ) + + +def _counter_cache() -> Redis: + return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + + +def _counter(cache: Redis, end_user: str) -> float | None: + raw: Final = cache.get(f"spend:end_user:{end_user}") + return None if raw is None else float(raw) + + +def _lose_counter(cache: Redis, end_user: str) -> None: + cache.delete(f"spend:end_user:{end_user}") + + +def _set_end_user_row_spend(end_user: str, spend: float) -> None: + write_rows('UPDATE "LiteLLM_EndUserTable" SET spend = %s WHERE user_id = %s', (str(spend), end_user)) + + +def _await_row_spend(end_user: str, spend: float) -> None: + eventually( + lambda: _end_user_row(end_user), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(spend), + seconds=70, + ) + + +def _refused_spend(denied: httpx.Response, end_user: str) -> float: + error: Final = _denial_error(denied) + assert error["type"] == "budget_exceeded", denied.text + matched: Final = _REFUSAL_MESSAGE.match(string_value(error["message"])) + assert matched is not None, denied.text + assert matched.group("user") == end_user + return float(matched.group("spend")) + + +@pytest.mark.parametrize("spelling", ("user", "responses_user", "messages_metadata_user_id")) +def test_overlapping_cold_starts_count_the_end_user_spend_once(gateway: Gateway, spelling: str) -> None: + chosen: Final = SPELLINGS[spelling] + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + _counter_cache() as cache, + ): + deployment: Final = _deployment(scenario, chosen) + key: Final = scenario.key(models=[deployment.model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + assert _customer_info(gateway, end_user).spend == _COLD_SEED_SPEND + assert _counter(cache, end_user) is None + upstream.delete("/__observations").raise_for_status() + + with ThreadPoolExecutor(max_workers=_COLD_BURST) as pool: + burst: Final = tuple( + pool.map( + lambda index: _send(gateway, chosen, deployment, key, end_user, f"cold {index}"), + range(_COLD_BURST), + ) + ) + refused: Final = [response.text for response in burst if response.status_code != 200] + assert refused == [], refused + follow_up: Final = _send(gateway, chosen, deployment, key, end_user, "cold follow-up") + _assert_served(chosen, follow_up) + assert len(_observed_requests(upstream)) == _COLD_BURST + 1 + + expected_spend: Final = _COLD_SEED_SPEND + (_COLD_BURST + 1) * _CHAT_COST + _await_row_spend(end_user, expected_spend) + assert _customer_info(gateway, end_user).spend == pytest.approx(expected_spend) + counted: Final = eventually( + lambda: _counter(cache, end_user), + lambda value: value is not None and value == pytest.approx(expected_spend), + seconds=30, + return_last_on_timeout=True, + ) + assert counted == pytest.approx(expected_spend) + + +def test_a_lost_counter_keeps_spend_this_proxy_has_not_written_yet(gateway: Gateway) -> None: + with gateway.scenario() as scenario, _counter_cache() as cache: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.25) + key: Final = scenario.key(models=[model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + call_cost: Final = 20 * 0.001 + 20 * 0.25 + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + served: Final = _chat(gateway, model, key, end_user) + assert served.status_code == 200, served.text + _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost) + _set_end_user_row_spend(end_user, _COLD_SEED_SPEND) + _lose_counter(cache, end_user) + + denied: Final = _chat(gateway, model, key, end_user) + + assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + call_cost) + + +def test_a_lost_counter_keeps_spend_another_proxy_has_not_written_yet(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario, _counter_cache() as cache: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.149) + eventually(lambda: _serves_model(peer, model), bool, seconds=30) + key: Final = scenario.key(models=[model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + call_cost: Final = 20 * 0.001 + 20 * 0.149 + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + first: Final = _chat(peer, model, key, end_user) + assert first.status_code == 200, first.text + _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost) + _set_end_user_row_spend(end_user, _COLD_SEED_SPEND) + _lose_counter(cache, end_user) + second: Final = _chat(peer, model, key, end_user) + assert second.status_code == 200, second.text + + denied: Final = _chat(gateway, model, key, end_user) + + assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + 2 * call_cost) 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 new file mode 100644 index 00000000000..d51c36a9c64 --- /dev/null +++ b/tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py @@ -0,0 +1,179 @@ +"""Overlapping cold starts of one spend counter inside one process add that process's cached spend once.""" + +import asyncio +from typing import Final, cast + +import pytest + +import litellm.proxy.proxy_server as ps +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch, spend_counter_batch_scope + +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.""" + + def __init__(self) -> None: + self.values: Final[dict[str, float]] = {} # mutable-ok: stands in for the Redis keyspace + self.failing = False + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + await asyncio.sleep(0) + return {key: self.values.get(key) for key in key_list} + + async def async_get_cache(self, key: str, **_: object) -> float | None: + await asyncio.sleep(0) + return self.values.get(key) + + async def async_increment(self, key: str, value: float, **_: object) -> float: + await asyncio.sleep(0) + if self.failing: + raise ConnectionError("redis down") + self.values[key] = self.values.get(key, 0.0) + value + return self.values[key] + + async def async_delete_cache(self, key: str) -> None: + await asyncio.sleep(0) + if not self.failing: + self.values.pop(key, None) + + +def _process(monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | 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=SOURCE, value={"spend": cached_spend}) + monkeypatch.setattr(ps, "spend_counter_cache", counters) + monkeypatch.setattr(ps, "user_api_key_cache", objects) + monkeypatch.setattr(ps, "prisma_client", object()) + return counters + + +async def _cold_start(redis: _SharedRedis | 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=SOURCE) + + +async def _read_cold(counter: str) -> None: + batch: Final = active_spend_counter_batch() + assert batch is not None + assert await batch.read(counter) == (None, True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("with_redis", (True, False)) +async def test_overlapping_cold_starts_in_one_process_add_its_cached_spend_once( + monkeypatch: pytest.MonkeyPatch, with_redis: bool +) -> None: + counter: Final = f"spend:end_user:once-{with_redis}" + redis: Final = _SharedRedis() if with_redis else None + counters: Final = _process(monkeypatch, redis, cached_spend=6.0) + + await asyncio.gather(*(_cold_start(redis, counter) for _ in range(5))) + + seeded: Final = redis.values.get(counter) if redis is not None else counters.in_memory_cache.get_cache(key=counter) + assert seeded == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_before_this_process_seeded_it_does_not_add_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:stale-reader" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: 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() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE) + + seeded: Final = asyncio.Event() + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_counter_lost_again_after_this_process_seeded_it_is_seeded_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:lost-again" + redis: Final = _SharedRedis() + _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] + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_another_process_seeding_after_this_request_read_the_counter_still_sums( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:two-processes" + redis: Final = _SharedRedis() + read: Final = asyncio.Event() + a_seeded: Final = asyncio.Event() + + async def _process_b() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await a_seeded.wait() + _process(monkeypatch, redis, cached_spend=5.0) + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE) + + async def _process_a() -> None: + await read.wait() + _process(monkeypatch, redis, cached_spend=3.0) + await _cold_start(redis, counter) + a_seeded.set() + + await asyncio.gather(_process_b(), _process_a()) + + assert redis.values[counter] == 8.0 + + +@pytest.mark.asyncio +async def test_a_failed_seed_fails_the_requests_that_joined_it_and_the_next_cold_start_seeds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:failed-seed" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + redis.failing = True + + outcomes: Final = await asyncio.gather(*(_cold_start(redis, counter) for _ in range(3)), return_exceptions=True) + redis.failing = False + await _cold_start(redis, counter) + + assert [type(outcome) for outcome in outcomes] == [ConnectionError] * 3 + assert redis.values[counter] == 6.0 From 5a031c45b5a3a5439cc89524ac5efdc92deb5bd6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 10:08:28 -0700 Subject: [PATCH 02/18] 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 | 21 ++++-- .../test_spend_counter_cold_seed.py | 65 +++++++++++++++++++ 2 files changed, 81 insertions(+), 5 deletions(-) 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 From 5264872c2016a58085253ae8709588078fe615b4 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 10:20:41 -0700 Subject: [PATCH 03/18] test(proxy): read the counter directly in the another-proxy lost-counter test The refusal there can come from budget reservation, whose message carries the request's own reservation, so the test now checks the counter in Redis. --- .../spend/test_end_user_budget_enforcement.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py index 9ba53f133ec..6c129015494 100644 --- a/tests/integration/spend/test_end_user_budget_enforcement.py +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -768,7 +768,15 @@ def test_a_lost_counter_keeps_spend_another_proxy_has_not_written_yet(gateway: G _lose_counter(cache, end_user) second: Final = _chat(peer, model, key, end_user) assert second.status_code == 200, second.text + expected_spend: Final = _COLD_SEED_SPEND + 2 * call_cost + counted: Final = eventually( + lambda: _counter(cache, end_user), + lambda value: value is not None and value == pytest.approx(expected_spend), + seconds=30, + return_last_on_timeout=True, + ) denied: Final = _chat(gateway, model, key, end_user) - assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + 2 * call_cost) + assert counted == pytest.approx(expected_spend) + assert _denial_error(denied)["type"] == "budget_exceeded", denied.text From 8a6d572b59164b1e5167a356cdb40ba360289ea0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 10:24:54 -0700 Subject: [PATCH 04/18] test(proxy): move the cold seed tests into test_spend_counters.py --- .../test_spend_counter_cold_seed.py | 244 ----------------- .../proxy/proxy_server/test_spend_counters.py | 247 +++++++++++++++++- 2 files changed, 242 insertions(+), 249 deletions(-) delete mode 100644 tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py 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 deleted file mode 100644 index 2662518e941..00000000000 --- a/tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py +++ /dev/null @@ -1,244 +0,0 @@ -"""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 - -import litellm.proxy.proxy_server as ps -from litellm.caching.dual_cache import DualCache -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache -from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch, spend_counter_batch_scope - -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.""" - - def __init__(self) -> None: - self.values: Final[dict[str, float]] = {} # mutable-ok: stands in for the Redis keyspace - self.failing = False - - async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: - await asyncio.sleep(0) - return {key: self.values.get(key) for key in key_list} - - async def async_get_cache(self, key: str, **_: object) -> float | None: - await asyncio.sleep(0) - return self.values.get(key) - - async def async_increment(self, key: str, value: float, **_: object) -> float: - await asyncio.sleep(0) - if self.failing: - raise ConnectionError("redis down") - self.values[key] = self.values.get(key, 0.0) + value - return self.values[key] - - async def async_delete_cache(self, key: str) -> None: - await asyncio.sleep(0) - if not self.failing: - self.values.pop(key, None) - - -def _process(monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | 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=SOURCE, value={"spend": cached_spend}) - monkeypatch.setattr(ps, "spend_counter_cache", counters) - monkeypatch.setattr(ps, "user_api_key_cache", objects) - monkeypatch.setattr(ps, "prisma_client", object()) - return counters - - -async def _cold_start(redis: _SharedRedis | 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=SOURCE) - - -async def _read_cold(counter: str) -> None: - batch: Final = active_spend_counter_batch() - assert batch is not None - assert await batch.read(counter) == (None, True) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("with_redis", (True, False)) -async def test_overlapping_cold_starts_in_one_process_add_its_cached_spend_once( - monkeypatch: pytest.MonkeyPatch, with_redis: bool -) -> None: - counter: Final = f"spend:end_user:once-{with_redis}" - redis: Final = _SharedRedis() if with_redis else None - counters: Final = _process(monkeypatch, redis, cached_spend=6.0) - - await asyncio.gather(*(_cold_start(redis, counter) for _ in range(5))) - - seeded: Final = redis.values.get(counter) if redis is not None else counters.in_memory_cache.get_cache(key=counter) - assert seeded == 6.0 - - -@pytest.mark.asyncio -async def test_a_request_that_read_the_counter_before_this_process_seeded_it_does_not_add_again( - monkeypatch: pytest.MonkeyPatch, -) -> None: - counter: Final = "spend:end_user:stale-reader" - redis: Final = _SharedRedis() - _process(monkeypatch, redis, cached_spend=6.0) - read: 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() - await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE) - - seeded: Final = asyncio.Event() - - async def _first_request() -> None: - await read.wait() - await _cold_start(redis, counter) - seeded.set() - - await asyncio.gather(_stale_request(), _first_request()) - - assert redis.values[counter] == 6.0 - - -@pytest.mark.asyncio -async def test_a_counter_lost_again_after_this_process_seeded_it_is_seeded_again( - monkeypatch: pytest.MonkeyPatch, -) -> None: - counter: Final = "spend:end_user:lost-again" - redis: Final = _SharedRedis() - _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] - await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE) - - async def _first_request() -> None: - await read.wait() - await _cold_start(redis, counter) - seeded.set() - - await asyncio.gather(_stale_request(), _first_request()) - - assert redis.values[counter] == 6.0 - - -@pytest.mark.asyncio -async def test_another_process_seeding_after_this_request_read_the_counter_still_sums( - monkeypatch: pytest.MonkeyPatch, -) -> None: - counter: Final = "spend:end_user:two-processes" - redis: Final = _SharedRedis() - read: Final = asyncio.Event() - a_seeded: Final = asyncio.Event() - - async def _process_b() -> None: - with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): - await _read_cold(counter) - read.set() - await a_seeded.wait() - _process(monkeypatch, redis, cached_spend=5.0) - await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE) - - async def _process_a() -> None: - await read.wait() - _process(monkeypatch, redis, cached_spend=3.0) - await _cold_start(redis, counter) - a_seeded.set() - - await asyncio.gather(_process_b(), _process_a()) - - assert redis.values[counter] == 8.0 - - -@pytest.mark.asyncio -async def test_a_failed_seed_fails_the_requests_that_joined_it_and_the_next_cold_start_seeds( - monkeypatch: pytest.MonkeyPatch, -) -> None: - counter: Final = "spend:end_user:failed-seed" - redis: Final = _SharedRedis() - _process(monkeypatch, redis, cached_spend=6.0) - redis.failing = True - - outcomes: Final = await asyncio.gather(*(_cold_start(redis, counter) for _ in range(3)), return_exceptions=True) - redis.failing = False - await _cold_start(redis, counter) - - 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 diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..ac5e3ff2388 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -22,6 +22,7 @@ Pins covered: from __future__ import annotations import asyncio +import threading from collections.abc import Awaitable, Callable from datetime import datetime from typing import Final, Literal, cast @@ -36,8 +37,9 @@ from opentelemetry.trace import StatusCode import litellm import litellm.proxy.proxy_server as ps -from litellm._service_logger import ServiceLogging 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.constants import GLOBAL_PROXY_SPEND_CACHE_KEY @@ -50,6 +52,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_loca from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team +from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch, spend_counter_batch_scope from .conftest import normalize @@ -1817,7 +1820,9 @@ class _RedisCacheOnClient(RedisCache): class _FakeRedisHoldingOneReply(FakeAsyncRedis): - def __init__(self, server: FakeServer, held_increment: float, applied: asyncio.Event, released: asyncio.Event) -> None: + def __init__( + self, server: FakeServer, held_increment: float, applied: asyncio.Event, released: asyncio.Event + ) -> None: super().__init__(server=server) self._held_increment: Final = held_increment self._applied: Final = applied @@ -1970,9 +1975,7 @@ async def test_update_cache_traces_the_global_spend_redis_call_under_the_request for span in exporter.get_finished_spans() if span.name == "redis" and span.attributes is not None and span.parent is not None ) - assert redis_spans == ( - ("async_increment_if_present", expected_status, request_span.get_span_context().span_id), - ) + assert redis_spans == (("async_increment_if_present", expected_status, request_span.get_span_context().span_id),) class _UnwritableKeyUserApiKeyCache(UserApiKeyCache): @@ -2013,3 +2016,237 @@ async def test_update_cache_still_adds_spend_to_the_others_when_one_cached_objec 1.25, 10.25, ) + + +# Overlapping cold starts of one spend counter inside one process add that process's cached spend once. + +_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.""" + + def __init__(self) -> None: + self.values: Final[dict[str, float]] = {} # mutable-ok: stands in for the Redis keyspace + self.failing = False + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + await asyncio.sleep(0) + return {key: self.values.get(key) for key in key_list} + + async def async_get_cache(self, key: str, **_: object) -> float | None: + await asyncio.sleep(0) + return self.values.get(key) + + async def async_increment(self, key: str, value: float, **_: object) -> float: + await asyncio.sleep(0) + if self.failing: + raise ConnectionError("redis down") + self.values[key] = self.values.get(key, 0.0) + value + return self.values[key] + + async def async_delete_cache(self, key: str) -> None: + await asyncio.sleep(0) + if not self.failing: + self.values.pop(key, None) + + +def _process(monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | 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}) + monkeypatch.setattr(ps, "spend_counter_cache", counters) + monkeypatch.setattr(ps, "user_api_key_cache", objects) + monkeypatch.setattr(ps, "prisma_client", object()) + return counters + + +async def _cold_start(redis: _SharedRedis | 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) + + +async def _read_cold(counter: str) -> None: + batch: Final = active_spend_counter_batch() + assert batch is not None + assert await batch.read(counter) == (None, True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("with_redis", (True, False)) +async def test_overlapping_cold_starts_in_one_process_add_its_cached_spend_once( + monkeypatch: pytest.MonkeyPatch, with_redis: bool +) -> None: + counter: Final = f"spend:end_user:once-{with_redis}" + redis: Final = _SharedRedis() if with_redis else None + counters: Final = _process(monkeypatch, redis, cached_spend=6.0) + + await asyncio.gather(*(_cold_start(redis, counter) for _ in range(5))) + + seeded: Final = redis.values.get(counter) if redis is not None else counters.in_memory_cache.get_cache(key=counter) + assert seeded == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_before_this_process_seeded_it_does_not_add_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:stale-reader" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: 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() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + seeded: Final = asyncio.Event() + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_counter_lost_again_after_this_process_seeded_it_is_seeded_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:lost-again" + redis: Final = _SharedRedis() + _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] + 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(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_another_process_seeding_after_this_request_read_the_counter_still_sums( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:two-processes" + redis: Final = _SharedRedis() + read: Final = asyncio.Event() + a_seeded: Final = asyncio.Event() + + async def _process_b() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await a_seeded.wait() + _process(monkeypatch, redis, cached_spend=5.0) + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _process_a() -> None: + await read.wait() + _process(monkeypatch, redis, cached_spend=3.0) + await _cold_start(redis, counter) + a_seeded.set() + + await asyncio.gather(_process_b(), _process_a()) + + assert redis.values[counter] == 8.0 + + +@pytest.mark.asyncio +async def test_a_failed_seed_fails_the_requests_that_joined_it_and_the_next_cold_start_seeds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:failed-seed" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + redis.failing = True + + outcomes: Final = await asyncio.gather(*(_cold_start(redis, counter) for _ in range(3)), return_exceptions=True) + redis.failing = False + await _cold_start(redis, counter) + + 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 From d2122a36bc8e6c9cba748e278bc48730e54bbc0b Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 15:11:58 -0700 Subject: [PATCH 05/18] 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 | 5 + litellm/constants.py | 1 + litellm/proxy/db/spend_counter_reseed.py | 51 +++-- litellm/proxy/proxy_server.py | 2 +- .../spend_tracking/spend_counter_batch.py | 42 +++-- .../proxy/proxy_server/test_spend_counters.py | 174 +++++++++++++++++- 6 files changed, 239 insertions(+), 36 deletions(-) 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 From 7bf40641a8ed72c25692299eb9ee89e0b753f9c9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 15:38:47 -0700 Subject: [PATCH 06/18] 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 | 26 +-- litellm/proxy/proxy_server.py | 8 +- .../spend_tracking/spend_counter_batch.py | 55 ++++-- .../proxy/proxy_server/test_spend_counters.py | 160 ++++++++++++++++++ 4 files changed, 217 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 2da5f0128e1..4faca60779d 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c29ff0bee4a..5a6e9756abb 100644 --- a/litellm/proxy/proxy_server.py +++ b/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), ) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 4937fea039b..e15d3a62122 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/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: diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 06c8444f2d8..f8e0dd16ea1 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/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 From 1944a38928a4ef5bfbc4fe498a5d0baf4922b430 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 15:45:10 -0700 Subject: [PATCH 07/18] test(proxy): drop the banner comment above the cold seed tests --- tests/unit/proxy/proxy_server/test_spend_counters.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index f8e0dd16ea1..a361fe36051 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -2019,8 +2019,6 @@ async def test_update_cache_still_adds_spend_to_the_others_when_one_cached_objec ) -# Overlapping cold starts of one spend counter inside one process add that process's cached spend once. - _COLD_SEED_SOURCE: Final = "end_user_id:customer-1" From 58b3e1b6d8cd920eacb0b8e78c5583498b4d1a55 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 11:03:11 -0800 Subject: [PATCH 08/18] fix(proxy): run a cold seed in the request's own task and bound its re-read A request whose batch MGET failed took the read tick from before that MGET, so its later cold GET counted as sent before this process's add landed and it could skip the add. A read outside the batch now counts as the moment it came back, and the reseed read that follows it covers an add landing in between A re-read could cost two Redis socket timeouts on top of main's add when Redis stalled. It is now one GET outside the request pipeline, bound by the 50 ms the request can wait in total, and when that time runs out the request adds as main does The add ran in its own shielded task, so it outlived a cancelled request and ran main's failed add path later, which deletes the counter. The add now runs in the request's task and requests waiting on it share a future, so a cancel stops the add as on main and the waiters add on their own --- litellm/proxy/db/spend_counter_reseed.py | 47 +-- litellm/proxy/proxy_server.py | 32 +- .../spend_tracking/spend_counter_batch.py | 25 +- .../proxy/proxy_server/test_spend_counters.py | 340 +++++++++++++++++- 4 files changed, 387 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4faca60779d..0c9e85dd788 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -80,23 +80,26 @@ class _CachedSpendSeed: started_at: int = 0 landed_at: int = 0 - in_flight: asyncio.Task[float | None] | None = None + in_flight: asyncio.Future[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]: + async def run(self, add: Callable[[], Awaitable[float | None]]) -> float | None: + """Run ``add`` in the calling task, so cancelling that task stops it, and publish its result to the requests + waiting on ``in_flight``, which ends cancelled when ``add`` fails or is cancelled.""" + in_flight: Final[asyncio.Future[float | None]] = asyncio.get_running_loop().create_future() 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 - return in_flight - - def _settle(self, done: asyncio.Future[float | None]) -> None: - if self.in_flight is done: - self.in_flight = None - if not done.cancelled() and done.exception() is None: + try: + added: Final = await add() self.landed_at = spend_counter_tick() + in_flight.set_result(added) + return added + finally: + if self.in_flight is in_flight: + self.in_flight = None + in_flight.cancel() @functools.lru_cache(maxsize=SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE) @@ -163,12 +166,12 @@ class SpendCounterReseed: ) -> 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 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.""" + came back and gets its result, and starts or joins another add when that one fails or is cancelled. 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 and on a + re-read 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) loop: Final = asyncio.get_running_loop() join_by: Final = loop.time() + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS @@ -179,17 +182,21 @@ class SpendCounterReseed: return await add() if in_flight is None: if seed.landed_at < cold.sent_at: - return await asyncio.shield(seed.start(add)) + return await seed.run(add) forget_spend_counter(counter_key) sent_at: Final = spend_counter_tick() - if await is_warm(): + try: + warm: Final = await asyncio.wait_for(is_warm(), timeout=max(0.0, join_by - loop.time())) + except asyncio.TimeoutError: + return await seed.run(add) + if 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 await seed.run(add) + if covers and not in_flight.cancelled(): return in_flight.result() return await add_unless_covered(cold) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5a6e9756abb..741f7d58f1d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -930,6 +930,7 @@ from litellm.proxy.spend_tracking.spend_capture_rate import ( from litellm.proxy.spend_tracking.spend_counter_batch import ( SPEND_COUNTERS_TARGET, PendingSpendIncrement, + SpendCounterRead, active_spend_counter_batch, forget_spend_counter, post_call_counter_keys, @@ -937,7 +938,6 @@ from litellm.proxy.spend_tracking.spend_counter_batch import ( record_spend_counter_value, spend_counter_batch_scope, spend_counter_read, - spend_counter_tick, ) from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, @@ -4087,10 +4087,8 @@ 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: Final = spend_counter_read(counter_key, sent_at) - if is_warm is False: + read: Final = await _cold_spend_counter_read(counter_key) + if read is not None: # Shares the per-counter lock with get_current_spend. db_spend: Final = await SpendCounterReseed.coalesced( prisma_client=prisma_client, @@ -4107,7 +4105,7 @@ async def ensure_spend_counter_initialized( counter_key=counter_key, 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), + is_warm=lambda: _is_spend_counter_warm_outside_batch(counter_key), ) @@ -4162,14 +4160,22 @@ _ensure_window_spend_counter_initialized: Final = ensure_window_spend_counter_in @with_service_target(SPEND_COUNTERS_TARGET) async def _is_spend_counter_cache_warm(counter_key: str) -> bool: - batched: Final = await read_batched_spend_counter(counter_key) - if batched is not None: - batched_value, _ = batched - if batched_value is None: - return False - spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched_value) - return True + return await _cold_spend_counter_read(counter_key) is None + +async def _cold_spend_counter_read(counter_key: str) -> SpendCounterRead | None: + """When the read that found the counter missing went out and came back, None when the counter has a value. A read + outside the batch counts as now, since ``SpendCounterReseed.coalesced`` reads the counter again after it.""" + batched: Final = await read_batched_spend_counter(counter_key) + if batched is not None and batched[0] is not None: + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched[0]) + return None + if batched is None and await _is_spend_counter_warm_outside_batch(counter_key): + return None + return spend_counter_read(counter_key) + + +async def _is_spend_counter_warm_outside_batch(counter_key: str) -> bool: if spend_counter_cache.redis_cache is not None: try: current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache( diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index e15d3a62122..55c4fab6893 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -40,7 +40,7 @@ class SpendCounterRead: _NO_READS: Final[Mapping[str, SpendCounterRead]] = MappingProxyType({}) -_PendingMGet: TypeAlias = tuple[BatchResult[Mapping[str, object]], int, "asyncio.Future[int]"] +_PendingMGet: TypeAlias = tuple[BatchResult[Mapping[str, object]], int] @dataclass(frozen=True, slots=True) @@ -157,25 +157,18 @@ class SpendCounterBatch: 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, sent_at, returned_at)) + self._inflight.append((read, sent_at)) async def _collect_inflight(self) -> None: reads: Final = tuple(self._inflight) self._inflight.clear() - for read, sent_at, returned_at in reads: + for read, sent_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, - SpendCounterRead( - sent_at=sent_at, returned_at=returned_at.result() if returned_at.done() else spend_counter_tick() - ), - ) + self._merge(fetched, SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick())) async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]: try: @@ -194,12 +187,14 @@ def active_spend_counter_batch() -> SpendCounterBatch | None: return _active_batch.get() -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.""" +def spend_counter_read(counter_key: str) -> SpendCounterRead: + """When this context's read of the counter went out and came back: the open batch's MGET for it, else now.""" batch: Final = _active_batch.get() 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()) + if batched is not None: + return batched + now: Final = spend_counter_tick() + return SpendCounterRead(sent_at=now, returned_at=now) 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 a361fe36051..424168de3e6 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -52,8 +52,13 @@ from litellm.models.user import LiteLLM_UserTable from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key +from litellm.proxy.db import spend_counter_reseed from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team -from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch, spend_counter_batch_scope +from litellm.proxy.spend_tracking.spend_counter_batch import ( + active_spend_counter_batch, + read_batched_spend_counter, + spend_counter_batch_scope, +) from .conftest import normalize @@ -2105,6 +2110,7 @@ async def test_a_request_that_read_the_counter_before_this_process_seeded_it_doe read.set() await seeded.wait() await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + assert await read_batched_spend_counter(counter) == (6.0, True) seeded: Final = asyncio.Event() @@ -2229,6 +2235,7 @@ async def test_a_counter_deleted_before_an_earlier_add_returns_is_seeded_by_a_la def test_a_cold_start_on_another_event_loop_adds_on_its_own(monkeypatch: pytest.MonkeyPatch) -> None: counter: Final = "spend:end_user:two-loops" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) held: Final = threading.Event() release: Final = threading.Event() @@ -2447,20 +2454,26 @@ async def test_a_seed_that_lands_while_a_cold_read_is_on_its_way_back_is_not_add class _HeldReadRedis(_SharedRedis): - """Answers an MGET with the values its keys had when it arrived, holding the reply while ``hold_reads`` is set.""" + """Answers a GET with the value its key 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() + self.reads_after_hold = 0 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} + self.reads_after_hold += self.hold_reads + return await super().async_batch_get_cache(key_list) + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.reads_after_hold += self.hold_reads + value: Final = self.values.get(key) if self.hold_reads: self.read_arrived.set() await self.release_read.wait() - return values + return value @pytest.mark.asyncio @@ -2536,29 +2549,41 @@ async def test_a_seed_that_lands_while_an_unbatched_cold_read_is_on_its_way_back class _QueuedReadRedis(_HeldReplyRedis): - """Holds the first MGET before Redis runs it, while increments apply at once and hold their replies.""" + """Holds the first read 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]: + async def _queue_the_first_read(self) -> None: if not self.read_sent.is_set(): self.read_sent.set() await self.run_read.wait() + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + await self._queue_the_first_read() return {key: self.values.get(key) for key in key_list} + async def async_get_cache(self, key: str, **_: object) -> float | None: + await self._queue_the_first_read() + return self.values.get(key) + @pytest.mark.asyncio +@pytest.mark.parametrize("batched", (True, False)) 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, + monkeypatch: pytest.MonkeyPatch, batched: bool ) -> None: - counter: Final = "spend:end_user:deleted-while-read-queued" + counter: Final = f"spend:end_user:deleted-while-read-queued-{batched}" redis: Final = _QueuedReadRedis() _process(monkeypatch, redis, cached_spend=6.0) - reader: Final = asyncio.ensure_future(_cold_start(redis, counter)) + reader: Final = asyncio.ensure_future( + _cold_start(redis, counter) + if batched + else ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + ) await redis.read_sent.wait() first: Final = asyncio.ensure_future(_cold_start(redis, counter)) await redis.applied.wait() @@ -2570,3 +2595,300 @@ async def test_a_cold_read_sent_before_an_add_started_does_not_join_it_when_the_ await asyncio.gather(reader, first) assert redis.values.get(counter) == 6.0 + + +@pytest.mark.asyncio +async def test_a_re_read_redis_does_not_answer_stops_at_the_wait_budget_and_the_request_adds_its_cached_spend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:stalled-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() + + await asyncio.wait_for(asyncio.gather(_stale_request(), _first_request()), timeout=2.0) + + assert redis.values[counter] == 6.0 + assert redis.reads_after_hold == 1 + + +class _LostAndRecreatedAfterFailedMGetRedis(_SharedRedis): + """Holds the first MGET until released, then loses the counter and resets that MGET. Once the request has read the + counter cold with its two GETs, another process's post-call increment lands before the request's next command.""" + + def __init__(self, counter: str) -> None: + super().__init__() + self.counter: Final = counter + self.mget_sent: Final = asyncio.Event() + self.release_mget: Final = asyncio.Event() + self.cold_gets = 0 + + def _other_process_increments_after_the_cold_gets(self) -> None: + if self.cold_gets == 2: + self.cold_gets += 1 + self.values[self.counter] = self.values.get(self.counter, 0.0) + 0.3 + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + if not self.mget_sent.is_set(): + self.mget_sent.set() + await self.release_mget.wait() + del self.values[self.counter] + raise ConnectionError("connection reset") + return await super().async_batch_get_cache(key_list) + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self._other_process_increments_after_the_cold_gets() + if self.release_mget.is_set() and self.cold_gets < 2: + self.cold_gets += 1 + return await super().async_get_cache(key) + + async def async_increment(self, key: str, value: float, **_: object) -> float: + self._other_process_increments_after_the_cold_gets() + return await super().async_increment(key, value) + + +@pytest.mark.asyncio +async def test_a_request_whose_mget_failed_and_whose_gets_went_out_after_this_process_seeded_adds_its_cached_spend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:failed-mget" + redis: Final = _LostAndRecreatedAfterFailedMGetRedis(counter) + _process(monkeypatch, redis, cached_spend=6.0) + + request: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.mget_sent.wait() + await _cold_start(redis, counter) + redis.release_mget.set() + await request + + assert redis.values[counter] == pytest.approx(6.3) + + +class _AppliedHeldThenResetRedis(_HeldReplyRedis): + """Applies an increment at once, holds its reply until released, then resets the connection instead of replying.""" + + def __init__(self) -> None: + super().__init__() + self.deleted = False + + async def async_increment(self, key: str, value: float, **_: object) -> float: + await super().async_increment(key, value) + raise ConnectionError("connection reset") + + async def async_delete_cache(self, key: str) -> None: + self.deleted = True + await super().async_delete_cache(key) + + +@pytest.mark.asyncio +async def test_a_cold_seed_whose_request_was_cancelled_does_not_delete_the_counter_later( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-seed" + redis: Final = _AppliedHeldThenResetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + request: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.applied.wait() + request.cancel() + await asyncio.wait((request,)) + redis.release.set() + for _ in range(20): + await asyncio.sleep(0) + + assert request.cancelled() + assert redis.deleted is False + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_requests_waiting_on_a_cold_seed_whose_request_was_cancelled_add_their_cached_spend_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-seed-with-waiters" + redis: Final = _StalledFirstIncrementRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + started: Final = asyncio.Event() + + async def _read_before_the_seed_started() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await started.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + joined_before_start: Final = asyncio.ensure_future(_read_before_the_seed_started()) + await read.wait() + seed: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.stalled.wait() + started.set() + joined_after_start: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(20): + await asyncio.sleep(0) + seed.cancel() + await asyncio.gather(joined_before_start, joined_after_start) + + assert seed.cancelled() + assert redis.values[counter] == 6.0 + + +class _HeldIncrementRedis(_SharedRedis): + """Holds each increment before Redis runs it until its release event in ``releases`` is set, counting reads.""" + + def __init__(self, held_increments: int) -> None: + super().__init__() + self.reads = 0 + self.arrived: Final = tuple(asyncio.Event() for _ in range(held_increments)) + self.releases: Final = tuple(asyncio.Event() for _ in range(held_increments)) + self.increments = 0 + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + self.reads += 1 + return await super().async_batch_get_cache(key_list) + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.reads += 1 + return await super().async_get_cache(key) + + async def async_increment(self, key: str, value: float, **_: object) -> float: + index: Final = self.increments + self.increments += 1 + if index < len(self.arrived): + self.arrived[index].set() + await self.releases[index].wait() + return await super().async_increment(key, value) + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_before_a_seed_started_takes_its_value_without_reading_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:joined-seed-value" + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + started: Final = asyncio.Event() + + async def _joined_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await started.wait() + reads_before_the_seed_landed: Final = redis.reads + ensuring: Final = asyncio.ensure_future( + ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + ) + for _ in range(10): + await asyncio.sleep(0) + redis.releases[0].set() + await ensuring + assert await read_batched_spend_counter(counter) == (6.0, True) + assert redis.reads == reads_before_the_seed_landed + + joined: Final = asyncio.ensure_future(_joined_request()) + await read.wait() + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + started.set() + await asyncio.gather(joined, first) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_while_a_seed_was_in_flight_does_not_add_again_after_it_lands( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:read-during-seed" + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + during: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(10): + await asyncio.sleep(0) + redis.releases[0].set() + await asyncio.gather(first, during) + + assert redis.values[counter] == 6.0 + + +class _HeldSecondGetRedis(_HeldIncrementRedis): + """Also answers the second GET with the value its key had when it arrived, holding the reply until released.""" + + def __init__(self) -> None: + super().__init__(held_increments=1) + self.gets = 0 + self.second_get_arrived: Final = asyncio.Event() + self.release_second_get: Final = asyncio.Event() + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.gets += 1 + value: Final = await super().async_get_cache(key) + if self.gets == 2: + self.second_get_arrived.set() + await self.release_second_get.wait() + return value + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_while_the_reseed_read_of_an_unbatched_request_is_out_is_not_added_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:seed-during-reseed-read" + redis: Final = _HeldSecondGetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + unbatched: Final = asyncio.ensure_future( + ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + ) + await redis.second_get_arrived.wait() + redis.releases[0].set() + await first + redis.release_second_get.set() + await unbatched + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_stalled_seed_finishing_does_not_hide_the_newer_seed_still_in_flight( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:older-seed-finishes" + redis: Final = _HeldIncrementRedis(held_increments=2) + _process(monkeypatch, redis, cached_spend=6.0) + + older: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + newer: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[1].wait() + redis.releases[0].set() + await older + del redis.values[counter] + late: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(10): + await asyncio.sleep(0) + redis.releases[1].set() + await asyncio.gather(newer, late) + + assert redis.values[counter] == 6.0 From 2e83346cb21fdfa8348aba47d1f01b9885748676 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 11:14:51 -0800 Subject: [PATCH 09/18] fix(proxy): give a cold seed's re-read its own time after a late add A request that waited most of the 50 ms on an add that then landed had almost no time left to re-read the counter, so the re-read timed out and the request added the spend again. The re-read now has 50 ms past the end of that wait --- litellm/proxy/db/spend_counter_reseed.py | 11 ++++--- .../proxy/proxy_server/test_spend_counters.py | 31 ++++++++++++++++++- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 0c9e85dd788..f47dc17bbd6 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -168,13 +168,14 @@ class SpendCounterReseed: 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 or is cancelled. 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 and on a - re-read 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.""" + 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, has that long again to finish a re-read, 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) loop: Final = asyncio.get_running_loop() join_by: Final = loop.time() + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS + re_read_by: Final = join_by + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS async def add_unless_covered(cold: SpendCounterRead) -> float | None: in_flight: Final = seed.in_flight @@ -186,7 +187,7 @@ class SpendCounterReseed: forget_spend_counter(counter_key) sent_at: Final = spend_counter_tick() try: - warm: Final = await asyncio.wait_for(is_warm(), timeout=max(0.0, join_by - loop.time())) + warm: Final = await asyncio.wait_for(is_warm(), timeout=max(0.0, re_read_by - loop.time())) except asyncio.TimeoutError: return await seed.run(add) if warm: diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 424168de3e6..c9342a1b895 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -25,7 +25,7 @@ import asyncio import threading from collections.abc import Awaitable, Callable from datetime import datetime -from typing import Final, Literal, cast +from typing import ClassVar, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call import pytest @@ -2830,6 +2830,35 @@ async def test_a_request_that_read_the_counter_while_a_seed_was_in_flight_does_n assert redis.values[counter] == 6.0 +class _SlowGetRedis(_HeldIncrementRedis): + """Also takes ``SECONDS`` to answer a GET.""" + + SECONDS: ClassVar[float] = 0.2 + + async def async_get_cache(self, key: str, **_: object) -> float | None: + await asyncio.sleep(self.SECONDS) + return await super().async_get_cache(key) + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_late_in_the_wait_leaves_time_to_re_read_the_counter_instead_of_adding_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:late-landing-seed" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.5) + redis: Final = _SlowGetRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + during: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await asyncio.sleep(0.4) + redis.releases[0].set() + await asyncio.gather(first, during) + + assert redis.values[counter] == 6.0 + + class _HeldSecondGetRedis(_HeldIncrementRedis): """Also answers the second GET with the value its key had when it arrived, holding the reply until released.""" From 0ad56125d279ba0d63ea41714acbb3e22c8a099a Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 15:51:37 -0800 Subject: [PATCH 10/18] fix(proxy): set a missing cold counter with SET NX instead of a timed re-read A request whose cold read may predate this process's add now sends SET NX, EXPIRE and GET in one pipeline, the same shape as main's INCRBYFLOAT and EXPIRE, so it sends one write and no read, cancels nothing, and needs no second time limit It adds on top only when the counter holds less than what the landed add should have left, which keeps main's sum when a counter was lost and re-created below that The wait on other requests' adds now shares one limit across a request's counters, and the cold check read goes back to main's decorated helper so its span and service target stay spend_counters --- litellm/caching/redis_cache.py | 24 +- litellm/proxy/db/spend_counter_reseed.py | 91 ++-- litellm/proxy/proxy_server.py | 57 ++- .../spend_tracking/spend_counter_batch.py | 3 +- tests/unit/caching/test_redis_cache.py | 50 +++ .../proxy/proxy_server/test_spend_counters.py | 418 +++++++++++++----- 6 files changed, 475 insertions(+), 168 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 3b90cce211d..1e469e26990 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1473,15 +1473,29 @@ class RedisCache(BaseCache): await _redis_client.expire(key, ttl) return float(result) + async def _set_if_missing_with_ttl(self, key: str, value: float, ttl: int | None) -> float: + """SET NX, its TTL command and a GET in one round trip; the value the key holds afterwards.""" + async with self._async_commands().pipeline(transaction=False) as pipe: + pipe.set(name=key, value=value, nx=True) + if ttl is not None: + pipe.expire(key, ttl) + pipe.get(key) + replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics + current: Final = _LUA_OPTIONAL_TOTAL.validate_python(replies[-1]) + return value if current is None else current + @_redis_circuit_breaker_guard async def async_increment( self, - key, + key: str, value: float, ttl: int | None = None, parent_otel_span: Span | None = None, refresh_ttl: bool = False, + only_if_missing: bool = False, ) -> float: + """Add ``value`` to the key and return its new value. With ``only_if_missing``, a key that exists keeps its + value, which is returned, and only the TTL is refreshed.""" from redis.asyncio import Redis _redis_client: Final[Redis] = self.init_async_client() @@ -1489,8 +1503,12 @@ class RedisCache(BaseCache): _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) try: - result: Final = await self._incrbyfloat_with_ttl( - _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl + result: Final = await ( + self._set_if_missing_with_ttl(key=key, value=value, ttl=_used_ttl) + if only_if_missing + else self._incrbyfloat_with_ttl( + _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl + ) ) ## LOGGING ## diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index f47dc17bbd6..8a0d413bdbc 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -16,6 +16,7 @@ import asyncio import functools from collections import OrderedDict from collections.abc import Awaitable, Callable, Mapping +from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType @@ -30,7 +31,6 @@ 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, spend_counter_tick, @@ -75,11 +75,12 @@ _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, the tick it started at, and the - tick the last one landed at.""" + """This process adding its cached spend to one cold counter: the add in flight, the tick it started at, the tick + the last one landed at, and the least the counter held when that one landed.""" started_at: int = 0 landed_at: int = 0 + landed_spend: float = 0.0 in_flight: asyncio.Future[float | None] | None = None def on_this_loop(self) -> bool: @@ -107,6 +108,18 @@ def _cached_spend_seed(spend_counter_cache: "DualCache", counter_key: str) -> _C return _CachedSpendSeed() +_cold_seed_waited: Final[ContextVar[float]] = ContextVar("cold_spend_seed_waited", default=0.0) + + +async def _wait_for_cold_seed(in_flight: asyncio.Future[float | None]) -> None: + """Wait on ``in_flight`` for what this context has left of ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS``.""" + loop: Final = asyncio.get_running_loop() + waited: Final = _cold_seed_waited.get() + started: Final = loop.time() + await asyncio.wait((in_flight,), timeout=max(0.0, SPEND_COUNTER_COLD_SEED_JOIN_SECONDS - waited)) + _cold_seed_waited.set(waited + loop.time() - started) + + 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 @@ -161,56 +174,66 @@ class SpendCounterReseed: spend_counter_cache: "DualCache", counter_key: str, read: SpendCounterRead, + spend: float, add: Callable[[], Awaitable[float | None]], - is_warm: Callable[[], Awaitable[bool]], + set_if_missing: Callable[[], Awaitable[float | None]], ) -> 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 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 or is cancelled. 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, has that long again to finish a re-read, 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.""" + """Run ``add``, this process adding ``spend`` to a cold counter, once for all of this process's 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 or is cancelled. A request whose read went + out before an add landed, and one still waiting on an add after ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across + its counters, runs ``set_if_missing`` in place of ``add``, then runs ``add`` too if the counter holds less than + the spend it should already hold. 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 run ``add``.""" seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) - loop: Final = asyncio.get_running_loop() - join_by: Final = loop.time() + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS - re_read_by: Final = join_by + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS - async def add_unless_covered(cold: SpendCounterRead) -> float | None: + async def add_spend() -> float | None: + added: Final = await add() + seed.landed_spend = spend + return added + + async def add_if_short() -> float | None: + least: Final = min(spend, seed.landed_spend) if seed.landed_at >= read.sent_at else spend + current: Final = await set_if_missing() + if current is not None and current < least: + return await add_spend() + seed.landed_spend = least + return current + + async def add_unless_covered() -> 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.sent_at: - return await seed.run(add) - forget_spend_counter(counter_key) - sent_at: Final = spend_counter_tick() - try: - warm: Final = await asyncio.wait_for(is_warm(), timeout=max(0.0, re_read_by - loop.time())) - except asyncio.TimeoutError: - return await seed.run(add) - if 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())) + return await seed.run(add_spend if seed.landed_at < read.sent_at else add_if_short) + covers: Final = seed.started_at > read.returned_at + await _wait_for_cold_seed(in_flight) if not in_flight.done(): - return await seed.run(add) + return await seed.run(add_if_short) if covers and not in_flight.cancelled(): return in_flight.result() - return await add_unless_covered(cold) + return await add_unless_covered() - added: Final = await add_unless_covered(read) + added: Final = await add_unless_covered() if added is not None: record_spend_counter_value(counter_key, added) @staticmethod @with_service_target(SPEND_COUNTERS_TARGET) - async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None: - """Apply local deltas after an in-flight reseed establishes the spend balance.""" + async def increment_in_memory( + spend_counter_cache: "DualCache", counter_key: str, increment: float, only_if_missing: bool = False + ) -> float | None: + """Apply local deltas after an in-flight reseed establishes the spend balance. With ``only_if_missing``, a + counter that exists keeps its value, which is returned.""" lock: Final = await SpendCounterReseed._get_lock(counter_key) async with lock: + current: Final[float | None] = ( # pyright: ignore[reportUnknownVariableType] # untyped in-memory cache + spend_counter_cache.in_memory_cache.get_cache(key=counter_key) # pyright: ignore[reportUnknownMemberType] # untyped in-memory cache + if only_if_missing + else None + ) + if current is not None: + return float(current) return await spend_counter_cache.async_increment_cache( key=counter_key, value=increment, local_only=True, refresh_ttl=True ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 741f7d58f1d..7c1676e22c1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -930,7 +930,6 @@ from litellm.proxy.spend_tracking.spend_capture_rate import ( from litellm.proxy.spend_tracking.spend_counter_batch import ( SPEND_COUNTERS_TARGET, PendingSpendIncrement, - SpendCounterRead, active_spend_counter_batch, forget_spend_counter, post_call_counter_keys, @@ -4087,8 +4086,9 @@ async def ensure_spend_counter_initialized( counter_key: str, source_cache_key: str | list[str], ) -> None: - read: Final = await _cold_spend_counter_read(counter_key) - if read is not None: + is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) + if is_warm is False: + read: Final = spend_counter_read(counter_key) # Shares the per-counter lock with get_current_spend. db_spend: Final = await SpendCounterReseed.coalesced( prisma_client=prisma_client, @@ -4104,8 +4104,11 @@ async def ensure_spend_counter_initialized( spend_counter_cache=spend_counter_cache, counter_key=counter_key, read=read, + spend=base_spend, add=lambda: increment_spend_counter_cache(counter_key=counter_key, increment=base_spend), - is_warm=lambda: _is_spend_counter_warm_outside_batch(counter_key), + set_if_missing=lambda: increment_spend_counter_cache( + counter_key=counter_key, increment=base_spend, only_if_missing=True + ), ) @@ -4160,22 +4163,14 @@ _ensure_window_spend_counter_initialized: Final = ensure_window_spend_counter_in @with_service_target(SPEND_COUNTERS_TARGET) async def _is_spend_counter_cache_warm(counter_key: str) -> bool: - return await _cold_spend_counter_read(counter_key) is None - - -async def _cold_spend_counter_read(counter_key: str) -> SpendCounterRead | None: - """When the read that found the counter missing went out and came back, None when the counter has a value. A read - outside the batch counts as now, since ``SpendCounterReseed.coalesced`` reads the counter again after it.""" batched: Final = await read_batched_spend_counter(counter_key) - if batched is not None and batched[0] is not None: - spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched[0]) - return None - if batched is None and await _is_spend_counter_warm_outside_batch(counter_key): - return None - return spend_counter_read(counter_key) + if batched is not None: + batched_value, _ = batched + if batched_value is None: + return False + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched_value) + return True - -async def _is_spend_counter_warm_outside_batch(counter_key: str) -> bool: if spend_counter_cache.redis_cache is not None: try: current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache( @@ -4217,13 +4212,24 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool: @with_service_target(SPEND_COUNTERS_TARGET) -async def increment_spend_counter_cache(counter_key: str, increment: float) -> float | None: +async def increment_spend_counter_cache( + counter_key: str, increment: float, only_if_missing: bool = False +) -> float | None: if spend_counter_cache.redis_cache is not None: try: - current_value: Final = await spend_counter_cache.redis_cache.async_increment( - key=counter_key, - value=increment, - refresh_ttl=True, + current_value: Final = await ( + spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=increment, + refresh_ttl=True, + only_if_missing=True, + ) + if only_if_missing + else spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=increment, + refresh_ttl=True, + ) ) except Exception: await invalidate_spend_counter(counter_key=counter_key) @@ -4236,7 +4242,10 @@ async def increment_spend_counter_cache(counter_key: str, increment: float) -> f return current_value return await SpendCounterReseed.increment_in_memory( - spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment + spend_counter_cache=spend_counter_cache, + counter_key=counter_key, + increment=increment, + only_if_missing=only_if_missing, ) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 55c4fab6893..8ff709086bb 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -188,7 +188,8 @@ def active_spend_counter_batch() -> SpendCounterBatch | None: def spend_counter_read(counter_key: str) -> SpendCounterRead: - """When this context's read of the counter went out and came back: the open batch's MGET for it, else now.""" + """When this context's read of the counter went out and came back: the open batch's MGET for it, else now, since + ``SpendCounterReseed.coalesced`` reads the counter again after a read outside the batch.""" batch: Final = _active_batch.get() batched: Final = batch.counter_read(counter_key) if batch is not None else None if batched is not None: diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..356a1f34985 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -8,6 +8,7 @@ from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeAsyncRedis, FakeServer import litellm from litellm._service_logger import ServiceLogging @@ -1474,6 +1475,55 @@ async def test_async_increment_pipelines_the_ttl_command( assert client.round_trips == expected_round_trips +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("existing", "expected"), + [ + pytest.param(None, 1.5, id="missing key: set to the value"), + pytest.param(0.25, 0.25, id="key below the value: kept"), + pytest.param(4.0, 4.0, id="key above the value: kept"), + ], +) +async def test_async_increment_only_if_missing_sets_a_missing_key_and_refreshes_the_ttl( + monkeypatch, redis_no_ping, existing, expected +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = FakeAsyncRedis(server=FakeServer()) + if existing is not None: + await client.set("ns:spend:key:k", existing, ex=5) + + with patch.object(redis_cache, "init_async_client", return_value=client): + current = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + ) + after_add = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=True) + + assert current == expected + assert after_add == expected + 2.0 + assert await client.ttl("ns:spend:key:k") == 60 + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = MagicMock() + client.pipeline.side_effect = ConnectionError("connection reset") + failure_hook = AsyncMock() + + with ( + patch.object(redis_cache, "init_async_client", return_value=client), + patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), + pytest.raises(ConnectionError), + ): + await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True) + await asyncio.sleep(0) + + assert failure_hook.await_args.kwargs["call_type"] == "async_increment" + assert redis_cache._circuit_breaker._hard_failure_count == 1 + + class _SetRecordingPipeline: def __init__(self) -> None: self.sets: list[tuple[str, str, timedelta | None]] = [] diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index c9342a1b895..e4a37721bba 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -25,7 +25,7 @@ import asyncio import threading from collections.abc import Awaitable, Callable from datetime import datetime -from typing import ClassVar, Final, Literal, cast +from typing import Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call import pytest @@ -37,13 +37,13 @@ from opentelemetry.trace import StatusCode import litellm import litellm.proxy.proxy_server as ps -from litellm._internal_context import current_service_target +from litellm._internal_context import current_service_target, 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_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.constants import AUTH_OBJECTS_TARGET, GLOBAL_PROXY_SPEND_CACHE_KEY from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.tag import LiteLLM_TagTable @@ -55,6 +55,7 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_u from litellm.proxy.db import spend_counter_reseed from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team from litellm.proxy.spend_tracking.spend_counter_batch import ( + SPEND_COUNTERS_TARGET, active_spend_counter_batch, read_batched_spend_counter, spend_counter_batch_scope, @@ -2044,12 +2045,16 @@ class _SharedRedis: await asyncio.sleep(0) return self.values.get(key) - async def async_increment(self, key: str, value: float, **_: object) -> float: + def _apply(self, key: str, value: float, only_if_missing: bool = False, **_: object) -> float: + if not only_if_missing or key not in self.values: + self.values[key] = self.values.get(key, 0.0) + value + return self.values[key] + + async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: await asyncio.sleep(0) if self.failing: raise ConnectionError("redis down") - self.values[key] = self.values.get(key, 0.0) + value - return self.values[key] + return self._apply(key, value, **kwargs) async def async_delete_cache(self, key: str) -> None: await asyncio.sleep(0) @@ -2091,7 +2096,9 @@ async def test_overlapping_cold_starts_in_one_process_add_its_cached_spend_once( await asyncio.gather(*(_cold_start(redis, counter) for _ in range(5))) - seeded: Final = redis.values.get(counter) if redis is not None else counters.in_memory_cache.get_cache(key=counter) + seeded: Final[float | None] = ( + redis.values.get(counter) if redis is not None else counters.in_memory_cache.get_cache(key=counter) + ) assert seeded == 6.0 @@ -2205,9 +2212,8 @@ class _HeldReplyRedis(_SharedRedis): 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] + async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + total: Final = self._apply(key, value, **kwargs) self.applied.set() await self.release.wait() return total @@ -2240,12 +2246,12 @@ def test_a_cold_start_on_another_event_loop_adds_on_its_own(monkeypatch: pytest. release: Final = threading.Event() class _ThreadHeldRedis(_SharedRedis): - async def async_increment(self, key: str, value: float, **_: object) -> float: + async def async_increment(self, key: str, value: float, **kwargs: bool) -> 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) + return await super().async_increment(key, value, **kwargs) redis: Final = _ThreadHeldRedis() _process(monkeypatch, redis, cached_spend=6.0) @@ -2326,12 +2332,12 @@ class _DroppedFirstIncrementRedis(_SharedRedis): super().__init__() self.dropped = False - async def async_increment(self, key: str, value: float, **_: object) -> float: + async def async_increment(self, key: str, value: float, **kwargs: bool) -> 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) + return await super().async_increment(key, value, **kwargs) @pytest.mark.asyncio @@ -2357,16 +2363,18 @@ class _StalledFirstIncrementRedis(_SharedRedis): self.stalled: Final = asyncio.Event() self.release: Final = asyncio.Event() - async def async_increment(self, key: str, value: float, **_: object) -> float: + async def async_increment(self, key: str, value: float, **kwargs: bool) -> 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) + return await super().async_increment(key, value, **kwargs) @pytest.mark.asyncio -async def test_cold_starts_stop_waiting_on_a_stalled_seed_and_add_their_own(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_cold_starts_stop_waiting_on_a_stalled_seed_and_set_the_missing_counter_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: counter: Final = "spend:end_user:stalled-seed" redis: Final = _StalledFirstIncrementRedis() _process(monkeypatch, redis, cached_spend=6.0) @@ -2381,7 +2389,7 @@ async def test_cold_starts_stop_waiting_on_a_stalled_seed_and_add_their_own(monk with pytest.raises(TimeoutError): await stalled - assert seeded_while_stalled == 18.0 + assert seeded_while_stalled == 6.0 class _FakeRedisHoldingIncrements(FakeAsyncRedis): @@ -2405,11 +2413,22 @@ async def test_a_cold_start_after_the_redis_breaker_opened_fails_fast_instead_of client: Final = _FakeRedisHoldingIncrements(FakeServer()) redis: Final = _RedisCacheOnClient(client, breaker) _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + opened: Final = asyncio.Event() + + async def _read_before_the_seed_started() -> None: + with spend_counter_batch_scope(redis, frozenset((counter,))): + await _read_cold(counter) + read.set() + await opened.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + later: Final = asyncio.ensure_future(_read_before_the_seed_started()) + await read.wait() 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)) + opened.set() for _ in range(100): await asyncio.sleep(0) client.release.set() @@ -2453,64 +2472,6 @@ async def test_a_seed_that_lands_while_a_cold_read_is_on_its_way_back_is_not_add assert redis.values[counter] == 6.0 -class _HeldReadRedis(_SharedRedis): - """Answers a GET with the value its key 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() - self.reads_after_hold = 0 - - async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: - self.reads_after_hold += self.hold_reads - return await super().async_batch_get_cache(key_list) - - async def async_get_cache(self, key: str, **_: object) -> float | None: - self.reads_after_hold += self.hold_reads - value: Final = self.values.get(key) - if self.hold_reads: - self.read_arrived.set() - await self.release_read.wait() - return value - - -@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.""" @@ -2598,11 +2559,11 @@ async def test_a_cold_read_sent_before_an_add_started_does_not_join_it_when_the_ @pytest.mark.asyncio -async def test_a_re_read_redis_does_not_answer_stops_at_the_wait_budget_and_the_request_adds_its_cached_spend( +async def test_a_counter_lost_and_re_created_below_the_seed_gets_the_cached_spend_on_top( monkeypatch: pytest.MonkeyPatch, ) -> None: - counter: Final = "spend:end_user:stalled-reread" - redis: Final = _HeldReadRedis() + counter: Final = "spend:end_user:re-created-below" + redis: Final = _SharedRedis() _process(monkeypatch, redis, cached_spend=6.0) read: Final = asyncio.Event() seeded: Final = asyncio.Event() @@ -2612,8 +2573,36 @@ async def test_a_re_read_redis_does_not_answer_stops_at_the_wait_budget_and_the_ await _read_cold(counter) read.set() await seeded.wait() - del redis.values[counter] - redis.hold_reads = True + redis.values[counter] = 0.5 + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + assert await read_batched_spend_counter(counter) == (6.5, True) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.5 + + +@pytest.mark.asyncio +async def test_a_counter_holding_the_seed_is_left_alone_by_a_request_whose_source_cache_moved_past_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:source-moved" + redis: Final = _SharedRedis() + counters: Final = _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() + ps.user_api_key_cache.in_memory_cache.set_cache(key=_COLD_SEED_SOURCE, value={"spend": 6.25}) await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) async def _first_request() -> None: @@ -2621,10 +2610,153 @@ async def test_a_re_read_redis_does_not_answer_stops_at_the_wait_budget_and_the_ await _cold_start(redis, counter) seeded.set() - await asyncio.wait_for(asyncio.gather(_stale_request(), _first_request()), timeout=2.0) + await asyncio.gather(_stale_request(), _first_request()) assert redis.values[counter] == 6.0 - assert redis.reads_after_hold == 1 + assert counters.in_memory_cache.get_cache(key=counter) == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_done_waiting_on_a_stalled_seed_adds_when_the_counter_holds_less_than_its_cached_spend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:below-own-spend" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((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) + while_stalled: Final = redis.values[counter] + redis.releases[0].set() + await stalled + + assert while_stalled == 6.5 + + +@pytest.mark.asyncio +async def test_a_counter_without_redis_set_only_if_missing_keeps_the_value_it_holds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:no-redis-set-if-missing" + monkeypatch.setattr(ps, "spend_counter_cache", DualCache(in_memory_cache=InMemoryCache())) + + first: Final = await ps.increment_spend_counter_cache(counter, 6.0, only_if_missing=True) + second: Final = await ps.increment_spend_counter_cache(counter, 4.0, only_if_missing=True) + added: Final = await ps.increment_spend_counter_cache(counter, 1.0) + + assert (first, second, added) == (6.0, 6.0, 7.0) + + +@pytest.mark.asyncio +async def test_a_counter_set_by_a_request_done_waiting_counts_as_the_seed_for_a_later_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:set-counts-as-seed" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + set_by_a_waiter: 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 set_by_a_waiter.wait() + redis.values[counter] = 0.5 + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + stale: Final = asyncio.ensure_future(_stale_request()) + await read.wait() + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + await _cold_start(redis, counter) + set_by_a_waiter.set() + await stale + after_the_stale_request: Final = redis.values[counter] + redis.releases[0].set() + await stalled + + assert after_the_stale_request == 6.5 + + +class _CommandCountingRedis(_SharedRedis): + def __init__(self) -> None: + super().__init__() + self.commands: Final[list[str]] = [] # mutable-ok: the commands the test asserts on + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + self.commands.append("mget") + return await super().async_batch_get_cache(key_list) + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.commands.append("get") + return await super().async_get_cache(key) + + async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + self.commands.append("set_if_missing" if kwargs.get("only_if_missing") else "add") + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +async def test_a_request_that_read_before_a_seed_landed_sends_one_write_and_keeps_its_value_for_later_reads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:one-write" + redis: Final = _CommandCountingRedis() + _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() + redis.commands.clear() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + assert await read_batched_spend_counter(counter) == (6.0, True) + assert await ps.get_current_spend(counter_key=counter, fallback_spend=6.0) == 6.0 + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.commands == ["set_if_missing"] + + +@pytest.mark.asyncio +async def test_a_request_waits_on_other_requests_adds_for_the_limit_in_total_across_its_counters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + limit: Final = 0.3 + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", limit) + keys: Final = tuple(f"spend:end_user:shared-wait-{index}" for index in range(4)) + redis: Final = _HeldIncrementRedis(held_increments=len(keys)) + _process(monkeypatch, redis, cached_spend=6.0) + stalled: Final = [asyncio.ensure_future(_cold_start(redis, key)) for key in keys] + for arrived in redis.arrived: + await arrived.wait() + + loop: Final = asyncio.get_running_loop() + started: Final = loop.time() + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset(keys)): + for key in keys: + await ps.ensure_spend_counter_initialized(counter_key=key, source_cache_key=_COLD_SEED_SOURCE) + waited: Final = loop.time() - started + for release in redis.releases: + release.set() + await asyncio.gather(*stalled) + + assert limit <= waited < 2 * limit class _LostAndRecreatedAfterFailedMGetRedis(_SharedRedis): @@ -2657,9 +2789,9 @@ class _LostAndRecreatedAfterFailedMGetRedis(_SharedRedis): self.cold_gets += 1 return await super().async_get_cache(key) - async def async_increment(self, key: str, value: float, **_: object) -> float: + async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: self._other_process_increments_after_the_cold_gets() - return await super().async_increment(key, value) + return await super().async_increment(key, value, **kwargs) @pytest.mark.asyncio @@ -2686,8 +2818,8 @@ class _AppliedHeldThenResetRedis(_HeldReplyRedis): super().__init__() self.deleted = False - async def async_increment(self, key: str, value: float, **_: object) -> float: - await super().async_increment(key, value) + async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + await super().async_increment(key, value, **kwargs) raise ConnectionError("connection reset") async def async_delete_cache(self, key: str) -> None: @@ -2721,6 +2853,7 @@ async def test_requests_waiting_on_a_cold_seed_whose_request_was_cancelled_add_t monkeypatch: pytest.MonkeyPatch, ) -> None: counter: Final = "spend:end_user:cancelled-seed-with-waiters" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) redis: Final = _StalledFirstIncrementRedis() _process(monkeypatch, redis, cached_spend=6.0) read: Final = asyncio.Event() @@ -2742,7 +2875,7 @@ async def test_requests_waiting_on_a_cold_seed_whose_request_was_cancelled_add_t for _ in range(20): await asyncio.sleep(0) seed.cancel() - await asyncio.gather(joined_before_start, joined_after_start) + await asyncio.wait_for(asyncio.gather(joined_before_start, joined_after_start), timeout=5) assert seed.cancelled() assert redis.values[counter] == 6.0 @@ -2766,13 +2899,13 @@ class _HeldIncrementRedis(_SharedRedis): self.reads += 1 return await super().async_get_cache(key) - async def async_increment(self, key: str, value: float, **_: object) -> float: + async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: index: Final = self.increments self.increments += 1 if index < len(self.arrived): self.arrived[index].set() await self.releases[index].wait() - return await super().async_increment(key, value) + return await super().async_increment(key, value, **kwargs) @pytest.mark.asyncio @@ -2830,23 +2963,11 @@ async def test_a_request_that_read_the_counter_while_a_seed_was_in_flight_does_n assert redis.values[counter] == 6.0 -class _SlowGetRedis(_HeldIncrementRedis): - """Also takes ``SECONDS`` to answer a GET.""" - - SECONDS: ClassVar[float] = 0.2 - - async def async_get_cache(self, key: str, **_: object) -> float | None: - await asyncio.sleep(self.SECONDS) - return await super().async_get_cache(key) - - @pytest.mark.asyncio -async def test_a_seed_that_lands_late_in_the_wait_leaves_time_to_re_read_the_counter_instead_of_adding_again( - monkeypatch: pytest.MonkeyPatch, -) -> None: +async def test_a_seed_that_lands_late_in_the_wait_is_not_added_again(monkeypatch: pytest.MonkeyPatch) -> None: counter: Final = "spend:end_user:late-landing-seed" monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.5) - redis: Final = _SlowGetRedis(held_increments=1) + redis: Final = _HeldIncrementRedis(held_increments=1) _process(monkeypatch, redis, cached_spend=6.0) first: Final = asyncio.ensure_future(_cold_start(redis, counter)) @@ -2859,6 +2980,35 @@ async def test_a_seed_that_lands_late_in_the_wait_leaves_time_to_re_read_the_cou assert redis.values[counter] == 6.0 +@pytest.mark.asyncio +async def test_a_request_joining_an_add_that_started_after_its_read_came_back_sends_no_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:joined-no-write" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + + async def _joining_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await redis.arrived[0].wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + joining: Final = asyncio.ensure_future(_joining_request()) + await read.wait() + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + for _ in range(100): + await asyncio.sleep(0) + redis.releases[0].set() + await asyncio.gather(joining, first) + + assert (redis.values[counter], redis.increments) == (6.0, 1) + + class _HeldSecondGetRedis(_HeldIncrementRedis): """Also answers the second GET with the value its key had when it arrived, holding the reply until released.""" @@ -2917,7 +3067,63 @@ async def test_a_stalled_seed_finishing_does_not_hide_the_newer_seed_still_in_fl late: Final = asyncio.ensure_future(_cold_start(redis, counter)) for _ in range(10): await asyncio.sleep(0) + commands_while_the_newer_seed_was_in_flight: Final = redis.increments redis.releases[1].set() await asyncio.gather(newer, late) + assert commands_while_the_newer_seed_was_in_flight == 2 assert redis.values[counter] == 6.0 + + +class _TargetRecordingRedis(_SharedRedis): + def __init__(self) -> None: + super().__init__() + self.targets: Final[list[str | None]] = [] # mutable-ok: the targets the test asserts on + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.targets.append(current_service_target()) + return 1.0 + + +@pytest.mark.asyncio +async def test_the_cold_check_read_reports_under_spend_counters_inside_another_phases_target( + monkeypatch: pytest.MonkeyPatch, +) -> None: + redis: Final = _TargetRecordingRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + with service_target(AUTH_OBJECTS_TARGET): + await ps.ensure_spend_counter_initialized(counter_key="spend:key:target", source_cache_key=_COLD_SEED_SOURCE) + + assert redis.targets == [SPEND_COUNTERS_TARGET] + + +@pytest.mark.asyncio +async def test_a_stale_read_on_a_real_client_seeds_a_missing_counter_with_its_ttl_and_keeps_an_existing_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:real-client" + client: Final = FakeAsyncRedis(server=FakeServer()) + redis: Final = _RedisCacheOnClient(client) + redis.default_ttl = 120 + _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(redis, frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + await client.delete(counter) + 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(_stale_request(), _first_request()) + + assert float(await client.get(counter) or 0) == 6.0 + assert 0 < await client.ttl(counter) <= 120 From 838467fea6a25d9920c595303bc268502cf1f4c1 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 2 Nov 2026 17:43:51 -0800 Subject: [PATCH 11/18] test(proxy): check the shared cold seed wait limit with a controlled clock The test measured real elapsed time, which tests/unit/AGENTS.md forbids, so it now reports each wait's limit under a clock the test controls --- .../proxy/proxy_server/test_spend_counters.py | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index e4a37721bba..95615f1810e 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -23,7 +23,7 @@ from __future__ import annotations import asyncio import threading -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Collection from datetime import datetime from typing import Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call @@ -2737,8 +2737,7 @@ async def test_a_request_that_read_before_a_seed_landed_sends_one_write_and_keep async def test_a_request_waits_on_other_requests_adds_for_the_limit_in_total_across_its_counters( monkeypatch: pytest.MonkeyPatch, ) -> None: - limit: Final = 0.3 - monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", limit) + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.3) keys: Final = tuple(f"spend:end_user:shared-wait-{index}" for index in range(4)) redis: Final = _HeldIncrementRedis(held_increments=len(keys)) _process(monkeypatch, redis, cached_spend=6.0) @@ -2746,17 +2745,31 @@ async def test_a_request_waits_on_other_requests_adds_for_the_limit_in_total_acr for arrived in redis.arrived: await arrived.wait() - loop: Final = asyncio.get_running_loop() - started: Final = loop.time() + clock: Final = [0.0] # mutable-ok: the time the event loop reports to the wait + limits: Final[list[float]] = [] # mutable-ok: the limit of each wait the test asserts on + + real_wait: Final = asyncio.wait + + async def _waits_on_an_add_take_a_tenth_of_a_second( + waited_on: Collection[asyncio.Future[object]], timeout: float + ) -> tuple[set[asyncio.Future[object]], set[asyncio.Future[object]]]: + if any(isinstance(future, asyncio.Task) for future in waited_on): + return await real_wait(waited_on, timeout=timeout) + limits.append(timeout) + clock[0] += 0.1 + return set(), set(waited_on) + + monkeypatch.setattr(asyncio.get_running_loop(), "time", lambda: clock[0]) + monkeypatch.setattr(asyncio, "wait", _waits_on_an_add_take_a_tenth_of_a_second) with spend_counter_batch_scope(cast(RedisCache, redis), frozenset(keys)): for key in keys: await ps.ensure_spend_counter_initialized(counter_key=key, source_cache_key=_COLD_SEED_SOURCE) - waited: Final = loop.time() - started + monkeypatch.undo() for release in redis.releases: release.set() await asyncio.gather(*stalled) - assert limit <= waited < 2 * limit + assert limits == pytest.approx([0.3, 0.2, 0.1, 0.0]) class _LostAndRecreatedAfterFailedMGetRedis(_SharedRedis): From 4360ef7a8e1444fd9707eb0d10a8e1b578554ba0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 03:26:09 -0800 Subject: [PATCH 12/18] fix(proxy): keep a late cold counter write to what a plain add can send A Redis user denied SET, GET or @slow still runs main's INCRBYFLOAT and EXPIRE, but the SET NX and GET pipeline failed there, so bursts main served got 503s or lost their reservation and counter The pipeline is now SETNX, INCRBYFLOAT k 0 and EXPIRE. The read back and TTL are the commands a plain add sends, and a user denied only SETNX reads back what the key holds, so a short counter takes main's add. A counter set this way is stored the way an add stores it A request that left a peer's counter alone also moved the landed stamp, so a later fresh read looked stale and lost its add. Only a plain add moves that stamp now The read back rounds to 17 decimal places, so a counter just under the spend by less than that no longer counts as short --- litellm/caching/redis_cache.py | 19 ++-- litellm/constants.py | 1 + litellm/proxy/db/spend_counter_reseed.py | 15 +-- tests/unit/caching/test_redis_cache.py | 100 +++++++++++++++++- .../proxy/proxy_server/test_spend_counters.py | 79 +++++++++++++- 5 files changed, 198 insertions(+), 16 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 1e469e26990..d9e43af77b5 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -139,6 +139,7 @@ _INCREMENT_IF_PRESENT_LUA: Final = ( ) _LUA_COUNT: Final = TypeAdapter(int) +_INCRBYFLOAT_REPLY: Final = TypeAdapter(float) _LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) @@ -1474,15 +1475,21 @@ class RedisCache(BaseCache): return float(result) async def _set_if_missing_with_ttl(self, key: str, value: float, ttl: int | None) -> float: - """SET NX, its TTL command and a GET in one round trip; the value the key holds afterwards.""" + """SETNX, a zero INCRBYFLOAT that reads the key back, and its TTL command in one round trip; the value the key + holds afterwards. Those last two are the commands a plain add sends, so a user denied only SETNX reads back what + the key holds, which is 0 for a key that was missing.""" + from redis.exceptions import NoPermissionError + async with self._async_commands().pipeline(transaction=False) as pipe: - pipe.set(name=key, value=value, nx=True) + pipe.setnx(name=key, value=value) + pipe.incrbyfloat(name=key, amount=0) if ttl is not None: pipe.expire(key, ttl) - pipe.get(key) - replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics - current: Final = _LUA_OPTIONAL_TOTAL.validate_python(replies[-1]) - return value if current is None else current + replies: Final[Sequence[object]] = await pipe.execute(raise_on_error=False) # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics + for index, reply in enumerate(replies): # pyright: ignore[reportUnknownArgumentType] # redis generics + if isinstance(reply, Exception) and not (index == 0 and isinstance(reply, NoPermissionError)): + raise reply + return _INCRBYFLOAT_REPLY.validate_python(replies[1]) @_redis_circuit_breaker_guard async def async_increment( diff --git a/litellm/constants.py b/litellm/constants.py index 199d7bc59bc..3152c11aff1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1925,6 +1925,7 @@ RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SES 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 +SPEND_COUNTER_INCRBYFLOAT_ROUNDING: Final = 1e-16 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 8a0d413bdbc..efdac62dd60 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -24,7 +24,11 @@ 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_COLD_SEED_JOIN_SECONDS, SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE +from litellm.constants import ( + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS, + SPEND_COUNTER_INCRBYFLOAT_ROUNDING, + 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 @@ -76,7 +80,7 @@ _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, the tick it started at, the tick - the last one landed at, and the least the counter held when that one landed.""" + the last plain add landed at, and the spend that add added.""" started_at: int = 0 landed_at: int = 0 @@ -94,7 +98,6 @@ class _CachedSpendSeed: self.in_flight = in_flight try: added: Final = await add() - self.landed_at = spend_counter_tick() in_flight.set_result(added) return added finally: @@ -183,21 +186,21 @@ class SpendCounterReseed: gets its result, and starts or joins another add when that one fails or is cancelled. A request whose read went out before an add landed, and one still waiting on an add after ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across its counters, runs ``set_if_missing`` in place of ``add``, then runs ``add`` too if the counter holds less than - the spend it should already hold. A request on another event loop than the add in flight, one that finds the + the spend it should already hold, beyond what INCRBYFLOAT's 17 decimal places can round away. 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 run ``add``.""" seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) async def add_spend() -> float | None: added: Final = await add() + seed.landed_at = spend_counter_tick() seed.landed_spend = spend return added async def add_if_short() -> float | None: least: Final = min(spend, seed.landed_spend) if seed.landed_at >= read.sent_at else spend current: Final = await set_if_missing() - if current is not None and current < least: + if current is not None and least - current > SPEND_COUNTER_INCRBYFLOAT_ROUNDING: return await add_spend() - seed.landed_spend = least return current async def add_unless_covered() -> float | None: diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 356a1f34985..dd24a8b6691 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fakeredis import FakeAsyncRedis, FakeServer +from redis.exceptions import NoPermissionError, ResponseError import litellm from litellm._service_logger import ServiceLogging @@ -1034,8 +1035,9 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) gc.collect() ServiceLogging() - assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == ( - redis_breaker_state_counts()["closed"] + assert ( + REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) + == (redis_breaker_state_counts()["closed"]) ) @@ -1504,6 +1506,100 @@ async def test_async_increment_only_if_missing_sets_a_missing_key_and_refreshes_ assert await client.ttl("ns:spend:key:k") == 60 +@pytest.mark.asyncio +@pytest.mark.parametrize(("value", "stored"), [(6.0, b"6"), (1e-05, b"0.00001")]) +async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_add_stores_it( + monkeypatch, redis_no_ping, value, stored +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = FakeAsyncRedis(server=FakeServer()) + + with patch.object(redis_cache, "init_async_client", return_value=client): + await redis_cache.async_increment(key="spend:key:k", value=value, ttl=60, only_if_missing=True) + + assert await client.get("ns:spend:key:k") == stored + + +class _RepliesPipeline: + """A pipeline whose execute returns the given replies the way ``raise_on_error=False`` does.""" + + def __init__(self, replies: list[object]) -> None: + self.replies = replies + self.commands: list[tuple[object, ...]] = [] + + async def __aenter__(self) -> "_RepliesPipeline": + return self + + async def __aexit__(self, *exc_info: object) -> None: + return None + + def setnx(self, name: str, value: float) -> None: + self.commands.append(("SETNX", name, value)) + + def incrbyfloat(self, name: str, amount: float) -> None: + self.commands.append(("INCRBYFLOAT", name, amount)) + + def expire(self, name: str, ttl: int) -> None: + self.commands.append(("EXPIRE", name, ttl)) + + async def execute(self, raise_on_error: bool = True) -> list[object]: + assert raise_on_error is False + return self.replies + + +@pytest.mark.asyncio +@pytest.mark.parametrize("read_back", [pytest.param(0.0, id="missing"), pytest.param(7.0, id="held")]) +async def test_async_increment_only_if_missing_reads_the_key_back_when_the_user_is_denied_setnx( + monkeypatch, redis_no_ping, read_back +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + pipeline = _RepliesPipeline([NoPermissionError("no permissions to run the 'setnx' command"), read_back, True]) + client = MagicMock() + client.pipeline.return_value = pipeline + failure_hook = AsyncMock() + + with ( + patch.object(redis_cache, "init_async_client", return_value=client), + patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), + ): + current = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + ) + await asyncio.sleep(0) + + assert current == read_back + assert pipeline.commands == [ + ("SETNX", "ns:spend:key:k", 1.5), + ("INCRBYFLOAT", "ns:spend:key:k", 0), + ("EXPIRE", "ns:spend:key:k", 60), + ] + failure_hook.assert_not_awaited() + assert redis_cache._circuit_breaker._hard_failure_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("replies", "raised"), + [ + pytest.param([ResponseError("LOADING"), 7.0, True], ResponseError, id="setnx-failed-for-another-reason"), + pytest.param([True, NoPermissionError("denied"), True], NoPermissionError, id="incrbyfloat-denied"), + pytest.param([True, 7.0, NoPermissionError("denied")], NoPermissionError, id="expire-denied"), + ], +) +async def test_async_increment_only_if_missing_raises_any_other_error_in_the_replies( + monkeypatch, redis_no_ping, replies, raised +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = MagicMock() + client.pipeline.return_value = _RepliesPipeline(replies) + + with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(raised): + await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True) + + @pytest.mark.asyncio async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 95615f1810e..3371c3a257b 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -2616,6 +2616,81 @@ async def test_a_counter_holding_the_seed_is_left_alone_by_a_request_whose_sourc assert counters.in_memory_cache.get_cache(key=counter) == 6.0 +@pytest.mark.asyncio +async def test_a_request_that_left_a_peers_counter_alone_does_not_cost_a_later_read_its_add( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:left-alone-then-fresh" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + stale_read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + fresh_read: Final = asyncio.Event() + left_alone: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + stale_read.set() + await fresh_read.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + left_alone.set() + + async def _first_request() -> None: + await stale_read.wait() + await _cold_start(redis, counter) + seeded.set() + + async def _fresh_request() -> None: + await seeded.wait() + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + del redis.values[counter] + await _read_cold(counter) + redis.values[counter] = 7.0 + fresh_read.set() + await left_alone.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + await asyncio.gather(_stale_request(), _first_request(), _fresh_request()) + + assert redis.values[counter] == 13.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("read_back", "expected"), + [ + pytest.param(0.00123, 0.00123, id="rounded-by-incrbyfloat-is-left-alone"), + pytest.param(0.0012299, 0.0024599000000000002, id="short-by-more-than-rounding-gets-the-add"), + ], +) +async def test_a_counter_reading_back_just_under_the_spend_is_only_topped_up_when_really_short( + monkeypatch: pytest.MonkeyPatch, read_back: float, expected: float +) -> None: + counter: Final = "spend:end_user:read-back-rounding" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=0.0012300000000000002) + 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() + redis.values[counter] = read_back + 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(_stale_request(), _first_request()) + + assert redis.values[counter] == pytest.approx(expected, abs=1e-15) + + @pytest.mark.asyncio async def test_a_request_done_waiting_on_a_stalled_seed_adds_when_the_counter_holds_less_than_its_cached_spend( monkeypatch: pytest.MonkeyPatch, @@ -2653,7 +2728,7 @@ async def test_a_counter_without_redis_set_only_if_missing_keeps_the_value_it_ho @pytest.mark.asyncio -async def test_a_counter_set_by_a_request_done_waiting_counts_as_the_seed_for_a_later_request( +async def test_a_counter_lost_after_a_request_done_waiting_set_it_gets_the_add_of_a_later_request( monkeypatch: pytest.MonkeyPatch, ) -> None: counter: Final = "spend:end_user:set-counts-as-seed" @@ -3085,7 +3160,7 @@ async def test_a_stalled_seed_finishing_does_not_hide_the_newer_seed_still_in_fl await asyncio.gather(newer, late) assert commands_while_the_newer_seed_was_in_flight == 2 - assert redis.values[counter] == 6.0 + assert redis.values[counter] == 12.0 class _TargetRecordingRedis(_SharedRedis): From c14135b99fd0d47966fa67dba216ce11d2206b66 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 04:07:49 -0800 Subject: [PATCH 13/18] fix(proxy): count a late set that created the counter as this process's add A request that set a missing counter and finished left no trace, so another request that already read cold added on top, and a stalled add that was then cancelled left the counter at twice the spend A set that created the counter now counts like an add landing. One that left a counter alone still does not --- litellm/caching/redis_cache.py | 13 +++- litellm/proxy/db/spend_counter_reseed.py | 31 +++++--- litellm/proxy/proxy_server.py | 11 ++- tests/unit/caching/test_redis_cache.py | 32 +++++++- .../proxy/proxy_server/test_spend_counters.py | 76 +++++++++++++++---- 5 files changed, 128 insertions(+), 35 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index d9e43af77b5..cb847ab2876 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1474,10 +1474,12 @@ class RedisCache(BaseCache): await _redis_client.expire(key, ttl) return float(result) - async def _set_if_missing_with_ttl(self, key: str, value: float, ttl: int | None) -> float: + async def _set_if_missing_with_ttl( + self, key: str, value: float, ttl: int | None, on_created: Callable[[], None] | None + ) -> float: """SETNX, a zero INCRBYFLOAT that reads the key back, and its TTL command in one round trip; the value the key holds afterwards. Those last two are the commands a plain add sends, so a user denied only SETNX reads back what - the key holds, which is 0 for a key that was missing.""" + the key holds, which is 0 for a key that was missing. ``on_created`` runs when the SETNX created the key.""" from redis.exceptions import NoPermissionError async with self._async_commands().pipeline(transaction=False) as pipe: @@ -1489,6 +1491,8 @@ class RedisCache(BaseCache): for index, reply in enumerate(replies): # pyright: ignore[reportUnknownArgumentType] # redis generics if isinstance(reply, Exception) and not (index == 0 and isinstance(reply, NoPermissionError)): raise reply + if on_created is not None and replies[0] == 1: + on_created() return _INCRBYFLOAT_REPLY.validate_python(replies[1]) @_redis_circuit_breaker_guard @@ -1500,9 +1504,10 @@ class RedisCache(BaseCache): parent_otel_span: Span | None = None, refresh_ttl: bool = False, only_if_missing: bool = False, + on_created: Callable[[], None] | None = None, ) -> float: """Add ``value`` to the key and return its new value. With ``only_if_missing``, a key that exists keeps its - value, which is returned, and only the TTL is refreshed.""" + value, which is returned, and only the TTL is refreshed, and ``on_created`` runs if this call created the key.""" from redis.asyncio import Redis _redis_client: Final[Redis] = self.init_async_client() @@ -1511,7 +1516,7 @@ class RedisCache(BaseCache): key = self.check_and_fix_namespace(key=key) try: result: Final = await ( - self._set_if_missing_with_ttl(key=key, value=value, ttl=_used_ttl) + self._set_if_missing_with_ttl(key=key, value=value, ttl=_used_ttl, on_created=on_created) if only_if_missing else self._incrbyfloat_with_ttl( _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index efdac62dd60..c901a4c2e8f 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -80,7 +80,7 @@ _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, the tick it started at, the tick - the last plain add landed at, and the spend that add added.""" + the last write of its spend landed at, and that spend.""" started_at: int = 0 landed_at: int = 0 @@ -179,26 +179,30 @@ class SpendCounterReseed: read: SpendCounterRead, spend: float, add: Callable[[], Awaitable[float | None]], - set_if_missing: Callable[[], Awaitable[float | None]], + set_if_missing: Callable[[Callable[[], None]], Awaitable[float | None]], ) -> None: """Run ``add``, this process adding ``spend`` to a cold counter, once for all of this process's 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 or is cancelled. A request whose read went out before an add landed, and one still waiting on an add after ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across - its counters, runs ``set_if_missing`` in place of ``add``, then runs ``add`` too if the counter holds less than - the spend it should already hold, beyond what INCRBYFLOAT's 17 decimal places can round away. A request on another event loop than the add in flight, one that finds the + its counters, runs ``set_if_missing``, which calls its argument if it created the counter, in place of ``add``, + then runs ``add`` too if the counter holds less than the spend it should already hold, beyond what + INCRBYFLOAT's 17 decimal places can round away. 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 run ``add``.""" seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) - async def add_spend() -> float | None: - added: Final = await add() + def spend_landed() -> None: seed.landed_at = spend_counter_tick() seed.landed_spend = spend + + async def add_spend() -> float | None: + added: Final = await add() + spend_landed() return added async def add_if_short() -> float | None: least: Final = min(spend, seed.landed_spend) if seed.landed_at >= read.sent_at else spend - current: Final = await set_if_missing() + current: Final = await set_if_missing(spend_landed) if current is not None and least - current > SPEND_COUNTER_INCRBYFLOAT_ROUNDING: return await add_spend() return current @@ -224,10 +228,14 @@ class SpendCounterReseed: @staticmethod @with_service_target(SPEND_COUNTERS_TARGET) async def increment_in_memory( - spend_counter_cache: "DualCache", counter_key: str, increment: float, only_if_missing: bool = False + spend_counter_cache: "DualCache", + counter_key: str, + increment: float, + only_if_missing: bool = False, + on_created: Callable[[], None] | None = None, ) -> float | None: """Apply local deltas after an in-flight reseed establishes the spend balance. With ``only_if_missing``, a - counter that exists keeps its value, which is returned.""" + counter that exists keeps its value, which is returned, and ``on_created`` runs if this call created it.""" lock: Final = await SpendCounterReseed._get_lock(counter_key) async with lock: current: Final[float | None] = ( # pyright: ignore[reportUnknownVariableType] # untyped in-memory cache @@ -237,9 +245,12 @@ class SpendCounterReseed: ) if current is not None: return float(current) - return await spend_counter_cache.async_increment_cache( + added: Final = await spend_counter_cache.async_increment_cache( key=counter_key, value=increment, local_only=True, refresh_ttl=True ) + if on_created is not None: + on_created() + return added @staticmethod async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7c1676e22c1..5e4c1801955 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4106,8 +4106,8 @@ async def ensure_spend_counter_initialized( read=read, spend=base_spend, add=lambda: increment_spend_counter_cache(counter_key=counter_key, increment=base_spend), - set_if_missing=lambda: increment_spend_counter_cache( - counter_key=counter_key, increment=base_spend, only_if_missing=True + set_if_missing=lambda on_created: increment_spend_counter_cache( + counter_key=counter_key, increment=base_spend, only_if_missing=True, on_created=on_created ), ) @@ -4213,7 +4213,10 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool: @with_service_target(SPEND_COUNTERS_TARGET) async def increment_spend_counter_cache( - counter_key: str, increment: float, only_if_missing: bool = False + counter_key: str, + increment: float, + only_if_missing: bool = False, + on_created: Callable[[], None] | None = None, ) -> float | None: if spend_counter_cache.redis_cache is not None: try: @@ -4223,6 +4226,7 @@ async def increment_spend_counter_cache( value=increment, refresh_ttl=True, only_if_missing=True, + on_created=on_created, ) if only_if_missing else spend_counter_cache.redis_cache.async_increment( @@ -4246,6 +4250,7 @@ async def increment_spend_counter_cache( counter_key=counter_key, increment=increment, only_if_missing=only_if_missing, + on_created=on_created, ) diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index dd24a8b6691..49e4979e04a 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1509,7 +1509,7 @@ async def test_async_increment_only_if_missing_sets_a_missing_key_and_refreshes_ @pytest.mark.asyncio @pytest.mark.parametrize(("value", "stored"), [(6.0, b"6"), (1e-05, b"0.00001")]) async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_add_stores_it( - monkeypatch, redis_no_ping, value, stored + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float, stored: bytes ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache = RedisCache(namespace="ns") @@ -1521,6 +1521,28 @@ async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_a assert await client.get("ns:spend:key:k") == stored +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("existing", "reported"), [pytest.param(None, ["created"], id="created"), pytest.param(7.0, [], id="kept")] +) +async def test_async_increment_only_if_missing_reports_only_a_key_it_created( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, existing: float | None, reported: list[str] +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = FakeAsyncRedis(server=FakeServer()) + if existing is not None: + await client.set("ns:spend:key:k", existing) + created: list[str] = [] + + with patch.object(redis_cache, "init_async_client", return_value=client): + await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, only_if_missing=True, on_created=lambda: created.append("created") + ) + + assert created == reported + + class _RepliesPipeline: """A pipeline whose execute returns the given replies the way ``raise_on_error=False`` does.""" @@ -1551,7 +1573,7 @@ class _RepliesPipeline: @pytest.mark.asyncio @pytest.mark.parametrize("read_back", [pytest.param(0.0, id="missing"), pytest.param(7.0, id="held")]) async def test_async_increment_only_if_missing_reads_the_key_back_when_the_user_is_denied_setnx( - monkeypatch, redis_no_ping, read_back + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, read_back: float ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache = RedisCache(namespace="ns") @@ -1559,17 +1581,19 @@ async def test_async_increment_only_if_missing_reads_the_key_back_when_the_user_ client = MagicMock() client.pipeline.return_value = pipeline failure_hook = AsyncMock() + created = MagicMock() with ( patch.object(redis_cache, "init_async_client", return_value=client), patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), ): current = await redis_cache.async_increment( - key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_created=created ) await asyncio.sleep(0) assert current == read_back + created.assert_not_called() assert pipeline.commands == [ ("SETNX", "ns:spend:key:k", 1.5), ("INCRBYFLOAT", "ns:spend:key:k", 0), @@ -1589,7 +1613,7 @@ async def test_async_increment_only_if_missing_reads_the_key_back_when_the_user_ ], ) async def test_async_increment_only_if_missing_raises_any_other_error_in_the_replies( - monkeypatch, redis_no_ping, replies, raised + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, replies: list[object], raised: type[Exception] ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache = RedisCache(namespace="ns") diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 3371c3a257b..246012b8bf8 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -22,6 +22,7 @@ Pins covered: from __future__ import annotations import asyncio +import contextlib import threading from collections.abc import Awaitable, Callable, Collection from datetime import datetime @@ -2045,12 +2046,21 @@ class _SharedRedis: await asyncio.sleep(0) return self.values.get(key) - def _apply(self, key: str, value: float, only_if_missing: bool = False, **_: object) -> float: + def _apply( + self, + key: str, + value: float, + only_if_missing: bool = False, + on_created: Callable[[], None] | None = None, + **_: object, + ) -> float: if not only_if_missing or key not in self.values: self.values[key] = self.values.get(key, 0.0) + value + if on_created is not None: + on_created() return self.values[key] - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: await asyncio.sleep(0) if self.failing: raise ConnectionError("redis down") @@ -2212,7 +2222,7 @@ class _HeldReplyRedis(_SharedRedis): self.applied: Final = asyncio.Event() self.release: Final = asyncio.Event() - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: total: Final = self._apply(key, value, **kwargs) self.applied.set() await self.release.wait() @@ -2246,7 +2256,7 @@ def test_a_cold_start_on_another_event_loop_adds_on_its_own(monkeypatch: pytest. release: Final = threading.Event() class _ThreadHeldRedis(_SharedRedis): - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: if threading.current_thread() is not threading.main_thread(): held.set() while not release.is_set(): @@ -2332,7 +2342,7 @@ class _DroppedFirstIncrementRedis(_SharedRedis): super().__init__() self.dropped = False - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: if not self.dropped: self.dropped = True await asyncio.sleep(0) @@ -2363,7 +2373,7 @@ class _StalledFirstIncrementRedis(_SharedRedis): self.stalled: Final = asyncio.Event() self.release: Final = asyncio.Event() - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: if not self.stalled.is_set(): self.stalled.set() await self.release.wait() @@ -2691,6 +2701,38 @@ async def test_a_counter_reading_back_just_under_the_spend_is_only_topped_up_whe assert redis.values[counter] == pytest.approx(expected, abs=1e-15) +@pytest.mark.asyncio +async def test_a_counter_set_while_the_stalled_seed_is_cancelled_is_not_added_to_by_a_request_that_read_it_cold( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:set-then-stalled-cancelled" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + set_by_a_waiter: Final = asyncio.Event() + + async def _read_before_the_set() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await set_by_a_waiter.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + reader: Final = asyncio.ensure_future(_read_before_the_set()) + await read.wait() + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + await _cold_start(redis, counter) + stalled.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stalled + set_by_a_waiter.set() + await reader + + assert redis.values[counter] == 6.0 + + @pytest.mark.asyncio async def test_a_request_done_waiting_on_a_stalled_seed_adds_when_the_counter_holds_less_than_its_cached_spend( monkeypatch: pytest.MonkeyPatch, @@ -2720,15 +2762,21 @@ async def test_a_counter_without_redis_set_only_if_missing_keeps_the_value_it_ho counter: Final = "spend:end_user:no-redis-set-if-missing" monkeypatch.setattr(ps, "spend_counter_cache", DualCache(in_memory_cache=InMemoryCache())) - first: Final = await ps.increment_spend_counter_cache(counter, 6.0, only_if_missing=True) - second: Final = await ps.increment_spend_counter_cache(counter, 4.0, only_if_missing=True) + created: Final[list[str]] = [] # mutable-ok: the created calls the test asserts on + first: Final = await ps.increment_spend_counter_cache( + counter, 6.0, only_if_missing=True, on_created=lambda: created.append("first") + ) + second: Final = await ps.increment_spend_counter_cache( + counter, 4.0, only_if_missing=True, on_created=lambda: created.append("second") + ) added: Final = await ps.increment_spend_counter_cache(counter, 1.0) assert (first, second, added) == (6.0, 6.0, 7.0) + assert created == ["first"] @pytest.mark.asyncio -async def test_a_counter_lost_after_a_request_done_waiting_set_it_gets_the_add_of_a_later_request( +async def test_a_counter_set_by_a_request_done_waiting_counts_as_the_seed_for_a_later_request( monkeypatch: pytest.MonkeyPatch, ) -> None: counter: Final = "spend:end_user:set-counts-as-seed" @@ -2773,7 +2821,7 @@ class _CommandCountingRedis(_SharedRedis): self.commands.append("get") return await super().async_get_cache(key) - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: self.commands.append("set_if_missing" if kwargs.get("only_if_missing") else "add") return await super().async_increment(key, value, **kwargs) @@ -2877,7 +2925,7 @@ class _LostAndRecreatedAfterFailedMGetRedis(_SharedRedis): self.cold_gets += 1 return await super().async_get_cache(key) - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: self._other_process_increments_after_the_cold_gets() return await super().async_increment(key, value, **kwargs) @@ -2906,7 +2954,7 @@ class _AppliedHeldThenResetRedis(_HeldReplyRedis): super().__init__() self.deleted = False - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: await super().async_increment(key, value, **kwargs) raise ConnectionError("connection reset") @@ -2987,7 +3035,7 @@ class _HeldIncrementRedis(_SharedRedis): self.reads += 1 return await super().async_get_cache(key) - async def async_increment(self, key: str, value: float, **kwargs: bool) -> float: + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: index: Final = self.increments self.increments += 1 if index < len(self.arrived): @@ -3160,7 +3208,7 @@ async def test_a_stalled_seed_finishing_does_not_hide_the_newer_seed_still_in_fl await asyncio.gather(newer, late) assert commands_while_the_newer_seed_was_in_flight == 2 - assert redis.values[counter] == 12.0 + assert redis.values[counter] == 6.0 class _TargetRecordingRedis(_SharedRedis): From 0a44bbbf6158627b55fe33159e548251f9ce5411 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 04:31:50 -0800 Subject: [PATCH 14/18] test(proxy): annotate the only-if-missing tests the way AGENTS.md asks Their locals carry Final and their parameters carry types --- tests/unit/caching/test_redis_cache.py | 48 ++++++++++++++------------ 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 49e4979e04a..5ebb5cab9c4 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1487,19 +1487,19 @@ async def test_async_increment_pipelines_the_ttl_command( ], ) async def test_async_increment_only_if_missing_sets_a_missing_key_and_refreshes_the_ttl( - monkeypatch, redis_no_ping, existing, expected + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, existing: float | None, expected: float ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace="ns") - client = FakeAsyncRedis(server=FakeServer()) + redis_cache: Final = RedisCache(namespace="ns") + client: Final = FakeAsyncRedis(server=FakeServer()) if existing is not None: await client.set("ns:spend:key:k", existing, ex=5) with patch.object(redis_cache, "init_async_client", return_value=client): - current = await redis_cache.async_increment( + current: Final = await redis_cache.async_increment( key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True ) - after_add = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=True) + after_add: Final = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=True) assert current == expected assert after_add == expected + 2.0 @@ -1512,8 +1512,8 @@ async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_a monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float, stored: bytes ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace="ns") - client = FakeAsyncRedis(server=FakeServer()) + redis_cache: Final = RedisCache(namespace="ns") + client: Final = FakeAsyncRedis(server=FakeServer()) with patch.object(redis_cache, "init_async_client", return_value=client): await redis_cache.async_increment(key="spend:key:k", value=value, ttl=60, only_if_missing=True) @@ -1529,11 +1529,11 @@ async def test_async_increment_only_if_missing_reports_only_a_key_it_created( monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, existing: float | None, reported: list[str] ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace="ns") - client = FakeAsyncRedis(server=FakeServer()) + redis_cache: Final = RedisCache(namespace="ns") + client: Final = FakeAsyncRedis(server=FakeServer()) if existing is not None: await client.set("ns:spend:key:k", existing) - created: list[str] = [] + created: Final[list[str]] = [] with patch.object(redis_cache, "init_async_client", return_value=client): await redis_cache.async_increment( @@ -1576,18 +1576,20 @@ async def test_async_increment_only_if_missing_reads_the_key_back_when_the_user_ monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, read_back: float ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace="ns") - pipeline = _RepliesPipeline([NoPermissionError("no permissions to run the 'setnx' command"), read_back, True]) - client = MagicMock() + redis_cache: Final = RedisCache(namespace="ns") + pipeline: Final = _RepliesPipeline( + [NoPermissionError("no permissions to run the 'setnx' command"), read_back, True] + ) + client: Final = MagicMock() client.pipeline.return_value = pipeline - failure_hook = AsyncMock() - created = MagicMock() + failure_hook: Final = AsyncMock() + created: Final = MagicMock() with ( patch.object(redis_cache, "init_async_client", return_value=client), patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), ): - current = await redis_cache.async_increment( + current: Final = await redis_cache.async_increment( key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_created=created ) await asyncio.sleep(0) @@ -1616,8 +1618,8 @@ async def test_async_increment_only_if_missing_raises_any_other_error_in_the_rep monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, replies: list[object], raised: type[Exception] ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace="ns") - client = MagicMock() + redis_cache: Final = RedisCache(namespace="ns") + client: Final = MagicMock() client.pipeline.return_value = _RepliesPipeline(replies) with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(raised): @@ -1625,12 +1627,14 @@ async def test_async_increment_only_if_missing_raises_any_other_error_in_the_rep @pytest.mark.asyncio -async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment(monkeypatch, redis_no_ping): +async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") - redis_cache = RedisCache(namespace="ns") - client = MagicMock() + redis_cache: Final = RedisCache(namespace="ns") + client: Final = MagicMock() client.pipeline.side_effect = ConnectionError("connection reset") - failure_hook = AsyncMock() + failure_hook: Final = AsyncMock() with ( patch.object(redis_cache, "init_async_client", return_value=client), From 5c640ea870e5e26b234f3816751f9149c3075cc1 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 10:37:28 -0800 Subject: [PATCH 15/18] fix(proxy): keep the cluster retry and main's add on a late cold counter write The late SETNX pipeline ran with raise_on_error=False, which skips ClusterPipeline's retry, so a dropped connection on Redis Cluster failed a request that main served, and the error lost redis-py's command text. It now uses the same raising execute as a plain add A server that refuses a command in that pipeline, such as a user denied SETNX or a renamed SETNX, now gets main's plain add instead of a read back, so the failure and its log line are main's A fresh cold read that joined a set which left a peer's counter alone sent no add. It now adds unless a write landed after its read went out A request cancelled while its set was on the wire lost the add a short counter needed. The set and that add now finish in their own task --- litellm/caching/redis_cache.py | 71 ++++-- litellm/proxy/db/spend_counter_reseed.py | 61 +++-- litellm/proxy/proxy_server.py | 10 +- tests/unit/caching/test_redis_cache.py | 125 +++++++--- .../proxy/proxy_server/test_spend_counters.py | 224 +++++++++++++++++- 5 files changed, 412 insertions(+), 79 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index cb847ab2876..870fa8f532c 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1474,26 +1474,46 @@ class RedisCache(BaseCache): await _redis_client.expire(key, ttl) return float(result) - async def _set_if_missing_with_ttl( - self, key: str, value: float, ttl: int | None, on_created: Callable[[], None] | None - ) -> float: - """SETNX, a zero INCRBYFLOAT that reads the key back, and its TTL command in one round trip; the value the key - holds afterwards. Those last two are the commands a plain add sends, so a user denied only SETNX reads back what - the key holds, which is 0 for a key that was missing. ``on_created`` runs when the SETNX created the key.""" - from redis.exceptions import NoPermissionError + async def _setnx_and_read_back(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None: + """SETNX, a zero INCRBYFLOAT that reads the key back, and its TTL command in one round trip: whether the key was + created and the value it holds, or None when the server refused a command rather than failed to answer.""" + from redis.exceptions import ClusterError, RedisError + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError - async with self._async_commands().pipeline(transaction=False) as pipe: - pipe.setnx(name=key, value=value) - pipe.incrbyfloat(name=key, amount=0) - if ttl is not None: - pipe.expire(key, ttl) - replies: Final[Sequence[object]] = await pipe.execute(raise_on_error=False) # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics - for index, reply in enumerate(replies): # pyright: ignore[reportUnknownArgumentType] # redis generics - if isinstance(reply, Exception) and not (index == 0 and isinstance(reply, NoPermissionError)): - raise reply - if on_created is not None and replies[0] == 1: - on_created() - return _INCRBYFLOAT_REPLY.validate_python(replies[1]) + try: + async with self._async_commands().pipeline(transaction=False) as pipe: + pipe.setnx(name=key, value=value) + pipe.incrbyfloat(name=key, amount=0) + if ttl is not None: + pipe.expire(key, ttl) + replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics + except RedisError as e: + if isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)): + raise + return None + return replies[0] == 1, _INCRBYFLOAT_REPLY.validate_python(replies[1]) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # redis generics + + async def _set_if_missing_with_ttl( + self, + _redis_client: "Redis[bytes]", + key: str, + value: float, + ttl: int | None, + refresh_ttl: bool, + on_added: Callable[[], None] | None, + ) -> float: + """The value the key holds after setting it if it was missing. A server that refuses the commands this sends, + such as a user denied SETNX, gets a plain add instead. ``on_added`` runs when this call added the value.""" + written: Final = await self._setnx_and_read_back(key=key, value=value, ttl=ttl) + current: Final = ( + await self._incrbyfloat_with_ttl(_redis_client, key=key, value=value, ttl=ttl, refresh_ttl=refresh_ttl) + if written is None + else written[1] + ) + if on_added is not None and (written is None or written[0]): + on_added() + return current @_redis_circuit_breaker_guard async def async_increment( @@ -1504,10 +1524,10 @@ class RedisCache(BaseCache): parent_otel_span: Span | None = None, refresh_ttl: bool = False, only_if_missing: bool = False, - on_created: Callable[[], None] | None = None, + on_added: Callable[[], None] | None = None, ) -> float: """Add ``value`` to the key and return its new value. With ``only_if_missing``, a key that exists keeps its - value, which is returned, and only the TTL is refreshed, and ``on_created`` runs if this call created the key.""" + value, which is returned, and only the TTL is refreshed, and ``on_added`` runs if this call added the value.""" from redis.asyncio import Redis _redis_client: Final[Redis] = self.init_async_client() @@ -1516,7 +1536,14 @@ class RedisCache(BaseCache): key = self.check_and_fix_namespace(key=key) try: result: Final = await ( - self._set_if_missing_with_ttl(key=key, value=value, ttl=_used_ttl, on_created=on_created) + self._set_if_missing_with_ttl( + _redis_client, # pyright: ignore[reportUnknownArgumentType] # redis generics + key=key, + value=value, + ttl=_used_ttl, + refresh_ttl=refresh_ttl, + on_added=on_added, + ) if only_if_missing else self._incrbyfloat_with_ttl( _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index c901a4c2e8f..8e3a7e8311c 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -15,7 +15,7 @@ memory in long-lived deployments. import asyncio import functools from collections import OrderedDict -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Coroutine, Mapping from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone @@ -123,6 +123,23 @@ async def _wait_for_cold_seed(in_flight: asyncio.Future[float | None]) -> None: _cold_seed_waited.set(waited + loop.time() - started) +_late_writes: Final[set[asyncio.Task[float | None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks + + +def _late_write_finished(task: asyncio.Task[float | None]) -> None: + _late_writes.discard(task) + if not task.cancelled(): + task.exception() + + +async def _landing_even_if_cancelled(write: Coroutine[object, object, float | None]) -> float | None: + """Await ``write`` in a task of its own, so cancelling the caller leaves it to land as a command already sent does.""" + task: Final = asyncio.ensure_future(write) + _late_writes.add(task) + task.add_done_callback(_late_write_finished) + return await asyncio.shield(task) + + 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 @@ -183,12 +200,13 @@ class SpendCounterReseed: ) -> None: """Run ``add``, this process adding ``spend`` to a cold counter, once for all of this process's 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 or is cancelled. A request whose read went - out before an add landed, and one still waiting on an add after ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across - its counters, runs ``set_if_missing``, which calls its argument if it created the counter, in place of ``add``, - then runs ``add`` too if the counter holds less than the spend it should already hold, beyond what - INCRBYFLOAT's 17 decimal places can round away. 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 run ``add``.""" + gets its result, and starts or joins another add when that one fails, is cancelled or wrote nothing after its + read went out. A request whose read went out before an add landed, and one still waiting on an add after + ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across its counters, runs ``set_if_missing``, which calls its argument + if it added the spend, in place of ``add``, then runs ``add`` too if it did not and the counter holds less than + the spend it should already hold, beyond what INCRBYFLOAT's 17 decimal places can round away. + Those two steps finish even if the request is cancelled. 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 run ``add``.""" seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) def spend_landed() -> None: @@ -200,12 +218,21 @@ class SpendCounterReseed: spend_landed() return added - async def add_if_short() -> float | None: + async def write_if_short() -> float | None: least: Final = min(spend, seed.landed_spend) if seed.landed_at >= read.sent_at else spend - current: Final = await set_if_missing(spend_landed) - if current is not None and least - current > SPEND_COUNTER_INCRBYFLOAT_ROUNDING: - return await add_spend() - return current + set_added: Final = asyncio.Event() + + def on_added() -> None: + set_added.set() + spend_landed() + + current: Final = await set_if_missing(on_added) + if current is None or set_added.is_set(): + return current + return await add_spend() if least - current > SPEND_COUNTER_INCRBYFLOAT_ROUNDING else current + + async def add_if_short() -> float | None: + return await _landing_even_if_cancelled(write_if_short()) async def add_unless_covered() -> float | None: in_flight: Final = seed.in_flight @@ -217,7 +244,7 @@ class SpendCounterReseed: await _wait_for_cold_seed(in_flight) if not in_flight.done(): return await seed.run(add_if_short) - if covers and not in_flight.cancelled(): + if covers and not in_flight.cancelled() and seed.landed_at >= read.sent_at: return in_flight.result() return await add_unless_covered() @@ -232,10 +259,10 @@ class SpendCounterReseed: counter_key: str, increment: float, only_if_missing: bool = False, - on_created: Callable[[], None] | None = None, + on_added: Callable[[], None] | None = None, ) -> float | None: """Apply local deltas after an in-flight reseed establishes the spend balance. With ``only_if_missing``, a - counter that exists keeps its value, which is returned, and ``on_created`` runs if this call created it.""" + counter that exists keeps its value, which is returned, and ``on_added`` runs if this call added to it.""" lock: Final = await SpendCounterReseed._get_lock(counter_key) async with lock: current: Final[float | None] = ( # pyright: ignore[reportUnknownVariableType] # untyped in-memory cache @@ -248,8 +275,8 @@ class SpendCounterReseed: added: Final = await spend_counter_cache.async_increment_cache( key=counter_key, value=increment, local_only=True, refresh_ttl=True ) - if on_created is not None: - on_created() + if on_added is not None: + on_added() return added @staticmethod diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5e4c1801955..dd81f836023 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4106,8 +4106,8 @@ async def ensure_spend_counter_initialized( read=read, spend=base_spend, add=lambda: increment_spend_counter_cache(counter_key=counter_key, increment=base_spend), - set_if_missing=lambda on_created: increment_spend_counter_cache( - counter_key=counter_key, increment=base_spend, only_if_missing=True, on_created=on_created + set_if_missing=lambda on_added: increment_spend_counter_cache( + counter_key=counter_key, increment=base_spend, only_if_missing=True, on_added=on_added ), ) @@ -4216,7 +4216,7 @@ async def increment_spend_counter_cache( counter_key: str, increment: float, only_if_missing: bool = False, - on_created: Callable[[], None] | None = None, + on_added: Callable[[], None] | None = None, ) -> float | None: if spend_counter_cache.redis_cache is not None: try: @@ -4226,7 +4226,7 @@ async def increment_spend_counter_cache( value=increment, refresh_ttl=True, only_if_missing=True, - on_created=on_created, + on_added=on_added, ) if only_if_missing else spend_counter_cache.redis_cache.async_increment( @@ -4250,7 +4250,7 @@ async def increment_spend_counter_cache( counter_key=counter_key, increment=increment, only_if_missing=only_if_missing, - on_created=on_created, + on_added=on_added, ) diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 5ebb5cab9c4..5cce440fdbc 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -9,7 +9,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fakeredis import FakeAsyncRedis, FakeServer -from redis.exceptions import NoPermissionError, ResponseError +from redis.exceptions import ClusterDownError, NoPermissionError, RedisError, ResponseError +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import TimeoutError as RedisTimeoutError import litellm from litellm._service_logger import ServiceLogging @@ -1537,20 +1539,20 @@ async def test_async_increment_only_if_missing_reports_only_a_key_it_created( with patch.object(redis_cache, "init_async_client", return_value=client): await redis_cache.async_increment( - key="spend:key:k", value=1.5, ttl=60, only_if_missing=True, on_created=lambda: created.append("created") + key="spend:key:k", value=1.5, ttl=60, only_if_missing=True, on_added=lambda: created.append("created") ) assert created == reported -class _RepliesPipeline: - """A pipeline whose execute returns the given replies the way ``raise_on_error=False`` does.""" +class _RaisingPipeline: + """A pipeline whose execute raises the first error in its replies, as redis-py's does by default.""" def __init__(self, replies: list[object]) -> None: self.replies = replies self.commands: list[tuple[object, ...]] = [] - async def __aenter__(self) -> "_RepliesPipeline": + async def __aenter__(self) -> "_RaisingPipeline": return self async def __aexit__(self, *exc_info: object) -> None: @@ -1565,66 +1567,129 @@ class _RepliesPipeline: def expire(self, name: str, ttl: int) -> None: self.commands.append(("EXPIRE", name, ttl)) + async def execute(self) -> list[object]: + first_error: Final = next((reply for reply in self.replies if isinstance(reply, Exception)), None) + if first_error is not None: + raise first_error + return self.replies + + +class _ClusterPipelineAfterADroppedConnection(_RaisingPipeline): + """The first attempt reads a dropped connection. As in redis-py 5.3.1's ClusterPipeline.execute, raising is what + drives the retry, and ``raise_on_error=False`` hands the error back as a reply instead.""" + async def execute(self, raise_on_error: bool = True) -> list[object]: - assert raise_on_error is False + if not raise_on_error: + return [RedisConnectionError("Connection closed by server."), self.replies[1], self.replies[2]] return self.replies @pytest.mark.asyncio -@pytest.mark.parametrize("read_back", [pytest.param(0.0, id="missing"), pytest.param(7.0, id="held")]) -async def test_async_increment_only_if_missing_reads_the_key_back_when_the_user_is_denied_setnx( - monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, read_back: float +async def test_async_increment_only_if_missing_keeps_the_cluster_retry_of_a_dropped_connection( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache: Final = RedisCache(namespace="ns") - pipeline: Final = _RepliesPipeline( - [NoPermissionError("no permissions to run the 'setnx' command"), read_back, True] - ) client: Final = MagicMock() - client.pipeline.return_value = pipeline + client.pipeline.return_value = _ClusterPipelineAfterADroppedConnection([True, 1.5, True]) + + with patch.object(redis_cache, "init_async_client", return_value=client): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + ) + + assert current == 1.5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "refusal", + [ + pytest.param(NoPermissionError("no permissions to run the 'setnx' command"), id="denied"), + pytest.param(ResponseError("unknown command 'SETNX'"), id="renamed"), + pytest.param(RedisError("SETNX command doesn't exist in Redis commands"), id="renamed-on-a-cluster-client"), + ], +) +async def test_async_increment_only_if_missing_adds_the_value_when_the_server_refuses_the_set( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, refusal: Exception +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add: Final = _RaisingPipeline([7.5, True]) + client: Final = MagicMock() + client.pipeline.side_effect = [_RaisingPipeline([refusal, 6.0, True]), add] failure_hook: Final = AsyncMock() - created: Final = MagicMock() + added: Final = MagicMock() with ( patch.object(redis_cache, "init_async_client", return_value=client), patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), ): current: Final = await redis_cache.async_increment( - key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_created=created + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added ) await asyncio.sleep(0) - assert current == read_back - created.assert_not_called() - assert pipeline.commands == [ - ("SETNX", "ns:spend:key:k", 1.5), - ("INCRBYFLOAT", "ns:spend:key:k", 0), - ("EXPIRE", "ns:spend:key:k", 60), - ] + assert current == 7.5 + added.assert_called_once_with() + assert add.commands == [("INCRBYFLOAT", "ns:spend:key:k", 1.5), ("EXPIRE", "ns:spend:key:k", 60)] failure_hook.assert_not_awaited() assert redis_cache._circuit_breaker._hard_failure_count == 0 +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_reports_a_write_redis_refuses_with_the_add_and_its_error_text( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add_error: Final = ResponseError("Command # 1 (INCRBYFLOAT ns:spend:key:k 1.5) of pipeline caused error") + client: Final = MagicMock() + client.pipeline.side_effect = [ + _RaisingPipeline([ResponseError("Command # 1 (SETNX) of pipeline caused error"), 0.0, True]), + _RaisingPipeline([add_error, True]), + ] + failure_hook: Final = AsyncMock() + added: Final = MagicMock() + + with ( + patch.object(redis_cache, "init_async_client", return_value=client), + patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), + pytest.raises(ResponseError) as raised, + ): + await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + await asyncio.sleep(0) + + assert raised.value is add_error + added.assert_not_called() + failure_hook.assert_awaited_once() + assert failure_hook.await_args.kwargs["error"] is add_error + + @pytest.mark.asyncio @pytest.mark.parametrize( - ("replies", "raised"), + "failure", [ - pytest.param([ResponseError("LOADING"), 7.0, True], ResponseError, id="setnx-failed-for-another-reason"), - pytest.param([True, NoPermissionError("denied"), True], NoPermissionError, id="incrbyfloat-denied"), - pytest.param([True, 7.0, NoPermissionError("denied")], NoPermissionError, id="expire-denied"), + pytest.param(RedisConnectionError("Connection closed by server."), id="dropped-connection"), + pytest.param(RedisTimeoutError("Timeout reading from socket"), id="stalled-node"), + pytest.param(ClusterDownError("CLUSTERDOWN The cluster is down"), id="cluster-down"), ], ) -async def test_async_increment_only_if_missing_raises_any_other_error_in_the_replies( - monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, replies: list[object], raised: type[Exception] +async def test_async_increment_only_if_missing_sends_no_add_after_a_failure_the_client_already_retried( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, failure: Exception ): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache: Final = RedisCache(namespace="ns") client: Final = MagicMock() - client.pipeline.return_value = _RepliesPipeline(replies) + client.pipeline.return_value = _RaisingPipeline([failure, 0.0, True]) - with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(raised): + with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(type(failure)): await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True) + assert client.pipeline.call_count == 1 + @pytest.mark.asyncio async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment( diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 246012b8bf8..3595a70bd37 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -23,6 +23,7 @@ from __future__ import annotations import asyncio import contextlib +import gc import threading from collections.abc import Awaitable, Callable, Collection from datetime import datetime @@ -2051,13 +2052,13 @@ class _SharedRedis: key: str, value: float, only_if_missing: bool = False, - on_created: Callable[[], None] | None = None, + on_added: Callable[[], None] | None = None, **_: object, ) -> float: if not only_if_missing or key not in self.values: self.values[key] = self.values.get(key, 0.0) + value - if on_created is not None: - on_created() + if on_added is not None: + on_added() return self.values[key] async def async_increment(self, key: str, value: float, **kwargs: object) -> float: @@ -2764,10 +2765,10 @@ async def test_a_counter_without_redis_set_only_if_missing_keeps_the_value_it_ho created: Final[list[str]] = [] # mutable-ok: the created calls the test asserts on first: Final = await ps.increment_spend_counter_cache( - counter, 6.0, only_if_missing=True, on_created=lambda: created.append("first") + counter, 6.0, only_if_missing=True, on_added=lambda: created.append("first") ) second: Final = await ps.increment_spend_counter_cache( - counter, 4.0, only_if_missing=True, on_created=lambda: created.append("second") + counter, 4.0, only_if_missing=True, on_added=lambda: created.append("second") ) added: Final = await ps.increment_spend_counter_cache(counter, 1.0) @@ -3234,6 +3235,219 @@ async def test_the_cold_check_read_reports_under_spend_counters_inside_another_p assert redis.targets == [SPEND_COUNTERS_TARGET] +class _HeldSetRedis(_SharedRedis): + """Holds each set on the wire, before Redis runs it, until ``release_set`` is set.""" + + def __init__(self) -> None: + super().__init__() + self.set_sent: Final = asyncio.Event() + self.release_set: Final = asyncio.Event() + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if kwargs.get("only_if_missing"): + self.set_sent.set() + await self.release_set.wait() + return await super().async_increment(key, value, **kwargs) + + +async def _a_request_read_cold_before_this_process_seeded( + redis: _SharedRedis, counter: str, seeded: asyncio.Event, go: asyncio.Event +) -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + seeded.set() + await go.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("peer_value", "expected"), + [ + pytest.param(7.0, 13.0, id="the-set-left-a-peers-counter-alone-so-the-fresh-read-adds"), + pytest.param(None, 6.0, id="the-set-created-the-counter-so-the-fresh-read-joins-it"), + ], +) +async def test_a_fresh_cold_read_joining_a_set_only_skips_its_add_when_that_set_wrote( + monkeypatch: pytest.MonkeyPatch, peer_value: float | None, expected: float +) -> None: + counter: Final = "spend:end_user:fresh-read-joins-a-set" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _HeldSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + del redis.values[counter] + + async def _fresh_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + if peer_value is not None: + redis.values[counter] = peer_value + go.set() + await redis.set_sent.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + fresh: Final = asyncio.ensure_future(_fresh_request()) + await redis.set_sent.wait() + for _ in range(20): + await asyncio.sleep(0) + redis.release_set.set() + await asyncio.wait_for(asyncio.gather(stale, fresh), timeout=5) + + assert redis.values[counter] == expected + + +@pytest.mark.asyncio +async def test_a_request_cancelled_with_its_set_on_the_wire_still_adds_to_a_short_counter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-with-the-set-out" + redis: Final = _HeldSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + redis.values[counter] = 0.25 + go.set() + await redis.set_sent.wait() + + stale.cancel() + await asyncio.wait((stale,)) + redis.release_set.set() + for _ in range(20): + await asyncio.sleep(0) + + assert stale.cancelled() + assert redis.values[counter] == 6.25 + + +@pytest.mark.asyncio +async def test_a_set_failing_after_its_request_was_cancelled_deletes_the_counter_and_logs_no_stray_task_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-set-fails" + redis: Final = _HeldSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + async def _cancel_a_request_with_its_set_on_the_wire() -> None: + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + go.set() + await redis.set_sent.wait() + stale.cancel() + await asyncio.wait((stale,)) + + await _cancel_a_request_with_its_set_on_the_wire() + stray: Final[list[object]] = [] # mutable-ok: collects what the loop's exception handler is given + asyncio.get_running_loop().set_exception_handler(lambda _, context: stray.append(context)) + redis.failing = True + redis.release_set.set() + for _ in range(20): + await asyncio.sleep(0) + redis.failing = False + gc.collect() + await asyncio.sleep(0) + + assert stray == [] + assert ps.spend_counter_cache.in_memory_cache.get_cache(key=counter) is None + + +@pytest.mark.asyncio +async def test_a_request_joining_an_add_that_is_cancelled_writes_the_counter_itself( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:joined-add-cancelled" + redis: Final = _HeldIncrementRedis(held_increments=2) + redis.releases[0].set() + _process(monkeypatch, redis, cached_spend=6.0) + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + older: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + del redis.values[counter] + fresh: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[1].wait() + go.set() + for _ in range(20): + await asyncio.sleep(0) + + fresh.cancel() + await asyncio.wait((fresh,)) + await older + + assert fresh.cancelled() + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_set_finding_a_short_counter_adds_even_if_another_requests_add_landed_during_that_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:add-landed-during-set" + redis: Final = _HeldIncrementRedis(held_increments=2) + _process(monkeypatch, redis, cached_spend=6.0) + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + second: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[1].wait() + redis.releases[0].set() + await first + redis.values[counter] = 0.3 + + redis.releases[1].set() + await second + + assert redis.values[counter] == pytest.approx(6.3) + + +class _RefusesTheSetRedis(_SharedRedis): + """A server that refuses SETNX, where ``async_increment`` falls back to a plain add.""" + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + return await super().async_increment(key, value, **{**kwargs, "only_if_missing": False}) + + +@pytest.mark.asyncio +async def test_a_late_request_whose_set_became_a_plain_add_does_not_add_again_to_a_counter_below_zero( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:refused-set-negative-counter" + redis: Final = _RefusesTheSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + redis.values[counter] = -1.5 + go.set() + await asyncio.wait_for(stale, timeout=5) + + assert redis.values[counter] == 4.5 + + @pytest.mark.asyncio async def test_a_stale_read_on_a_real_client_seeds_a_missing_counter_with_its_ttl_and_keeps_an_existing_one( monkeypatch: pytest.MonkeyPatch, From f4a3998979154408816c6d2752e3a35524635232 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 13:45:10 -0800 Subject: [PATCH 16/18] fix(proxy): keep a cancelled request's late write as the add in flight until it lands A request cancelled with its late set on the wire cleared the add in flight while the set kept running, so a newer cold read could add on top of it and turn 6 into 12. The shielded write now owns the in flight slot, and the slot is taken before the task starts so a request resuming in the same tick joins it --- litellm/proxy/db/spend_counter_reseed.py | 37 ++++---- .../proxy/proxy_server/test_spend_counters.py | 88 +++++++++++++++++++ 2 files changed, 109 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 8e3a7e8311c..fdcc924cac9 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -90,20 +90,24 @@ class _CachedSpendSeed: def on_this_loop(self) -> bool: return self.in_flight is None or self.in_flight.get_loop() is asyncio.get_running_loop() - async def run(self, add: Callable[[], Awaitable[float | None]]) -> float | None: - """Run ``add`` in the calling task, so cancelling that task stops it, and publish its result to the requests - waiting on ``in_flight``, which ends cancelled when ``add`` fails or is cancelled.""" + def run(self, add: Callable[[], Awaitable[float | None]]) -> Coroutine[object, object, float | None]: + """Make ``add`` the add in flight from this call on, and return the coroutine that runs it and publishes its + result to the requests waiting on ``in_flight``, which ends cancelled when ``add`` fails or is cancelled.""" in_flight: Final[asyncio.Future[float | None]] = asyncio.get_running_loop().create_future() self.started_at = spend_counter_tick() self.in_flight = in_flight - try: - added: Final = await add() - in_flight.set_result(added) - return added - finally: - if self.in_flight is in_flight: - self.in_flight = None - in_flight.cancel() + + async def add_and_publish() -> float | None: + try: + added: Final = await add() + in_flight.set_result(added) + return added + finally: + if self.in_flight is in_flight: + self.in_flight = None + in_flight.cancel() + + return add_and_publish() @functools.lru_cache(maxsize=SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE) @@ -205,8 +209,9 @@ class SpendCounterReseed: ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across its counters, runs ``set_if_missing``, which calls its argument if it added the spend, in place of ``add``, then runs ``add`` too if it did not and the counter holds less than the spend it should already hold, beyond what INCRBYFLOAT's 17 decimal places can round away. - Those two steps finish even if the request is cancelled. 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 run ``add``.""" + Those two steps finish even if the request is cancelled, and stay the add in flight until they do. 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 run ``add``.""" seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) def spend_landed() -> None: @@ -232,18 +237,18 @@ class SpendCounterReseed: return await add_spend() if least - current > SPEND_COUNTER_INCRBYFLOAT_ROUNDING else current async def add_if_short() -> float | None: - return await _landing_even_if_cancelled(write_if_short()) + return await _landing_even_if_cancelled(seed.run(write_if_short)) async def add_unless_covered() -> 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: - return await seed.run(add_spend if seed.landed_at < read.sent_at else add_if_short) + return await (seed.run(add_spend) if seed.landed_at < read.sent_at else add_if_short()) covers: Final = seed.started_at > read.returned_at await _wait_for_cold_seed(in_flight) if not in_flight.done(): - return await seed.run(add_if_short) + return await add_if_short() if covers and not in_flight.cancelled() and seed.landed_at >= read.sent_at: return in_flight.result() return await add_unless_covered() diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 3595a70bd37..b6a0b50c9cd 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -3331,6 +3331,94 @@ async def test_a_request_cancelled_with_its_set_on_the_wire_still_adds_to_a_shor assert redis.values[counter] == 6.25 +class _AddLandsAfterSetRedis(_HeldSetRedis): + """Also holds a plain add on the wire while ``set_landed`` is clear, so a set still in flight lands first.""" + + def __init__(self) -> None: + super().__init__() + self.set_landed: Final = asyncio.Event() + self.set_landed.set() + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if not kwargs.get("only_if_missing"): + await self.set_landed.wait() + added: Final = await super().async_increment(key, value, **kwargs) + if kwargs.get("only_if_missing"): + self.set_landed.set() + return added + + +@pytest.mark.asyncio +async def test_a_set_still_on_the_wire_after_its_request_was_cancelled_is_joined_by_a_newer_cold_read( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-set-joined" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _AddLandsAfterSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + del redis.values[counter] + redis.set_landed.clear() + go.set() + await redis.set_sent.wait() + stale.cancel() + await asyncio.wait((stale,)) + + newer: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(20): + await asyncio.sleep(0) + redis.release_set.set() + await asyncio.wait_for(newer, timeout=5) + + assert stale.cancelled() + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_fresh_read_released_with_a_stale_read_joins_the_set_the_stale_read_starts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:released-together" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _AddLandsAfterSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + stale_read: Final = asyncio.Event() + fresh_read: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, stale_read, go) + ) + await stale_read.wait() + await _cold_start(redis, counter) + del redis.values[counter] + redis.set_landed.clear() + fresh: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, fresh_read, go) + ) + await fresh_read.wait() + owner: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(20): + await asyncio.sleep(0) + go.set() + for _ in range(20): + await asyncio.sleep(0) + + owner.cancel() + await redis.set_sent.wait() + for _ in range(20): + await asyncio.sleep(0) + redis.release_set.set() + await asyncio.wait_for(asyncio.gather(stale, fresh), timeout=5) + + assert redis.values[counter] == 6.0 + + @pytest.mark.asyncio async def test_a_set_failing_after_its_request_was_cancelled_deletes_the_counter_and_logs_no_stray_task_error( monkeypatch: pytest.MonkeyPatch, From deef69c78623684f1c2723b3f3bf3228e8fb83f7 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 4 Nov 2026 17:22:44 -0800 Subject: [PATCH 17/18] fix(proxy): set a late cold counter with one SET NX GET and keep it after a cancelled failure A late request sent SETNX, INCRBYFLOAT k 0 and EXPIRE. When Redis refused SETNX the zero add still created the key at 0, so a lost add after it left a warm 0 counter, and when Redis ran out of memory after SETNX the add after it doubled the counter It now sends SET NX GET and EXPIRE. SET NX GET is one command, so a refused one wrote nothing and main's add after it leaves what main's add leaves. The value is written the way INCRBYFLOAT writes it. Redis before 7.0, a key holding a value INCRBYFLOAT might read differently and a spend that is not finite get main's add A late write that fails after its request was cancelled no longer deletes the counter, since main's cancelled add never reaches that DEL --- litellm/caching/redis_cache.py | 43 ++++- litellm/proxy/db/spend_counter_reseed.py | 26 ++- litellm/proxy/proxy_server.py | 5 +- tests/unit/caching/test_redis_cache.py | 174 ++++++++++++++++-- .../proxy/proxy_server/test_spend_counters.py | 95 ++++++++-- 5 files changed, 296 insertions(+), 47 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 870fa8f532c..f6223fb8d94 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -16,6 +16,8 @@ import inspect import itertools import json import logging +import math +import re import threading import time import weakref @@ -24,6 +26,7 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta +from decimal import Decimal from types import FrameType, MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast @@ -139,11 +142,26 @@ _INCREMENT_IF_PRESENT_LUA: Final = ( ) _LUA_COUNT: Final = TypeAdapter(int) -_INCRBYFLOAT_REPLY: Final = TypeAdapter(float) +_PLAIN_DECIMAL: Final = re.compile(r"-?[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?") _LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) +def _as_incrbyfloat_writes(value: float) -> str: + """``value`` as INCRBYFLOAT writes it into a missing key: 17 decimal places with the trailing zeros cut.""" + text: Final = format(Decimal(repr(value)), ".17f").rstrip("0").rstrip(".") + return "0" if text == "-0" else text + + +def _plain_float(reply: object) -> float | None: + """The finite float a stored value spells in plain decimal or exponent form, or None for anything else.""" + text: Final = reply.decode(errors="replace") if isinstance(reply, bytes) else reply + if not isinstance(text, str) or _PLAIN_DECIMAL.fullmatch(text) is None: + return None + number: Final = float(text) + return number if math.isfinite(number) else None + + def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: return _OPTIONAL_COUNTS.validate_python( tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) @@ -1474,17 +1492,20 @@ class RedisCache(BaseCache): await _redis_client.expire(key, ttl) return float(result) - async def _setnx_and_read_back(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None: - """SETNX, a zero INCRBYFLOAT that reads the key back, and its TTL command in one round trip: whether the key was - created and the value it holds, or None when the server refused a command rather than failed to answer.""" + async def _set_if_missing_and_get(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None: + """SET NX GET and its TTL command in one round trip: whether the key was created and the value it holds, or + None when the server refused a command rather than failed to answer, or the key holds a value this does not + read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing.""" from redis.exceptions import ClusterError, RedisError from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError + if not math.isfinite(value): + return None + stored: Final = _as_incrbyfloat_writes(value) try: async with self._async_commands().pipeline(transaction=False) as pipe: - pipe.setnx(name=key, value=value) - pipe.incrbyfloat(name=key, amount=0) + pipe.set(name=key, value=stored, nx=True, get=True) if ttl is not None: pipe.expire(key, ttl) replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics @@ -1492,7 +1513,10 @@ class RedisCache(BaseCache): if isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)): raise return None - return replies[0] == 1, _INCRBYFLOAT_REPLY.validate_python(replies[1]) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # redis generics + if replies[0] is None: + return True, float(stored) + held: Final = _plain_float(replies[0]) # pyright: ignore[reportUnknownArgumentType] # redis generics + return None if held is None else (False, held) async def _set_if_missing_with_ttl( self, @@ -1504,8 +1528,9 @@ class RedisCache(BaseCache): on_added: Callable[[], None] | None, ) -> float: """The value the key holds after setting it if it was missing. A server that refuses the commands this sends, - such as a user denied SETNX, gets a plain add instead. ``on_added`` runs when this call added the value.""" - written: Final = await self._setnx_and_read_back(key=key, value=value, ttl=ttl) + such as a user denied SET or a Redis before 7.0, and a key holding a value INCRBYFLOAT might read differently, + get a plain add instead. ``on_added`` runs when this call added the value.""" + written: Final = await self._set_if_missing_and_get(key=key, value=value, ttl=ttl) current: Final = ( await self._incrbyfloat_with_ttl(_redis_client, key=key, value=value, ttl=ttl, refresh_ttl=refresh_ttl) if written is None diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index fdcc924cac9..f27c67497f9 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -16,7 +16,7 @@ import asyncio import functools from collections import OrderedDict from collections.abc import Awaitable, Callable, Coroutine, Mapping -from contextvars import ContextVar +from contextvars import ContextVar, copy_context from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType @@ -128,6 +128,13 @@ async def _wait_for_cold_seed(in_flight: asyncio.Future[float | None]) -> None: _late_writes: Final[set[asyncio.Task[float | None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks +_late_write_abandoned: Final[ContextVar[asyncio.Event | None]] = ContextVar("late_spend_write_abandoned", default=None) + + +def late_write_abandoned() -> bool: + """Whether this runs in a late write whose request was cancelled, which main's add would not have outlived.""" + abandoned: Final = _late_write_abandoned.get() + return abandoned is not None and abandoned.is_set() def _late_write_finished(task: asyncio.Task[float | None]) -> None: @@ -137,11 +144,19 @@ def _late_write_finished(task: asyncio.Task[float | None]) -> None: async def _landing_even_if_cancelled(write: Coroutine[object, object, float | None]) -> float | None: - """Await ``write`` in a task of its own, so cancelling the caller leaves it to land as a command already sent does.""" - task: Final = asyncio.ensure_future(write) + """Await ``write`` in a task of its own, so cancelling the caller leaves it to land as a command already sent does, + and marks it abandoned from then on, so ``late_write_abandoned`` holds inside it.""" + abandoned: Final = asyncio.Event() + context: Final = copy_context() + context.run(_late_write_abandoned.set, abandoned) + task: Final = context.run(asyncio.ensure_future, write) _late_writes.add(task) task.add_done_callback(_late_write_finished) - return await asyncio.shield(task) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + abandoned.set() + raise def _redis_breaker_closed(spend_counter_cache: "DualCache") -> bool: @@ -209,7 +224,8 @@ class SpendCounterReseed: ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across its counters, runs ``set_if_missing``, which calls its argument if it added the spend, in place of ``add``, then runs ``add`` too if it did not and the counter holds less than the spend it should already hold, beyond what INCRBYFLOAT's 17 decimal places can round away. - Those two steps finish even if the request is cancelled, and stay the add in flight until they do. A request on + Those two steps finish even if the request is cancelled, and stay the add in flight until they do, but once it is + cancelled a failure in them leaves the counter as main's cancelled add does, without the DEL. 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 run ``add``.""" seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dd81f836023..1ef0fba452f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -639,7 +639,7 @@ from litellm.proxy.db.request_error_tracking import ( flush_request_errors, request_error_accumulator, ) -from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed +from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed, late_write_abandoned from litellm.proxy.discovery_endpoints import ( agent_skills_discovery_router, ui_discovery_endpoints_router, @@ -4236,7 +4236,8 @@ async def increment_spend_counter_cache( ) ) except Exception: - await invalidate_spend_counter(counter_key=counter_key) + if not late_write_abandoned(): + await invalidate_spend_counter(counter_key=counter_key) raise spend_counter_cache.in_memory_cache.set_cache( key=counter_key, diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 5cce440fdbc..2a2ac5f7063 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1509,7 +1509,10 @@ async def test_async_increment_only_if_missing_sets_a_missing_key_and_refreshes_ @pytest.mark.asyncio -@pytest.mark.parametrize(("value", "stored"), [(6.0, b"6"), (1e-05, b"0.00001")]) +@pytest.mark.parametrize( + ("value", "stored"), + [(6.0, b"6"), (1e-05, b"0.00001"), (6.000000002, b"6.000000002"), (0.1, b"0.1"), (1e-20, b"0"), (-0.0, b"0")], +) async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_add_stores_it( monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float, stored: bytes ): @@ -1518,9 +1521,10 @@ async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_a client: Final = FakeAsyncRedis(server=FakeServer()) with patch.object(redis_cache, "init_async_client", return_value=client): - await redis_cache.async_increment(key="spend:key:k", value=value, ttl=60, only_if_missing=True) + current: Final = await redis_cache.async_increment(key="spend:key:k", value=value, ttl=60, only_if_missing=True) assert await client.get("ns:spend:key:k") == stored + assert current == float(stored) @pytest.mark.asyncio @@ -1558,8 +1562,8 @@ class _RaisingPipeline: async def __aexit__(self, *exc_info: object) -> None: return None - def setnx(self, name: str, value: float) -> None: - self.commands.append(("SETNX", name, value)) + def set(self, name: str, value: str, nx: bool = False, get: bool = False) -> None: + self.commands.append(("SET", name, value, nx, get)) def incrbyfloat(self, name: str, amount: float) -> None: self.commands.append(("INCRBYFLOAT", name, amount)) @@ -1580,7 +1584,7 @@ class _ClusterPipelineAfterADroppedConnection(_RaisingPipeline): async def execute(self, raise_on_error: bool = True) -> list[object]: if not raise_on_error: - return [RedisConnectionError("Connection closed by server."), self.replies[1], self.replies[2]] + return [RedisConnectionError("Connection closed by server."), self.replies[1]] return self.replies @@ -1591,7 +1595,7 @@ async def test_async_increment_only_if_missing_keeps_the_cluster_retry_of_a_drop monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache: Final = RedisCache(namespace="ns") client: Final = MagicMock() - client.pipeline.return_value = _ClusterPipelineAfterADroppedConnection([True, 1.5, True]) + client.pipeline.return_value = _ClusterPipelineAfterADroppedConnection([None, True]) with patch.object(redis_cache, "init_async_client", return_value=client): current: Final = await redis_cache.async_increment( @@ -1605,9 +1609,10 @@ async def test_async_increment_only_if_missing_keeps_the_cluster_retry_of_a_drop @pytest.mark.parametrize( "refusal", [ - pytest.param(NoPermissionError("no permissions to run the 'setnx' command"), id="denied"), - pytest.param(ResponseError("unknown command 'SETNX'"), id="renamed"), - pytest.param(RedisError("SETNX command doesn't exist in Redis commands"), id="renamed-on-a-cluster-client"), + pytest.param(NoPermissionError("no permissions to run the 'set' command"), id="denied"), + pytest.param(ResponseError("syntax error"), id="redis-before-7-has-no-set-nx-get"), + pytest.param(ResponseError("unknown command 'SET'"), id="renamed"), + pytest.param(RedisError("SET command doesn't exist in Redis commands"), id="renamed-on-a-cluster-client"), ], ) async def test_async_increment_only_if_missing_adds_the_value_when_the_server_refuses_the_set( @@ -1617,7 +1622,7 @@ async def test_async_increment_only_if_missing_adds_the_value_when_the_server_re redis_cache: Final = RedisCache(namespace="ns") add: Final = _RaisingPipeline([7.5, True]) client: Final = MagicMock() - client.pipeline.side_effect = [_RaisingPipeline([refusal, 6.0, True]), add] + client.pipeline.side_effect = [_RaisingPipeline([refusal, True]), add] failure_hook: Final = AsyncMock() added: Final = MagicMock() @@ -1646,7 +1651,7 @@ async def test_async_increment_only_if_missing_reports_a_write_redis_refuses_wit add_error: Final = ResponseError("Command # 1 (INCRBYFLOAT ns:spend:key:k 1.5) of pipeline caused error") client: Final = MagicMock() client.pipeline.side_effect = [ - _RaisingPipeline([ResponseError("Command # 1 (SETNX) of pipeline caused error"), 0.0, True]), + _RaisingPipeline([ResponseError("Command # 1 (SET) of pipeline caused error"), True]), _RaisingPipeline([add_error, True]), ] failure_hook: Final = AsyncMock() @@ -1683,7 +1688,7 @@ async def test_async_increment_only_if_missing_sends_no_add_after_a_failure_the_ monkeypatch.setenv("REDIS_HOST", "https://my-test-host") redis_cache: Final = RedisCache(namespace="ns") client: Final = MagicMock() - client.pipeline.return_value = _RaisingPipeline([failure, 0.0, True]) + client.pipeline.return_value = _RaisingPipeline([failure, True]) with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(type(failure)): await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True) @@ -1691,6 +1696,151 @@ async def test_async_increment_only_if_missing_sends_no_add_after_a_failure_the_ assert client.pipeline.call_count == 1 +class _Keyspace: + """String keys as Redis holds them, for pipelines that run each command in order and raise the first error after + running the rest, as a non-transaction pipeline does. Commands named in ``refused`` are refused, a pipeline that + has run ``writes_before_full`` write commands has the rest refused as Redis does past ``maxmemory``, and pipelines + after the first ``reached`` never reach Redis, as when the connection drops.""" + + def __init__( + self, refused: frozenset[str] = frozenset(), writes_before_full: int | None = None, reached: int | None = None + ) -> None: + self.values: Final[dict[str, str]] = {} # mutable-ok: stands in for the Redis keyspace + self.refused = refused + self.writes_before_full = writes_before_full + self.reached = reached + self.pipelines = 0 + + def pipeline(self, **_: object) -> "_KeyspacePipeline": + self.pipelines += 1 + return _KeyspacePipeline(self, lost=self.reached is not None and self.pipelines > self.reached) + + +class _KeyspacePipeline(_RaisingPipeline): + def __init__(self, keyspace: _Keyspace, lost: bool) -> None: + super().__init__([]) + self.keyspace = keyspace + self.lost = lost + self.writes_left = keyspace.writes_before_full + + def setnx(self, name: str, value: float) -> None: + self.commands.append(("SETNX", name, value)) + + def _write(self, command: str) -> None: + if command in self.keyspace.refused: + raise NoPermissionError(f"this user has no permissions to run the '{command.lower()}' command") + if self.writes_left is not None: + if self.writes_left == 0: + raise ResponseError("OOM command not allowed when used memory > 'maxmemory'.") + self.writes_left -= 1 + + def _run(self, command: tuple[object, ...]) -> object: + name: Final = str(command[1]) + held: Final = self.keyspace.values.get(name) + if command[0] == "EXPIRE": + return held is not None + self._write(str(command[0])) + if command[0] == "SET": + if held is None: + self.keyspace.values[name] = str(command[2]) + return None if held is None else held.encode() + if command[0] == "SETNX": + if held is None: + self.keyspace.values[name] = repr(command[2]) + return held is None + total: Final = float(held or 0) + float(cast(float, command[2])) + self.keyspace.values[name] = repr(total) + return total + + async def execute(self) -> list[object]: + if self.lost: + raise RedisConnectionError("Connection closed by server.") + replies: Final[list[object]] = [] + for command in self.commands: + try: + replies.append(self._run(command)) + except RedisError as e: + replies.append(e) + first_error: Final = next((reply for reply in replies if isinstance(reply, Exception)), None) + if first_error is not None: + raise first_error + return replies + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_leaves_a_missing_key_missing_when_the_set_is_refused_and_the_add_lost( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(refused=frozenset(("SET", "SETNX")), reached=1) + + with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(RedisConnectionError): + await redis_cache.async_increment(key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True) + + assert keyspace.values == {} + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_sets_a_missing_key_once_when_redis_runs_out_of_memory_after_the_set( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(writes_before_full=1) + added: Final = MagicMock() + + with patch.object(redis_cache, "init_async_client", return_value=keyspace): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + + assert current == 6.0 + assert float(keyspace.values["ns:spend:key:k"]) == 6.0 + assert keyspace.pipelines == 1 + added.assert_called_once_with() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("held", ["not-a-number", "0x10", "6_0", " 6", "inf", "nan", "1e999"]) +async def test_async_increment_only_if_missing_leaves_a_value_it_cannot_read_like_incrbyfloat_to_mains_add( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, held: str +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add: Final = _RaisingPipeline([7.5, True]) + client: Final = MagicMock() + client.pipeline.side_effect = [_RaisingPipeline([held.encode(), True]), add] + + with patch.object(redis_cache, "init_async_client", return_value=client): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + ) + + assert current == 7.5 + assert add.commands == [("INCRBYFLOAT", "ns:spend:key:k", 1.5), ("EXPIRE", "ns:spend:key:k", 60)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [float("inf"), float("nan")]) +async def test_async_increment_only_if_missing_sends_a_value_incrbyfloat_refuses_as_mains_add( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add: Final = _RaisingPipeline([ResponseError("increment would produce NaN or Infinity"), True]) + client: Final = MagicMock() + client.pipeline.return_value = add + + with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(ResponseError): + await redis_cache.async_increment( + key="spend:key:k", value=value, ttl=60, refresh_ttl=True, only_if_missing=True + ) + + assert client.pipeline.call_count == 1 + assert add.commands[0][0] == "INCRBYFLOAT" + + @pytest.mark.asyncio async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment( monkeypatch: pytest.MonkeyPatch, redis_no_ping: object diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index b6a0b50c9cd..d63950ebbdf 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -3419,40 +3419,97 @@ async def test_a_fresh_read_released_with_a_stale_read_joins_the_set_the_stale_r assert redis.values[counter] == 6.0 +class _SetFailsOnReleaseRedis(_HeldSetRedis): + """Holds each set on the wire and fails it once released, while a DEL still reaches Redis.""" + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if kwargs.get("only_if_missing"): + self.set_sent.set() + await self.release_set.wait() + raise ConnectionError("Timeout reading from socket") + return await super().async_increment(key, value, **kwargs) + + @pytest.mark.asyncio -async def test_a_set_failing_after_its_request_was_cancelled_deletes_the_counter_and_logs_no_stray_task_error( - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize( + ("cancelled", "counter_left", "in_memory_left"), + [ + pytest.param(True, 9.0, 6.0, id="cancelled-keeps-the-counter-as-mains-cancelled-add-does"), + pytest.param(False, None, None, id="still-waiting-deletes-it-as-mains-failed-add-does"), + ], +) +async def test_a_set_failing_after_its_request_was_cancelled_keeps_the_counter_and_logs_no_stray_task_error( + monkeypatch: pytest.MonkeyPatch, cancelled: bool, counter_left: float | None, in_memory_left: float | None ) -> None: counter: Final = "spend:end_user:cancelled-set-fails" - redis: Final = _HeldSetRedis() + redis: Final = _SetFailsOnReleaseRedis() _process(monkeypatch, redis, cached_spend=6.0) - - async def _cancel_a_request_with_its_set_on_the_wire() -> None: - read_before_seed: Final = asyncio.Event() - go: Final = asyncio.Event() - stale: Final = asyncio.ensure_future( - _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) - ) - await read_before_seed.wait() - await _cold_start(redis, counter) - go.set() - await redis.set_sent.wait() + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + go.set() + await redis.set_sent.wait() + if cancelled: stale.cancel() await asyncio.wait((stale,)) - - await _cancel_a_request_with_its_set_on_the_wire() stray: Final[list[object]] = [] # mutable-ok: collects what the loop's exception handler is given asyncio.get_running_loop().set_exception_handler(lambda _, context: stray.append(context)) - redis.failing = True + redis.values[counter] = 9.0 redis.release_set.set() + await asyncio.wait((stale,)) for _ in range(20): await asyncio.sleep(0) - redis.failing = False gc.collect() await asyncio.sleep(0) + assert redis.values.get(counter) == counter_left + assert ps.spend_counter_cache.in_memory_cache.get_cache(key=counter) == in_memory_left + assert stale.cancelled() or isinstance(stale.exception(), ConnectionError) assert stray == [] - assert ps.spend_counter_cache.in_memory_cache.get_cache(key=counter) is None + + +@pytest.mark.asyncio +async def test_only_the_late_write_sees_that_its_request_was_cancelled(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = "spend:end_user:abandoned-stays-in-the-write" + redis: Final = _HeldSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + seen: Final[list[tuple[str, bool]]] = [] # mutable-ok: records who saw the request as gone + + original_increment: Final = redis.async_increment + + async def increment_noting_abandoned(key: str, value: float, **kwargs: object) -> float: + current: Final = await original_increment(key, value, **kwargs) + if kwargs.get("only_if_missing"): + seen.append(("write", spend_counter_reseed.late_write_abandoned())) + return current + + monkeypatch.setattr(redis, "async_increment", increment_noting_abandoned) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + + async def request() -> None: + try: + await _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + except asyncio.CancelledError: + seen.append(("request", spend_counter_reseed.late_write_abandoned())) + raise + + stale: Final = asyncio.ensure_future(request()) + await read_before_seed.wait() + await _cold_start(redis, counter) + go.set() + await redis.set_sent.wait() + stale.cancel() + await asyncio.wait((stale,)) + redis.release_set.set() + for _ in range(20): + await asyncio.sleep(0) + + assert seen == [("request", False), ("write", True)] @pytest.mark.asyncio From a35cf6bcb0af701410f425cd9eeee04252e6b767 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 4 Nov 2026 18:07:39 -0800 Subject: [PATCH 18/18] fix(proxy): raise a refused TTL command after a late set as main's add does When Redis applied SET NX GET but refused the EXPIRE after it, the late write fell back to main's add, which added the spend on top of the value the set had just written A refusal that redis-py names after the first pipeline command now raises, as main's add raises on the same EXPIRE, and the caller deletes the counter as on main --- litellm/caching/redis_cache.py | 16 ++++++++++++--- tests/unit/caching/test_redis_cache.py | 27 +++++++++++++++++++++++--- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index f6223fb8d94..6ef99f5ffd1 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -143,6 +143,7 @@ _INCREMENT_IF_PRESENT_LUA: Final = ( _LUA_COUNT: Final = TypeAdapter(int) _PLAIN_DECIMAL: Final = re.compile(r"-?[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?") +_PIPELINE_COMMAND_NUMBER: Final = re.compile(r"Command # ([0-9]+) \(") _LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) @@ -153,6 +154,13 @@ def _as_incrbyfloat_writes(value: float) -> str: return "0" if text == "-0" else text +def _failed_after_the_first_command(error: Exception) -> bool: + """Whether redis-py names a pipeline command after the first as the one that failed, so the first one applied.""" + message: Final = error.args[0] if error.args else None + number: Final = _PIPELINE_COMMAND_NUMBER.match(message) if isinstance(message, str) else None + return number is not None and int(number.group(1)) > 1 + + def _plain_float(reply: object) -> float | None: """The finite float a stored value spells in plain decimal or exponent form, or None for anything else.""" text: Final = reply.decode(errors="replace") if isinstance(reply, bytes) else reply @@ -1494,8 +1502,9 @@ class RedisCache(BaseCache): async def _set_if_missing_and_get(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None: """SET NX GET and its TTL command in one round trip: whether the key was created and the value it holds, or - None when the server refused a command rather than failed to answer, or the key holds a value this does not - read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing.""" + None when the server refused SET NX GET rather than failed to answer, or the key holds a value this does not + read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing. A refused TTL command + raises, as it does in main's add.""" from redis.exceptions import ClusterError, RedisError from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError @@ -1510,7 +1519,8 @@ class RedisCache(BaseCache): pipe.expire(key, ttl) replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics except RedisError as e: - if isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)): + failed_to_answer: Final = isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)) + if failed_to_answer or _failed_after_the_first_command(e): raise return None if replies[0] is None: diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2a2ac5f7063..e673f16e5a6 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1727,8 +1727,6 @@ class _KeyspacePipeline(_RaisingPipeline): self.commands.append(("SETNX", name, value)) def _write(self, command: str) -> None: - if command in self.keyspace.refused: - raise NoPermissionError(f"this user has no permissions to run the '{command.lower()}' command") if self.writes_left is not None: if self.writes_left == 0: raise ResponseError("OOM command not allowed when used memory > 'maxmemory'.") @@ -1737,6 +1735,8 @@ class _KeyspacePipeline(_RaisingPipeline): def _run(self, command: tuple[object, ...]) -> object: name: Final = str(command[1]) held: Final = self.keyspace.values.get(name) + if command[0] in self.keyspace.refused: + raise NoPermissionError(f"this user has no permissions to run the '{str(command[0]).lower()}' command") if command[0] == "EXPIRE": return held is not None self._write(str(command[0])) @@ -1756,10 +1756,12 @@ class _KeyspacePipeline(_RaisingPipeline): if self.lost: raise RedisConnectionError("Connection closed by server.") replies: Final[list[object]] = [] - for command in self.commands: + for number, command in enumerate(self.commands, start=1): try: replies.append(self._run(command)) except RedisError as e: + words: str = " ".join(str(part) for part in command) + e.args = (f"Command # {number} ({words}) of pipeline caused error: {e.args}",) replies.append(e) first_error: Final = next((reply for reply in replies if isinstance(reply, Exception)), None) if first_error is not None: @@ -1801,6 +1803,25 @@ async def test_async_increment_only_if_missing_sets_a_missing_key_once_when_redi added.assert_called_once_with() +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_raises_a_refused_ttl_command_as_mains_add_does_without_adding_again( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(refused=frozenset(("EXPIRE",))) + added: Final = MagicMock() + + with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(NoPermissionError): + await redis_cache.async_increment( + key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + + assert keyspace.values == {"ns:spend:key:k": "6"} + assert keyspace.pipelines == 1 + added.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.parametrize("held", ["not-a-number", "0x10", "6_0", " 6", "inf", "nan", "1e999"]) async def test_async_increment_only_if_missing_leaves_a_value_it_cannot_read_like_incrbyfloat_to_mains_add(