diff --git a/litellm/constants.py b/litellm/constants.py index 1012ca831f2e..8d241886b51d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1879,6 +1879,7 @@ SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 +SPEND_COUNTER_CACHE_MAX_SIZE: Final = 10000 SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6fa7a7d761fd..cd5a5c70cefe 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -256,7 +256,7 @@ def generate_feedback_box(): import litellm._redis from litellm import Router from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger -from litellm.caching.caching import DualCache, RedisCache +from litellm.caching.caching import DualCache, InMemoryCache, RedisCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( @@ -285,6 +285,7 @@ def generate_feedback_box(): REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, + SPEND_COUNTER_CACHE_MAX_SIZE, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, ) @@ -2506,7 +2507,10 @@ def general_settings_view() -> Mapping[str, object]: user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) -spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) +spend_counter_cache: Final = DualCache( + in_memory_cache=InMemoryCache(max_size_in_memory=SPEND_COUNTER_CACHE_MAX_SIZE), + default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, +) cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index e1b5940935fb..d455c8b1da3f 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -107,6 +107,9 @@ "tests/integration/spend/test_cache_and_quota.py::test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores": [ "quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores" ], + "tests/integration/spend/test_spend_counter_scope_churn.py::test_over_budget_key_stays_blocked_while_hundreds_of_end_users_take_traffic": [ + "quota_management.budget.spend_counter.survives_scope_churn" + ], "tests/integration/spend/test_cache_and_quota.py::test_different_system_messages_do_not_share_a_cached_response": [ "quota_management.response_cache.system_messages_partition_cache_identity" ], diff --git a/tests/integration/spend/test_spend_counter_scope_churn.py b/tests/integration/spend/test_spend_counter_scope_churn.py new file mode 100644 index 000000000000..1dd0577b4286 --- /dev/null +++ b/tests/integration/spend/test_spend_counter_scope_churn.py @@ -0,0 +1,70 @@ +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from pydantic import TypeAdapter + +GENERIC_CACHE_SIZE: Final = 200 +CHURN_SCOPES: Final = 250 + + +def _chat(candidate: Gateway, model: str, key: str, user: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"scope churn {uuid.uuid4().hex}"}], "user": user}, + key=key, + ) + + +@pytest.mark.covers("quota_management.budget.spend_counter.survives_scope_churn") +@pytest.mark.timeout(180) +def test_over_budget_key_stays_blocked_while_hundreds_of_end_users_take_traffic( + gateway: Gateway, tmp_path: Path +) -> None: + shared: Final = TypeAdapter(dict[str, object]).validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + general_settings: Final = TypeAdapter(dict[str, object]).validate_python(shared["general_settings"]) + path: Final = tmp_path / "in-memory-counters.yaml" + path.write_text( + yaml.safe_dump( + { + "model_list": [], + "general_settings": {**general_settings, "proxy_batch_write_at": 600}, + "router_settings": shared["router_settings"], + } + ) + ) + run: Final = uuid.uuid4().hex + with ( + owned_proxy(gateway, tmp_path, {}, config=path, remove_environment=("REDIS_HOST", "REDIS_PORT")) as candidate, + candidate.scenario() as scenario, + httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as upstream, + ): + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + probe: Final = scenario.key(models=[model], max_budget=0.06) + churn: Final = scenario.key(models=[model]) + + def _churn_status(index: int) -> int: + return _chat(candidate, model, churn, f"churn-{run}-{index}").status_code + + first: Final = _chat(candidate, model, probe, f"probe-{run}") + assert first.status_code == 200 and first.json()["usage"]["total_tokens"] == 40, first.text + blocked: Final = _chat(candidate, model, probe, f"probe-{run}") + assert blocked.status_code == 422 and blocked.json()["error"]["type"] == "budget_exceeded", blocked.text + with ThreadPoolExecutor(max_workers=8) as pool: + statuses: Final = tuple(pool.map(_churn_status, range(CHURN_SCOPES))) + assert statuses.count(200) > GENERIC_CACHE_SIZE, statuses + upstream.get("/__observations").raise_for_status() + after: Final = _chat(candidate, model, probe, f"probe-{run}") + assert after.status_code == 422 and after.json()["error"]["type"] == "budget_exceeded", ( + f"over-budget key admitted after {statuses.count(200)} other scopes took traffic: {after.status_code} {after.text}" + ) + assert upstream.get("/__observations").json()["requests"] == [] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 26bfd5c52bd8..647bcb082846 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8553,6 +8553,31 @@ async def test_get_current_spend_fallback_to_in_memory(): ps.spend_counter_cache = original +@pytest.mark.asyncio +async def test_over_budget_key_spend_survives_hundreds_of_other_scopes_incrementing(): + from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY + from litellm.proxy.proxy_server import get_current_spend, increment_spend_counters, spend_counter_cache + + token = "unit-40221-probe" + other_scopes = DEFAULT_MAX_SIZE_IN_MEMORY + 100 + end_users = [f"unit-40221-{index}" for index in range(other_scopes)] + try: + await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.25) + for end_user in end_users: + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.01, end_user_id=end_user + ) + spend = await get_current_spend(counter_key=f"spend:key:{token}", fallback_spend=0.0) + finally: + for counter_key in [f"spend:key:{token}", *(f"spend:end_user:{end_user}" for end_user in end_users)]: + spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) + + assert spend == 0.25, ( + f"key spend read back as {spend} after {other_scopes} other scopes were charged; " + "the counter was evicted and the budget check fell back to the stale value (#40221)" + ) + + @pytest.mark.asyncio async def test_increment_spend_counters_initializes_and_increments(): """Counter should initialize from cached object spend, then increment.