From 3bc6b5d11927d166e6590e2be293fe12e91be68b Mon Sep 17 00:00:00 2001 From: sumanth Date: Fri, 3 May 2024 10:50:45 +0530 Subject: [PATCH] usage-based-routing-ttl-on-cache --- litellm/router_strategy/lowest_tpm_rpm.py | 8 ++-- litellm/router_strategy/lowest_tpm_rpm_v2.py | 8 ++-- litellm/tests/test_router_caching.py | 49 ++++++++++++++++++++ 3 files changed, 57 insertions(+), 8 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 0a7773a84b..a58eb9e02e 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -59,13 +59,13 @@ class LowestTPMLoggingHandler(CustomLogger): request_count_dict = self.router_cache.get_cache(key=tpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens - self.router_cache.set_cache(key=tpm_key, value=request_count_dict) + self.router_cache.set_cache(key=tpm_key, value=request_count_dict, ttl= 60) ## RPM request_count_dict = self.router_cache.get_cache(key=rpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + 1 - self.router_cache.set_cache(key=rpm_key, value=request_count_dict) + self.router_cache.set_cache(key=rpm_key, value=request_count_dict, ttl= 60) ### TESTING ### if self.test_flag: @@ -110,13 +110,13 @@ class LowestTPMLoggingHandler(CustomLogger): request_count_dict = self.router_cache.get_cache(key=tpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens - self.router_cache.set_cache(key=tpm_key, value=request_count_dict) + self.router_cache.set_cache(key=tpm_key, value=request_count_dict, ttl= 60) ## RPM request_count_dict = self.router_cache.get_cache(key=rpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + 1 - self.router_cache.set_cache(key=rpm_key, value=request_count_dict) + self.router_cache.set_cache(key=rpm_key, value=request_count_dict, ttl= 60) ### TESTING ### if self.test_flag: diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index f7a55d9709..59004c22b9 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -91,7 +91,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger): ) else: # if local result below limit, check redis ## prevent unnecessary redis checks - result = self.router_cache.increment_cache(key=rpm_key, value=1) + result = self.router_cache.increment_cache(key=rpm_key, value=1, ttl = 60) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message="Deployment over defined rpm limit={}. current usage={}".format( @@ -170,7 +170,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger): else: # if local result below limit, check redis ## prevent unnecessary redis checks result = await self.router_cache.async_increment_cache( - key=rpm_key, value=1 + key=rpm_key, value=1, ttl = 60 ) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( @@ -231,7 +231,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger): # update cache ## TPM - self.router_cache.increment_cache(key=tpm_key, value=total_tokens) + self.router_cache.increment_cache(key=tpm_key, value=total_tokens, ttl = 60) ### TESTING ### if self.test_flag: self.logged_success += 1 @@ -275,7 +275,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger): ## TPM await self.router_cache.async_increment_cache( - key=tpm_key, value=total_tokens + key=tpm_key, value=total_tokens, ttl = 60 ) ### TESTING ### diff --git a/litellm/tests/test_router_caching.py b/litellm/tests/test_router_caching.py index ebace161c9..ed1ddac6ee 100644 --- a/litellm/tests/test_router_caching.py +++ b/litellm/tests/test_router_caching.py @@ -134,6 +134,55 @@ async def test_acompletion_caching_on_router(): traceback.print_exc() pytest.fail(f"Error occurred: {e}") +@pytest.mark.asyncio +async def test_completion_caching_on_router(): + # tests completion + caching on router + try: + litellm.set_verbose = True + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + "tpm": 1000, + "rpm": 1, + }, + ] + + messages = [ + {"role": "user", "content": f"write a one sentence poem {time.time()}?"} + ] + router = Router( + model_list=model_list, + redis_host=os.environ["REDIS_HOST"], + redis_password=os.environ["REDIS_PASSWORD"], + redis_port=os.environ["REDIS_PORT"], + cache_responses=True, + timeout=30, + routing_strategy="usage-based-routing", + ) + response1 = await router.completion( + model="gpt-3.5-turbo", messages=messages, temperature=1 + ) + print(f"response1: {response1}") + await asyncio.sleep(60) + response2 = await router.completion( + model="gpt-3.5-turbo", messages=messages, temperature=1 + ) + print(f"response2: {response2}") + assert len(response1.choices[0].message.content) > 0 + assert len(response2.choices[0].message.content) > 0 + + router.reset() + except litellm.Timeout as e: + end_time = time.time() + print(f"timeout error occurred: {end_time - start_time}") + pass + except Exception as e: + traceback.print_exc() + pytest.fail(f"Error occurred: {e}") @pytest.mark.asyncio async def test_acompletion_caching_with_ttl_on_router():