diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index a11c6d87237..3300e56ff70 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -333,7 +333,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger): tpm_dict[tpm_key] = 0 all_deployments = tpm_dict - deployment = None + potential_deployments = [] # if multiple deployments have the same low value for item, item_tpm in all_deployments.items(): ## get the item from model list _deployment = None @@ -369,11 +369,17 @@ class LowestTPMLoggingHandler_v2(CustomLogger): rpm_dict[item] + 1 > _deployment_rpm ): continue + elif item_tpm == lowest_tpm: + potential_deployments.append(_deployment) elif item_tpm < lowest_tpm: lowest_tpm = item_tpm - deployment = _deployment + potential_deployments = [_deployment] print_verbose("returning picked lowest tpm/rpm deployment.") - return deployment + + if len(potential_deployments) > 0: + return random.choice(potential_deployments) + else: + return None async def async_get_available_deployments( self, diff --git a/litellm/tests/test_tpm_rpm_routing_v2.py b/litellm/tests/test_tpm_rpm_routing_v2.py index 9a43ae3ca10..78620a6e552 100644 --- a/litellm/tests/test_tpm_rpm_routing_v2.py +++ b/litellm/tests/test_tpm_rpm_routing_v2.py @@ -282,6 +282,64 @@ def test_router_skip_rate_limited_deployments(): print(f"An exception occurred! {str(e)}") +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_multiple_potential_deployments(sync_mode): + """ + If multiple deployments have the same tpm value + + call 5 times, test if deployments are shuffled. + + -> prevents single deployment from being overloaded in high-concurrency scenario + """ + + model_list = [ + { + "model_name": "azure-model", + "litellm_params": { + "model": "azure/gpt-turbo", + "api_key": "os.environ/AZURE_FRANCE_API_KEY", + "api_base": "https://openai-france-1234.openai.azure.com", + "tpm": 1440, + }, + }, + { + "model_name": "azure-model", + "litellm_params": { + "model": "azure/gpt-turbo-2", + "api_key": "os.environ/AZURE_FRANCE_API_KEY", + "api_base": "https://openai-france-1234.openai.azure.com", + "tpm": 1440, + }, + }, + ] + router = Router( + model_list=model_list, + routing_strategy="usage-based-routing-v2", + set_verbose=False, + num_retries=3, + ) # type: ignore + + model_ids = set() + for _ in range(5): + if sync_mode: + deployment = router.get_available_deployment( + model="azure-model", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + ) + else: + deployment = await router.async_get_available_deployment( + model="azure-model", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + ) + + ## get id ## + id = deployment.get("model_info", {}).get("id") + model_ids.add(id) + + assert len(model_ids) == 2 + + def test_single_deployment_tpm_zero(): import litellm import os