diff --git a/.gitignore b/.gitignore index f5591b7f2c..59812ed6ed 100644 --- a/.gitignore +++ b/.gitignore @@ -99,4 +99,5 @@ STABILIZATION_TODO.md **/test-results **/playwright-report **/*.storageState.json -**/coverage \ No newline at end of file +**/coverage +test-config \ No newline at end of file diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index be819a04ca..7187690375 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -307,9 +307,70 @@ router_settings: | token_rate_limit_type | string | Rate limit counting method: "total", "output", or "input" tokens | | use_redis_transaction_buffer | boolean | If true, buffers database transactions in Redis before writing | | use_shared_health_check | boolean | If true, uses Redis-backed shared health check state across multiple proxy instances | +| enable_redis_auth_cache | boolean | **[Beta]** When `true`, attaches Redis to the virtual-key auth cache (`user_api_key_cache`) so all proxy workers/pods share the same cache instead of each pod resolving keys independently against the database. Requires `litellm_settings.cache: true` with a Redis backend. Significantly reduces database load in multi-worker deployments by eliminating per-pod cache misses on the `combined_view` query. Off by default for a safe phased rollout — enable once your Redis cluster is healthy. Will become the default in a future release. See [Redis Auth Cache](#redis-auth-cache-multi-worker-db-load-reduction). | | user_header_mappings | dict | Map custom request headers to user IDs using lookup rules | | user_header_name | string | HTTP header name to extract user identity from requests | +## Redis Auth Cache — Multi-Worker DB Load Reduction + +### Problem + +In multi-worker or multi-pod deployments each worker process keeps its own **in-memory** virtual-key cache (`user_api_key_cache`). When a request is routed to a worker that has not seen a key before, it runs a `combined_view` SQL query against the database (8+ JOINs). With many workers and a large number of unique keys this causes: + +- Redundant DB queries on every pod's cold-start +- Sustained CPU spikes on the database (especially visible in RDS Performance Insights) +- Increased p99 latency for requests that miss the local cache + +### Solution + +Setting `enable_redis_auth_cache: true` attaches Redis to `user_api_key_cache` so the resolved key object is stored in a **shared** Redis cache. A cache hit on any worker prevents the DB query entirely. + +```yaml +# config.yaml + +litellm_settings: + cache: true + cache_params: + type: redis + host: os.environ/REDIS_HOST + port: os.environ/REDIS_PORT + +general_settings: + master_key: sk-1234 + enable_redis_auth_cache: true # ← share the auth cache across workers +``` + +### Requirements + +| Requirement | Notes | +|---|---| +| `litellm_settings.cache: true` | Redis must be configured as the cache backend | +| Redis cluster is healthy | Auth lookups now depend on Redis availability — monitor it | +| LiteLLM ≥ version with CacheCodec support | All read/write paths now use `CacheCodec` for safe serialisation across Redis round-trips | + +### Rollout recommendation + +1. Deploy with `enable_redis_auth_cache: false` (default) and baseline DB CPU. +2. Enable on a canary pod first; confirm cache-hit rate via proxy debug logs (`LITELLM_LOG=DEBUG`). +3. Roll out to all pods; monitor DB CPU drop. +4. The flag will be removed and the behaviour made permanent in a future LiteLLM release. + +### Debugging + +Set `LITELLM_LOG=DEBUG` and look for: + +``` +enable_redis_auth_cache=True: attached Redis to user_api_key_cache — virtual-key lookups are now shared across all proxy workers. +``` + +If the flag is off you will see: + +``` +enable_redis_auth_cache is not set: user_api_key_cache remains in-memory only (per-worker). Set general_settings.enable_redis_auth_cache: true to share the auth cache across workers and reduce DB load. +``` + +--- + ### router_settings - Reference :::info diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a592a81ed8..960261d5cf 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2358,6 +2358,19 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) + enable_redis_auth_cache: Optional[bool] = Field( + None, + description=( + "When True, attaches Redis to user_api_key_cache so virtual-key lookups " + "are shared across all proxy workers/pods instead of being resolved " + "per-process. Requires a Redis cache to be configured under litellm_settings. " + "Reduces DB load significantly in multi-worker deployments by eliminating " + "redundant combined_view SQL queries caused by per-pod cache misses. " + "Off by default for a safe phased rollout — set to True once your Redis " + "cluster is healthy and the CacheCodec serialisation has been validated in " + "your environment. Will be enabled by default in a future release." + ), + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 760c0fb12d..02d41041e8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -823,7 +823,9 @@ async def get_default_end_user_budget( # Check cache first cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) if cached_budget is not None: - return LiteLLM_BudgetTable(**cached_budget) + deserialized = CacheCodec.deserialize(cached_budget, LiteLLM_BudgetTable) + if deserialized is not None: + return deserialized # Fetch from database try: @@ -837,14 +839,15 @@ async def get_default_end_user_budget( ) return None + _budget_obj = LiteLLM_BudgetTable(**budget_record.dict()) # Cache the budget for 60 seconds await user_api_key_cache.async_set_cache( key=cache_key, - value=budget_record.dict(), + value=CacheCodec.serialize(_budget_obj, model_type=LiteLLM_BudgetTable), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) - return LiteLLM_BudgetTable(**budget_record.dict()) + return _budget_obj except Exception as e: verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}") @@ -960,20 +963,20 @@ async def get_end_user_object( # Check cache first cached_user_obj = await user_api_key_cache.async_get_cache(key=_key) if cached_user_obj is not None: - return_obj = LiteLLM_EndUserTable(**cached_user_obj) + return_obj = CacheCodec.deserialize(cached_user_obj, LiteLLM_EndUserTable) + if return_obj is not None: + # Apply default budget if needed + return_obj = await _apply_default_budget_to_end_user( + end_user_obj=return_obj, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) - # Apply default budget if needed - return_obj = await _apply_default_budget_to_end_user( - end_user_obj=return_obj, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - ) + # Check budget limits + _check_end_user_budget(end_user_obj=return_obj, route=route) - # Check budget limits - _check_end_user_budget(end_user_obj=return_obj, route=route) - - return return_obj + return return_obj # Fetch from database try: @@ -996,9 +999,10 @@ async def get_end_user_object( parent_otel_span=parent_otel_span, ) - # Save to cache (always store as dict for consistency) + # Save to cache await user_api_key_cache.async_set_cache( - key="end_user_id:{}".format(end_user_id), value=_response.dict() + key="end_user_id:{}".format(end_user_id), + value=CacheCodec.serialize(_response, model_type=LiteLLM_EndUserTable), ) # Check budget limits @@ -1051,10 +1055,11 @@ async def get_tag_objects_batch( cache_key = f"tag:{tag_name}" cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) if cached_tag is not None: - if isinstance(cached_tag, dict): - tag_objects[tag_name] = LiteLLM_TagTable(**cached_tag) + deserialized_tag = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable) + if deserialized_tag is not None: + tag_objects[tag_name] = deserialized_tag else: - tag_objects[tag_name] = cached_tag + uncached_tags.append(tag_name) else: uncached_tags.append(tag_name) @@ -1070,11 +1075,12 @@ async def get_tag_objects_batch( for db_tag in db_tags: tag_name = db_tag.tag_name cache_key = f"tag:{tag_name}" - # Cache with default TTL (same as end_user objects) + _tag_obj = LiteLLM_TagTable(**db_tag.dict()) await user_api_key_cache.async_set_cache( - key=cache_key, value=db_tag.dict() + key=cache_key, + value=CacheCodec.serialize(_tag_obj, model_type=LiteLLM_TagTable), ) - tag_objects[tag_name] = LiteLLM_TagTable(**db_tag.dict()) + tag_objects[tag_name] = _tag_obj except Exception as e: verbose_proxy_logger.debug(f"Error batch fetching tags from database: {e}") @@ -1146,7 +1152,11 @@ async def get_team_membership( # check if in cache cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key) if cached_membership_obj is not None: - return LiteLLM_TeamMembership(**cached_membership_obj) + deserialized_membership = CacheCodec.deserialize( + cached_membership_obj, LiteLLM_TeamMembership + ) + if deserialized_membership is not None: + return deserialized_membership # else, check db try: @@ -1158,10 +1168,11 @@ async def get_team_membership( if response is None: return None - # save the team membership object to cache (store as dict) - await user_api_key_cache.async_set_cache(key=_key, value=response.dict()) - _response = LiteLLM_TeamMembership(**response.dict()) + await user_api_key_cache.async_set_cache( + key=_key, + value=CacheCodec.serialize(_response, model_type=LiteLLM_TeamMembership), + ) return _response except Exception: @@ -1350,10 +1361,11 @@ async def get_user_object( if not check_db_only: cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id) if cached_user_obj is not None: - if isinstance(cached_user_obj, dict): - return LiteLLM_UserTable(**cached_user_obj) - elif isinstance(cached_user_obj, LiteLLM_UserTable): - return cached_user_obj + deserialized_user = CacheCodec.deserialize( + cached_user_obj, LiteLLM_UserTable + ) + if deserialized_user is not None: + return deserialized_user # else, check db if prisma_client is None: raise Exception("No db connected") @@ -1415,7 +1427,7 @@ async def get_user_object( # save the user object to cache await user_api_key_cache.async_set_cache( key=user_id, - value=response_dict, + value=CacheCodec.serialize(_response, model_type=LiteLLM_UserTable), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1886,16 +1898,19 @@ async def get_team_object_by_alias( ) # Cache the result by both alias and team_id + _serialized_team = CacheCodec.serialize( + team_obj, model_type=LiteLLM_TeamTableCachedObj + ) await user_api_key_cache.async_set_cache( key=cache_key, - value=team_obj, + value=_serialized_team, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by team_id for consistency team_id_cache_key = "team_id:{}".format(team_obj.team_id) await user_api_key_cache.async_set_cache( key=team_id_cache_key, - value=team_obj, + value=_serialized_team, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -1946,10 +1961,11 @@ async def get_org_object_by_alias( cache_key = "org_alias:{}".format(org_alias) cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) if cached_org_obj is not None: - if isinstance(cached_org_obj, dict): - return LiteLLM_OrganizationTable(**cached_org_obj) - elif isinstance(cached_org_obj, LiteLLM_OrganizationTable): - return cached_org_obj + deserialized_org = CacheCodec.deserialize( + cached_org_obj, LiteLLM_OrganizationTable + ) + if deserialized_org is not None: + return deserialized_org # Query database by organization_alias try: @@ -1976,16 +1992,19 @@ async def get_org_object_by_alias( org = orgs[0] org_obj = LiteLLM_OrganizationTable(**org.model_dump()) + _serialized_org = CacheCodec.serialize( + org_obj, model_type=LiteLLM_OrganizationTable + ) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=org_obj.model_dump(), + value=_serialized_org, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by org_id for consistency await user_api_key_cache.async_set_cache( key="org_id:{}".format(org_obj.organization_id), - value=org_obj.model_dump(), + value=_serialized_org, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2287,10 +2306,11 @@ async def get_object_permission( key = "object_permission_id:{}".format(object_permission_id) cached_obj_permission = await user_api_key_cache.async_get_cache(key=key) if cached_obj_permission is not None: - if isinstance(cached_obj_permission, dict): - return LiteLLM_ObjectPermissionTable(**cached_obj_permission) - elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable): - return cached_obj_permission + deserialized_perm = CacheCodec.deserialize( + cached_obj_permission, LiteLLM_ObjectPermissionTable + ) + if deserialized_perm is not None: + return deserialized_perm # else, check db try: @@ -2301,14 +2321,16 @@ async def get_object_permission( if response is None: return None - # save the object permission to cache + _perm_obj = LiteLLM_ObjectPermissionTable(**response.dict()) await user_api_key_cache.async_set_cache( key=key, - value=response.model_dump(), + value=CacheCodec.serialize( + _perm_obj, model_type=LiteLLM_ObjectPermissionTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) - return LiteLLM_ObjectPermissionTable(**response.dict()) + return _perm_obj except Exception: return None @@ -2339,10 +2361,11 @@ async def get_managed_vector_store_rows_by_uuids( key = "managed_vector_store_id:{}".format(uuid) cached = await user_api_key_cache.async_get_cache(key=key) if cached is not None: - if isinstance(cached, dict): - result.append(LiteLLM_ManagedVectorStoresTable(**cached)) - elif isinstance(cached, LiteLLM_ManagedVectorStoresTable): - result.append(cached) + deserialized_vs = CacheCodec.deserialize( + cached, LiteLLM_ManagedVectorStoresTable + ) + if deserialized_vs is not None: + result.append(deserialized_vs) else: cache_misses.append(uuid) else: @@ -2370,7 +2393,9 @@ async def get_managed_vector_store_rows_by_uuids( key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id) await user_api_key_cache.async_set_cache( key=key, - value=row_dict, + value=CacheCodec.serialize( + cached_obj, model_type=LiteLLM_ManagedVectorStoresTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) result.append(cached_obj) @@ -2413,12 +2438,13 @@ async def get_org_object( cache_key = "org_id:{}:with_budget".format(org_id) # check if in cache - cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key) + cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) if cached_org_obj is not None: - if isinstance(cached_org_obj, dict): - return LiteLLM_OrganizationTable(**cached_org_obj) - elif isinstance(cached_org_obj, LiteLLM_OrganizationTable): - return cached_org_obj + deserialized_org = CacheCodec.deserialize( + cached_org_obj, LiteLLM_OrganizationTable + ) + if deserialized_org is not None: + return deserialized_org # else, check db try: query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}} @@ -2432,16 +2458,15 @@ async def get_org_object( if response is None: raise Exception + _org_obj = LiteLLM_OrganizationTable(**response.model_dump()) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=( - response.model_dump() if hasattr(response, "model_dump") else response - ), + value=CacheCodec.serialize(_org_obj, model_type=LiteLLM_OrganizationTable), ttl=DEFAULT_IN_MEMORY_TTL, ) - return response + return _org_obj except Exception: raise Exception( f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call." @@ -3451,10 +3476,11 @@ async def get_project_object( cache_key = "project_id:{}".format(project_id) cached_obj = await user_api_key_cache.async_get_cache(key=cache_key) if cached_obj is not None: - if isinstance(cached_obj, dict): - return LiteLLM_ProjectTableCachedObj(**cached_obj) - elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj): - return cached_obj + deserialized_project = CacheCodec.deserialize( + cached_obj, LiteLLM_ProjectTableCachedObj + ) + if deserialized_project is not None: + return deserialized_project # Fetch from DB project_row = await prisma_client.db.litellm_projecttable.find_unique( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 4a6856b6d1..07d944ae93 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -45,6 +45,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import can_team_access_model +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.utils import PrismaClient, ProxyLogging from .auth_checks import ( @@ -1376,7 +1377,9 @@ class JWTAuthManager: if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=user_object.model_dump(), + value=CacheCodec.serialize( + user_object, model_type=LiteLLM_UserTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1399,7 +1402,9 @@ class JWTAuthManager: if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=user_object.model_dump(), + value=CacheCodec.serialize( + user_object, model_type=LiteLLM_UserTable + ), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 855ce50ec0..3fe8863e26 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1297,27 +1297,39 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if prisma_client is not None: _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" - team_member_info = await user_api_key_cache.async_get_cache( + _cached_member = await user_api_key_cache.async_get_cache( key=_cache_key ) + team_member_info: Optional[LiteLLM_TeamMembership] = ( + CacheCodec.deserialize(_cached_member, LiteLLM_TeamMembership) + if _cached_member is not None + else None + ) if team_member_info is None: # read from DB _user_id = valid_token.user_id _team_id = valid_token.team_id if _user_id is not None and _team_id is not None: - team_member_info = await prisma_client.db.litellm_teammembership.find_first( + _db_member = await prisma_client.db.litellm_teammembership.find_first( where={ "user_id": _user_id, "team_id": _team_id, }, # type: ignore include={"litellm_budget_table": True}, ) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - ttl=5, - ) + if _db_member is not None: + team_member_info = LiteLLM_TeamMembership( + **_db_member.dict() + ) + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=CacheCodec.serialize( + team_member_info, + model_type=LiteLLM_TeamMembership, + ), + ttl=5, + ) if ( team_member_info is not None diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py index af0c3eef18..80a8d6281a 100644 --- a/litellm/proxy/common_utils/cache_pydantic_utils.py +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -44,10 +44,17 @@ class CacheCodec: If ``model_type`` is set, the payload is validated with that model, then ``model_dump(mode="json", exclude_none=True)`` — symmetric with ``deserialize``. + If the value is already an instance of ``model_type`` (or a subclass), + ``model_validate`` is skipped to avoid an unnecessary Pydantic copy — the + value is dumped directly. + If ``model_type`` is omitted, any ``BaseModel`` is dumped as above; other values (e.g. plain ``dict``) are returned unchanged. """ if model_type is not None: + if isinstance(value, model_type): + # Already the right type: dump directly, skip re-validation copy. + return value.model_dump(mode="json", exclude_none=True) if isinstance(value, (dict, BaseModel)): return model_type.model_validate(value).model_dump( mode="json", exclude_none=True diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index caaec12f7a..30fed09623 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -9,6 +9,7 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.auth.auth_checks import ( _cache_access_object, _cache_key_object, @@ -236,12 +237,11 @@ async def _patch_key_caches_add_access_group( ) -> None: """Patch cached key objects to include access_group_id.""" for token in key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) - if cached_key is None: + raw = await user_api_key_cache.async_get_cache(key=token) + if raw is None: continue - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if not isinstance(cached_key, UserAPIKeyAuth): + cached_key = CacheCodec.deserialize(raw, UserAPIKeyAuth) + if cached_key is None: continue if cached_key.access_group_ids is None: cached_key.access_group_ids = [access_group_id] @@ -267,12 +267,11 @@ async def _patch_key_caches_remove_access_group( ) -> None: """Patch cached key objects to remove access_group_id.""" for token in key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) - if cached_key is None: + raw = await user_api_key_cache.async_get_cache(key=token) + if raw is None: continue - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids: + cached_key = CacheCodec.deserialize(raw, UserAPIKeyAuth) + if cached_key is not None and cached_key.access_group_ids: cached_key.access_group_ids = [ ag for ag in cached_key.access_group_ids if ag != access_group_id ] diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 46e7963da7..0362c3ff6b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -70,6 +70,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.auth.auth_utils import _has_user_setup_sso from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -1233,11 +1234,7 @@ async def _sync_user_role_from_jwt_role_map( user_info.user_role = mapped_role.value await user_api_key_cache.async_set_cache( key=user_info.user_id, - value=( - user_info.model_dump() - if hasattr(user_info, "model_dump") - else dict(user_info) - ), + value=CacheCodec.serialize(user_info, model_type=LiteLLM_UserTable), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 72ec41ef97..f8c52a5289 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -78,8 +78,11 @@ from litellm.proxy._types import ( InvitationNew, InvitationUpdate, Litellm_EntityType, + LiteLLM_EndUserTable, LiteLLM_JWTAuth, + LiteLLM_TagTable, LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, PassThroughGenericEndpoint, @@ -2029,34 +2032,38 @@ async def update_cache( # noqa: PLR0915 ### UPDATE USER SPEND ### async def _update_user_cache(): ## UPDATE CACHE FOR USER ID + GLOBAL PROXY + if response_cost is None: + return user_ids = [user_id] try: for _id in user_ids: # Fetch the existing cost for the given user if _id is None: continue - existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id) - if existing_spend_obj is None: + cached_user = await user_api_key_cache.async_get_cache(key=_id) + if cached_user is None: # do nothing if there is no cache value return + existing_spend_obj = CacheCodec.deserialize(cached_user, LiteLLM_UserTable) + if existing_spend_obj is None: + return verbose_proxy_logger.debug( f"_update_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj.json())) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_UserTable + ), + ) + ) ## UPDATE GLOBAL PROXY ## global_proxy_spend = await user_api_key_cache.async_get_cache( key="{}:spend".format(litellm_proxy_admin_name) @@ -2088,31 +2095,33 @@ async def update_cache( # noqa: PLR0915 _id = "end_user_id:{}".format(end_user_id) try: # Fetch the existing cost for the given user - existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id) - if existing_spend_obj is None: + cached_end_user = await user_api_key_cache.async_get_cache(key=_id) + if cached_end_user is None: # if user does not exist in LiteLLM_UserTable, create a new user # do nothing if end-user not in api key cache return + existing_spend_obj = CacheCodec.deserialize( + cached_end_user, LiteLLM_EndUserTable + ) + if existing_spend_obj is None: + return verbose_proxy_logger.debug( f"_update_end_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if existing_spend_obj is None: - existing_spend = 0 - else: - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend + + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj.json())) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_EndUserTable + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update end user spend in cache. " @@ -2131,36 +2140,32 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: - # Fetch the existing cost for the given user - existing_spend_obj: Optional[LiteLLM_TeamTable] = ( - await user_api_key_cache.async_get_cache(key=_id) + cached_team = await user_api_key_cache.async_get_cache(key=_id) + if cached_team is None: + # do nothing if team not in api key cache + return + existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = ( + CacheCodec.deserialize(cached_team, LiteLLM_TeamTableCachedObj) ) if existing_spend_obj is None: - # do nothing if team not in api key cache return verbose_proxy_logger.debug( f"_update_team_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if existing_spend_obj is None: - existing_spend: Optional[float] = 0.0 - else: - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend - if existing_spend is None: - existing_spend = 0.0 + existing_spend: float = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update team spend in cache. " @@ -2187,32 +2192,32 @@ async def update_cache( # noqa: PLR0915 cache_key = f"tag:{tag_name}" # Fetch the existing tag object from cache - existing_tag_obj = await user_api_key_cache.async_get_cache( - key=cache_key - ) - if existing_tag_obj is None: + cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) + if cached_tag is None: # do nothing if tag not in api key cache continue + existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable) + if existing_tag_obj is None: + continue + verbose_proxy_logger.debug( f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}" ) - if isinstance(existing_tag_obj, dict): - existing_spend = existing_tag_obj.get("spend", 0) or 0 - else: - existing_spend = getattr(existing_tag_obj, "spend", 0) or 0 - + existing_spend = existing_tag_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the spend column for the given tag - if isinstance(existing_tag_obj, dict): - existing_tag_obj["spend"] = new_spend - values_to_update_in_cache.append((cache_key, existing_tag_obj)) - else: - existing_tag_obj.spend = new_spend - values_to_update_in_cache.append((cache_key, existing_tag_obj)) + existing_tag_obj.spend = new_spend + values_to_update_in_cache.append( + ( + cache_key, + CacheCodec.serialize( + existing_tag_obj, model_type=LiteLLM_TagTable + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update tag spend in cache. " @@ -2849,7 +2854,7 @@ class ProxyConfig: self, cache_params: dict, ): - global redis_usage_cache, llm_router + global redis_usage_cache, llm_router, general_settings from litellm import Cache if "default_in_memory_ttl" in cache_params: @@ -2871,16 +2876,23 @@ class ProxyConfig: ) # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. - # Share the same Redis client for virtual-key lookups (same DualCache as - # model_max_budget_limiter). attach_redis_cache is a no-op if Redis is - # already set (e.g. config reload). - user_api_key_cache.attach_redis_cache( - redis_usage_cache, - default_redis_ttl=litellm.default_redis_ttl, - ) - verbose_proxy_logger.debug( - "Attached redis_usage_cache Redis client to user_api_key_cache" - ) + if general_settings.get("enable_redis_auth_cache") is True: + user_api_key_cache.attach_redis_cache( + redis_usage_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) + verbose_proxy_logger.info( + "enable_redis_auth_cache=True: attached Redis to " + "user_api_key_cache — virtual-key lookups are now " + "shared across all proxy workers." + ) + else: + verbose_proxy_logger.info( + "enable_redis_auth_cache is not set: user_api_key_cache " + "remains in-memory only (per-worker). Set " + "general_settings.enable_redis_auth_cache: true to share " + "the auth cache across workers and reduce DB load." + ) def switch_on_llm_response_caching(self): """ diff --git a/tests/test_litellm/proxy/common_utils/test_cache_codec.py b/tests/test_litellm/proxy/common_utils/test_cache_codec.py index e99f22a75d..ef52eab4b5 100644 --- a/tests/test_litellm/proxy/common_utils/test_cache_codec.py +++ b/tests/test_litellm/proxy/common_utils/test_cache_codec.py @@ -1,5 +1,6 @@ import logging from typing import Optional +from unittest.mock import patch import pytest from pydantic import BaseModel, ValidationError @@ -50,6 +51,41 @@ class TestCacheCodecSerialize: with pytest.raises(ValidationError): CacheCodec.serialize({"count": 1}, model_type=_SampleModel) + def test_with_model_type_already_correct_instance_skips_revalidation(self): + """Fast-path: value is already model_type — model_validate must NOT be called.""" + m = _SampleModel(name="fast", count=7) + with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate: + out = CacheCodec.serialize(m, model_type=_SampleModel) + assert out == {"name": "fast", "count": 7} + mock_validate.assert_not_called() + + def test_with_model_type_subclass_instance_skips_revalidation(self): + """Subclass is isinstance of base → should also take the fast path.""" + sub = _SampleSubModel(name="sub", count=2) + with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate: + out = CacheCodec.serialize(sub, model_type=_SampleModel) + assert out == {"name": "sub", "count": 2} + mock_validate.assert_not_called() + + def test_with_model_type_dict_input_goes_through_model_validate(self): + """A dict value (not yet an instance) must still go through model_validate.""" + raw = {"name": "via-dict", "count": 5} + with patch.object( + _SampleModel, "model_validate", wraps=_SampleModel.model_validate + ) as mock_validate: + out = CacheCodec.serialize(raw, model_type=_SampleModel) + assert out == {"name": "via-dict", "count": 5} + mock_validate.assert_called_once() + + def test_with_model_type_incompatible_model_raises_validation_error(self): + """Passing an instance of a completely different model is a caller error and raises.""" + + class _IncompatibleModel(BaseModel): + name: str + + with pytest.raises(Exception): + CacheCodec.serialize(_IncompatibleModel(name="x"), model_type=_SampleModel) + class TestCacheCodecDeserialize: def test_none_returns_none(self): diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py new file mode 100644 index 0000000000..f824fd2f9c --- /dev/null +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -0,0 +1,130 @@ +""" +Tests for the enable_redis_auth_cache general_settings flag. + +Verifies that _init_cache attaches Redis to user_api_key_cache only when +the flag is explicitly set to True, and leaves it in-memory-only otherwise. +""" + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +import litellm.proxy.proxy_server as ps +from litellm.caching.caching import RedisCache +from litellm.caching.dual_cache import DualCache + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeRedisCache(RedisCache): + """ + Minimal RedisCache subclass that passes isinstance checks without + requiring a real Redis connection. __init__ is bypassed so no + network calls are made. + """ + + def __init__(self): # noqa: super().__init__ skipped intentionally + pass + + +@contextmanager +def _patched_init_cache(general_settings: dict, cache_params: dict): + """ + Context manager that: + 1. Replaces the module-level globals with fresh DualCache instances. + 2. Patches ``litellm.Cache`` (locally imported inside _init_cache) so + it returns a fake cache whose ``.cache`` attribute is a + _FakeRedisCache (passes the isinstance guard in _init_cache). + 3. Yields (user_api_key_cache, spend_counter_cache) after calling + _init_cache, then restores everything. + """ + fake_redis = _FakeRedisCache() + + mock_litellm_cache = MagicMock() + mock_litellm_cache.cache = fake_redis + + fresh_user_cache = DualCache() + fresh_spend_cache = DualCache() + + with ( + patch.object(ps, "general_settings", general_settings), + patch.object(ps, "user_api_key_cache", fresh_user_cache), + patch.object(ps, "spend_counter_cache", fresh_spend_cache), + patch.object(ps, "llm_router", None), + # Cache is locally imported inside _init_cache: patch it at source. + patch("litellm.Cache", return_value=mock_litellm_cache), + ): + litellm.cache = None + ps.ProxyConfig()._init_cache(cache_params) + yield fresh_user_cache, fresh_spend_cache + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestRedisAuthCacheFlag: + def test_flag_true_attaches_redis_to_user_api_key_cache(self): + """When enable_redis_auth_cache=True, user_api_key_cache.redis_cache must be set.""" + with _patched_init_cache( + general_settings={"enable_redis_auth_cache": True}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is not None, ( + "Redis should be attached to user_api_key_cache when " + "enable_redis_auth_cache=True" + ) + + def test_flag_false_leaves_user_api_key_cache_in_memory_only(self): + """When enable_redis_auth_cache=False, user_api_key_cache must stay in-memory.""" + with _patched_init_cache( + general_settings={"enable_redis_auth_cache": False}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is None, ( + "user_api_key_cache must remain in-memory-only when " + "enable_redis_auth_cache=False" + ) + + def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self): + """When enable_redis_auth_cache is not set at all, default is in-memory-only.""" + with _patched_init_cache( + general_settings={}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is None, ( + "user_api_key_cache must remain in-memory-only when " + "enable_redis_auth_cache is absent from general_settings" + ) + + def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self): + """spend_counter_cache must receive Redis regardless of the auth-cache flag.""" + for flag_value in (True, False, None): + gs = ( + {"enable_redis_auth_cache": flag_value} + if flag_value is not None + else {} + ) + with _patched_init_cache( + general_settings=gs, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (_, spend_cache): + assert spend_cache.redis_cache is not None, ( + f"spend_counter_cache must always get Redis " + f"(enable_redis_auth_cache={flag_value!r})" + ) + + def test_flag_false_spend_gets_redis_but_user_cache_does_not(self): + """Explicit False: spend cache wired, auth cache left in-memory.""" + with _patched_init_cache( + general_settings={"enable_redis_auth_cache": False}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, spend_cache): + assert spend_cache.redis_cache is not None + assert user_cache.redis_cache is None