diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index 65bcce0e532..24e96a3bf2a 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -41,7 +41,18 @@ def simple_shuffle( ############## Check if 'weight' or 'rpm' or 'tpm' param set for a weighted pick ################# for weight_by in ["weight", "rpm", "tpm"]: - weight = healthy_deployments[0].get("litellm_params").get(weight_by, None) + # Check if ANY deployment has this metric set, not just the first one. + # Deployment order can change after cooldown/filtering, so checking only + # index 0 caused weighted routing to be silently skipped whenever the + # first healthy deployment omitted the metric but a later one set it. + weight = next( + ( + m["litellm_params"].get(weight_by) + for m in healthy_deployments + if m["litellm_params"].get(weight_by) is not None + ), + None, + ) if weight is not None: weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] verbose_router_logger.debug("\nweight %s", weights) diff --git a/tests/test_litellm/router_strategy/test_simple_shuffle.py b/tests/test_litellm/router_strategy/test_simple_shuffle.py new file mode 100644 index 00000000000..1b88c7086f1 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_simple_shuffle.py @@ -0,0 +1,172 @@ +""" +Tests for simple_shuffle router strategy. + +Covers: https://github.com/BerriAI/litellm/issues/33329 +Bug: simple_shuffle only checked healthy_deployments[0] to decide whether +weighted routing is enabled. If the first deployment omitted the metric but +a later deployment set it, weighted routing was silently skipped. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from litellm.router_strategy.simple_shuffle import simple_shuffle + +# A minimal router stub — simple_shuffle only calls print_deployment on it. +_ROUTER_STUB = SimpleNamespace(print_deployment=lambda d: "deployment") + + +def _make_deployments(first_has_weight: bool): + """Return two deployments; weight is on first or second depending on flag.""" + if first_has_weight: + return [ + { + "model_name": "chat", + "litellm_params": {"model": "openai/first", "weight": 1}, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/second"}, + }, + ] + else: + return [ + { + "model_name": "chat", + "litellm_params": {"model": "openai/first"}, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/second", "weight": 1}, + }, + ] + + +class TestSimpleShuffleWeightOnFirstDeployment: + """Baseline: weight on first deployment — weighted routing must be used.""" + + def test_uses_weighted_pick_when_first_deployment_has_weight(self): + deployments = _make_deployments(first_has_weight=True) + + with ( + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + ) as uniform_pick, + patch( + "litellm.router_strategy.simple_shuffle.random.choices", + return_value=[0], + ) as weighted_pick, + ): + simple_shuffle(_ROUTER_STUB, deployments, "chat") + + assert weighted_pick.call_count == 1, "Should use weighted pick" + assert uniform_pick.call_count == 0, "Should NOT use uniform pick" + + +class TestSimpleShuffleWeightOnSecondDeployment: + """ + Regression for #33329: + weight is only on the second deployment — weighted routing must still be used. + Before the fix, simple_shuffle checked only index 0 and fell through to + random.choice (uniform) because index 0 had no weight. + """ + + def test_uses_weighted_pick_when_only_second_deployment_has_weight(self): + deployments = _make_deployments(first_has_weight=False) + + with ( + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + ) as uniform_pick, + patch( + "litellm.router_strategy.simple_shuffle.random.choices", + return_value=[1], + ) as weighted_pick, + ): + simple_shuffle(_ROUTER_STUB, deployments, "chat") + + assert weighted_pick.call_count == 1, ( + "Should use weighted pick — weight exists on second deployment" + ) + assert uniform_pick.call_count == 0, ( + "Should NOT use uniform pick when any deployment has weight" + ) + + def test_order_independence(self): + """ + Reversing the deployment list must not change routing strategy. + Both orderings should use weighted routing since weight is configured. + """ + for first_has_weight in [True, False]: + deployments = _make_deployments(first_has_weight=first_has_weight) + + with ( + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + ) as uniform_pick, + patch( + "litellm.router_strategy.simple_shuffle.random.choices", + return_value=[0], + ) as weighted_pick, + ): + simple_shuffle(_ROUTER_STUB, deployments, "chat") + + assert weighted_pick.call_count == 1, ( + f"Expected weighted pick when first_has_weight={first_has_weight}" + ) + assert uniform_pick.call_count == 0, ( + f"Expected no uniform pick when first_has_weight={first_has_weight}" + ) + + +class TestSimpleShuffleNoWeights: + """When no deployment has any weight/rpm/tpm, uniform random pick is used.""" + + def test_uses_uniform_pick_when_no_weights(self): + deployments = [ + {"model_name": "chat", "litellm_params": {"model": "openai/first"}}, + {"model_name": "chat", "litellm_params": {"model": "openai/second"}}, + ] + + with ( + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + return_value=deployments[0], + ) as uniform_pick, + patch( + "litellm.router_strategy.simple_shuffle.random.choices", + ) as weighted_pick, + ): + simple_shuffle(_ROUTER_STUB, deployments, "chat") + + assert uniform_pick.call_count == 1, "Should use uniform pick when no weights" + assert weighted_pick.call_count == 0, "Should NOT use weighted pick" + + +class TestSimpleShuffleRpmTpm: + """Same order-independence fix applies to rpm and tpm metrics.""" + + @pytest.mark.parametrize("metric", ["rpm", "tpm"]) + def test_uses_weighted_pick_when_only_second_deployment_has_metric(self, metric): + deployments = [ + {"model_name": "chat", "litellm_params": {"model": "openai/first"}}, + {"model_name": "chat", "litellm_params": {"model": "openai/second", metric: 100}}, + ] + + with ( + patch( + "litellm.router_strategy.simple_shuffle.random.choice", + ) as uniform_pick, + patch( + "litellm.router_strategy.simple_shuffle.random.choices", + return_value=[1], + ) as weighted_pick, + ): + simple_shuffle(_ROUTER_STUB, deployments, "chat") + + assert weighted_pick.call_count == 1, ( + f"Should use weighted pick — {metric} exists on second deployment" + ) + assert uniform_pick.call_count == 0