diff --git a/litellm/tests/test_router_get_deployments.py b/litellm/tests/test_router_get_deployments.py index 1e642babae..d943364588 100644 --- a/litellm/tests/test_router_get_deployments.py +++ b/litellm/tests/test_router_get_deployments.py @@ -297,7 +297,7 @@ def test_model_group_aliases(): "litellm_params": { "model": "gpt-3.5-turbo-0613", "api_key": os.getenv("OPENAI_API_KEY"), - "rpm": 6, + "tpm": 1, }, }, { @@ -307,14 +307,14 @@ def test_model_group_aliases(): "api_key": os.getenv("AZURE_API_KEY"), "api_base": os.getenv("AZURE_API_BASE"), "api_version": os.getenv("AZURE_API_VERSION"), - "rpm": 1440, + "tpm": 99, }, }, { "model_name": "claude-1", "litellm_params": { "model": "bedrock/claude1.2", - "rpm": 1440, + "tpm": 1, }, } ] @@ -330,9 +330,24 @@ def test_model_group_aliases(): selected_model_name = selected_model.get("model_name") if selected_model_name != "gpt-3.5-turbo": pytest.fail(f"Selected model {selected_model_name} is not gpt-3.5-turbo") - + + # test that + # call get_available_deployment 1k times, it should pick azure/chatgpt-v-2 about 90% of the time + selection_counts = defaultdict(int) + for _ in range(1000): + selected_model = router.get_available_deployment("gpt-3.5-turbo") + selected_model_id = selected_model["litellm_params"]["model"] + selected_model_name = litellm.utils.remove_model_id(selected_model_id) + selection_counts[selected_model_name] +=1 + print(selection_counts) + + total_requests = sum(selection_counts.values()) + + # Assert that 'azure/chatgpt-v-2' has about 90% of the total requests + assert selection_counts['azure/chatgpt-v-2'] / total_requests > 0.89, f"Assertion failed: 'azure/chatgpt-v-2' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}" + router.reset() except Exception as e: traceback.print_exc() pytest.fail(f"Error occurred: {e}") -test_model_group_aliases() \ No newline at end of file +# test_model_group_aliases() \ No newline at end of file