mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(simple_shuffle): honor weights when first deployment has none
This commit is contained in:
parent
69a491e168
commit
f77948b762
2 changed files with 37 additions and 3 deletions
|
|
@ -41,9 +41,9 @@ 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]
|
||||
raw_weights = [m["litellm_params"].get(weight_by, None) for m in healthy_deployments]
|
||||
if any(weight is not None for weight in raw_weights):
|
||||
weights = [weight if weight is not None else 0 for weight in raw_weights]
|
||||
verbose_router_logger.debug(f"\nweight {weights}")
|
||||
total_weight = sum(weights)
|
||||
if total_weight <= 0:
|
||||
|
|
|
|||
34
tests/test_litellm/router_strategy/test_simple_shuffle.py
Normal file
34
tests/test_litellm/router_strategy/test_simple_shuffle.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
from collections import Counter
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.router_strategy.simple_shuffle import simple_shuffle
|
||||
|
||||
|
||||
def _deployment(deployment_id: str, **litellm_params) -> dict:
|
||||
return {
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "gpt-4o", **litellm_params},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
|
||||
|
||||
def test_weights_respected_when_first_deployment_has_no_weight():
|
||||
"""Regression for #33329: `simple_shuffle` decided whether weighted routing
|
||||
was enabled by inspecting only the first healthy deployment. When the first
|
||||
deployment had no `weight`, every deployment's weight was ignored and the
|
||||
pick fell back to uniform random. The weighted pick must be driven by any
|
||||
deployment declaring a weight, regardless of ordering.
|
||||
"""
|
||||
healthy_deployments = [
|
||||
_deployment("A"),
|
||||
_deployment("B", weight=1),
|
||||
_deployment("C", weight=99),
|
||||
]
|
||||
|
||||
counts: Counter = Counter()
|
||||
for _ in range(2000):
|
||||
deployment = simple_shuffle(MagicMock(), healthy_deployments, "test-model")
|
||||
counts[deployment["model_info"]["id"]] += 1
|
||||
|
||||
assert counts["A"] == 0
|
||||
assert counts["C"] > counts["B"] * 2
|
||||
Loading…
Add table
Reference in a new issue