Spaces:
Configuration error
Configuration error
File size: 2,689 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import sys
import os
from dotenv import load_dotenv
load_dotenv()
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
from litellm import Router, Timeout
import time
### Test calling router async
async def call_acompletion(semaphore, router: Router, input_data):
async with semaphore:
try:
# Use asyncio.wait_for to set a timeout for the task
response = await router.acompletion(**input_data)
# Handle the response as needed
print(response)
return response
except Timeout:
print(f"Task timed out: {input_data}")
return None # You may choose to return something else or raise an exception
async def main():
# Initialize the Router
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"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"),
},
},
]
router = Router(model_list=model_list, num_retries=3, timeout=10)
# Create a semaphore with a capacity of 100
semaphore = asyncio.Semaphore(100)
# List to hold all task references
tasks = []
start_time_all_tasks = time.time()
# Launch 1000 tasks
for _ in range(500):
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)
# Wait for all tasks to complete
responses = await asyncio.gather(*tasks)
# Process responses as needed
# Record the end time for all tasks
end_time_all_tasks = time.time()
# Calculate the total time for all tasks
total_time_all_tasks = end_time_all_tasks - start_time_all_tasks
print(f"Total time for all tasks: {total_time_all_tasks} seconds")
# Calculate the average time per response
average_time_per_response = total_time_all_tasks / len(responses)
print(f"Average time per response: {average_time_per_response} seconds")
print(f"NUMBER OF COMPLETED TASKS: {len(responses)}")
# Run the main function
asyncio.run(main())
|