diff --git a/litellm/router.py b/litellm/router.py index 6cbad19dbb3..d7a3e54fac0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -14079,21 +14079,19 @@ class Router: # Safe to check statically since `rpm` is a direct litellm_params setting, unlike # max_input_tokens which get_router_model_info() can also derive from the cost map. - _any_deployment_has_rpm = self.routing_strategy != "usage-based-routing-v2" and any( + any_deployment_has_rpm: Final = self.routing_strategy != "usage-based-routing-v2" and any( isinstance(litellm_params := d.get("litellm_params"), dict) and litellm_params.get("rpm") is not None for d in _returned_deployments ) ## get model group RPM ## - model_group_cache: dict[str, int] | None = None # mutable-ok: incremented below - if _any_deployment_has_rpm: - dt = get_utc_datetime() - current_minute = dt.strftime("%H-%M") - rpm_key = f"{model}:rpm:{current_minute}" - cached_model_group: Final = self.cache.get_cache( - key=rpm_key, local_only=True, parent_otel_span=parent_otel_span - ) - model_group_cache = cached_model_group or {} # mutable-ok: updated below + rpm_key: Final = f"{model}:rpm:{get_utc_datetime().strftime('%H-%M')}" + cached_model_group: Final = ( + self.cache.get_cache(key=rpm_key, local_only=True, parent_otel_span=parent_otel_span) + if any_deployment_has_rpm + else None + ) + model_group_cache: Final = cached_model_group if isinstance(cached_model_group, dict) else {} for idx, deployment in enumerate(_returned_deployments): # Cache nested dict access to avoid repeated temporary dict allocations _litellm_params = deployment.get("litellm_params", {}) @@ -14117,7 +14115,7 @@ class Router: continue ## RPM CHECK ## - if _any_deployment_has_rpm and isinstance(model_group_cache, dict): + if any_deployment_has_rpm: model_id = _model_info.get("id", "") ### get local router cache ### current_request_cache_local = ( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4a9dd199c08..40d60214b4c 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -24748,6 +24748,158 @@ class TestPreCallChecksOptimization: assert deployments[0].get("model_info", {}).get("id") == "small", "First deployment ID changed!" assert deployments[1].get("model_info", {}).get("id") == "large", "Second deployment ID changed!" + +def _rpm_cache_keys_read_by(router: Router, model: str, deployments: list) -> list: + get_cache_calls: list = [] + original_get_cache = router.cache.get_cache + + def _tracking_get_cache(*args, **kwargs): + get_cache_calls.append(kwargs.get("key", args[0] if args else None)) + return original_get_cache(*args, **kwargs) + + router.cache.get_cache = _tracking_get_cache + router._pre_call_checks( + model=model, + healthy_deployments=deployments, + messages=[{"role": "user", "content": "hi"}], + ) + return get_cache_calls + + +def test_no_rpm_cache_read_when_no_deployment_has_rpm(): + router = Router( + model_list=[ + { + "model_name": "test", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-test"}, + "model_info": {"id": "no-rpm-1"}, + }, + { + "model_name": "test", + "litellm_params": {"model": "gpt-5.5", "api_key": "sk-test"}, + "model_info": {"id": "no-rpm-2"}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + ) + deployments = router.get_model_list(model_name="test") + assert deployments is not None + + get_cache_calls = _rpm_cache_keys_read_by(router, "test", deployments) + + assert get_cache_calls == [], f"Expected zero cache reads, got {get_cache_calls}" + + +def test_rpm_cache_still_read_and_enforced_when_a_deployment_has_rpm(): + router = Router( + model_list=[ + { + "model_name": "test", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-test", "rpm": 1}, + "model_info": {"id": "low-rpm"}, + }, + { + "model_name": "test", + "litellm_params": {"model": "gpt-5.5", "api_key": "sk-test"}, + "model_info": {"id": "no-rpm"}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + ) + deployments = router.get_model_list(model_name="test") + assert deployments is not None + router.cache.set_cache(key="low-rpm", value=1, local_only=True) + + filtered = router._pre_call_checks( + model="test", + healthy_deployments=deployments, + messages=[{"role": "user", "content": "hi"}], + ) + + assert [d["model_info"]["id"] for d in filtered] == ["no-rpm"] + + +def test_no_rpm_cache_read_for_usage_based_routing_v2(): + router = Router( + model_list=[ + { + "model_name": "test", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-test", "rpm": 1}, + "model_info": {"id": "low-rpm"}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + routing_strategy="usage-based-routing-v2", + ) + deployments = router.get_model_list(model_name="test") + assert deployments is not None + + get_cache_calls = _rpm_cache_keys_read_by(router, "test", deployments) + + assert get_cache_calls == [], f"Expected zero cache reads under usage-based-routing-v2, got {get_cache_calls}" + + +def test_token_counting_failure_returns_original_deployments(monkeypatch): + router = Router( + model_list=[ + { + "model_name": "test", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-test"}, + "model_info": {"id": "dep-1", "max_input_tokens": 100}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + ) + + def _raise_token_counting(*args, **kwargs): + raise ValueError("boom") + + monkeypatch.setattr(router, "_count_pre_call_check_tokens", _raise_token_counting) + deployments = router.get_model_list(model_name="test") + assert deployments is not None + + result = router._pre_call_checks( + model="test", + healthy_deployments=deployments, + messages=[{"role": "user", "content": "hi"}], + ) + + assert result == deployments + + +def test_model_info_lookup_failure_is_swallowed(monkeypatch): + router = Router( + model_list=[ + { + "model_name": "test", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-test"}, + "model_info": {"id": "dep-1"}, + }, + ], + set_verbose=False, + enable_pre_call_checks=True, + ) + + def _raise_model_info(*args, **kwargs): + raise ValueError("boom") + + monkeypatch.setattr(router, "get_router_model_info", _raise_model_info) + deployments = router.get_model_list(model_name="test") + assert deployments is not None + + result = router._pre_call_checks( + model="test", + healthy_deployments=deployments, + messages=[{"role": "user", "content": "hi"}], + ) + + assert [d["model_info"]["id"] for d in result] == ["dep-1"] + + if __name__ == "__main__": pytest.main([__file__, "-v"])