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

test: load test router

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-17 08:23:35 -08:00
parent c7aba49d83
commit 3d45d8a58c
1 changed files with 60 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

120
litellm/tests/test_loadtest_router.py
View file

@ -1,69 +1,69 @@
import sys, os # import sys, os
import traceback # import traceback
from dotenv import load_dotenv # from dotenv import load_dotenv
import copy # import copy
load_dotenv() # load_dotenv()
sys.path.insert( # sys.path.insert(
0, os.path.abspath("../..") # 0, os.path.abspath("../..")
) # Adds the parent directory to the system path # ) # Adds the parent directory to the system path
import asyncio # import asyncio
from litellm import Router, Timeout # from litellm import Router, Timeout
async def call_acompletion(semaphore, router: Router, input_data): # async def call_acompletion(semaphore, router: Router, input_data):
async with semaphore: # async with semaphore:
try: # try:
# Use asyncio.wait_for to set a timeout for the task # # Use asyncio.wait_for to set a timeout for the task
response = await router.acompletion(**input_data) # response = await router.acompletion(**input_data)
# Handle the response as needed # # Handle the response as needed
return response # return response
except Timeout: # except Timeout:
print(f"Task timed out: {input_data}") # print(f"Task timed out: {input_data}")
return None # You may choose to return something else or raise an exception # return None # You may choose to return something else or raise an exception
async def main(): # async def main():
# Initialize the Router # # Initialize the Router
model_list= [{ # model_list= [{
"model_name": "gpt-3.5-turbo", # "model_name": "gpt-3.5-turbo",
"litellm_params": { # "litellm_params": {
"model": "gpt-3.5-turbo", # "model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"), # "api_key": os.getenv("OPENAI_API_KEY"),
}, # },
}, { # }, {
"model_name": "gpt-3.5-turbo", # "model_name": "gpt-3.5-turbo",
"litellm_params": { # "litellm_params": {
"model": "azure/chatgpt-v-2", # "model": "azure/chatgpt-v-2",
"api_key": os.getenv("AZURE_API_KEY"), # "api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"), # "api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION") # "api_version": os.getenv("AZURE_API_VERSION")
}, # },
}, { # }, {
"model_name": "gpt-3.5-turbo", # "model_name": "gpt-3.5-turbo",
"litellm_params": { # "litellm_params": {
"model": "azure/chatgpt-functioncalling", # "model": "azure/chatgpt-functioncalling",
"api_key": os.getenv("AZURE_API_KEY"), # "api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"), # "api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION") # "api_version": os.getenv("AZURE_API_VERSION")
}, # },
}] # }]
router = Router(model_list=model_list, num_retries=3, timeout=10) # router = Router(model_list=model_list, num_retries=3, timeout=10)
# Create a semaphore with a capacity of 100 # # Create a semaphore with a capacity of 100
semaphore = asyncio.Semaphore(100) # semaphore = asyncio.Semaphore(100)
# List to hold all task references # # List to hold all task references
tasks = [] # tasks = []
# Launch 1000 tasks # # Launch 1000 tasks
for _ in range(1000): # for _ in range(1000):
task = asyncio.create_task(call_acompletion(semaphore, router, {"model": "gpt-3.5-turbo", "messages": [{"role":"user", "content": "Hey, how's it going?"}]})) # task = asyncio.create_task(call_acompletion(semaphore, router, {"model": "gpt-3.5-turbo", "messages": [{"role":"user", "content": "Hey, how's it going?"}]}))
tasks.append(task) # tasks.append(task)
# Wait for all tasks to complete # # Wait for all tasks to complete
responses = await asyncio.gather(*tasks) # responses = await asyncio.gather(*tasks)
# Process responses as needed # # Process responses as needed
print(f"NUMBER OF COMPLETED TASKS: {len(responses)}") # print(f"NUMBER OF COMPLETED TASKS: {len(responses)}")
# Run the main function # # Run the main function
asyncio.run(main()) # asyncio.run(main())