From db3338b2068b6fa6448d9f82540208fae499fcad Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 09:55:30 -0700 Subject: [PATCH] feat(proxy): make the in-memory management cache capacity configurable (#40725) * feat(proxy): make the in-memory management cache capacity configurable Add general_settings.user_api_key_cache_max_size (positive int, default 200) to resize the in-memory tier of the shared user_api_key_cache at startup and on DB config reloads, expose it in the Admin UI general settings, and cover it with behavioral tests. Prior art: #34726 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): resize the in-memory tier from DualCache so any cache instance honours the cap Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): wrap the cache capacity field description to the 120 col limit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/dual_cache.py | 5 +- litellm/caching/in_memory_cache.py | 6 +- litellm/proxy/_types.py | 9 ++ litellm/proxy/proxy_server.py | 28 ++++ tests/test_litellm/proxy/test_proxy_server.py | 145 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 6 files changed, 195 insertions(+), 3 deletions(-) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 26ad89a70cc..81e2af45686 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -22,7 +22,7 @@ from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from .base_cache import BaseCache -from .in_memory_cache import InMemoryCache +from .in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure if TYPE_CHECKING: @@ -83,6 +83,9 @@ class DualCache(BaseCache): if default_redis_ttl is not None: self.default_redis_ttl = default_redis_ttl + def update_in_memory_max_size(self, max_size: int | None) -> None: + self.in_memory_cache.max_size_in_memory = DEFAULT_MAX_SIZE_IN_MEMORY if max_size is None else max_size + def attach_redis_cache( self, redis_cache: RedisCache | None = None, diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 38a9966f9f9..4058a3d72dd 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -24,11 +24,13 @@ from litellm.constants import MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB from .base_cache import BaseCache +DEFAULT_MAX_SIZE_IN_MEMORY: Final = 200 + class InMemoryCache(BaseCache): def __init__( self, - max_size_in_memory: int | None = 200, + max_size_in_memory: int | None = DEFAULT_MAX_SIZE_IN_MEMORY, default_ttl: int | None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute max_size_per_item: int | None = 1024, # 1MB = 1024KB @@ -37,7 +39,7 @@ class InMemoryCache(BaseCache): max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default """ self.max_size_in_memory = ( - max_size_in_memory if max_size_in_memory is not None else 200 + max_size_in_memory if max_size_in_memory is not None else DEFAULT_MAX_SIZE_IN_MEMORY ) # set an upper bound of 200 items in-memory self.default_ttl = default_ttl or 600 self.max_size_per_item = max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB # 1MB = 1024KB diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2cbff128635..43af18cbef4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2602,6 +2602,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): global_max_parallel_requests: int | None = Field( None, description="global max parallel requests to allow for a proxy instance." ) + user_api_key_cache_max_size: int | None = Field( + None, + gt=0, + description=( + "max number of entries (virtual keys, teams, users, end users, memberships, ...) each worker keeps in " + "its in-memory auth cache. Defaults to 200. Raise this if you have more active keys than that or auth " + "lookups keep hitting the DB" + ), + ) max_request_size_mb: int | None = Field( None, description="max request size in MB, if a request is larger than this size it will be rejected", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5d94b65ccba..687fbc0cb9b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5808,6 +5808,16 @@ class ProxyConfig: default_redis_ttl=ttl, ) + ### USER API KEY CACHE MAX SIZE (in-memory tier shared by keys, teams, users, end users, ...) ### + if "user_api_key_cache_max_size" in general_settings: + user_api_key_cache.update_in_memory_max_size( + ConfigGeneralSettings.model_validate( + MappingProxyType( + {"user_api_key_cache_max_size": general_settings["user_api_key_cache_max_size"]} + ) + ).user_api_key_cache_max_size + ) + ### PKCE MULTI-INSTANCE PREREQUISITE CHECK ### # PKCE verifiers are stored in redis_usage_cache when available so they can # be read back by any instance (not just the one that started the auth flow). @@ -7058,6 +7068,23 @@ class ProxyConfig: "enable_openai_websocket_passthrough" ) + if "user_api_key_cache_max_size" not in self._yaml_general_settings_keys: + db_cache_max_size: Final = _general_settings.get("user_api_key_cache_max_size") + try: + cache_max_size: Final = ConfigGeneralSettings.model_validate( + MappingProxyType({"user_api_key_cache_max_size": db_cache_max_size}) + ).user_api_key_cache_max_size + except ValidationError: + verbose_proxy_logger.warning( + "Ignoring invalid general_settings.user_api_key_cache_max_size=%r from the DB", db_cache_max_size + ) + else: + if cache_max_size is None: + general_settings.pop("user_api_key_cache_max_size", None) + else: + general_settings["user_api_key_cache_max_size"] = cache_max_size + user_api_key_cache.update_in_memory_max_size(cache_max_size) + ## STORE MODEL IN DB ## if "store_model_in_db" in _general_settings: value = _general_settings["store_model_in_db"] @@ -16970,6 +16997,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "cancel_on_disconnect": "Boolean", "disable_auto_add_proxy_admin_to_teams": "Boolean", "apply_user_budget_to_team_keys": "Boolean", + "user_api_key_cache_max_size": "Integer", } ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 4f4a9e87d18..3b3031647dc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7313,6 +7313,88 @@ async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins( assert ps.general_settings["apply_user_budget_to_team_keys"] is True +def _fill_user_api_key_cache(cache: DualCache, count: int) -> None: + for index in range(count): + cache.set_cache(key=f"key-{index}", value={"token": f"key-{index}"}, local_only=True) + + +@pytest.mark.asyncio +async def test_update_general_settings_user_api_key_cache_max_size_resizes_the_running_cache(monkeypatch): + """The Admin UI writes the capacity to the DB config, so the running cache has + to pick it up on reload; otherwise the knob only works after a restart.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + cache = UserApiKeyCache() + monkeypatch.setattr(proxy_server_module, "general_settings", {}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig()._update_general_settings(db_general_settings={"user_api_key_cache_max_size": 300}) + + assert proxy_server_module.general_settings["user_api_key_cache_max_size"] == 300 + + _fill_user_api_key_cache(cache, 250) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + +@pytest.mark.asyncio +async def test_update_general_settings_clearing_user_api_key_cache_max_size_restores_the_default(monkeypatch): + """Blanking the field in the dashboard deletes the key, so the cache must fall + back to the default capacity rather than keep the last configured size.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + cache = UserApiKeyCache() + cache.update_in_memory_max_size(5000) + monkeypatch.setattr(proxy_server_module, "general_settings", {"user_api_key_cache_max_size": 5000}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig()._update_general_settings(db_general_settings={"store_model_in_db": True}) + + assert "user_api_key_cache_max_size" not in proxy_server_module.general_settings + + _fill_user_api_key_cache(cache, 201) + assert cache.get_cache(key="key-0", local_only=True) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("db_value", [0, -5, "lots"]) +async def test_update_general_settings_ignores_an_invalid_user_api_key_cache_max_size(db_value, monkeypatch): + """A non-positive capacity would make the eviction loop pop an empty heap on the + next write, so a bad DB value must leave the running cache untouched.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + cache = UserApiKeyCache() + cache.update_in_memory_max_size(300) + monkeypatch.setattr(proxy_server_module, "general_settings", {}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig()._update_general_settings(db_general_settings={"user_api_key_cache_max_size": db_value}) + + assert "user_api_key_cache_max_size" not in proxy_server_module.general_settings + + _fill_user_api_key_cache(cache, 250) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + +@pytest.mark.asyncio +async def test_update_general_settings_user_api_key_cache_max_size_yaml_wins(monkeypatch): + """A DB value must not silently override an explicit YAML capacity on reload.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"user_api_key_cache_max_size"} + cache = UserApiKeyCache() + cache.update_in_memory_max_size(300) + monkeypatch.setattr(proxy_server_module, "general_settings", {"user_api_key_cache_max_size": 300}) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await proxy_config._update_general_settings(db_general_settings={"user_api_key_cache_max_size": 10}) + + assert proxy_server_module.general_settings["user_api_key_cache_max_size"] == 300 + + _fill_user_api_key_cache(cache, 250) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "db_value,expected", @@ -10344,6 +10426,27 @@ def test_get_config_list_includes_apply_user_budget_to_team_keys(monkeypatch): app.dependency_overrides.clear() +def test_get_config_list_includes_user_api_key_cache_max_size(monkeypatch): + """The Admin UI General Settings table renders whatever /config/list returns, + so the cache capacity has to be exposed there as an Integer to be editable.""" + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(proxy_server_module, "prisma_client", mock_prisma) + app.dependency_overrides[proxy_server_module.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + assert fields["user_api_key_cache_max_size"]["field_type"] == "Integer" + finally: + app.dependency_overrides.clear() + + def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch): """The throttle fraction is a litellm_settings scalar surfaced on the General Settings table as a Float field so it sits with the other global limits; it @@ -12946,6 +13049,48 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin assert router.fallback_access_check is router_fallback_access_check +@pytest.mark.asyncio +async def test_load_config_user_api_key_cache_max_size_keeps_more_than_200_entries(tmp_path, monkeypatch): + """The auth cache used to be pinned at InMemoryCache's 200 entry default, so a + deployment with more keys than that evicted constantly and every request + fell through to the DB. The YAML knob has to raise the cap on the live cache.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump({"general_settings": {"user_api_key_cache_max_size": "1000"}})) + + cache = UserApiKeyCache() + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + _fill_user_api_key_cache(cache, 999) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_value", [0, -1, "unbounded"]) +async def test_load_config_rejects_a_non_positive_user_api_key_cache_max_size(tmp_path, bad_value, monkeypatch): + """InMemoryCache treats 0 as 'cache nothing' and a negative cap makes eviction + pop an empty heap, so the proxy must refuse to boot with such a value instead + of silently disabling auth caching.""" + from pydantic import ValidationError + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump({"general_settings": {"user_api_key_cache_max_size": bad_value}})) + + cache = UserApiKeyCache() + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", cache) + with pytest.raises(ValidationError): + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + _fill_user_api_key_cache(cache, 150) + assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} + + def test_docs_redoc_openapi_are_reachable_by_default(): """ LIT-6745: the interactive/machine-readable docs surfaces are on by diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 1b34c6f6a51..7fd4f8e413d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26007,6 +26007,11 @@ export interface components { * @description If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False. */ use_spend_logs_partitioning?: boolean | null; + /** + * User Api Key Cache Max Size + * @description max number of entries (virtual keys, teams, users, end users, memberships, ...) each worker keeps in its in-memory auth cache. Defaults to 200. Raise this if you have more active keys than that or auth lookups keep hitting the DB + */ + user_api_key_cache_max_size?: number | null; /** User Header Mappings */ user_header_mappings?: components["schemas"]["UserHeaderMapping"][] | null; /**