From f6cdb4ca0d23a1b8d5679d21db5f2aa1b17bba27 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 25 Sep 2024 19:56:17 -0700 Subject: [PATCH] [Perf improvement Proxy] Use Dual Cache for getting key and team objects (#5903) * use dual cache - perf * fix auth checks * fix budget checks for keys * fix get / set team tests --- litellm/proxy/auth/auth_checks.py | 28 ++++++---------------------- litellm/tests/test_proxy_server.py | 12 ++++++------ 2 files changed, 12 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 404b3ffc9..80113fff3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -434,15 +434,6 @@ async def _cache_management_object( ): await user_api_key_cache.async_set_cache(key=key, value=value) - ## UPDATE REDIS CACHE ## - if proxy_logging_obj is not None: - _value = value.model_dump_json( - exclude_unset=True, exclude={"parent_otel_span": True} - ) - await proxy_logging_obj.internal_usage_cache.async_set_cache( - key=key, value=_value, litellm_parent_otel_span=None - ) - async def _cache_team_object( team_id: str, @@ -471,7 +462,7 @@ async def _cache_key_object( ): key = hashed_token - ## CACHE REFRESH TIME! + ## CACHE REFRESH TIME user_api_key_obj.last_refreshed_at = time.time() await _cache_management_object( @@ -524,10 +515,12 @@ async def get_team_object( ## CHECK REDIS CACHE ## if ( proxy_logging_obj is not None - and proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache is not None + and proxy_logging_obj.internal_usage_cache.dual_cache ): - cached_team_obj = await proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache.async_get_cache( - key=key + cached_team_obj = ( + await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( + key=key + ) ) if cached_team_obj is None: @@ -592,15 +585,6 @@ async def get_key_object( key = hashed_token cached_team_obj: Optional[UserAPIKeyAuth] = None - ## CHECK REDIS CACHE ## - if ( - proxy_logging_obj is not None - and proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache is not None - ): - cached_team_obj = await proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache.async_get_cache( - key=key - ) - if cached_team_obj is None: cached_team_obj = await user_api_key_cache.async_get_cache(key=key) diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 7357785e3..9271549d0 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -753,17 +753,17 @@ async def test_team_update_redis(): litellm.proxy.proxy_server, "proxy_logging_obj" ) - proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache = RedisCache() + redis_cache = RedisCache() with patch.object( - proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, + redis_cache, "async_set_cache", new=AsyncMock(), ) as mock_client: await _cache_team_object( team_id="1234", team_table=LiteLLM_TeamTableCachedObj(team_id="1234"), - user_api_key_cache=DualCache(), + user_api_key_cache=DualCache(redis_cache=redis_cache), proxy_logging_obj=proxy_logging_obj, ) @@ -782,17 +782,17 @@ async def test_get_team_redis(client_no_auth): litellm.proxy.proxy_server, "proxy_logging_obj" ) - proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache = RedisCache() + redis_cache = RedisCache() with patch.object( - proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, + redis_cache, "async_get_cache", new=AsyncMock(), ) as mock_client: try: await get_team_object( team_id="1234", - user_api_key_cache=DualCache(), + user_api_key_cache=DualCache(redis_cache=redis_cache), parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, prisma_client=AsyncMock(),