test(test_lowest_latency_routing.py): add more tests

This commit is contained in:
Krrish Dholakia 2023-12-30 17:41:42 +05:30
parent f2d0d5584a
commit 027218c3f0
3 changed files with 146 additions and 139 deletions

View file

@ -314,6 +314,7 @@ class Router:
raise e raise e
async def _acompletion(self, model: str, messages: List[Dict[str, str]], **kwargs): async def _acompletion(self, model: str, messages: List[Dict[str, str]], **kwargs):
model_name = None
try: try:
self.print_verbose( self.print_verbose(
f"Inside _acompletion()- model: {model}; kwargs: {kwargs}" f"Inside _acompletion()- model: {model}; kwargs: {kwargs}"

View file

@ -3,7 +3,7 @@
import dotenv, os, requests, random import dotenv, os, requests, random
from typing import Optional from typing import Optional
from datetime import datetime from datetime import datetime, timedelta
dotenv.load_dotenv() # Loading env variables using dotenv dotenv.load_dotenv() # Loading env variables using dotenv
import traceback import traceback
@ -135,6 +135,9 @@ class LowestLatencyLoggingHandler(CustomLogger):
if _deployment is None: if _deployment is None:
continue # skip to next one continue # skip to next one
if isinstance(item_latency, timedelta):
item_latency = float(item_latency.total_seconds())
if item_latency == 0: if item_latency == 0:
deployment = _deployment deployment = _deployment
break break

View file

@ -69,7 +69,7 @@ def test_get_available_deployments():
"model_info": {"id": "5678"}, "model_info": {"id": "5678"},
}, },
] ]
lowest_tpm_logger = LowestLatencyLoggingHandler( lowest_latency_logger = LowestLatencyLoggingHandler(
router_cache=test_cache, model_list=model_list router_cache=test_cache, model_list=model_list
) )
model_group = "gpt-3.5-turbo" model_group = "gpt-3.5-turbo"
@ -86,8 +86,9 @@ def test_get_available_deployments():
} }
start_time = time.time() start_time = time.time()
response_obj = {"usage": {"total_tokens": 50}} response_obj = {"usage": {"total_tokens": 50}}
time.sleep(3)
end_time = time.time() end_time = time.time()
lowest_tpm_logger.log_success_event( lowest_latency_logger.log_success_event(
response_obj=response_obj, response_obj=response_obj,
kwargs=kwargs, kwargs=kwargs,
start_time=start_time, start_time=start_time,
@ -106,8 +107,9 @@ def test_get_available_deployments():
} }
start_time = time.time() start_time = time.time()
response_obj = {"usage": {"total_tokens": 20}} response_obj = {"usage": {"total_tokens": 20}}
time.sleep(2)
end_time = time.time() end_time = time.time()
lowest_tpm_logger.log_success_event( lowest_latency_logger.log_success_event(
response_obj=response_obj, response_obj=response_obj,
kwargs=kwargs, kwargs=kwargs,
start_time=start_time, start_time=start_time,
@ -116,12 +118,12 @@ def test_get_available_deployments():
## CHECK WHAT'S SELECTED ## ## CHECK WHAT'S SELECTED ##
print( print(
lowest_tpm_logger.get_available_deployments( lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list model_group=model_group, healthy_deployments=model_list
) )
) )
assert ( assert (
lowest_tpm_logger.get_available_deployments( lowest_latency_logger.get_available_deployments(
model_group=model_group, healthy_deployments=model_list model_group=model_group, healthy_deployments=model_list
)["model_info"]["id"] )["model_info"]["id"]
== "5678" == "5678"
@ -131,149 +133,150 @@ def test_get_available_deployments():
# test_get_available_deployments() # test_get_available_deployments()
# def test_router_get_available_deployments(): def test_router_get_available_deployments():
# """ """
# Test if routers 'get_available_deployments' returns the least busy deployment Test if routers 'get_available_deployments' returns the fastest deployment
# """ """
# model_list = [ model_list = [
# { {
# "model_name": "azure-model", "model_name": "azure-model",
# "litellm_params": { "litellm_params": {
# "model": "azure/gpt-turbo", "model": "azure/gpt-turbo",
# "api_key": "os.environ/AZURE_FRANCE_API_KEY", "api_key": "os.environ/AZURE_FRANCE_API_KEY",
# "api_base": "https://openai-france-1234.openai.azure.com", "api_base": "https://openai-france-1234.openai.azure.com",
# "rpm": 1440, "rpm": 1440,
# }, },
# "model_info": {"id": 1}, "model_info": {"id": 1},
# }, },
# { {
# "model_name": "azure-model", "model_name": "azure-model",
# "litellm_params": { "litellm_params": {
# "model": "azure/gpt-35-turbo", "model": "azure/gpt-35-turbo",
# "api_key": "os.environ/AZURE_EUROPE_API_KEY", "api_key": "os.environ/AZURE_EUROPE_API_KEY",
# "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com", "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
# "rpm": 6, "rpm": 6,
# }, },
# "model_info": {"id": 2}, "model_info": {"id": 2},
# }, },
# ] ]
# router = Router( router = Router(
# model_list=model_list, model_list=model_list,
# routing_strategy="usage-based-routing", routing_strategy="latency-based-routing",
# set_verbose=False, set_verbose=False,
# num_retries=3, num_retries=3,
# ) # type: ignore ) # type: ignore
# ## DEPLOYMENT 1 ## ## DEPLOYMENT 1 ##
# deployment_id = 1 deployment_id = 1
# kwargs = { kwargs = {
# "litellm_params": { "litellm_params": {
# "metadata": { "metadata": {
# "model_group": "azure-model", "model_group": "azure-model",
# }, },
# "model_info": {"id": 1}, "model_info": {"id": 1},
# } }
# } }
# start_time = time.time() start_time = time.time()
# response_obj = {"usage": {"total_tokens": 50}} response_obj = {"usage": {"total_tokens": 50}}
# end_time = time.time() time.sleep(3)
# router.lowesttpm_logger.log_success_event( end_time = time.time()
# response_obj=response_obj, router.lowestlatency_logger.log_success_event(
# kwargs=kwargs, response_obj=response_obj,
# start_time=start_time, kwargs=kwargs,
# end_time=end_time, start_time=start_time,
# ) end_time=end_time,
# ## DEPLOYMENT 2 ## )
# deployment_id = 2 ## DEPLOYMENT 2 ##
# kwargs = { deployment_id = 2
# "litellm_params": { kwargs = {
# "metadata": { "litellm_params": {
# "model_group": "azure-model", "metadata": {
# }, "model_group": "azure-model",
# "model_info": {"id": 2}, },
# } "model_info": {"id": 2},
# } }
# start_time = time.time() }
# response_obj = {"usage": {"total_tokens": 20}} start_time = time.time()
# end_time = time.time() response_obj = {"usage": {"total_tokens": 20}}
# router.lowesttpm_logger.log_success_event( time.sleep(2)
# response_obj=response_obj, end_time = time.time()
# kwargs=kwargs, router.lowestlatency_logger.log_success_event(
# start_time=start_time, response_obj=response_obj,
# end_time=end_time, kwargs=kwargs,
# ) start_time=start_time,
end_time=end_time,
)
# ## CHECK WHAT'S SELECTED ## ## CHECK WHAT'S SELECTED ##
# # print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model")) # print(router.lowesttpm_logger.get_available_deployments(model_group="azure-model"))
# print(router.get_available_deployment(model="azure-model")) print(router.get_available_deployment(model="azure-model"))
# assert router.get_available_deployment(model="azure-model")["model_info"]["id"] == 2 assert router.get_available_deployment(model="azure-model")["model_info"]["id"] == 2
# # test_get_available_deployments() # test_get_available_deployments()
# # test_router_get_available_deployments() # test_router_get_available_deployments()
# @pytest.mark.asyncio @pytest.mark.asyncio
# async def test_router_completion_streaming(): async def test_router_completion_streaming():
# messages = [ messages = [
# {"role": "user", "content": "Hello, can you generate a 500 words poem?"} {"role": "user", "content": "Hello, can you generate a 500 words poem?"}
# ] ]
# model = "azure-model" model = "azure-model"
# model_list = [ model_list = [
# { {
# "model_name": "azure-model", "model_name": "azure-model",
# "litellm_params": { "litellm_params": {
# "model": "azure/gpt-turbo", "model": "azure/gpt-turbo",
# "api_key": "os.environ/AZURE_FRANCE_API_KEY", "api_key": "os.environ/AZURE_FRANCE_API_KEY",
# "api_base": "https://openai-france-1234.openai.azure.com", "api_base": "https://openai-france-1234.openai.azure.com",
# "rpm": 1440, "rpm": 1440,
# }, },
# "model_info": {"id": 1}, "model_info": {"id": 1},
# }, },
# { {
# "model_name": "azure-model", "model_name": "azure-model",
# "litellm_params": { "litellm_params": {
# "model": "azure/gpt-35-turbo", "model": "azure/gpt-35-turbo",
# "api_key": "os.environ/AZURE_EUROPE_API_KEY", "api_key": "os.environ/AZURE_EUROPE_API_KEY",
# "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com", "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
# "rpm": 6, "rpm": 6,
# }, },
# "model_info": {"id": 2}, "model_info": {"id": 2},
# }, },
# ] ]
# router = Router( router = Router(
# model_list=model_list, model_list=model_list,
# routing_strategy="usage-based-routing", routing_strategy="latency-based-routing",
# set_verbose=False, set_verbose=False,
# num_retries=3, num_retries=3,
# ) # type: ignore ) # type: ignore
# ### Make 3 calls, test if 3rd call goes to lowest tpm deployment ### Make 3 calls, test if 3rd call goes to fastest deployment
# ## CALL 1+2 ## CALL 1+2
# tasks = [] tasks = []
# response = None response = None
# final_response = None final_response = None
# for _ in range(2): for _ in range(2):
# tasks.append(router.acompletion(model=model, messages=messages)) tasks.append(router.acompletion(model=model, messages=messages))
# response = await asyncio.gather(*tasks) response = await asyncio.gather(*tasks)
# if response is not None: if response is not None:
# ## CALL 3 ## CALL 3
# await asyncio.sleep(1) # let the token update happen await asyncio.sleep(1) # let the cache update happen
# current_minute = datetime.now().strftime("%H-%M") picked_deployment = router.lowestlatency_logger.get_available_deployments(
# picked_deployment = router.lowesttpm_logger.get_available_deployments( model_group=model, healthy_deployments=router.healthy_deployments
# model_group=model, healthy_deployments=router.healthy_deployments )
# ) final_response = await router.acompletion(model=model, messages=messages)
# final_response = await router.acompletion(model=model, messages=messages) print(f"min deployment id: {picked_deployment}")
# print(f"min deployment id: {picked_deployment}") print(f"model id: {final_response._hidden_params['model_id']}")
# print(f"model id: {final_response._hidden_params['model_id']}") assert (
# assert ( final_response._hidden_params["model_id"]
# final_response._hidden_params["model_id"] == picked_deployment["model_info"]["id"]
# == picked_deployment["model_info"]["id"] )
# )
# # asyncio.run(test_router_completion_streaming()) # asyncio.run(test_router_completion_streaming())