fix(main.py): multiple deployments fix - run in parallel

This commit is contained in:
Krrish Dholakia 2023-10-21 14:28:34 -07:00
parent 34136a2c0d
commit cd0e699bcf
2 changed files with 8 additions and 9 deletions

View file

@ -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

View file

@ -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}")
traceback.print_exc()
pytest.fail(f"An exception occurred: {e}")
test_multiple_deployments()