From b313fc0b63aedfe6521584d6070bc3df3ece51b7 Mon Sep 17 00:00:00 2001 From: Aloys Jehwin Date: Wed, 15 Jul 2026 16:47:24 +0530 Subject: [PATCH] fix: simple_shuffle ignores weights when first healthy deployment has none (fixes #33329) The weighted-routing guard checked only healthy_deployments[0] for each metric (weight/rpm/tpm). If the first deployment lacked the field, the entire metric was skipped even when later deployments had it, causing silent fallback to uniform random.choice. Fix: replace the single-element probe with any(...) over all healthy deployments so the metric is activated whenever at least one deployment carries it. The existing zero-total-weight guard already handles the case where the first-element value is 0 (numeric falsy) correctly. Add four unit tests covering: - weight present only on second deployment (was broken) - same after reversing the list (was broken in opposite order) - no weights -> falls through to uniform pick - all weights present -> weighted pick (regression) Signed-off-by: Aloys Jehwin --- litellm/router_strategy/simple_shuffle.py | 47 ++++++----- .../router_unit_tests/test_simple_shuffle.py | 84 +++++++++++++++++++ 2 files changed, 109 insertions(+), 22 deletions(-) create mode 100644 tests/router_unit_tests/test_simple_shuffle.py diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index d3349ee29ce..13e21c7a4e6 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -41,28 +41,31 @@ 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) - if weight is not None: - weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] - verbose_router_logger.debug(f"\nweight {weights}") - total_weight = sum(weights) - if total_weight <= 0: - # All remaining candidates have weight 0 for this metric (e.g. - # after a weighted-failover exclusion left only zero-weight - # backups). Skip to the next metric (rpm/tpm) which may still - # provide a meaningful weighted pick; if none do, we fall - # through to the uniform random pick at the end. - continue - weights = [weight / total_weight for weight in weights] - verbose_router_logger.debug(f"\n weights {weights} by {weight_by}") - # Perform weighted random pick - selected_index = random.choices(range(len(weights)), weights=weights)[0] - verbose_router_logger.debug(f"\n selected index, {selected_index}") - deployment = healthy_deployments[selected_index] - verbose_router_logger.info( - f"get_available_deployment for model: {model}, Selected deployment: {llm_router_instance.print_deployment(deployment) or deployment[0]} for model: {model}" - ) - return deployment or deployment[0] + if not any( + m.get("litellm_params", {}).get(weight_by) is not None + for m in healthy_deployments + ): + continue + weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] + verbose_router_logger.debug(f"\nweight {weights}") + total_weight = sum(weights) + if total_weight <= 0: + # All remaining candidates have weight 0 for this metric (e.g. + # after a weighted-failover exclusion left only zero-weight + # backups). Skip to the next metric (rpm/tpm) which may still + # provide a meaningful weighted pick; if none do, we fall + # through to the uniform random pick at the end. + continue + weights = [weight / total_weight for weight in weights] + verbose_router_logger.debug(f"\n weights {weights} by {weight_by}") + # Perform weighted random pick + selected_index = random.choices(range(len(weights)), weights=weights)[0] + verbose_router_logger.debug(f"\n selected index, {selected_index}") + deployment = healthy_deployments[selected_index] + verbose_router_logger.info( + f"get_available_deployment for model: {model}, Selected deployment: {llm_router_instance.print_deployment(deployment) or deployment[0]} for model: {model}" + ) + return deployment or deployment[0] ############## No RPM/TPM passed, we do a random pick ################# item = random.choice(healthy_deployments) diff --git a/tests/router_unit_tests/test_simple_shuffle.py b/tests/router_unit_tests/test_simple_shuffle.py new file mode 100644 index 00000000000..1a5981441a1 --- /dev/null +++ b/tests/router_unit_tests/test_simple_shuffle.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: MIT +# Tests for fix: simple_shuffle ignored weights when the first healthy +# deployment had none, falling through to uniform random.choice (fixes #33329). + +from types import SimpleNamespace +from unittest.mock import patch + +from litellm.router_strategy.simple_shuffle import simple_shuffle + +_router = SimpleNamespace(print_deployment=lambda d: d) + + +def _deployments(first_has_weight: bool): + first = {"model_name": "chat", "litellm_params": {"model": "openai/first"}} + second = {"model_name": "chat", "litellm_params": {"model": "openai/second", "weight": 1}} + if first_has_weight: + first["litellm_params"]["weight"] = 0 + return [first, second] + + +def test_weight_on_second_deployment_triggers_weighted_pick(): + """Weight present only on second deployment must use random.choices, not random.choice.""" + deployments = _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, deployments, "chat") + assert weighted_pick.call_count == 1, "weighted pick not used" + assert uniform_pick.call_count == 0, "uniform pick used unexpectedly" + + +def test_reversed_deployment_order_same_result(): + """Reversing the list must not change whether weighted routing is used.""" + deployments = list(reversed(_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=[0], + ) as weighted_pick, + ): + simple_shuffle(_router, deployments, "chat") + assert weighted_pick.call_count == 1, "weighted pick not used after reversing" + assert uniform_pick.call_count == 0, "uniform pick used after reversing" + + +def test_no_weights_falls_through_to_uniform(): + """When no deployment has any weight, random.choice must be used.""" + deployments = [ + {"model_name": "chat", "litellm_params": {"model": "openai/a"}}, + {"model_name": "chat", "litellm_params": {"model": "openai/b"}}, + ] + 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, deployments, "chat") + assert uniform_pick.call_count == 1 + assert weighted_pick.call_count == 0 + + +def test_all_weights_present_uses_weighted_pick(): + """Sanity check: original behaviour with all weights present is preserved.""" + deployments = [ + {"model_name": "chat", "litellm_params": {"model": "openai/a", "weight": 2}}, + {"model_name": "chat", "litellm_params": {"model": "openai/b", "weight": 1}}, + ] + 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, deployments, "chat") + assert weighted_pick.call_count == 1 + assert uniform_pick.call_count == 0