From 1beba93cc8934afa0b456fa0034f5166c1ae2887 Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Thu, 14 Aug 2025 11:17:49 -0700 Subject: [PATCH] Fix - add safe divide by 0 for most places to prevent crash (#13624) --- litellm/litellm_core_utils/core_helpers.py | 21 ++++++ litellm/router_strategy/simple_shuffle.py | 7 +- .../litellm_core_utils/test_core_helpers.py | 72 ++++++++++++++++++- 3 files changed, 96 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 13a2e554f12..4aeb9d4d640 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -37,6 +37,27 @@ def safe_divide_seconds( return float(seconds / denominator) +def safe_divide( + numerator: Union[int, float], + denominator: Union[int, float], + default: Union[int, float] = 0 +) -> Union[int, float]: + """ + Safely divide two numbers, returning a default value if denominator is zero. + + Args: + numerator: The number to divide + denominator: The number to divide by + default: Value to return if denominator is zero (defaults to 0) + + Returns: + The result of numerator/denominator, or default if denominator is zero + """ + if denominator == 0: + return default + return numerator / denominator + + def map_finish_reason( finish_reason: str, ): # openai supports 5 stop sequences - 'stop', 'length', 'function_call', 'content_filter', 'null' diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index da24c02f2e3..bef145c44ba 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -9,6 +9,7 @@ import random from typing import TYPE_CHECKING, Any, Dict, List, Union from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.core_helpers import safe_divide if TYPE_CHECKING: from litellm.router import Router as _Router @@ -46,7 +47,7 @@ def simple_shuffle( 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] + weights = [safe_divide(weight, total_weight, 0) 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] @@ -63,7 +64,7 @@ def simple_shuffle( 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] + weights = [safe_divide(rpm, total_rpm, 0) 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] @@ -80,7 +81,7 @@ def simple_shuffle( 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] + weights = [safe_divide(tpm, total_tpm, 0) 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] diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index d7869e6b800..32f3ad3f55c 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -9,7 +9,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, safe_divide def test_get_litellm_metadata_from_kwargs(): @@ -57,3 +57,73 @@ def test_preserve_upstream_non_openai_attributes(): ) assert model_response.test_key == "test_value" + + +def test_safe_divide_basic(): + """Test basic safe division functionality""" + # Normal division + result = safe_divide(10, 2) + assert result == 5.0, f"Expected 5.0, got {result}" + + # Division with float + result = safe_divide(7.5, 2.5) + assert result == 3.0, f"Expected 3.0, got {result}" + + # Division by zero with default + result = safe_divide(10, 0) + assert result == 0, f"Expected 0, got {result}" + + # Division by zero with custom default + result = safe_divide(10, 0, default=1) + assert result == 1, f"Expected 1, got {result}" + + # Division by zero with custom default as float + result = safe_divide(10, 0, default=0.5) + assert result == 0.5, f"Expected 0.5, got {result}" + + +def test_safe_divide_edge_cases(): + """Test edge cases for safe division""" + # Zero numerator + result = safe_divide(0, 5) + assert result == 0.0, f"Expected 0.0, got {result}" + + # Negative numbers + result = safe_divide(-10, 2) + assert result == -5.0, f"Expected -5.0, got {result}" + + # Negative denominator + result = safe_divide(10, -2) + assert result == -5.0, f"Expected -5.0, got {result}" + + # Both negative + result = safe_divide(-10, -2) + assert result == 5.0, f"Expected 5.0, got {result}" + + # Float division + result = safe_divide(1, 3) + assert abs(result - 0.3333333333333333) < 1e-10, f"Expected ~0.333..., got {result}" + + +def test_safe_divide_weight_scenario(): + """Test safe division in the context of weight calculations""" + # Simulate weight calculation scenario + weights = [3, 7, 0, 2] + total_weight = sum(weights) # 12 + + # Normal case + normalized_weights = [safe_divide(w, total_weight) for w in weights] + expected = [0.25, 7/12, 0.0, 1/6] + + for i, (actual, exp) in enumerate(zip(normalized_weights, expected)): + assert abs(actual - exp) < 1e-10, f"Weight {i}: Expected {exp}, got {actual}" + + # Zero total weight scenario (division by zero) + zero_weights = [0, 0, 0] + zero_total = sum(zero_weights) # 0 + + # Should return default values (0) for all weights + normalized_zero_weights = [safe_divide(w, zero_total) for w in zero_weights] + expected_zero = [0, 0, 0] + + assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}"