From 902e5e71e115322d8d8c38ca69dc2cd28b207272 Mon Sep 17 00:00:00 2001
From: sudu <teric@outlook.com>
Date: Wed, 13 Aug 2025 12:40:34 +0800
Subject: [PATCH] Update simple_shuffle.py, choose weights by 'weight', 'rpm',
 'tpm' in one loop

---
 litellm/router_strategy/simple_shuffle.py | 68 ++++++-----------------
 1 file changed, 17 insertions(+), 51 deletions(-)

diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py
index da24c02f2e..e35ecfdcd2 100644
--- a/litellm/router_strategy/simple_shuffle.py
+++ b/litellm/router_strategy/simple_shuffle.py
@@ -39,57 +39,23 @@ def simple_shuffle(
         Dict: A single healthy deployment
     """
 
-    ############## Check if 'weight' param set for a weighted pick #################
-    weight = healthy_deployments[0].get("litellm_params").get("weight", None)
-    if weight is not None:
-        # use weight-random pick if rpms provided
-        weights = [m["litellm_params"].get("weight", 0) for m in healthy_deployments]
-        verbose_router_logger.debug(f"\nweight {weights}")
-        total_weight = sum(weights)
-        weights = [weight / total_weight for weight in weights]
-        verbose_router_logger.debug(f"\n weights {weights}")
-        # 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]
-    ############## Check if we can do a RPM/TPM based weighted pick #################
-    rpm = healthy_deployments[0].get("litellm_params").get("rpm", None)
-    if rpm is not None:
-        # use weight-random pick if rpms provided
-        rpms = [m["litellm_params"].get("rpm", 0) for m in healthy_deployments]
-        verbose_router_logger.debug(f"\nrpms {rpms}")
-        total_rpm = sum(rpms)
-        weights = [rpm / total_rpm for rpm in rpms]
-        verbose_router_logger.debug(f"\n weights {weights}")
-        # Perform weighted random pick
-        selected_index = random.choices(range(len(rpms)), 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]
-    ############## Check if we can do a RPM/TPM based weighted pick #################
-    tpm = healthy_deployments[0].get("litellm_params").get("tpm", None)
-    if tpm is not None:
-        # use weight-random pick if rpms provided
-        tpms = [m["litellm_params"].get("tpm", 0) for m in healthy_deployments]
-        verbose_router_logger.debug(f"\ntpms {tpms}")
-        total_tpm = sum(tpms)
-        weights = [tpm / total_tpm for tpm in tpms]
-        verbose_router_logger.debug(f"\n weights {weights}")
-        # Perform weighted random pick
-        selected_index = random.choices(range(len(tpms)), 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]
+    ############## 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)
+            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)