From d076743146c3ce3b23fd7cfeb8bb6192ef05c31d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Tue, 15 Sep 2026 13:35:44 -0700 Subject: [PATCH] fix(router): reject non-finite selective head outputs --- .../complexity_router/selective_policy.py | 5 ++++- .../router_strategy/test_selective_policy.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/router_strategy/complexity_router/selective_policy.py b/litellm/router_strategy/complexity_router/selective_policy.py index 0e8f217f34c..afe0bf93b07 100644 --- a/litellm/router_strategy/complexity_router/selective_policy.py +++ b/litellm/router_strategy/complexity_router/selective_policy.py @@ -76,10 +76,13 @@ class SelectiveHead(BaseModel): return self def linear(self, values: Mapping[str, float]) -> tuple[float, ...]: - return tuple( + outputs: Final = tuple( sum(coefficient * values.get(feature, 0.0) for feature, coefficient in zip(self.features, row)) + intercept for row, intercept in zip(self.coefficients, self.intercept) ) + if not all(math.isfinite(value) for value in outputs): + raise OverflowError("Selective head produced a non-finite value") + return outputs def probabilities(self, values: Mapping[str, float]) -> tuple[float, ...]: if self.constant is not None: diff --git a/tests/test_litellm/router_strategy/test_selective_policy.py b/tests/test_litellm/router_strategy/test_selective_policy.py index 429ec8e1179..3075b7b0112 100644 --- a/tests/test_litellm/router_strategy/test_selective_policy.py +++ b/tests/test_litellm/router_strategy/test_selective_policy.py @@ -72,6 +72,21 @@ def test_overflow_falls_back_to_the_capable_model() -> None: assert not policy.evaluate({}, 0.5, 0.9).use_efficient +@pytest.mark.parametrize("coefficient", [1e308, -1e308]) +def test_overflowing_binary_logits_cannot_hide_as_equal_probabilities(coefficient: float) -> None: + head: Final = SelectiveHead( + features=("logit_p",), coefficients=((coefficient,),), intercept=(0.0,), classes=(0, 1) + ) + policy: Final = SelectivePolicy( + version="test", + feature_schema="v2-v1", + target="scalar_calibration", + threshold=0.05, + heads=(head, head), + ) + assert not policy.evaluate({}, 0.99999, 0.99999).use_efficient + + def test_per_model_policy_uses_both_fitted_heads() -> None: efficient: Final = SelectiveHead( features=("p_e",), coefficients=((0.0,),), intercept=(math.log(4.0),), classes=(0, 1)