From 6052ce1017aa27e7692da2d0664bfe91f659acfc Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Fri, 24 Apr 2026 16:44:50 -0700 Subject: [PATCH] cache LiteLLM_Config param reads in DualCache + batch scheduler-tick fetch --- litellm/proxy/proxy_server.py | 55 +++++++++--- litellm/proxy/utils.py | 89 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 20 +++++ 3 files changed, 150 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 007dbe5fa7..8f676df04c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -497,14 +497,18 @@ from litellm.proxy.utils import ( _get_redoc_url, _is_projected_spend_over_limit, _is_valid_team_configs, + get_config_param, get_custom_url, get_error_message_str, get_server_root_path, handle_exception_on_proxy, hash_password, hash_token, + invalidate_config_param, + litellm_config_cache, migrate_passwords_to_scrypt_async, model_dump_with_preserved_fields, + prefetch_config_params, update_spend, ) from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router @@ -2929,8 +2933,13 @@ class ProxyConfig: ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache spend_counter_cache.redis_cache = redis_usage_cache + litellm_config_cache.redis_cache = redis_usage_cache # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. + elif litellm_config_cache.redis_cache is None: + verbose_proxy_logger.info( + "litellm_config_cache: no Redis configured; cluster-wide cache sharing disabled." + ) def switch_on_llm_response_caching(self): """ @@ -4846,10 +4855,7 @@ class ProxyConfig: "environment_variables", ] for k in keys: - response = prisma_client.get_generic_data( - key="param_name", value=k, table_name="config" - ) - _tasks.append(response) + _tasks.append(get_config_param(prisma_client, k)) responses = await asyncio.gather(*_tasks) for response in responses: @@ -4931,6 +4937,19 @@ class ProxyConfig: global llm_router, llm_model_list, master_key, general_settings try: + # warm the config cache so the per-param reads below all hit + await prefetch_config_params( + prisma_client, + [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + "model_cost_map_reload_config", + "anthropic_beta_headers_reload_config", + ], + ) + # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) if self._should_load_db_object(object_type="models"): new_models = await self._get_models_from_db(prisma_client=prisma_client) @@ -4940,8 +4959,8 @@ class ProxyConfig: new_models=new_models, proxy_logging_obj=proxy_logging_obj ) - db_general_settings = await prisma_client.db.litellm_config.find_first( - where={"param_name": "general_settings"} + db_general_settings = await get_config_param( + prisma_client, "general_settings" ) # update general settings @@ -5034,10 +5053,7 @@ class ProxyConfig: from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook try: - # Load litellm_settings from DB - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "litellm_settings"} - ) + config_record = await get_config_param(prisma_client, "litellm_settings") if config_record is None or config_record.param_value is None: return @@ -5192,8 +5208,8 @@ class ProxyConfig: """ try: # Get model cost map reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "model_cost_map_reload_config"} + config_record = await get_config_param( + prisma_client, "model_cost_map_reload_config" ) if config_record is None or config_record.param_value is None: @@ -5288,6 +5304,7 @@ class ProxyConfig: }, }, ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}" @@ -5307,8 +5324,8 @@ class ProxyConfig: """ try: # Get anthropic beta headers reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "anthropic_beta_headers_reload_config"} + config_record = await get_config_param( + prisma_client, "anthropic_beta_headers_reload_config" ) if config_record is None or config_record.param_value is None: @@ -5396,6 +5413,7 @@ class ProxyConfig: }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") # Count providers in config provider_count = sum( @@ -12674,6 +12692,7 @@ async def update_config( # noqa: PLR0915 "update": {"param_value": v}, }, ) + await invalidate_config_param(k) ### OLD LOGIC [TODO] MOVE TO DB ### @@ -12861,6 +12880,7 @@ async def update_config_general_settings( "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) + await invalidate_config_param("general_settings") return response @@ -13144,6 +13164,7 @@ async def delete_config_general_settings( "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) + await invalidate_config_param("general_settings") return response @@ -13509,6 +13530,7 @@ async def reload_model_cost_map( }, }, ) + await invalidate_config_param("model_cost_map_reload_config") models_count = len(new_model_cost_map) if new_model_cost_map else 0 verbose_proxy_logger.info( @@ -13578,6 +13600,7 @@ async def schedule_model_cost_map_reload( }, }, ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reload scheduled for every {hours} hours" @@ -13631,6 +13654,7 @@ async def cancel_model_cost_map_reload( await prisma_client.db.litellm_config.delete( where={"param_name": "model_cost_map_reload_config"} ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info("Model cost map reload schedule cancelled") @@ -13861,6 +13885,7 @@ async def reload_anthropic_beta_headers( }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") provider_count = sum( 1 for k in new_config.keys() if k not in ["provider_aliases", "description"] @@ -13934,6 +13959,7 @@ async def schedule_anthropic_beta_headers_reload( }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info( f"Anthropic beta headers reload scheduled for every {hours} hours" @@ -13987,6 +14013,7 @@ async def cancel_anthropic_beta_headers_reload( await prisma_client.db.litellm_config.delete( where={"param_name": "anthropic_beta_headers_reload_config"} ) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 712853a33c..3a1184c434 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2442,6 +2442,92 @@ async def _lookup_deprecated_key( return None +# DualCache for LiteLLM_Config param_name reads. +# Redis layer is attached in proxy_server._init_cache. +LITELLM_CONFIG_CACHE_TTL_SECONDS: int = int( + os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60") +) +_CONFIG_CACHE_MISS: str = "__litellm_config_param_miss__" + +litellm_config_cache: DualCache = DualCache( + default_in_memory_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS, + default_redis_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS, +) + + +class _ConfigRow: + """Mimics the Prisma litellm_config row shape for cached entries.""" + + __slots__ = ("param_name", "param_value") + + def __init__(self, param_name: str, param_value: Any) -> None: + self.param_name = param_name + self.param_value = param_value + + +def _config_cache_key(param_name: str) -> str: + return f"litellm_config:param:{param_name}" + + +def _pack_config_row(row: Any) -> Dict[str, Any]: + return {"param_name": row.param_name, "param_value": row.param_value} + + +def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]: + if cached is None or cached == _CONFIG_CACHE_MISS: + return None + if isinstance(cached, dict): + return _ConfigRow(cached["param_name"], cached["param_value"]) + return None + + +async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any]: + """Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None.""" + cache_key = _config_cache_key(param_name) + cached = await litellm_config_cache.async_get_cache(cache_key) + if cached is not None: + return _unpack_config_row(cached) + + row = await prisma_client.get_generic_data( + key="param_name", value=param_name, table_name="config" + ) + cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + await litellm_config_cache.async_set_cache( + cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS + ) + return row + + +async def invalidate_config_param(param_name: str) -> None: + """Evict from both cache layers; call after every LiteLLM_Config write.""" + await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) + + +async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: + """Batch-load LiteLLM_Config rows into the cache with one find_many.""" + if not param_names: + return + try: + rows = await prisma_client.db.litellm_config.find_many( + where={"param_name": {"in": param_names}} # type: ignore + ) + except Exception as e: + verbose_proxy_logger.debug( + "prefetch_config_params failed, falling through to per-param queries: %s", + e, + ) + return + by_name = {row.param_name: row for row in rows} + for name in param_names: + row = by_name.get(name) + cache_value: Any = ( + _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + ) + await litellm_config_cache.async_set_cache( + _config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS + ) + + class PrismaClient: spend_log_transactions: List = [] _spend_log_transactions_lock = asyncio.Lock() @@ -3310,6 +3396,9 @@ class PrismaClient: tasks.append(updated_table_row) await asyncio.gather(*tasks) + # invalidate cache so other pods see writes from save_config + for k in data.keys(): + await invalidate_config_param(k) verbose_proxy_logger.info("Data Inserted into Config Table") elif table_name == "spend": db_data = self.jsonify_object(data=data) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 3349a138ee..1f4f82a64e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2544,6 +2544,14 @@ class TestPriceDataReloadAPI: class TestPriceDataReloadIntegration: """Integration tests for the complete price data reload feature""" + @pytest.fixture(autouse=True) + def _flush_litellm_config_cache(self): + from litellm.proxy.utils import litellm_config_cache + + litellm_config_cache.flush_cache() + yield + litellm_config_cache.flush_cache() + @pytest.fixture def client_with_auth(self): """Create a test client with authentication""" @@ -2601,6 +2609,7 @@ class TestPriceDataReloadIntegration: def test_distributed_reload_check_function(self): """Test the _check_and_reload_model_cost_map function""" from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import litellm_config_cache proxy_config = ProxyConfig() @@ -2609,14 +2618,19 @@ class TestPriceDataReloadIntegration: # Test case 1: No config in database mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + # _check_and_reload_model_cost_map routes through get_config_param, + # which calls prisma.get_generic_data on a cache miss. + mock_prisma.get_generic_data = AsyncMock(return_value=None) # Should return early without reloading asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) # Test case 2: Config with interval but not time to reload + litellm_config_cache.flush_cache() mock_config = MagicMock() mock_config.param_value = {"interval_hours": 6, "force_reload": False} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) # Mock current time and last reload time with patch( @@ -2632,8 +2646,10 @@ class TestPriceDataReloadIntegration: asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) # Test case 3: Config with force reload + litellm_config_cache.flush_cache() mock_config.param_value = {"interval_hours": 6, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) original_model_cost = litellm.model_cost.copy() @@ -2675,6 +2691,8 @@ class TestPriceDataReloadIntegration: mock_config = MagicMock() mock_config.param_value = {"interval_hours": 24, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + # _check_and_reload_model_cost_map now reads through get_generic_data. + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) original_model_cost = litellm.model_cost.copy() @@ -2770,6 +2788,8 @@ class TestPriceDataReloadIntegration: mock_config = MagicMock() mock_config.param_value = {"interval_hours": 12, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + # _check_and_reload_anthropic_beta_headers now reads through get_generic_data. + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) with patch(