From cd0e699bcfdde3be5795ba0312fbfbe6abdfc25f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 21 Oct 2023 14:28:34 -0700 Subject: [PATCH] fix(main.py): multiple deployments fix - run in parallel --- litellm/main.py | 11 +++-------- litellm/tests/test_multiple_deployments.py | 6 +++++- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index fa1d8e7757..f5618c0554 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1394,15 +1394,10 @@ def batch_completion_models(*args, **kwargs): kwargs = {**deployment, **nested_kwargs} futures[deployment["model"]] = executor.submit(completion, **kwargs) - print(f"futures: {futures}") - # done, not_done = concurrent.futures.wait(futures.values(), return_when=concurrent.futures.FIRST_COMPLETED) + done, not_done = concurrent.futures.wait(futures.values(), return_when=concurrent.futures.FIRST_COMPLETED) - # done is a set of futures that completed - for _, future in futures.items(): - if future.result() is not None: - return future.result() - # for future in done: - # return future.result() + for future in done: + return future.result() return None # If no response is received from any model diff --git a/litellm/tests/test_multiple_deployments.py b/litellm/tests/test_multiple_deployments.py index 792926ad2e..76b2685146 100644 --- a/litellm/tests/test_multiple_deployments.py +++ b/litellm/tests/test_multiple_deployments.py @@ -44,5 +44,9 @@ def test_multiple_deployments(): try: ## LiteLLM completion call ## returns first response response = completion(model="mistral-7b-instruct", messages=messages, model_list=model_list) + print(f"response: {response}") except Exception as e: - pytest.fail(f"An exception occurred: {e}") \ No newline at end of file + traceback.print_exc() + pytest.fail(f"An exception occurred: {e}") + +test_multiple_deployments() \ No newline at end of file