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