ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

usage-based-routing-ttl-on-cache

Browse Source
This commit is contained in:
sumanth 2024-05-03 10:50:45 +05:30
parent 91971fa9e0
commit 3bc6b5d119
3 changed files with 57 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
litellm/router_strategy/lowest_tpm_rpm.py
View File

@ -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:

8
litellm/router_strategy/lowest_tpm_rpm_v2.py
View File

@ -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 ###

49
litellm/tests/test_router_caching.py
View File

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