From 85607949638baf4d602e60211e6461f2687968f1 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 27 Nov 2023 16:37:57 -0800 Subject: [PATCH] (test) load test router --- litellm/proxy/tests/load_test_completion.py | 28 ++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/tests/load_test_completion.py b/litellm/proxy/tests/load_test_completion.py index 133b336646..dfb40850d1 100644 --- a/litellm/proxy/tests/load_test_completion.py +++ b/litellm/proxy/tests/load_test_completion.py @@ -1,6 +1,7 @@ import time, asyncio from openai import AsyncOpenAI import uuid +import traceback litellm_client = AsyncOpenAI( @@ -8,20 +9,41 @@ litellm_client = AsyncOpenAI( base_url="http://0.0.0.0:8000" ) + async def litellm_completion(): - return await litellm_client.chat.completions.create( + # Your existing code for litellm_completion goes here + try: + return await litellm_client.chat.completions.create( model="gpt-3.5-turbo", messages=[{"role": "user", "content": f"This is a test: {uuid.uuid4()}"}], ) + except Exception as e: + # If there's an exception, log the error message + with open("error_log.txt", "a") as error_log: + error_log.write(f"Error during completion: {str(e)}. Tracbeack {traceback.format_exc()}\n, load_test ") + pass + async def main(): start = time.time() - n = 1 # Number of concurrent tasks + n = 500 # Number of concurrent tasks tasks = [litellm_completion() for _ in range(n)] + chat_completions = await asyncio.gather(*tasks) + successful_completions = [c for c in chat_completions if c is not None] + + # Write errors to error_log.txt + with open("error_log.txt", "a") as error_log: + for completion in chat_completions: + if isinstance(completion, str): + error_log.write(completion + "\n") + print(n, time.time() - start, len(successful_completions)) if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + # Blank out contents of error_log.txt + open("error_log.txt", "w").close() + + asyncio.run(main())