phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

[Perf improvement Proxy] Use Dual Cache for getting key and team objects (#5903)

Browse source 
* use dual cache - perf

* fix auth checks

* fix budget checks for keys

* fix get / set team tests
This commit is contained in:
Ishaan Jaff 2024-09-25 19:56:17 -07:00 • committed by GitHub
parent 39c9150e97
commit f6cdb4ca0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 12 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

28
litellm/proxy/auth/auth_checks.py
View file

@ -434,15 +434,6 @@ async def _cache_management_object(
): ):
await user_api_key_cache.async_set_cache(key=key, value=value) 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( async def _cache_team_object(
team_id: str, team_id: str,
@ -471,7 +462,7 @@ async def _cache_key_object(
): ):
key = hashed_token key = hashed_token
## CACHE REFRESH TIME! ## CACHE REFRESH TIME
user_api_key_obj.last_refreshed_at = time.time() user_api_key_obj.last_refreshed_at = time.time()
await _cache_management_object( await _cache_management_object(
@ -524,10 +515,12 @@ async def get_team_object(
## CHECK REDIS CACHE ## ## CHECK REDIS CACHE ##
if ( if (
proxy_logging_obj is not None 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( cached_team_obj = (
key=key await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache(
key=key
)
) )
if cached_team_obj is None: if cached_team_obj is None:
@ -592,15 +585,6 @@ async def get_key_object(
key = hashed_token key = hashed_token
cached_team_obj: Optional[UserAPIKeyAuth] = None 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: if cached_team_obj is None:
cached_team_obj = await user_api_key_cache.async_get_cache(key=key) cached_team_obj = await user_api_key_cache.async_get_cache(key=key)

12
litellm/tests/test_proxy_server.py
View file

@ -753,17 +753,17 @@ async def test_team_update_redis():
litellm.proxy.proxy_server, "proxy_logging_obj" litellm.proxy.proxy_server, "proxy_logging_obj"
) )
proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache = RedisCache() redis_cache = RedisCache()
with patch.object( with patch.object(
proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, redis_cache,
"async_set_cache", "async_set_cache",
new=AsyncMock(), new=AsyncMock(),
) as mock_client: ) as mock_client:
await _cache_team_object( await _cache_team_object(
team_id="1234", team_id="1234",
team_table=LiteLLM_TeamTableCachedObj(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, 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" litellm.proxy.proxy_server, "proxy_logging_obj"
) )
proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache = RedisCache() redis_cache = RedisCache()
with patch.object( with patch.object(
proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache, redis_cache,
"async_get_cache", "async_get_cache",
new=AsyncMock(), new=AsyncMock(),
) as mock_client: ) as mock_client:
try: try:
await get_team_object( await get_team_object(
team_id="1234", team_id="1234",
user_api_key_cache=DualCache(), user_api_key_cache=DualCache(redis_cache=redis_cache),
parent_otel_span=None, parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj, proxy_logging_obj=proxy_logging_obj,
prisma_client=AsyncMock(), prisma_client=AsyncMock(),