From a34de5628976c9c826a0fa4d668486b5c0fb8774 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 30 Dec 2023 07:28:13 +0530 Subject: [PATCH] fix(router.py): handle initial scenario for tpm/rpm routing --- litellm/router.py | 7 ++++++- litellm/router_strategy/lowest_tpm_rpm.py | 3 +++ litellm/tests/test_least_busy_routing.py | 3 --- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index ebe0e877f..ecdc8f211 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1606,7 +1606,12 @@ class Router: self.routing_strategy == "usage-based-routing" and self.lowesttpm_logger is not None ): - return self.lowesttpm_logger.get_available_deployments(model_group=model) + min_deployment = self.lowesttpm_logger.get_available_deployments( + model_group=model + ) + if min_deployment is None: + min_deployment = random.choice(healthy_deployments) + return min_deployment raise ValueError("No models available.") diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 9217383bb..c531f3320 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -136,6 +136,9 @@ class LowestTPMLoggingHandler(CustomLogger): lowest_tpm = float("inf") deployment = None + if tpm_dict is None: # base case + return + for item, item_tpm in tpm_dict.items(): ## get the item from model list _deployment = None diff --git a/litellm/tests/test_least_busy_routing.py b/litellm/tests/test_least_busy_routing.py index 0bc125fe5..74e27f34a 100644 --- a/litellm/tests/test_least_busy_routing.py +++ b/litellm/tests/test_least_busy_routing.py @@ -128,6 +128,3 @@ def test_router_get_available_deployments(): assert return_dict[1] == 10 assert return_dict[2] == 54 assert return_dict[3] == 100 - - -# test_router_get_available_deployments()