fix(router): let simple-shuffle weight by any deployment's weight/rpm/tpm (#40222)

simple_shuffle only looked at healthy_deployments[0] to decide whether a
metric was configured, so a weight, rpm, or tpm on a later deployment was
ignored and the pick fell back to uniform random. Decide the metric from
all healthy deployments and keep the total_weight <= 0 fall-through.

Resolves LIT-7112

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/router_strategy/simple_shuffle.py

@@ -41,8 +41,7 @@ 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:
+        if any(m["litellm_params"].get(weight_by) is not None for m in healthy_deployments):
             weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments]
             verbose_router_logger.debug("\nweight %s", weights)
             total_weight = sum(weights)

tests/test_litellm/router_strategy/test_simple_shuffle.py

@@ -0,0 +1,54 @@
+from collections import Counter
+
+import pytest
+
+from litellm import Router
+from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict
+
+DRAWS = 200
+
+
+def _deployment(dep_id: str, metric: LiteLLMParamsTypedDict | None = None) -> DeploymentTypedDict:
+    params: LiteLLMParamsTypedDict = {"model": "gpt-4o", "api_key": "key", "mock_response": f"from {dep_id}"}
+    return {
+        "model_name": "test-model",
+        "litellm_params": {**params, **(metric or {})},
+        "model_info": {"id": dep_id},
+    }
+
+
+async def _draw_model_ids(router: Router) -> Counter[str]:
+    counts: Counter[str] = Counter()
+    for _ in range(DRAWS):
+        response = await router.acompletion(model="test-model", messages=[{"role": "user", "content": "hi"}])
+        counts[response._hidden_params["model_id"]] += 1
+    return counts
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("metric", [{"weight": 5}, {"rpm": 5}, {"tpm": 5}], ids=["weight", "rpm", "tpm"])
+async def test_weighted_pick_when_only_a_later_deployment_carries_the_metric(metric: LiteLLMParamsTypedDict):
+    router = Router(
+        model_list=[_deployment("unweighted"), _deployment("weighted", metric)],
+        routing_strategy="simple-shuffle",
+        num_retries=0,
+    )
+
+    counts = await _draw_model_ids(router)
+
+    assert counts["weighted"] == DRAWS
+    assert counts["unweighted"] == 0
+
+
+@pytest.mark.asyncio
+async def test_uniform_pick_when_every_configured_weight_is_zero():
+    router = Router(
+        model_list=[_deployment("unweighted"), _deployment("standby", {"weight": 0})],
+        routing_strategy="simple-shuffle",
+        num_retries=0,
+    )
+
+    counts = await _draw_model_ids(router)
+
+    assert counts["unweighted"] > 0
+    assert counts["standby"] > 0