fix(router): keep the rpm cache skip inside the type rules

The old pre-call test file was deleted upstream, so the regression coverage now sits with the other router pre-call tests.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
TobiMayr 2026-10-14 12:21:29 +01:00
parent 0e582644c4
commit 01e81369bb
2 changed files with 161 additions and 11 deletions

View file

@ -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 = (

View file

@ -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"])