cache LiteLLM_Config param reads in DualCache + batch scheduler-tick fetch

2026-07-15 18:21:44 +00:00 · 2026-04-28 16:29:50 -07:00
parent 89f0d4024e
commit 6052ce1017
3 changed files with 150 additions and 14 deletions
@@ -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")

@@ -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)
@@ -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(