phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 03:34:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(test) router:get_available_deployment

Browse source
This commit is contained in:
ishaan-jaff 2023-11-29 17:54:41 -08:00
parent e365aea138
commit 54678d5f9c
2 changed files with 288 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

53
litellm/tests/test_router.py
View file

@ -1,5 +1,5 @@
#### What this tests ####
# This tests calling batch_completions by running 100 messages together
#This tests litellm router
import sys, os, time
import traceback, asyncio
@ -299,54 +299,3 @@ def test_aembedding_on_router():
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
# test_aembedding_on_router()
def test_weighted_selection_router():
# this tests if load balancing works based on the provided rpms in the router
# it's fast test, only tests get_available_deployment
# users can pass rpms as a litellm_param
try:
litellm.set_verbose = False
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo-0613",
"api_key": os.getenv("OPENAI_API_KEY"),
"rpm": 6,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-2",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"rpm": 1440,
},
}
]
router = Router(
model_list=model_list,
)
selection_counts = defaultdict(int)
# call get_available_deployment 1k times, it should pick azure/chatgpt-v-2 about 90% of the time
for _ in range(1000):
selected_model = router.get_available_deployment("gpt-3.5-turbo")
selected_model_id = selected_model["litellm_params"]["model"]
selected_model_name = litellm.utils.remove_model_id(selected_model_id)
selection_counts[selected_model_name] +=1
print(selection_counts)
total_requests = sum(selection_counts.values())
# Assert that 'azure/chatgpt-v-2' has about 90% of the total requests
assert selection_counts['azure/chatgpt-v-2'] / total_requests > 0.89, f"Assertion failed: 'azure/chatgpt-v-2' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
router.reset()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
# test_weighted_selection_router()